From 96f550c5cd8e8648ae87ab2624416ef4317948f5 Mon Sep 17 00:00:00 2001
From: Hongze Zhang <hongze.zhang@intel.com>
Date: Wed, 14 Aug 2024 14:32:08 +0800
Subject: [PATCH] update test

---
 .../columnarbatch/ColumnarBatchTest.java      | 54 ++++++++++++-------
 1 file changed, 34 insertions(+), 20 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 3388e75beeebd..956c75e4d528b 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
@@ -37,7 +37,7 @@ public void testOffloadAndLoad() {
     TaskResources$.MODULE$.runUnsafe(
         () -> {
           final int numRows = 100;
-          final ColumnarBatch batch = newArrowBatch(numRows);
+          final ColumnarBatch batch = newArrowBatch("a boolean, b int", numRows);
           Assert.assertTrue(ColumnarBatches.isHeavyBatch(batch));
           final ColumnarBatch offloaded =
               ColumnarBatches.ensureOffloaded(ArrowBufferAllocators.contextInstance(), batch);
@@ -52,7 +52,6 @@ public void testOffloadAndLoad() {
                       false)
                   .count();
           Assert.assertEquals(numRows, cnt);
-          Assert.assertEquals(loaded.getRow(0).getInt(1), 15);
           loaded.close();
           return null;
         });
@@ -92,28 +91,43 @@ public void testCreateByHandle() {
         });
   }
 
-  private static ColumnarBatch newArrowBatch(String schema, int numRows) {
-    final ArrowWritableColumnVector[] columns =
-        ArrowWritableColumnVector.allocateColumns(numRows, StructType.fromDDL(schema));
-    for (ArrowWritableColumnVector col : columns) {
-      col.setValueCount(numRows);
-    }
-    final ColumnarBatch batch = new ColumnarBatch(columns);
-    batch.setNumRows(numRows);
-    return batch;
+  @Test
+  public void testOffloadAndLoadWithWrites() {
+    TaskResources$.MODULE$.runUnsafe(
+        () -> {
+          final int numRows = 100;
+          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);
+          Assert.assertTrue(ColumnarBatches.isLightBatch(offloaded));
+          final ColumnarBatch loaded =
+              ColumnarBatches.ensureLoaded(ArrowBufferAllocators.contextInstance(), offloaded);
+          Assert.assertTrue(ColumnarBatches.isHeavyBatch(loaded));
+          long cnt =
+              StreamSupport.stream(
+                      Spliterators.spliteratorUnknownSize(
+                          loaded.rowIterator(), Spliterator.ORDERED),
+                      false)
+                  .count();
+          Assert.assertEquals(numRows, cnt);
+          Assert.assertEquals(loaded.getRow(0).getInt(1), 15);
+          loaded.close();
+          return null;
+        });
   }
 
-  private static ColumnarBatch newArrowBatch(int numRows) {
-    String schema = "a boolean, b int";
+
+  private static ColumnarBatch newArrowBatch(String schema, int numRows) {
     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);
     }