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};
-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 = RT.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,36 @@ 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; if a name is registered twice,
+   * the later registration replaces the earlier one (deduplicated when the query is built).
+   *
+   * <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,9 +86,26 @@ 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 (ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) {
+      // A query with registered streams is one-shot: the native call below consumes them. Mark it consumed here,
+      // just before the native call, so a failure in the output-stream allocation above 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);
     }
   }
@@ -64,7 +116,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,130 @@ public void testToRecordBatches() throws IOException {
         reader.getVectorSchemaRoot().getSchema().toString());
     reader.close();
   }
+
+  @Test
+  public void testRegisterArrow() throws Exception {
+    // An additional 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 in the dataset and in both 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 registered 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 (would be unexecutable).
+      try (VectorSchemaRoot more = idTable(3);
+          ArrowArrayStream moreStream = toStream(more)) {
+        Assertions.assertThrows(
+            IllegalStateException.class, () -> q.registerArrow("more", moreStream));
+      }
+    }
+  }
+
+  /** 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;
+  }
 }

python/python/lance/dataset.py

@@ -5443,6 +5443,39 @@ def with_row_addr(self, with_row_addr: bool = True) -> "SqlQueryBuilder":
         self._builder = self._builder.with_row_addr(with_row_addr)
         return self
 
+    def register_arrow(
+        self, name: str, data: Union[pa.Table, pa.RecordBatchReader]
+    ) -> "SqlQueryBuilder":
+        """
+        Register an in-memory Arrow relation that the query can join.
+
+        The relation is exposed under ``name`` and can be referenced like any
+        other table in the SQL (for example joined or used in a subquery). It
+        must be non-empty, since the schema is derived from the provided data.
+
+        Parameters
+        ----------
+        name: str
+            The name the relation is registered under in the query.
+        data: pyarrow.Table or pyarrow.RecordBatchReader
+            The in-memory relation to register. Must be non-empty.
+
+        Examples
+        --------
+        Register an in-memory relation and semi-join it in the query::
+
+            ids = pa.table({"id": [1, 2, 3]})
+            batches = (
+                dataset.sql("SELECT * FROM d WHERE id IN (SELECT id FROM ids)")
+                .table_name("d")
+                .register_arrow("ids", ids)
+                .build()
+                .to_batch_records()
+            )
+        """
+        self._builder = self._builder.register_arrow(name, data)
+        return self
+
     def build(self) -> SqlQuery:
         """
         Build the query.