diff --git a/java/lance-jni/src/sql.rs b/java/lance-jni/src/sql.rs index f378577eedf..452e06754d3 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_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,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,15 @@ fn sql_builder( table_name: JObject, with_row_id: jboolean, with_row_addr: jboolean, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) -> Result { 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) } diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index cce6d939222..ea25c9c3c2a 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,6 +31,9 @@ public class SqlQuery { private String table = DEFAULT_TABLE_NAME; private boolean withRowId = false; 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; @@ -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). + * + *

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 tableName, boolean withRowId, boolean withRowAddr, - long streamAddress) + long streamAddress, + 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 2e3dcab9892..54a45c6fa33 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,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 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 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; + } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..bebd7063809 100644 --- a/python/python/lance/dataset.py +++ b/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. diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 898bdf93069..755925375af 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5381,6 +5381,41 @@ def test_dataset_sql(tmp_path: Path): assert pa.Table.from_batches(complex_result) == expected_complex +@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. 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", data) + .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_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/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..650314cdd59 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,108 @@ 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. + /// + /// # 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 + } + + /// 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 n_batches = batches.len(); + let schema = batches.first().map(|b| b.schema()).ok_or_else(|| { + 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])?))) + } + /// 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 +183,20 @@ impl SqlQueryBuilder { row_addr, )), )?; + // 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 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)?; + } + } register_functions(&ctx); let df = ctx.sql(&self.sql).await?; Ok(SqlQuery::new(df)) @@ -128,6 +248,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 +486,143 @@ 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["id"].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]["n"].as_primitive::().value(0), 4); + + // 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] + 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()); + + // 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}"); + } }