feat(sql): register additional tables for Dataset.sql queries

java/lance-jni/src/sql.rs

@@ -5,7 +5,7 @@ use crate::blocking_dataset::{BlockingDataset, NATIVE_DATASET};
 use crate::error::Result;
 use crate::traits::FromJString;
 use crate::{Error, JNIEnvExt, RT, block_on};
-use arrow::ffi_stream::FFI_ArrowArrayStream;
+use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
 use jni::JNIEnv;
 use jni::objects::{JClass, JObject, JString};
 use jni::sys::{JNI_TRUE, jboolean, jlong};
@@ -23,6 +23,8 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords(
     with_row_id: jboolean,
     with_row_addr: jboolean,
     stream_addr: jlong,
+    extra_table_names: JObject,
+    extra_stream_addrs: JObject,
 ) {
     ok_or_throw_without_return!(
         env,
@@ -34,11 +36,14 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords(
             with_row_id,
             with_row_addr,
             stream_addr,
+            extra_table_names,
+            extra_stream_addrs,
         )
         .map_err(|e| Error::input_error(e.to_string()))
     )
 }
 
+#[allow(clippy::too_many_arguments)]
 fn inner_into_batch_records(
     env: &mut JNIEnv,
     java_dataset: JObject,
@@ -47,6 +52,8 @@ fn inner_into_batch_records(
     with_row_id: jboolean,
     with_row_addr: jboolean,
     stream_addr: jlong,
+    extra_table_names: JObject,
+    extra_stream_addrs: JObject,
 ) -> Result<()> {
     let builder = sql_builder(
         env,
@@ -55,6 +62,8 @@ fn inner_into_batch_records(
         table_name,
         with_row_id,
         with_row_addr,
+        extra_table_names,
+        extra_stream_addrs,
     )?;
 
     let stream = block_on(async move {
@@ -69,16 +78,23 @@ fn inner_into_batch_records(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 fn sql_builder(
     env: &mut JNIEnv,
     java_dataset: JObject,
     sql: JString,
     table_name: JObject,
     with_row_id: jboolean,
     with_row_addr: jboolean,
+    extra_table_names: JObject,
+    extra_stream_addrs: JObject,
 ) -> Result<SqlQueryBuilder> {
     let sql_str = sql.extract(env)?;
     let table_str = env.get_string_opt(&table_name)?;
+    // Read every env-derived input before taking the dataset guard below, which holds a mutable borrow of `env`
+    // for its lifetime (a second `env` borrow while it is alive would not compile).
+    let names = env.get_strings(&extra_table_names)?;
+    let addrs = env.get_longs(&extra_stream_addrs)?;
 
     let dataset_guard =
         unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
@@ -93,5 +109,17 @@ fn sql_builder(
         builder = builder.table_name(table.as_str())
     }
 
+    // Register each caller-supplied relation (parallel name/address lists, read above). Each stream was exported
+    // to a C-Data stream on the Java side; import it to RecordBatches and register it under its name.
+    for (name, addr) in names.into_iter().zip(addrs) {
+        let stream_ptr = addr as *mut FFI_ArrowArrayStream;
+        let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr)? };
+        let mut batches = Vec::new();
+        for batch in reader {
+            batches.push(batch.map_err(|e| Error::input_error(e.to_string()))?);
+        }
+        builder = builder.register_arrow(name.as_str(), batches)?;
+    }
+
     Ok(builder)
 }

java/src/main/java/org/lance/SqlQuery.java

@@ -19,6 +19,8 @@
 import org.apache.arrow.vector.ipc.ArrowReader;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 
 public class SqlQuery {
@@ -29,6 +31,9 @@ public class SqlQuery {
   private String table = DEFAULT_TABLE_NAME;
   private boolean withRowId = false;
   private boolean withRowAddr = false;
+  private final List<String> extraTableNames = new ArrayList<>();
+  private final List<Long> extraStreamAddresses = new ArrayList<>();
+  private boolean consumed = false;
 
   public SqlQuery(Dataset dataset, String sql) {
     this.dataset = dataset;
@@ -40,6 +45,51 @@ public SqlQuery tableName(String tableName) {
     return this;
   }
 
+  /**
+   * Register an additional in-memory Arrow relation (exported to {@code stream} via the C Data
+   * Interface) as a table named {@code name}, joinable in the SQL alongside the dataset. {@link
+   * #intoBatchRecords()} consumes the stream during the native call (it takes ownership of the
+   * underlying C stream); the caller still owns the {@link ArrowArrayStream} handle and should
+   * close it afterwards (typically via try-with-resources), as with {@code MergeInsert}. May be
+   * called multiple times to register multiple tables.
+   *
+   * <p>The stream is fully materialized in memory before the query executes, so callers must keep
+   * the input bounded to available memory. For streaming or out-of-core relations, use a custom
+   * table provider through the Rust API instead of this Arrow convenience method.
+   *
+   * <p>{@code name} is parsed as a SQL identifier, so an unquoted name is lowercased ({@code IDs}
+   * registers the table {@code ids}) while a quoted one keeps its case ({@code "IDs"} registers the
+   * table {@code IDs}). If two names resolve to the same table, the later registration replaces the
+   * earlier one (deduplicated when the query is built). A name that resolves to the query's own
+   * {@link #tableName(String)} would hide the dataset, so {@link #intoBatchRecords()} rejects it
+   * with an {@link IllegalArgumentException}.
+   *
+   * <p>Because the registered stream is consumed on the first {@link #intoBatchRecords()} call, a
+   * query with registered relations is single-use; calling {@link #intoBatchRecords()} a second
+   * time throws.
+   *
+   * @throws IllegalArgumentException if {@code name} is null or blank, or {@code stream} is null
+   * @throws IllegalStateException if the query was already run via {@link #intoBatchRecords()}
+   */
+  public SqlQuery registerArrow(String name, ArrowArrayStream stream) {
+    if (consumed) {
+      throw new IllegalStateException(
+          "registerArrow cannot be called after intoBatchRecords(); build a new query and "
+              + "re-register relations.");
+    }
+    if (name == null || name.trim().isEmpty()) {
+      throw new IllegalArgumentException(
+          "registerArrow: table name must be non-empty, got: " + name);
+    }
+    if (stream == null) {
+      throw new IllegalArgumentException(
+          "registerArrow: stream must not be null (table name: " + name + ")");
+    }
+    this.extraTableNames.add(name);
+    this.extraStreamAddresses.add(stream.memoryAddress());
+    return this;
+  }
+
   public SqlQuery withRowId(boolean withRowId) {
     this.withRowId = withRowId;
     return this;
@@ -51,10 +101,29 @@ public SqlQuery withRowAddr(boolean withAddr) {
   }
 
   public ArrowReader intoBatchRecords() throws IOException {
+    if (consumed) {
+      throw new IllegalStateException(
+          "intoBatchRecords() was already called on this SqlQuery; a query with registered "
+              + "Arrow relations is single-use because the registered streams are consumed. "
+              + "Build a new query and re-register them.");
+    }
     try (LockManager.ReadLock readLock = dataset.acquireReadLock();
         ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) {
+      // A query with registered streams is one-shot: the native call below consumes them. Mark it
+      // consumed just before that call, so a failed output-stream allocation does not brick a
+      // retry.
+      if (!extraTableNames.isEmpty()) {
+        consumed = true;
+      }
       intoBatchRecords(
-          dataset, sql, Optional.ofNullable(table), withRowId, withRowAddr, s.memoryAddress());
+          dataset,
+          sql,
+          Optional.ofNullable(table),
+          withRowId,
+          withRowAddr,
+          s.memoryAddress(),
+          extraTableNames,
+          extraStreamAddresses);
       return Data.importArrayStream(dataset.allocator(), s);
     }
   }
@@ -65,7 +134,9 @@ private static native void intoBatchRecords(
       Optional<String> tableName,
       boolean withRowId,
       boolean withRowAddr,
-      long streamAddress)
+      long streamAddress,
+      List<String> extraTableNames,
+      List<Long> extraStreamAddresses)
       throws IOException;
 
   @Override

java/src/test/java/org/lance/SqlQueryTest.java

@@ -13,17 +13,32 @@
  */
 package org.lance;
 
+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.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.ipc.ArrowStreamReader;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 
 public class SqlQueryTest {
   private static final String NAME = "sqlquery_test_dataset";
@@ -108,4 +123,152 @@ public void testToRecordBatches() throws IOException {
         reader.getVectorSchemaRoot().getSchema().toString());
     reader.close();
   }
+
+  @Test
+  public void testRegisterArrow() throws Exception {
+    // An in-memory Arrow relation (column `id`) to semi-join against the dataset; 100 is absent.
+    Schema idSchema =
+        new Schema(Collections.singletonList(Field.nullable("id", new ArrowType.Int(32, true))));
+    try (VectorSchemaRoot ids = VectorSchemaRoot.create(idSchema, allocator)) {
+      ids.allocateNew();
+      IntVector idVector = (IntVector) ids.getVector("id");
+      int[] wanted = {1, 5, 39, 100};
+      for (int i = 0; i < wanted.length; i++) {
+        idVector.setSafe(i, wanted[i]);
+      }
+      ids.setRowCount(wanted.length);
+
+      // The native side consumes the stream; we allocate and close it (try-with-resources).
+      try (ArrowArrayStream stream = toStream(ids)) {
+        ArrowReader reader =
+            dataset
+                .sql(
+                    "select id from "
+                        + NAME
+                        + " where id in (select id from filter_ids) order by id")
+                .tableName(NAME)
+                .registerArrow("filter_ids", stream)
+                .intoBatchRecords();
+        List<Integer> got = new ArrayList<>();
+        while (reader.loadNextBatch()) {
+          VectorSchemaRoot root = reader.getVectorSchemaRoot();
+          for (int i = 0; i < root.getRowCount(); i++) {
+            got.add((Integer) root.getVector(0).getObject(i));
+          }
+        }
+        reader.close();
+        // 100 is not in the dataset; the rest are the intersection, in order.
+        Assertions.assertEquals(Arrays.asList(1, 5, 39), got);
+      }
+    }
+  }
+
+  @Test
+  public void testRegisterArrowMultiple() throws Exception {
+    // Register two relations in one query and join both; only ids present in all three survive.
+    try (VectorSchemaRoot a = idTable(1, 2, 3, 10);
+        VectorSchemaRoot b = idTable(2, 3, 4, 10);
+        ArrowArrayStream sa = toStream(a);
+        ArrowArrayStream sb = toStream(b)) {
+      ArrowReader reader =
+          dataset
+              .sql(
+                  "select id from "
+                      + NAME
+                      + " where id in (select id from a) and id in (select id from b) order by id")
+              .tableName(NAME)
+              .registerArrow("a", sa)
+              .registerArrow("b", sb)
+              .intoBatchRecords();
+      List<Integer> got = new ArrayList<>();
+      while (reader.loadNextBatch()) {
+        VectorSchemaRoot root = reader.getVectorSchemaRoot();
+        for (int i = 0; i < root.getRowCount(); i++) {
+          got.add((Integer) root.getVector(0).getObject(i));
+        }
+      }
+      reader.close();
+      // a = {1,2,3,10}, b = {2,3,4,10}, dataset ids 0..39, so the intersection is {2,3,10}.
+      Assertions.assertEquals(Arrays.asList(2, 3, 10), got);
+    }
+  }
+
+  @Test
+  public void testRegisterArrowValidationAndReuse() throws Exception {
+    // registerArrow rejects invalid inputs at the boundary.
+    try (VectorSchemaRoot ids = idTable(1, 2);
+        ArrowArrayStream stream = toStream(ids)) {
+      SqlQuery q = dataset.sql("select id from " + NAME).tableName(NAME);
+      Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow("", stream));
+      Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow(null, stream));
+      Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow("ids", null));
+    }
+
+    // A query with registered relations is single-use: the stream is consumed on the first call,
+    // so a second intoBatchRecords() throws rather than handing JNI a dead stream.
+    try (VectorSchemaRoot ids = idTable(1, 2);
+        ArrowArrayStream stream = toStream(ids)) {
+      SqlQuery q =
+          dataset
+              .sql("select id from " + NAME + " where id in (select id from ids)")
+              .tableName(NAME)
+              .registerArrow("ids", stream);
+      q.intoBatchRecords().close();
+      Assertions.assertThrows(IllegalStateException.class, q::intoBatchRecords);
+
+      // Registering another relation after the query is consumed is also rejected.
+      try (VectorSchemaRoot more = idTable(3);
+          ArrowArrayStream moreStream = toStream(more)) {
+        Assertions.assertThrows(
+            IllegalStateException.class, () -> q.registerArrow("more", moreStream));
+      }
+    }
+
+    // A relation whose name resolves to the query's own table name would hide the dataset, so the
+    // native build rejects it (the Rust InvalidInput surfaces as IllegalArgumentException). Assert
+    // on the message too: DataFusion's own "table already exists" error maps to the same exception.
+    try (VectorSchemaRoot ids = idTable(1, 2);
+        ArrowArrayStream stream = toStream(ids)) {
+      SqlQuery q =
+          dataset.sql("select id from " + NAME).tableName(NAME).registerArrow(NAME, stream);
+      IllegalArgumentException error =
+          Assertions.assertThrows(IllegalArgumentException.class, q::intoBatchRecords);
+      Assertions.assertTrue(
+          error.getMessage().contains("resolves to")
+              && error.getMessage().contains("the query's table name '" + NAME + "'"),
+          "unexpected message: " + error.getMessage());
+    }
+  }
+
+  /** Build a single-column (`id`: Int32) in-memory relation. */
+  private VectorSchemaRoot idTable(int... ids) {
+    Schema schema =
+        new Schema(Collections.singletonList(Field.nullable("id", new ArrowType.Int(32, true))));
+    VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
+    root.allocateNew();
+    IntVector idVector = (IntVector) root.getVector("id");
+    for (int i = 0; i < ids.length; i++) {
+      idVector.setSafe(i, ids[i]);
+    }
+    root.setRowCount(ids.length);
+    return root;
+  }
+
+  /**
+   * Serialize a single-batch root to a self-contained Arrow C-Data stream (mirrors
+   * MergeInsertTest).
+   */
+  private ArrowArrayStream toStream(VectorSchemaRoot root) throws IOException {
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) {
+      writer.start();
+      writer.writeBatch();
+      writer.end();
+    }
+    ArrowStreamReader reader =
+        new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray()), allocator);
+    ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator);
+    Data.exportArrayStream(allocator, reader, stream);
+    return stream;
+  }
 }