feat(java): support add columns via reader (#3456)

Author: yanghua

java/core/lance-jni/src/blocking_dataset.rs

@@ -1057,3 +1057,50 @@ fn inner_add_columns_by_sql_expressions(
     )?;
     Ok(())
 }
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeAddColumnsByReader(
+    mut env: JNIEnv,
+    java_dataset: JObject,
+    arrow_array_stream_addr: jlong,
+    batch_size: JObject, // Optional<Long>
+) {
+    ok_or_throw_without_return!(
+        env,
+        inner_add_columns_by_reader(&mut env, java_dataset, arrow_array_stream_addr, batch_size)
+    )
+}
+
+fn inner_add_columns_by_reader(
+    env: &mut JNIEnv,
+    java_dataset: JObject,
+    arrow_array_stream_addr: jlong,
+    batch_size: JObject, // Optional<Long>
+) -> Result<()> {
+    let stream_ptr = arrow_array_stream_addr as *mut FFI_ArrowArrayStream;
+
+    let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?;
+
+    let transform = NewColumnTransform::Reader(Box::new(reader));
+
+    let batch_size = if env.call_method(&batch_size, "isPresent", "()Z", &[])?.z()? {
+        let batch_size_value = env.get_long_opt(&batch_size)?;
+        match batch_size_value {
+            Some(value) => Some(
+                value
+                    .try_into()
+                    .map_err(|_| Error::input_error("Batch size conversion error".to_string()))?,
+            ),
+            None => None,
+        }
+    } else {
+        None
+    };
+
+    let mut dataset_guard =
+        unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+
+    RT.block_on(dataset_guard.inner.add_columns(transform, None, batch_size))?;
+
+    Ok(())
+}

java/core/src/main/java/com/lancedb/lance/Dataset.java

@@ -287,6 +287,23 @@ public void addColumns(SqlExpressions sqlExpressions, Optional<Long> batchSize)
   private native void nativeAddColumnsBySqlExpressions(
       SqlExpressions sqlExpressions, Optional<Long> batchSize);
 
+  /**
+   * Add columns to the dataset.
+   *
+   * @param stream The Arrow Array Stream generated by arrow reader to add columns.
+   * @param batchSize The number of rows to read at a time from the source dataset when applying the
+   *     transform.
+   */
+  public void addColumns(ArrowArrayStream stream, Optional<Long> batchSize) {
+    try (LockManager.WriteLock writeLock = lockManager.acquireWriteLock()) {
+      Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+      nativeAddColumnsByReader(stream.memoryAddress(), batchSize);
+    }
+  }
+
+  private native void nativeAddColumnsByReader(
+      long arrowStreamMemoryAddress, Optional<Long> batchSize);
+
   /**
    * Drop columns from the dataset.
    *

java/core/src/test/java/com/lancedb/lance/DatasetTest.java

@@ -13,11 +13,16 @@
  */
 package com.lancedb.lance;
 
+import com.lancedb.lance.ipc.LanceScanner;
 import com.lancedb.lance.schema.ColumnAlteration;
 import com.lancedb.lance.schema.SqlExpressions;
 
+import org.apache.arrow.c.ArrowArrayStream;
+import org.apache.arrow.c.Data;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowReader;
 import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -351,6 +356,107 @@ void testAddColumnBySqlExpressions() {
     }
   }
 
+  @Test
+  void testAddColumnsByStream() throws IOException {
+    String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName();
+    String datasetPath = tempDir.resolve(testMethodName).toString();
+    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      TestUtils.SimpleTestDataset testDataset =
+          new TestUtils.SimpleTestDataset(allocator, datasetPath);
+
+      try (Dataset initialDataset = testDataset.createEmptyDataset()) {
+        try (Dataset datasetV1 = testDataset.write(1, 3)) {
+          assertEquals(3, datasetV1.countRows());
+        }
+      }
+
+      dataset = Dataset.open(datasetPath, allocator);
+
+      Schema newColumnSchema =
+          new Schema(
+              Collections.singletonList(Field.nullable("age", new ArrowType.Int(32, true))), null);
+
+      try (VectorSchemaRoot vector = VectorSchemaRoot.create(newColumnSchema, allocator);
+          ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
+
+        IntVector ageVector = (IntVector) vector.getVector("age");
+        ageVector.allocateNew(3);
+        ageVector.set(0, 25);
+        ageVector.set(1, 30);
+        ageVector.set(2, 35);
+        vector.setRowCount(3);
+
+        class SimpleVectorReader extends ArrowReader {
+          private boolean batchLoaded = false;
+
+          protected SimpleVectorReader(BufferAllocator allocator) {
+            super(allocator);
+          }
+
+          @Override
+          public boolean loadNextBatch() {
+            if (!batchLoaded) {
+              batchLoaded = true;
+              return true;
+            }
+            return false;
+          }
+
+          @Override
+          public VectorSchemaRoot getVectorSchemaRoot() {
+            return vector;
+          }
+
+          @Override
+          public long bytesRead() {
+            return vector.getFieldVectors().stream().mapToLong(FieldVector::getBufferSize).sum();
+          }
+
+          @Override
+          protected void closeReadSource() {}
+
+          @Override
+          protected Schema readSchema() {
+            return newColumnSchema;
+          }
+        }
+
+        try (ArrowReader reader = new SimpleVectorReader(allocator)) {
+          Data.exportArrayStream(allocator, reader, stream);
+
+          dataset.addColumns(stream, Optional.of(3L));
+
+          Schema expectedSchema =
+              new Schema(
+                  Arrays.asList(
+                      Field.nullable("id", new ArrowType.Int(32, true)),
+                      Field.nullable("name", new ArrowType.Utf8()),
+                      Field.nullable("age", new ArrowType.Int(32, true))),
+                  null);
+          Schema actualSchema = dataset.getSchema();
+          assertEquals(expectedSchema.getFields(), actualSchema.getFields());
+
+          try (LanceScanner scanner = dataset.newScan()) {
+            try (ArrowReader resultReader = scanner.scanBatches()) {
+              assertTrue(resultReader.loadNextBatch());
+              VectorSchemaRoot root = resultReader.getVectorSchemaRoot();
+              assertEquals(3, root.getRowCount());
+
+              IntVector idVector = (IntVector) root.getVector("id");
+              IntVector ageVectorResult = (IntVector) root.getVector("age");
+              for (int i = 0; i < 3; i++) {
+                assertEquals(i, idVector.get(i));
+                assertEquals(25 + i * 5, ageVectorResult.get(i));
+              }
+            }
+          }
+        }
+      }
+    } catch (Exception e) {
+      fail("Exception occurred during test: " + e.getMessage(), e);
+    }
+  }
+
   @Test
   void testDropPath() {
     String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName();