From cb5a8505fa6970efc768d8088e78ed88db57810e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sat, 18 Jul 2026 13:01:00 +0200 Subject: [PATCH 1/5] feat(sql): register additional tables for Dataset.sql queries Add a way to register extra tables in the DataFusion SessionContext that backs Dataset.sql, so a query can join a caller-supplied relation (for example an id set to semi-join, or another dataset via LanceTableProvider) instead of scoping by a large literal IN list. rust: SqlQueryBuilder gains register_table(name, Arc) and a register_arrow(name, Vec) convenience that wraps the batches in a MemTable. build() registers each extra table before running the query. register_arrow errors on empty input (no schema to derive); register_table accepts an already-built provider, including an empty one. java: SqlQuery.registerArrow(name, ArrowArrayStream) imports the stream through the JNI and registers it as an in-memory table. python: SqlQueryBuilder.register_arrow(name, data) accepts a pyarrow Table or RecordBatchReader. tests: rust unit tests for register_arrow and register_table (semi-join, direct select, empty relation, empty input error); a java SqlQueryTest case; and a python test. --- java/lance-jni/src/sql.rs | 29 +++- java/src/main/java/org/lance/SqlQuery.java | 28 +++- .../src/test/java/org/lance/SqlQueryTest.java | 65 +++++++++ python/python/lance/dataset.py | 20 +++ python/python/tests/test_dataset.py | 20 +++ python/src/dataset.rs | 20 +++ rust/lance/src/dataset/sql.rs | 130 ++++++++++++++++++ 7 files changed, 309 insertions(+), 3 deletions(-) diff --git a/java/lance-jni/src/sql.rs b/java/lance-jni/src/sql.rs index f378577eedf..ca27329a054 100644 --- a/java/lance-jni/src/sql.rs +++ b/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_name: JObject, + extra_stream_addr: jlong, ) { 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_name, + extra_stream_addr, ) .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_name: JObject, + extra_stream_addr: jlong, ) -> 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_name, + extra_stream_addr, )?; let stream = RT.block_on(async move { @@ -69,6 +78,7 @@ fn inner_into_batch_records( Ok(()) } +#[allow(clippy::too_many_arguments)] fn sql_builder( env: &mut JNIEnv, java_dataset: JObject, @@ -76,9 +86,12 @@ fn sql_builder( table_name: JObject, with_row_id: jboolean, with_row_addr: jboolean, + extra_table_name: JObject, + extra_stream_addr: jlong, ) -> Result { let sql_str = sql.extract(env)?; let table_str = env.get_string_opt(&table_name)?; + let extra_name = env.get_string_opt(&extra_table_name)?; let dataset_guard = unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?; @@ -93,5 +106,19 @@ fn sql_builder( builder = builder.table_name(table.as_str()) } + // An additional in-memory Arrow relation (the caller exported it to a C-Data stream). Import it to + // RecordBatches and register it under `extra_name`, so the SQL can semi-join it natively. + if extra_stream_addr != 0 { + let name = extra_name + .ok_or_else(|| Error::input_error("extra table stream given without a name".to_string()))?; + let stream_ptr = extra_stream_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) } diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index cce6d939222..4de8d24dc98 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/java/src/main/java/org/lance/SqlQuery.java @@ -29,6 +29,8 @@ public class SqlQuery { private String table = DEFAULT_TABLE_NAME; private boolean withRowId = false; private boolean withRowAddr = false; + private String extraTableName = null; + private long extraStreamAddress = 0L; public SqlQuery(Dataset dataset, String sql) { this.dataset = dataset; @@ -40,6 +42,19 @@ 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}. Only one extra table is supported per query; the last call wins. + */ + public SqlQuery registerArrow(String name, ArrowArrayStream stream) { + this.extraTableName = name; + this.extraStreamAddress = stream.memoryAddress(); + return this; + } + public SqlQuery withRowId(boolean withRowId) { this.withRowId = withRowId; return this; @@ -53,7 +68,14 @@ public SqlQuery withRowAddr(boolean withAddr) { public ArrowReader intoBatchRecords() throws IOException { try (ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { intoBatchRecords( - dataset, sql, Optional.ofNullable(table), withRowId, withRowAddr, s.memoryAddress()); + dataset, + sql, + Optional.ofNullable(table), + withRowId, + withRowAddr, + s.memoryAddress(), + Optional.ofNullable(extraTableName), + extraStreamAddress); return Data.importArrayStream(dataset.allocator(), s); } } @@ -64,7 +86,9 @@ private static native void intoBatchRecords( Optional tableName, boolean withRowId, boolean withRowAddr, - long streamAddress) + long streamAddress, + Optional extraTableName, + long extraStreamAddress) throws IOException; @Override diff --git a/java/src/test/java/org/lance/SqlQueryTest.java b/java/src/test/java/org/lance/SqlQueryTest.java index 2e3dcab9892..d78beecf151 100644 --- a/java/src/test/java/org/lance/SqlQueryTest.java +++ b/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,54 @@ 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 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); + } + } + } + + /** 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; + } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..671f41f701e 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -5443,6 +5443,26 @@ 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. + """ + self._builder = self._builder.register_arrow(name, data) + return self + def build(self) -> SqlQuery: """ Build the query. diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 898bdf93069..13fffe14345 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5381,6 +5381,26 @@ def test_dataset_sql(tmp_path: Path): assert pa.Table.from_batches(complex_result) == expected_complex +def test_dataset_sql_register_arrow(tmp_path: Path): + table = pa.table({"id": [1, 2, 3, 4, 5], "value": ["a", "b", "c", "d", "e"]}) + ds = lance.write_dataset(table, tmp_path / "test") + + # A subset of ids to join against, including one id (99) not in the dataset. + ids = pa.table({"id": [4, 2, 99]}) + + result = ( + ds.sql("SELECT id FROM t WHERE id IN (SELECT id FROM ids) ORDER BY id") + .table_name("t") + .register_arrow("ids", ids) + .build() + .to_batch_records() + ) + + # Only the intersection, in ascending order. + expected = pa.table({"id": [2, 4]}) + assert pa.Table.from_batches(result) == expected + + def test_file_reader_options(tmp_path: Path): """Test cache_repetition_index and validate_on_decode options""" # Create a dataset with large repetitive strings to test cache_repetition_index diff --git a/python/src/dataset.rs b/python/src/dataset.rs index a361e0b63a8..0591e594f74 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -3954,6 +3954,26 @@ impl SqlQueryBuilder { } } + /// Register an in-memory Arrow relation that the SQL query can join. + /// + /// The relation is exposed under `name` and can be referenced like any other + /// table in the query. The `data` may be a pyarrow `Table` or + /// `RecordBatchReader` (ingested the same way as `add`/`merge_insert`), and + /// must be non-empty since the schema is derived from its batches. + #[pyo3(signature = (name, data))] + fn register_arrow(&self, name: &str, data: &Bound) -> PyResult { + let reader = convert_reader(data)?; + let batches = reader + .collect::, _>>() + .map_err(|err| PyValueError::new_err(err.to_string()))?; + let builder = self + .builder + .clone() + .register_arrow(name, batches) + .infer_error()?; + Ok(Self { builder }) + } + /// Build the SQL query. fn build(&self) -> PyResult { Ok(SqlQuery { diff --git a/rust/lance/src/dataset/sql.rs b/rust/lance/src/dataset/sql.rs index 8a1ccda2df6..568fa8d596e 100644 --- a/rust/lance/src/dataset/sql.rs +++ b/rust/lance/src/dataset/sql.rs @@ -6,6 +6,7 @@ use crate::datafusion::LanceTableProvider; use crate::dataset::utils::SchemaAdapter; use arrow_array::RecordBatch; use datafusion::dataframe::DataFrame; +use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::SendableRecordBatchStream; use datafusion::prelude::SessionContext; use futures::TryStreamExt; @@ -29,6 +30,12 @@ pub struct SqlQueryBuilder { /// If true, the query result will include the internal row address pub(crate) with_row_addr: bool, + + /// Additional tables to register in the DataFusion context before running the query (name to provider). + /// This lets a `Dataset.sql` query join a caller-supplied relation, such as an in-memory `MemTable` (for + /// example a ranked id set to semi-join), another Lance dataset (`LanceTableProvider`), or a custom + /// `TableProvider`. + pub(crate) extra_tables: Vec<(String, Arc)>, } impl SqlQueryBuilder { @@ -39,9 +46,32 @@ impl SqlQueryBuilder { table_name: "dataset".to_string(), with_row_id: false, with_row_addr: false, + extra_tables: Vec::new(), } } + /// Register an additional table named `name` in the query's DataFusion context, joinable alongside the + /// dataset. Accepts any `TableProvider`, such as a `MemTable` (an in-memory Arrow relation), another Lance + /// dataset's `LanceTableProvider`, or a view. May be called more than once; a later duplicate name wins. + pub fn register_table(mut self, name: &str, provider: Arc) -> Self { + self.extra_tables.push((name.to_string(), provider)); + self + } + + /// Convenience over [`Self::register_table`] that wraps in-memory Arrow `batches` in a `MemTable`. `batches` + /// must be non-empty (the schema is derived from the first batch); for an empty relation, build a `MemTable` + /// yourself and pass it to [`Self::register_table`]. + pub fn register_arrow(self, name: &str, batches: Vec) -> lance_core::Result { + let schema = batches.first().map(|b| b.schema()).ok_or_else(|| { + lance_core::Error::invalid_input( + "register_arrow requires a non-empty relation to derive a schema; \ + use register_table with a MemTable for an empty relation" + .to_string(), + ) + })?; + Ok(self.register_table(name, Arc::new(MemTable::try_new(schema, vec![batches])?))) + } + /// The table name to register in the datafusion context. /// This is used to specify a "table name" for the dataset. /// So that you can run SQL queries against it. @@ -77,6 +107,9 @@ impl SqlQueryBuilder { row_addr, )), )?; + for (name, provider) in self.extra_tables { + ctx.register_table(name, provider)?; + } register_functions(&ctx); let df = ctx.sql(&self.sql).await?; Ok(SqlQuery::new(df)) @@ -128,6 +161,7 @@ mod tests { use crate::Dataset; use all_asserts::assert_true; + use datafusion::datasource::MemTable; use arrow_array::cast::AsArray; use arrow_array::types::{Int32Type, Int64Type, UInt64Type}; use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray}; @@ -365,4 +399,100 @@ mod tests { pretty_assertions::assert_eq!(batch.num_columns(), 1); pretty_assertions::assert_eq!(batch.column(0).as_primitive::().value(0), 1); } + + /// A dataset with an `id` column = 0..=99 across 10 fragments (for the register-table tests below). + async fn ids_dataset(uri: &str) -> Dataset { + gen_batch() + .col("id", array::step::()) + .into_dataset(uri, FragmentCount::from(10), FragmentRowCount::from(10)) + .await + .unwrap() + } + + /// A one-batch relation `(id: Int32)` from `ids`. + fn ids_batch(ids: Vec) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![Field::new("id", DataType::Int32, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids))]).unwrap() + } + + fn collect_ids(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect() + } + + #[tokio::test] + async fn test_sql_register_arrow() { + let ds = ids_dataset("memory://test_sql_register_arrow").await; + + // Semi-join against a registered in-memory relation; 200 is absent from the dataset. + let results = ds + .sql("SELECT id FROM t WHERE id IN (SELECT id FROM ids) ORDER BY id") + .table_name("t") + .register_arrow("ids", vec![ids_batch(vec![5, 10, 42, 200])]) + .unwrap() + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + // 200 is dropped (not in the dataset); the rest are the intersection, in order. + pretty_assertions::assert_eq!(collect_ids(&results), vec![5, 10, 42]); + + // The registered relation is also selectable directly. + let counted = ds + .sql("SELECT count(*) AS n FROM ids") + .table_name("t") + .register_arrow("ids", vec![ids_batch(vec![5, 10, 42, 200])]) + .unwrap() + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + pretty_assertions::assert_eq!(counted[0].column(0).as_primitive::().value(0), 4); + + // register_arrow with no batches cannot derive a schema, so it returns an error (not a silent skip). + let err = ds.sql("SELECT 1").table_name("t").register_arrow("ids", vec![]); + assert_true!(err.is_err()); + } + + #[tokio::test] + async fn test_sql_register_table_provider() { + let ds = ids_dataset("memory://test_sql_register_table_provider").await; + let schema = ids_batch(vec![]).schema(); + + // The general path: register any TableProvider (here a MemTable) and join it. + let table = Arc::new(MemTable::try_new(schema.clone(), vec![vec![ids_batch(vec![7, 88])]]).unwrap()); + let results = ds + .sql("SELECT id FROM t WHERE id IN (SELECT id FROM ids) ORDER BY id") + .table_name("t") + .register_table("ids", table) + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + pretty_assertions::assert_eq!(collect_ids(&results), vec![7, 88]); + + // An empty registered relation is a valid table (schema, zero rows), so the semi-join returns zero rows + // rather than a "table not found" error. (`vec![vec![]]` is one partition with no batches; MemTable + // requires at least one partition.) + let empty = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap()); + let results = ds + .sql("SELECT id FROM t WHERE id IN (SELECT id FROM ids)") + .table_name("t") + .register_table("ids", empty) + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + pretty_assertions::assert_eq!(collect_ids(&results), Vec::::new()); + } } From 90084658d00e239033fc3c757d1f3ae93c429ea4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sat, 18 Jul 2026 13:54:24 +0200 Subject: [PATCH 2/5] feat(sql): register multiple tables from the Java binding SqlQuery.registerArrow now accumulates into a list rather than a single field, and intoBatchRecords passes the (name, stream) pairs as two parallel lists through the JNI, which loops and registers each. Rust and Python already accept multiple tables (they accumulate into a Vec). Add SqlQueryTest.testRegisterArrowMultiple covering two tables in one query. --- java/lance-jni/src/sql.rs | 35 ++++++++------- java/src/main/java/org/lance/SqlQuery.java | 21 +++++---- .../src/test/java/org/lance/SqlQueryTest.java | 44 +++++++++++++++++++ 3 files changed, 74 insertions(+), 26 deletions(-) diff --git a/java/lance-jni/src/sql.rs b/java/lance-jni/src/sql.rs index ca27329a054..452e06754d3 100644 --- a/java/lance-jni/src/sql.rs +++ b/java/lance-jni/src/sql.rs @@ -23,8 +23,8 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords( with_row_id: jboolean, with_row_addr: jboolean, stream_addr: jlong, - extra_table_name: JObject, - extra_stream_addr: jlong, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) { ok_or_throw_without_return!( env, @@ -36,8 +36,8 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords( with_row_id, with_row_addr, stream_addr, - extra_table_name, - extra_stream_addr, + extra_table_names, + extra_stream_addrs, ) .map_err(|e| Error::input_error(e.to_string())) ) @@ -52,8 +52,8 @@ fn inner_into_batch_records( with_row_id: jboolean, with_row_addr: jboolean, stream_addr: jlong, - extra_table_name: JObject, - extra_stream_addr: jlong, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) -> Result<()> { let builder = sql_builder( env, @@ -62,8 +62,8 @@ fn inner_into_batch_records( table_name, with_row_id, with_row_addr, - extra_table_name, - extra_stream_addr, + extra_table_names, + extra_stream_addrs, )?; let stream = RT.block_on(async move { @@ -86,12 +86,15 @@ fn sql_builder( table_name: JObject, with_row_id: jboolean, with_row_addr: jboolean, - extra_table_name: JObject, - extra_stream_addr: jlong, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) -> Result { let sql_str = sql.extract(env)?; let table_str = env.get_string_opt(&table_name)?; - let extra_name = env.get_string_opt(&extra_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) }?; @@ -106,12 +109,10 @@ fn sql_builder( builder = builder.table_name(table.as_str()) } - // An additional in-memory Arrow relation (the caller exported it to a C-Data stream). Import it to - // RecordBatches and register it under `extra_name`, so the SQL can semi-join it natively. - if extra_stream_addr != 0 { - let name = extra_name - .ok_or_else(|| Error::input_error("extra table stream given without a name".to_string()))?; - let stream_ptr = extra_stream_addr as *mut FFI_ArrowArrayStream; + // 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 { diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index 4de8d24dc98..37395a34f14 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/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,8 +31,8 @@ public class SqlQuery { private String table = DEFAULT_TABLE_NAME; private boolean withRowId = false; private boolean withRowAddr = false; - private String extraTableName = null; - private long extraStreamAddress = 0L; + private final List extraTableNames = new ArrayList<>(); + private final List extraStreamAddresses = new ArrayList<>(); public SqlQuery(Dataset dataset, String sql) { this.dataset = dataset; @@ -47,11 +49,12 @@ public SqlQuery tableName(String tableName) { * 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}. Only one extra table is supported per query; the last call wins. + * {@code MergeInsert}. May be called multiple times to register multiple tables; if a name is registered twice, + * the later registration replaces the earlier one at query time (DataFusion register semantics). */ public SqlQuery registerArrow(String name, ArrowArrayStream stream) { - this.extraTableName = name; - this.extraStreamAddress = stream.memoryAddress(); + this.extraTableNames.add(name); + this.extraStreamAddresses.add(stream.memoryAddress()); return this; } @@ -74,8 +77,8 @@ public ArrowReader intoBatchRecords() throws IOException { withRowId, withRowAddr, s.memoryAddress(), - Optional.ofNullable(extraTableName), - extraStreamAddress); + extraTableNames, + extraStreamAddresses); return Data.importArrayStream(dataset.allocator(), s); } } @@ -87,8 +90,8 @@ private static native void intoBatchRecords( boolean withRowId, boolean withRowAddr, long streamAddress, - Optional extraTableName, - long extraStreamAddress) + List extraTableNames, + List extraStreamAddresses) throws IOException; @Override diff --git a/java/src/test/java/org/lance/SqlQueryTest.java b/java/src/test/java/org/lance/SqlQueryTest.java index d78beecf151..e49183b0de5 100644 --- a/java/src/test/java/org/lance/SqlQueryTest.java +++ b/java/src/test/java/org/lance/SqlQueryTest.java @@ -160,6 +160,50 @@ public void testRegisterArrow() throws Exception { } } + @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 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); + } + } + + /** 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(); From e1cd1a0e1fefe7a00b45f498510179fad30e385c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sat, 18 Jul 2026 15:10:53 +0200 Subject: [PATCH 3/5] fix(sql): address review feedback on table registration - build(): a table name registered more than once now keeps the last provider (dedup), instead of erroring when the same name is registered twice; register_table/register_arrow docs already state last wins. - java: SqlQuery.registerArrow rejects null/blank names and null streams; a query with registered relations is single-use (intoBatchRecords consumes the streams, so a second call throws instead of reusing dead pointers). - docs: add runnable examples to register_table and register_arrow (Rust) and to the Python register_arrow. - tests: use the b["id"] accessor; assert the empty-input error variant and message; cover duplicate-name registration; Java input validation and single-use; Python RecordBatchReader input and empty-input ValueError. --- java/src/main/java/org/lance/SqlQuery.java | 22 ++++ .../src/test/java/org/lance/SqlQueryTest.java | 25 ++++ python/python/lance/dataset.py | 13 ++ python/python/tests/test_dataset.py | 21 ++- rust/lance/src/dataset/sql.rs | 123 +++++++++++++++++- 5 files changed, 194 insertions(+), 10 deletions(-) diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index 37395a34f14..aa3b94922c5 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/java/src/main/java/org/lance/SqlQuery.java @@ -33,6 +33,7 @@ public class SqlQuery { private boolean withRowAddr = false; private final List extraTableNames = new ArrayList<>(); private final List extraStreamAddresses = new ArrayList<>(); + private boolean consumed = false; public SqlQuery(Dataset dataset, String sql) { this.dataset = dataset; @@ -51,8 +52,19 @@ public SqlQuery tableName(String tableName) { * {@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 at query time (DataFusion register semantics). + * + *

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 */ public SqlQuery registerArrow(String name, ArrowArrayStream stream) { + if (name == null || name.trim().isEmpty()) { + throw new IllegalArgumentException("table name must be non-empty"); + } + if (stream == null) { + throw new IllegalArgumentException("stream must not be null"); + } this.extraTableNames.add(name); this.extraStreamAddresses.add(stream.memoryAddress()); return this; @@ -69,6 +81,16 @@ 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."); + } + // A query with registered streams is one-shot: the native call consumes the streams, so mark it consumed + // regardless of outcome to prevent a later call from handing JNI dead stream pointers. + if (!extraTableNames.isEmpty()) { + consumed = true; + } try (ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { intoBatchRecords( dataset, diff --git a/java/src/test/java/org/lance/SqlQueryTest.java b/java/src/test/java/org/lance/SqlQueryTest.java index e49183b0de5..c5d32f27d20 100644 --- a/java/src/test/java/org/lance/SqlQueryTest.java +++ b/java/src/test/java/org/lance/SqlQueryTest.java @@ -190,6 +190,31 @@ public void testRegisterArrowMultiple() throws Exception { } } + @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); + } + } + /** Build a single-column (`id`: Int32) in-memory relation. */ private VectorSchemaRoot idTable(int... ids) { Schema schema = diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 671f41f701e..bebd7063809 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -5459,6 +5459,19 @@ def register_arrow( 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 diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 13fffe14345..755925375af 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5381,17 +5381,20 @@ def test_dataset_sql(tmp_path: Path): assert pa.Table.from_batches(complex_result) == expected_complex -def test_dataset_sql_register_arrow(tmp_path: Path): +@pytest.mark.parametrize("as_reader", [False, True], ids=["table", "reader"]) +def test_dataset_sql_register_arrow(tmp_path: Path, as_reader: bool): table = pa.table({"id": [1, 2, 3, 4, 5], "value": ["a", "b", "c", "d", "e"]}) ds = lance.write_dataset(table, tmp_path / "test") - # A subset of ids to join against, including one id (99) not in the dataset. + # A subset of ids to join against, including one id (99) not in the dataset. Exercise both accepted input + # forms: a pyarrow Table and a RecordBatchReader. ids = pa.table({"id": [4, 2, 99]}) + data = ids.to_reader() if as_reader else ids result = ( ds.sql("SELECT id FROM t WHERE id IN (SELECT id FROM ids) ORDER BY id") .table_name("t") - .register_arrow("ids", ids) + .register_arrow("ids", data) .build() .to_batch_records() ) @@ -5401,6 +5404,18 @@ def test_dataset_sql_register_arrow(tmp_path: Path): assert pa.Table.from_batches(result) == expected +def test_dataset_sql_register_arrow_empty(tmp_path: Path): + table = pa.table({"id": [1, 2, 3]}) + ds = lance.write_dataset(table, tmp_path / "test") + + # An empty relation has no batches to derive a schema from, so registration raises (translated to ValueError + # by the PyO3 binding). + schema = pa.schema([("id", pa.int64())]) + empty = pa.RecordBatchReader.from_batches(schema, []) + with pytest.raises(ValueError): + ds.sql("SELECT id FROM ids").table_name("t").register_arrow("ids", empty) + + def test_file_reader_options(tmp_path: Path): """Test cache_repetition_index and validate_on_decode options""" # Create a dataset with large repetitive strings to test cache_repetition_index diff --git a/rust/lance/src/dataset/sql.rs b/rust/lance/src/dataset/sql.rs index 568fa8d596e..b6f0ea300e2 100644 --- a/rust/lance/src/dataset/sql.rs +++ b/rust/lance/src/dataset/sql.rs @@ -53,6 +53,43 @@ impl SqlQueryBuilder { /// Register an additional table named `name` in the query's DataFusion context, joinable alongside the /// dataset. Accepts any `TableProvider`, such as a `MemTable` (an in-memory Arrow relation), another Lance /// dataset's `LanceTableProvider`, or a view. May be called more than once; a later duplicate name wins. + /// + /// # Examples + /// + /// ```rust + /// # use std::sync::Arc; + /// # use tokio::runtime::Runtime; + /// # use arrow_array::{RecordBatch, RecordBatchIterator, Int32Array}; + /// # use arrow_schema::{Schema, Field, DataType}; + /// # use datafusion::datasource::MemTable; + /// # use lance::dataset::{WriteParams, Dataset}; + /// # + /// # let rt = Runtime::new().unwrap(); + /// # rt.block_on(async { + /// # let test_dir = tempfile::tempdir().unwrap(); + /// # let uri = test_dir.path().to_str().unwrap().to_string(); + /// # let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + /// # let batch = RecordBatch::try_new( + /// # schema.clone(), + /// # vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + /// # ).unwrap(); + /// # let reader = RecordBatchIterator::new(vec![batch.clone()].into_iter().map(Ok), schema.clone()); + /// # let dataset = Dataset::write(reader, &uri, Some(WriteParams::default())).await.unwrap(); + /// # + /// // Register any `TableProvider` (here a `MemTable`) and join it in the query. + /// let ids = Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch]]).unwrap()); + /// let batches = dataset + /// .sql("SELECT id FROM data WHERE id IN (SELECT id FROM ids)") + /// .table_name("data") + /// .register_table("ids", ids) + /// .build() + /// .await? + /// .into_batch_records() + /// .await?; + /// # assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + /// # Ok::<(), lance_core::Error>(()) + /// # }).unwrap(); + /// ``` pub fn register_table(mut self, name: &str, provider: Arc) -> Self { self.extra_tables.push((name.to_string(), provider)); self @@ -61,6 +98,45 @@ impl SqlQueryBuilder { /// Convenience over [`Self::register_table`] that wraps in-memory Arrow `batches` in a `MemTable`. `batches` /// must be non-empty (the schema is derived from the first batch); for an empty relation, build a `MemTable` /// yourself and pass it to [`Self::register_table`]. + /// + /// # Examples + /// + /// ```rust + /// # use std::sync::Arc; + /// # use tokio::runtime::Runtime; + /// # use arrow_array::{RecordBatch, RecordBatchIterator, Int32Array}; + /// # use arrow_schema::{Schema, Field, DataType}; + /// # use lance::dataset::{WriteParams, Dataset}; + /// # + /// # let rt = Runtime::new().unwrap(); + /// # rt.block_on(async { + /// # let test_dir = tempfile::tempdir().unwrap(); + /// # let uri = test_dir.path().to_str().unwrap().to_string(); + /// # let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + /// # let batch = RecordBatch::try_new( + /// # schema.clone(), + /// # vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4]))], + /// # ).unwrap(); + /// # let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + /// # let dataset = Dataset::write(reader, &uri, Some(WriteParams::default())).await.unwrap(); + /// # + /// // A caller-supplied relation to semi-join against the dataset. + /// let ids = RecordBatch::try_new( + /// Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + /// vec![Arc::new(Int32Array::from(vec![2, 3]))], + /// ).unwrap(); + /// let batches = dataset + /// .sql("SELECT id FROM data WHERE id IN (SELECT id FROM ids)") + /// .table_name("data") + /// .register_arrow("ids", vec![ids])? + /// .build() + /// .await? + /// .into_batch_records() + /// .await?; + /// # assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + /// # Ok::<(), lance_core::Error>(()) + /// # }).unwrap(); + /// ``` pub fn register_arrow(self, name: &str, batches: Vec) -> lance_core::Result { let schema = batches.first().map(|b| b.schema()).ok_or_else(|| { lance_core::Error::invalid_input( @@ -107,8 +183,14 @@ impl SqlQueryBuilder { row_addr, )), )?; - for (name, provider) in self.extra_tables { - ctx.register_table(name, provider)?; + // Register the extra tables, keeping the LAST provider for a name registered more than once (registering + // the same name into the context twice would otherwise error). Iterate in reverse and skip names already + // registered, so a later duplicate wins. + let mut registered = std::collections::HashSet::new(); + for (name, provider) in self.extra_tables.into_iter().rev() { + if registered.insert(name.clone()) { + ctx.register_table(name, provider)?; + } } register_functions(&ctx); let df = ctx.sql(&self.sql).await?; @@ -418,7 +500,7 @@ mod tests { fn collect_ids(batches: &[RecordBatch]) -> Vec { batches .iter() - .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .flat_map(|b| b["id"].as_primitive::().values().to_vec()) .collect() } @@ -453,11 +535,38 @@ mod tests { .into_batch_records() .await .unwrap(); - pretty_assertions::assert_eq!(counted[0].column(0).as_primitive::().value(0), 4); + pretty_assertions::assert_eq!(counted[0]["n"].as_primitive::().value(0), 4); - // register_arrow with no batches cannot derive a schema, so it returns an error (not a silent skip). - let err = ds.sql("SELECT 1").table_name("t").register_arrow("ids", vec![]); - assert_true!(err.is_err()); + // Registering the same name twice: the later relation replaces the earlier one. + let dup = ds + .sql("SELECT id FROM ids ORDER BY id") + .table_name("t") + .register_arrow("ids", vec![ids_batch(vec![1, 2])]) + .unwrap() + .register_arrow("ids", vec![ids_batch(vec![7, 8, 9])]) + .unwrap() + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + pretty_assertions::assert_eq!(collect_ids(&dup), vec![7, 8, 9]); + + // register_arrow with no batches cannot derive a schema, so it returns an InvalidInput error. + let err = ds + .sql("SELECT 1") + .table_name("t") + .register_arrow("ids", vec![]) + .unwrap_err(); + assert!( + matches!(err, lance_core::Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + assert!( + err.to_string().contains("schema"), + "message should mention the missing schema: {err}" + ); } #[tokio::test] From 468a8851c8eb153050762a82513b9197b3d895b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sat, 18 Jul 2026 17:05:10 +0200 Subject: [PATCH 4/5] fix(sql): richer registration error context; correct duplicate-name doc - register_arrow: the empty-input error now includes the table name and the batch count. - java registerArrow: the null/blank-name and null-stream errors now include the parameter context and rejected value. - java javadoc: describe last-registration-wins as build()-side dedup, not DataFusion register semantics (a MemorySchemaProvider errors on a duplicate name; the last-wins behavior is implemented by SqlQueryBuilder::build). --- java/src/main/java/org/lance/SqlQuery.java | 6 +++--- rust/lance/src/dataset/sql.rs | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index aa3b94922c5..3f21ad5fa54 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/java/src/main/java/org/lance/SqlQuery.java @@ -51,7 +51,7 @@ public SqlQuery tableName(String tableName) { * 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 at query time (DataFusion register semantics). + * the later registration replaces the earlier one (deduplicated when the query is built). * *

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. @@ -60,10 +60,10 @@ public SqlQuery tableName(String tableName) { */ public SqlQuery registerArrow(String name, ArrowArrayStream stream) { if (name == null || name.trim().isEmpty()) { - throw new IllegalArgumentException("table name must be non-empty"); + throw new IllegalArgumentException("registerArrow: table name must be non-empty, got: " + name); } if (stream == null) { - throw new IllegalArgumentException("stream must not be null"); + throw new IllegalArgumentException("registerArrow: stream must not be null (table name: " + name + ")"); } this.extraTableNames.add(name); this.extraStreamAddresses.add(stream.memoryAddress()); diff --git a/rust/lance/src/dataset/sql.rs b/rust/lance/src/dataset/sql.rs index b6f0ea300e2..1a3c39439ca 100644 --- a/rust/lance/src/dataset/sql.rs +++ b/rust/lance/src/dataset/sql.rs @@ -138,12 +138,12 @@ impl SqlQueryBuilder { /// # }).unwrap(); /// ``` pub fn register_arrow(self, name: &str, batches: Vec) -> lance_core::Result { + let n_batches = batches.len(); let schema = batches.first().map(|b| b.schema()).ok_or_else(|| { - lance_core::Error::invalid_input( - "register_arrow requires a non-empty relation to derive a schema; \ - use register_table with a MemTable for an empty relation" - .to_string(), - ) + lance_core::Error::invalid_input(format!( + "register_arrow for table '{name}' requires a non-empty relation to derive a schema \ + (received {n_batches} batches); use register_table with a MemTable for an empty relation" + )) })?; Ok(self.register_table(name, Arc::new(MemTable::try_new(schema, vec![batches])?))) } From ae6853f21c83b88999316497ffa4fcbf42aa6f9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sat, 18 Jul 2026 19:08:22 +0200 Subject: [PATCH 5/5] fix(sql): validate relation names in build; harden Java one-shot guards - build(): reject an empty or whitespace-only registered table name with a descriptive InvalidInput, so register_table/register_arrow both validate names (previously only the Java binding did). - java: reject registerArrow after the query was consumed, and set the consumed flag inside the try (just before the native call) so a failed output-stream allocation does not permanently block a retry. - tests: blank-name rejection (rust); register-after-consumed (java). --- java/src/main/java/org/lance/SqlQuery.java | 15 ++++++++----- .../src/test/java/org/lance/SqlQueryTest.java | 7 +++++++ rust/lance/src/dataset/sql.rs | 21 +++++++++++++++++++ 3 files changed, 38 insertions(+), 5 deletions(-) diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index 3f21ad5fa54..ea25c9c3c2a 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/java/src/main/java/org/lance/SqlQuery.java @@ -57,8 +57,13 @@ public SqlQuery tableName(String tableName) { * 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); } @@ -86,12 +91,12 @@ public ArrowReader intoBatchRecords() throws IOException { "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."); } - // A query with registered streams is one-shot: the native call consumes the streams, so mark it consumed - // regardless of outcome to prevent a later call from handing JNI dead stream pointers. - if (!extraTableNames.isEmpty()) { - consumed = true; - } 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, diff --git a/java/src/test/java/org/lance/SqlQueryTest.java b/java/src/test/java/org/lance/SqlQueryTest.java index c5d32f27d20..54a45c6fa33 100644 --- a/java/src/test/java/org/lance/SqlQueryTest.java +++ b/java/src/test/java/org/lance/SqlQueryTest.java @@ -212,6 +212,13 @@ public void testRegisterArrowValidationAndReuse() throws Exception { .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)); + } } } diff --git a/rust/lance/src/dataset/sql.rs b/rust/lance/src/dataset/sql.rs index 1a3c39439ca..650314cdd59 100644 --- a/rust/lance/src/dataset/sql.rs +++ b/rust/lance/src/dataset/sql.rs @@ -188,6 +188,11 @@ impl SqlQueryBuilder { // registered, so a later duplicate wins. let mut registered = std::collections::HashSet::new(); for (name, provider) in self.extra_tables.into_iter().rev() { + if name.trim().is_empty() { + return Err(lance_core::Error::invalid_input(format!( + "registered table name must be non-empty, got '{name}'" + ))); + } if registered.insert(name.clone()) { ctx.register_table(name, provider)?; } @@ -603,5 +608,21 @@ mod tests { .await .unwrap(); pretty_assertions::assert_eq!(collect_ids(&results), Vec::::new()); + + // A blank relation name is rejected when the query is built. + let blank = Arc::new(MemTable::try_new(ids_batch(vec![]).schema(), vec![vec![]]).unwrap()); + let err = ds + .sql("SELECT 1") + .table_name("t") + .register_table(" ", blank) + .build() + .await + .err() + .expect("expected an error for a blank table name"); + assert!( + matches!(&err, lance_core::Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + assert!(err.to_string().contains("non-empty"), "message should mention the empty name: {err}"); } }