Update ShortDecimalType to use LongArrayBlock

soscpd · Jul 27, 2016 · b7fcb29 · b7fcb29
1 parent 11fca35
commit b7fcb29
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 5 deletions.
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/type/DecimalType.java b/presto-spi/src/main/java/com/facebook/presto/spi/type/DecimalType.java
@@ -25,7 +25,8 @@
 import static java.util.Collections.unmodifiableList;
 
 public abstract class DecimalType
-        extends AbstractFixedWidthType
+        extends AbstractType
+        implements FixedWidthType
 {
     public static final int DEFAULT_SCALE = 0;
     public static final int DEFAULT_PRECISION = MAX_PRECISION;
@@ -53,9 +54,9 @@ public static DecimalType createDecimalType()
     private final int precision;
     private final int scale;
 
-    DecimalType(int precision, int scale, Class<?> javaType, int fixedSize)
+    DecimalType(int precision, int scale, Class<?> javaType)
     {
-        super(new TypeSignature(StandardTypes.DECIMAL, buildTypeParameters(precision, scale)), javaType, fixedSize);
+        super(new TypeSignature(StandardTypes.DECIMAL, buildTypeParameters(precision, scale)), javaType);
         this.precision = precision;
         this.scale = scale;
     }

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/type/LongDecimalType.java b/presto-spi/src/main/java/com/facebook/presto/spi/type/LongDecimalType.java
@@ -17,6 +17,8 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.block.Block;
 import com.facebook.presto.spi.block.BlockBuilder;
+import com.facebook.presto.spi.block.BlockBuilderStatus;
+import com.facebook.presto.spi.block.FixedWidthBlockBuilder;
 import io.airlift.slice.Slice;
 
 import static com.facebook.presto.spi.type.Decimals.MAX_PRECISION;
@@ -30,10 +32,37 @@ final class LongDecimalType
 
     LongDecimalType(int precision, int scale)
     {
-        super(precision, scale, Slice.class, SIZE_OF_LONG_DECIMAL);
+        super(precision, scale, Slice.class);
         validatePrecisionScale(precision, scale, MAX_PRECISION);
     }
 
+    @Override
+    public int getFixedSize()
+    {
+        return SIZE_OF_LONG_DECIMAL;
+    }
+
+    @Override
+    public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
+    {
+        return new FixedWidthBlockBuilder(
+                getFixedSize(),
+                blockBuilderStatus,
+                Math.min(expectedEntries, blockBuilderStatus.getMaxBlockSizeInBytes() / getFixedSize()));
+    }
+
+    @Override
+    public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return createBlockBuilder(blockBuilderStatus, expectedEntries, getFixedSize());
+    }
+
+    @Override
+    public BlockBuilder createFixedSizeBlockBuilder(int positionCount)
+    {
+        return new FixedWidthBlockBuilder(getFixedSize(), positionCount);
+    }
+
     @Override
     public Object getObjectValue(ConnectorSession session, Block block, int position)
     {
@@ -92,4 +121,10 @@ public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int l
     {
         blockBuilder.writeBytes(value, offset, length).closeEntry();
     }
+
+    @Override
+    public Slice getSlice(Block block, int position)
+    {
+        return block.getSlice(position, 0, getFixedSize());
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/type/ShortDecimalType.java b/presto-spi/src/main/java/com/facebook/presto/spi/type/ShortDecimalType.java
@@ -17,6 +17,8 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.block.Block;
 import com.facebook.presto.spi.block.BlockBuilder;
+import com.facebook.presto.spi.block.BlockBuilderStatus;
+import com.facebook.presto.spi.block.LongArrayBlockBuilder;
 
 import java.math.BigInteger;
 
@@ -28,10 +30,36 @@ final class ShortDecimalType
 {
     ShortDecimalType(int precision, int scale)
     {
-        super(precision, scale, long.class, SIZE_OF_LONG);
+        super(precision, scale, long.class);
         validatePrecisionScale(precision, scale, MAX_SHORT_PRECISION);
     }
 
+    @Override
+    public int getFixedSize()
+    {
+        return SIZE_OF_LONG;
+    }
+
+    @Override
+    public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
+    {
+        return new LongArrayBlockBuilder(
+                blockBuilderStatus,
+                Math.min(expectedEntries, blockBuilderStatus.getMaxBlockSizeInBytes() / getFixedSize()));
+    }
+
+    @Override
+    public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return createBlockBuilder(blockBuilderStatus, expectedEntries, getFixedSize());
+    }
+
+    @Override
+    public BlockBuilder createFixedSizeBlockBuilder(int positionCount)
+    {
+        return new LongArrayBlockBuilder(new BlockBuilderStatus(), positionCount);
+    }
+
     @Override
     public Object getObjectValue(ConnectorSession session, Block block, int position)
     {