From 96de7833855b6da0d93aae4ce5c6d157f2b24a02 Mon Sep 17 00:00:00 2001
From: Jin Chengcheng <chengcheng.jin@intel.com>
Date: Mon, 19 Aug 2024 14:41:15 +0800
Subject: [PATCH] Update ColumnarBatchTest.java to wrap function

---
 .../columnarbatch/ColumnarBatchTest.java      | 40 +++++++++----------
 1 file changed, 18 insertions(+), 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
index 298f11073bff..54994ccd4836 100644
--- 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;
-  }
 }