[REVIEW] Initial work for decimal type in Java/JNI [skip ci]

java/src/main/java/ai/rapids/cudf/DType.java

@@ -65,6 +65,8 @@ public enum DType {
 
   STRING(0, 23, "str"),
   LIST(0, 24, "list"),
+  DECIMAL32(4, 25, "decimal32"),
+  DECIMAL64(8, 26, "decimal64"),
   STRUCT(0, 27, "struct");
 
   private static final DType[] D_TYPES = DType.values();
@@ -119,6 +121,23 @@ public boolean isBackedByLong() {
     return LONGS.contains(this);
   }
 
+  /**
+   * Returns true if this type is backed by short type
+   * Namely this method will return true for the following types
+   *       DType.INT16,
+   *       DType.UINT16
+   */
+  public boolean isBackedByShort() { return SHORTS.contains(this); }
+
+  /**
+   * Returns true if this type is backed by byte type
+   * Namely this method will return true for the following types
+   *       DType.INT8,
+   *       DType.UINT8,
+   *       DType.BOOL8
+   */
+  public boolean isBackedByByte() { return BYTES.contains(this); }
+
   /**
    * Returns true for duration types
    */
@@ -196,6 +215,17 @@ public static DType fromNative(int nativeId) {
       DType.TIMESTAMP_DAYS
   );
 
+  private static final EnumSet<DType> SHORTS = EnumSet.of(
+      DType.INT16,
+      DType.UINT16
+  );
+
+  private static final EnumSet<DType> BYTES = EnumSet.of(
+          DType.INT8,
+          DType.UINT8,
+          DType.BOOL8
+  );
+
   private static final EnumSet<DType> NESTED_TYPE = EnumSet.of(
       DType.LIST,
       DType.STRUCT

java/src/main/java/ai/rapids/cudf/DataType.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ai.rapids.cudf;
+
+public class DataType {
+
+  final DType typeId;
+  int scale = 0;
+
+  public DataType(DType id) { typeId = id; }
+
+  public DataType(DType id, int decimalScale) {
+    typeId = id;
+    scale = decimalScale;
+  }
+
+  public boolean isTimestamp() { return typeId.isTimestamp();}
+
+  public boolean hasTimeResolution() { return typeId.hasTimeResolution(); }
+
+  public boolean isBackedByInt() { return typeId.isBackedByInt(); }
+
+  public boolean isBackedByLong() { return typeId.isBackedByLong(); }
+
+  public boolean isBackedByShort() { return typeId.isBackedByShort(); }
+
+  public boolean isBackedByByte() { return typeId.isBackedByByte(); }
+
+  public boolean isNestedType() { return typeId.isNestedType(); }
+
+  public int getNativeId() { return typeId.getNativeId(); }
+
+}

java/src/main/java/ai/rapids/cudf/HostColumnVector.java

@@ -50,6 +50,11 @@ public final class HostColumnVector extends HostColumnVectorCore {
     this(type, rows, nullCount, hostDataBuffer, hostValidityBuffer, null);
   }
 
+  HostColumnVector(ai.rapids.cudf.DataType type, long rows, Optional<Long> nullCount,
+                   HostMemoryBuffer hostDataBuffer, HostMemoryBuffer hostValidityBuffer) {
+    this(type, rows, nullCount, hostDataBuffer, hostValidityBuffer, null);
+  }
+
   /**
    * Create a new column vector with data populated on the host.
    * @param type               the type of the vector
@@ -89,6 +94,21 @@ public final class HostColumnVector extends HostColumnVectorCore {
     incRefCountInternal(true);
   }
 
+  HostColumnVector(ai.rapids.cudf.DataType type, long rows, Optional<Long> nullCount,
+                   HostMemoryBuffer hostDataBuffer, HostMemoryBuffer hostValidityBuffer,
+                   HostMemoryBuffer offsetBuffer) {
+    super(type, rows, nullCount, hostDataBuffer, hostValidityBuffer, offsetBuffer, new ArrayList<>());
+    assert type.typeId != DType.LIST : "This constructor should not be used for list type";
+    if (nullCount.isPresent() && nullCount.get() > 0 && hostValidityBuffer == null) {
+      throw new IllegalStateException("Buffer cannot have a nullCount without a validity buffer");
+    }
+    if (type.typeId != DType.STRING && type.typeId != DType.LIST) {
+      assert offsetBuffer == null : "offsets are only supported for STRING and LIST";
+    }
+    refCount = 0;
+    incRefCountInternal(true);
+  }
+
   /**
    * This is a really ugly API, but it is possible that the lifecycle of a column of
    * data may not have a clear lifecycle thanks to java and GC. This API informs the leak
@@ -175,8 +195,8 @@ public ColumnVector copyToDevice() {
       if (!type.isNestedType()) {
         HostMemoryBuffer hdata = this.offHeap.data;
         if (hdata != null) {
-          long dataLen = rows * type.sizeInBytes;
-          if (type == DType.STRING) {
+          long dataLen = rows * type.typeId.sizeInBytes;
+          if (type.typeId == DType.STRING) {
             // This needs a different type
             dataLen = getEndStringOffset(rows - 1);
             if (dataLen == 0 && getNullCount() == 0) {
@@ -239,6 +259,10 @@ public static Builder builder(DType type, int rows) {
     return new Builder(type, rows, 0);
   }
 
+  public static Builder builder(ai.rapids.cudf.DataType type, int rows) {
+    return new Builder(type, rows, 0);
+  }
+
   /**
    * Create a new Builder to hold the specified number of rows and with enough space to hold the
    * given amount of string data. Be sure to close the builder when done with it. Please try to
@@ -1074,7 +1098,7 @@ public String toString() {
 
   public static final class Builder implements AutoCloseable {
     private final long rows;
-    private final DType type;
+    private final ai.rapids.cudf.DataType type;
     private HostMemoryBuffer data;
     private HostMemoryBuffer valid;
     private HostMemoryBuffer offsets;
@@ -1091,7 +1115,7 @@ public static final class Builder implements AutoCloseable {
      *                         working with Strings.  It is ignored otherwise.
      */
     Builder(DType type, long rows, long stringBufferSize) {
-      this.type = type;
+      this.type = new ai.rapids.cudf.DataType(type);
       this.rows = rows;
       if (type == DType.STRING) {
         if (stringBufferSize <= 0) {
@@ -1108,6 +1132,24 @@ public static final class Builder implements AutoCloseable {
       }
     }
 
+    Builder(ai.rapids.cudf.DataType type, long rows, long stringBufferSize) {
+      this.type = type;
+      this.rows = rows;
+      if (type.typeId == DType.STRING) {
+        if (stringBufferSize <= 0) {
+          // We need at least one byte or we will get NULL back for data
+          stringBufferSize = 1;
+        }
+        this.data = HostMemoryBuffer.allocate(stringBufferSize);
+        // The offsets are ints and there is 1 more than the number of rows.
+        this.offsets = HostMemoryBuffer.allocate((rows + 1) * OFFSET_SIZE);
+        // The first offset is always 0
+        this.offsets.setInt(0, 0);
+      } else {
+        this.data = HostMemoryBuffer.allocate(rows * type.typeId.sizeInBytes);
+      }
+    }
+
     /**
      * Create a builder with a buffer of size rows (for testing ONLY).
      * @param type       datatype
@@ -1119,72 +1161,72 @@ public static final class Builder implements AutoCloseable {
      */
     Builder(DType type, long rows, HostMemoryBuffer testData,
             HostMemoryBuffer testValid, HostMemoryBuffer testOffsets) {
-      this.type = type;
+      this.type = new ai.rapids.cudf.DataType(type);
       this.rows = rows;
       this.data = testData;
       this.valid = testValid;
     }
 
     public final Builder append(boolean value) {
-      assert type == DType.BOOL8;
+      assert type.typeId == DType.BOOL8;
       assert currentIndex < rows;
-      data.setByte(currentIndex * type.sizeInBytes, value ? (byte)1 : (byte)0);
+      data.setByte(currentIndex * type.typeId.sizeInBytes, value ? (byte)1 : (byte)0);
       currentIndex++;
       return this;
     }
 
     public final Builder append(byte value) {
-      assert type == DType.INT8 || type == DType.UINT8 || type == DType.BOOL8;
+      assert type.typeId == DType.INT8 || type.typeId == DType.UINT8 || type.typeId == DType.BOOL8;
       assert currentIndex < rows;
-      data.setByte(currentIndex * type.sizeInBytes, value);
+      data.setByte(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
 
     public final Builder append(byte value, long count) {
       assert (count + currentIndex) <= rows;
-      assert type == DType.INT8 || type == DType.UINT8 || type == DType.BOOL8;
-      data.setMemory(currentIndex * type.sizeInBytes, count, value);
+      assert type.typeId == DType.INT8 || type.typeId == DType.UINT8 || type.typeId == DType.BOOL8;
+      data.setMemory(currentIndex * type.typeId.sizeInBytes, count, value);
       currentIndex += count;
       return this;
     }
 
     public final Builder append(short value) {
-      assert type == DType.INT16 || type == DType.UINT16;
+      assert type.typeId == DType.INT16 || type.typeId == DType.UINT16;
       assert currentIndex < rows;
-      data.setShort(currentIndex * type.sizeInBytes, value);
+      data.setShort(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
 
     public final Builder append(int value) {
       assert type.isBackedByInt();
       assert currentIndex < rows;
-      data.setInt(currentIndex * type.sizeInBytes, value);
+      data.setInt(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
 
     public final Builder append(long value) {
       assert type.isBackedByLong();
       assert currentIndex < rows;
-      data.setLong(currentIndex * type.sizeInBytes, value);
+      data.setLong(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
 
     public final Builder append(float value) {
-      assert type == DType.FLOAT32;
+      assert type.typeId == DType.FLOAT32;
       assert currentIndex < rows;
-      data.setFloat(currentIndex * type.sizeInBytes, value);
+      data.setFloat(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
 
     public final Builder append(double value) {
-      assert type == DType.FLOAT64;
+      assert type.typeId == DType.FLOAT64;
       assert currentIndex < rows;
-      data.setDouble(currentIndex * type.sizeInBytes, value);
+      data.setDouble(currentIndex * type.typeId.sizeInBytes, value);
       currentIndex++;
       return this;
     }
@@ -1203,7 +1245,7 @@ public Builder appendUTF8String(byte[] value, int offset, int length) {
       assert offset >= 0;
       assert length >= 0;
       assert value.length + offset <= length;
-      assert type == DType.STRING;
+      assert type.typeId == DType.STRING;
       assert currentIndex < rows;
       // just for strings we want to throw a real exception if we would overrun the buffer
       long oldLen = data.getLength();
@@ -1239,48 +1281,48 @@ public Builder appendUTF8String(byte[] value, int offset, int length) {
 
     public Builder appendArray(byte... values) {
       assert (values.length + currentIndex) <= rows;
-      assert type == DType.INT8 || type == DType.UINT8 || type == DType.BOOL8;
-      data.setBytes(currentIndex * type.sizeInBytes, values, 0, values.length);
+      assert type.isBackedByByte();
+      data.setBytes(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
 
     public Builder appendArray(short... values) {
-      assert type == DType.INT16 || type == DType.UINT16;
+      assert type.isBackedByShort();
       assert (values.length + currentIndex) <= rows;
-      data.setShorts(currentIndex * type.sizeInBytes, values, 0, values.length);
+      data.setShorts(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
 
     public Builder appendArray(int... values) {
       assert type.isBackedByInt();
       assert (values.length + currentIndex) <= rows;
-      data.setInts(currentIndex * type.sizeInBytes, values, 0, values.length);
+      data.setInts(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
 
     public Builder appendArray(long... values) {
       assert type.isBackedByLong();
       assert (values.length + currentIndex) <= rows;
-      data.setLongs(currentIndex * type.sizeInBytes, values, 0, values.length);
+      data.setLongs(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
 
     public Builder appendArray(float... values) {
-      assert type == DType.FLOAT32;
+      assert type.typeId == DType.FLOAT32;
       assert (values.length + currentIndex) <= rows;
-      data.setFloats(currentIndex * type.sizeInBytes, values, 0, values.length);
+      data.setFloats(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
 
     public Builder appendArray(double... values) {
-      assert type == DType.FLOAT64;
+      assert type.typeId == DType.FLOAT64;
       assert (values.length + currentIndex) <= rows;
-      data.setDoubles(currentIndex * type.sizeInBytes, values, 0, values.length);
+      data.setDoubles(currentIndex * type.typeId.sizeInBytes, values, 0, values.length);
       currentIndex += values.length;
       return this;
     }
@@ -1429,15 +1471,15 @@ public final Builder appendBoxed(String... values) throws IndexOutOfBoundsExcept
      */
     public final Builder append(HostColumnVector columnVector) {
       assert columnVector.rows <= (rows - currentIndex);
-      assert columnVector.type == type;
+      assert columnVector.type.typeId == type.typeId;
 
-      if (type == DType.STRING) {
+      if (type.typeId == DType.STRING) {
         throw new UnsupportedOperationException(
             "Appending a string column vector client side is not currently supported");
       } else {
-        data.copyFromHostBuffer(currentIndex * type.sizeInBytes, columnVector.offHeap.data,
+        data.copyFromHostBuffer(currentIndex * type.typeId.sizeInBytes, columnVector.offHeap.data,
             0L,
-            columnVector.getRowCount() * type.sizeInBytes);
+            columnVector.getRowCount() * type.typeId.sizeInBytes);
       }
 
       //As this is doing the append on the host assume that a null count is available
@@ -1467,7 +1509,7 @@ private void allocateBitmaskAndSetDefaultValues() {
     public final Builder appendNull() {
       setNullAt(currentIndex);
       currentIndex++;
-      if (type == DType.STRING) {
+      if (type.typeId == DType.STRING) {
         offsets.setInt(currentIndex * OFFSET_SIZE, currentStringByteIndex);
       }
       return this;