frame writer test data

LakshSingla · Sep 7, 2023 · a5f24c5 · a5f24c5
1 parent 357d2ff
commit a5f24c5
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 0 deletions.
diff --git a/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java b/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java
@@ -57,6 +57,8 @@ public static ColumnType ofType(TypeSignature<ValueType> type)
         switch (type.getElementType().getType()) {
           case LONG:
             return ColumnType.LONG_ARRAY;
+          case FLOAT:
+            return ColumnType.FLOAT_ARRAY;
           case DOUBLE:
             return ColumnType.DOUBLE_ARRAY;
           case STRING:

diff --git a/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java b/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTest.java
@@ -153,6 +153,30 @@ public void test_long()
     testWithDataset(FrameWriterTestData.TEST_LONGS);
   }
 
+  @Test
+  public void test_arrayLong()
+  {
+    // TODO(laksh): add explanation
+    Assume.assumeFalse(inputFrameType == FrameType.COLUMNAR || outputFrameType == FrameType.COLUMNAR);
+    testWithDataset(FrameWriterTestData.TEST_ARRAYS_LONG);
+  }
+
+  @Test
+  public void test_arrayFloat()
+  {
+    // TODO(laksh): add explanation
+    Assume.assumeFalse(inputFrameType == FrameType.COLUMNAR || outputFrameType == FrameType.COLUMNAR);
+    testWithDataset(FrameWriterTestData.TEST_ARRAYS_FLOAT);
+  }
+
+  @Test
+  public void test_arrayDouble()
+  {
+    // TODO(laksh): add explanation
+    Assume.assumeFalse(inputFrameType == FrameType.COLUMNAR || outputFrameType == FrameType.COLUMNAR);
+    testWithDataset(FrameWriterTestData.TEST_ARRAYS_DOUBLE);
+  }
+
   @Test
   public void test_float()
   {
@@ -226,6 +250,13 @@ public void test_typePairs()
     // Test all possible arrangements of two different types.
     for (final FrameWriterTestData.Dataset<?> dataset1 : FrameWriterTestData.DATASETS) {
       for (final FrameWriterTestData.Dataset<?> dataset2 : FrameWriterTestData.DATASETS) {
+        if (dataset1.getType().isArray() && dataset1.getType().getElementType().isNumeric()
+            || dataset2.getType().isArray() && dataset2.getType().getElementType().isNumeric()) {
+          if (inputFrameType == FrameType.COLUMNAR || outputFrameType == FrameType.COLUMNAR) {
+            // TODO(laksh)
+            continue;
+          }
+        }
         final RowSignature signature = makeSignature(Arrays.asList(dataset1, dataset2));
         final Sequence<List<Object>> rowSequence = unsortAndMakeRows(Arrays.asList(dataset1, dataset2));
 
@@ -265,6 +296,7 @@ public void test_typePairs()
   public void test_insufficientWriteCapacity()
   {
     // Test every possible capacity, up to the amount required to write all items from every list.
+    Assume.assumeFalse(inputFrameType == FrameType.COLUMNAR || outputFrameType == FrameType.COLUMNAR);
     final RowSignature signature = makeSignature(FrameWriterTestData.DATASETS);
     final Sequence<List<Object>> rowSequence = unsortAndMakeRows(FrameWriterTestData.DATASETS);
     final int totalRows = rowSequence.toList().size();

diff --git a/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTestData.java b/processing/src/test/java/org/apache/druid/frame/write/FrameWriterTestData.java
@@ -136,6 +136,24 @@ public class FrameWriterTestData
       ).sorted(Comparators.naturalNullsFirst()).collect(Collectors.toList())
   );
 
+  public static final Dataset<Object> TEST_ARRAYS_LONG = new Dataset<>(
+      ColumnType.LONG_ARRAY,
+      Arrays.asList(
+          null,
+          ObjectArrays.EMPTY_ARRAY,
+          new Object[]{null},
+          new Object[]{null, 6L, null, 5L, null},
+          new Object[]{null, 6L, null, 5L, NullHandling.defaultLongValue()},
+          new Object[]{null, 6L, null, 5L, 0L, -1L},
+          new Object[]{null, 6L, null, 5L, 0L, -1L, Long.MIN_VALUE},
+          new Object[]{null, 6L, null, 5L, 0L, -1L, Long.MAX_VALUE},
+          new Object[]{5L},
+          new Object[]{5L, 6L},
+          new Object[]{5L, 6L, null},
+          new Object[]{Long.MAX_VALUE, Long.MIN_VALUE}
+      )
+  );
+
   public static final Dataset<Float> TEST_FLOATS = new Dataset<>(
       ColumnType.FLOAT,
       Stream.of(
@@ -158,6 +176,28 @@ public class FrameWriterTestData
       ).sorted(Comparators.naturalNullsFirst()).collect(Collectors.toList())
   );
 
+  //CHECKSTYLE.OFF: Regexp
+  public static final Dataset<Object> TEST_ARRAYS_FLOAT = new Dataset<>(
+      ColumnType.FLOAT_ARRAY,
+      Arrays.asList(
+          null,
+          ObjectArrays.EMPTY_ARRAY,
+          new Object[]{null},
+          new Object[]{null, 6.2f, null, 5.1f, null},
+          new Object[]{null, 6.2f, null, 5.1f, NullHandling.defaultFloatValue()},
+          new Object[]{null, 6.2f, null, 5.7f, 0.0f, -1.0f},
+          new Object[]{null, 6.2f, null, 5.7f, 0.0f, -1.0f, Float.MIN_VALUE},
+          new Object[]{null, 6.2f, null, 5.7f, 0.0f, -1.0f, Float.MAX_VALUE},
+          new Object[]{Float.NEGATIVE_INFINITY, Float.MIN_VALUE},
+          new Object[]{5.7f},
+          new Object[]{5.7f, 6.2f},
+          new Object[]{5.7f, 6.2f, null},
+          new Object[]{Float.MAX_VALUE, Float.MIN_VALUE},
+          new Object[]{Float.POSITIVE_INFINITY, Float.MIN_VALUE}
+      )
+  );
+  //CHECKSTYLE.ON: Regexp
+
   public static final Dataset<Double> TEST_DOUBLES = new Dataset<>(
       ColumnType.DOUBLE,
       Stream.of(
@@ -180,6 +220,28 @@ public class FrameWriterTestData
       ).sorted(Comparators.naturalNullsFirst()).collect(Collectors.toList())
   );
 
+  //CHECKSTYLE.OFF: Regexp
+  public static final Dataset<Object> TEST_ARRAYS_DOUBLE = new Dataset<>(
+      ColumnType.DOUBLE_ARRAY,
+      Arrays.asList(
+          null,
+          ObjectArrays.EMPTY_ARRAY,
+          new Object[]{null},
+          new Object[]{null, 6.2d, null, 5.1d, null},
+          new Object[]{null, 6.2d, null, 5.1d, NullHandling.defaultDoubleValue()},
+          new Object[]{null, 6.2d, null, 5.7d, 0.0d, -1.0d},
+          new Object[]{null, 6.2d, null, 5.7d, 0.0d, -1.0d, Double.MIN_VALUE},
+          new Object[]{null, 6.2d, null, 5.7d, 0.0d, -1.0d, Double.MAX_VALUE},
+          new Object[]{Double.NEGATIVE_INFINITY, Double.MIN_VALUE},
+          new Object[]{5.7d},
+          new Object[]{5.7d, 6.2d},
+          new Object[]{5.7d, 6.2d, null},
+          new Object[]{Double.MAX_VALUE, Double.MIN_VALUE},
+          new Object[]{Double.POSITIVE_INFINITY, Double.MIN_VALUE}
+      )
+  );
+  //CHECKSTYLE.ON: Regexp
+
   public static final Dataset<HyperLogLogCollector> TEST_COMPLEX = new Dataset<>(
       HyperUniquesAggregatorFactory.TYPE,
       Arrays.asList(
@@ -200,6 +262,9 @@ public class FrameWriterTestData
                    .add(TEST_STRINGS_SINGLE_VALUE)
                    .add(TEST_STRINGS_MULTI_VALUE)
                    .add(TEST_ARRAYS_STRING)
+                   .add(TEST_ARRAYS_LONG)
+                   .add(TEST_ARRAYS_FLOAT)
+                   .add(TEST_ARRAYS_DOUBLE)
                    .add(TEST_COMPLEX)
                    .build();