Update ColumnarBatchTest.java to wrap function

apache · Aug 19, 2024 · 96de783 · 96de783
1 parent 7bafde5
commit 96de783
Showing 1 changed file with 18 additions and 22 deletions.
diff --git a/backends-velox/src/test/java/org/apache/gluten/columnarbatch/ColumnarBatchTest.java b/backends-velox/src/test/java/org/apache/gluten/columnarbatch/ColumnarBatchTest.java
@@ -100,7 +100,14 @@ public void testOffloadAndLoadReadRow() {
     TaskResources$.MODULE$.runUnsafe(
         () -> {
           final int numRows = 20;
-          final ColumnarBatch batch = newArrowBatch(numRows);
+          final ColumnarBatch batch = newArrowBatch("a boolean, b int", numRows);
+          final ArrowWritableColumnVector col0 = (ArrowWritableColumnVector) batch.column(0);
+          final ArrowWritableColumnVector col1 = (ArrowWritableColumnVector) batch.column(1);
+          for (int j = 0; j < numRows; j++) {
+            col0.putBoolean(j, j % 2 == 0);
+            col1.putInt(j, 15 - j);
+          }
+          col1.putNull(numRows - 1);
           Assert.assertTrue(ColumnarBatches.isHeavyBatch(batch));
           final ColumnarBatch offloaded =
               ColumnarBatches.ensureOffloaded(ArrowBufferAllocators.contextInstance(), batch);
@@ -125,13 +132,21 @@ public void testOffloadAndLoadReadRow() {
   public void testToString() {
     TaskResources$.MODULE$.runUnsafe(
         () -> {
-          final ColumnarBatch batch = newArrowBatch(20);
+          final int numRows = 20;
+          final ColumnarBatch batch = newArrowBatch("a boolean, b int", numRows);
+          final ArrowWritableColumnVector col0 = (ArrowWritableColumnVector) batch.column(0);
+          final ArrowWritableColumnVector col1 = (ArrowWritableColumnVector) batch.column(1);
+          for (int j = 0; j < numRows; j++) {
+            col0.putBoolean(j, j % 2 == 0);
+            col1.putInt(j, 15 - j);
+          }
+          col1.putNull(numRows - 1);
           StructType structType = new StructType();
           structType = structType.add("a", DataTypes.BooleanType, true);
           structType = structType.add("b", DataTypes.IntegerType, true);
           ColumnarBatch veloxBatch =
               RowToVeloxColumnarExec.toColumnarBatchIterator(
-                      JavaConverters.asScalaIterator(batch.rowIterator()), structType, 20)
+                      JavaConverters.asScalaIterator(batch.rowIterator()), structType, numRows)
                   .next();
           Assert.assertEquals("[true,15]\n[false,14]", ColumnarBatches.toString(veloxBatch, 0, 2));
           Assert.assertEquals(
@@ -151,23 +166,4 @@ private static ColumnarBatch newArrowBatch(String schema, int numRows) {
     batch.setNumRows(numRows);
     return batch;
   }
-
-  private static ColumnarBatch newArrowBatch(int numRows) {
-    String schema = "a boolean, b int";
-    final ArrowWritableColumnVector[] columns =
-        ArrowWritableColumnVector.allocateColumns(numRows, StructType.fromDDL(schema));
-    ArrowWritableColumnVector col1 = columns[0];
-    ArrowWritableColumnVector col2 = columns[1];
-    for (int j = 0; j < numRows; j++) {
-      col1.putBoolean(j, j % 2 == 0);
-      col2.putInt(j, 15 - j);
-    }
-    col2.putNull(numRows - 1);
-    for (ArrowWritableColumnVector col : columns) {
-      col.setValueCount(numRows);
-    }
-    final ColumnarBatch batch = new ColumnarBatch(columns);
-    batch.setNumRows(numRows);
-    return batch;
-  }
 }