Add test

rapidsai · Jun 14, 2023 · a40e667 · a40e667
1 parent 4f09f43
commit a40e667
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java b/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
@@ -2877,7 +2877,7 @@ void testListConcatByRowIgnoreNull() {
   }
 
   @Test
-  void testFlattenLists() {
+  void testFlattenListsNoNulls() {
     HostColumnVector.ListType listType = new HostColumnVector.ListType(true,
         new HostColumnVector.BasicType(true, DType.INT32));
     HostColumnVector.ListType listOfListsType = new HostColumnVector.ListType(true, listType);
@@ -2893,6 +2893,23 @@ void testFlattenLists() {
     }
   }
 
+  @Test
+  void testFlattenListsWithNulls() {
+    HostColumnVector.ListType listType = new HostColumnVector.ListType(true,
+        new HostColumnVector.BasicType(true, DType.INT32));
+    HostColumnVector.ListType listOfListsType = new HostColumnVector.ListType(true, listType);
+
+    try (ColumnVector input = ColumnVector.fromLists(listOfListsType,
+        Arrays.asList(null, Arrays.asList(3), Arrays.asList(4, 5, 6)),
+        Arrays.asList(Arrays.asList(null, 8, 9), Arrays.asList(10, 11, 12, 13, 14, null)));
+         ColumnVector result = input.flattenLists(true);
+         ColumnVector expected = ColumnVector.fromLists(listType,
+             null,
+             Arrays.asList(null, 8, 9, 10, 11, 12, 13, 14, null))) {
+      assertColumnsAreEqual(expected, result);
+    }
+  }
+
   @Test
   void testPrefixSum() {
     try (ColumnVector v1 = ColumnVector.fromLongs(1, 2, 3, 5, 8, 10);