diff --git a/java/lance-jni/src/sql.rs b/java/lance-jni/src/sql.rs index 2111d02c9fd..56d69279b0e 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, block_on}; -use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; use jni::JNIEnv; use jni::objects::{JClass, JObject, JString}; use jni::sys::{JNI_TRUE, jboolean, jlong}; @@ -23,6 +23,8 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords( with_row_id: jboolean, with_row_addr: jboolean, stream_addr: jlong, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) { ok_or_throw_without_return!( env, @@ -34,11 +36,14 @@ pub extern "system" fn Java_org_lance_SqlQuery_intoBatchRecords( with_row_id, with_row_addr, stream_addr, + extra_table_names, + extra_stream_addrs, ) .map_err(|e| Error::input_error(e.to_string())) ) } +#[allow(clippy::too_many_arguments)] fn inner_into_batch_records( env: &mut JNIEnv, java_dataset: JObject, @@ -47,6 +52,8 @@ fn inner_into_batch_records( with_row_id: jboolean, with_row_addr: jboolean, stream_addr: jlong, + extra_table_names: JObject, + extra_stream_addrs: JObject, ) -> Result<()> { let builder = sql_builder( env, @@ -55,6 +62,8 @@ fn inner_into_batch_records( table_name, with_row_id, with_row_addr, + extra_table_names, + extra_stream_addrs, )?; let stream = block_on(async move { @@ -69,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 cb149a7fa6f..4794d217343 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,51 @@ public SqlQuery tableName(String tableName) { return this; } + /** + * Register an additional in-memory Arrow relation (exported to {@code stream} via the C Data + * Interface) as a table named {@code name}, joinable in the SQL alongside the dataset. {@link + * #intoBatchRecords()} consumes the stream during the native call (it takes ownership of the + * underlying C stream); the caller still owns the {@link ArrowArrayStream} handle and should + * close it afterwards (typically via try-with-resources), as with {@code MergeInsert}. May be + * called multiple times to register multiple tables. + * + *

The stream is fully materialized in memory before the query executes, so callers must keep + * the input bounded to available memory. For streaming or out-of-core relations, use a custom + * table provider through the Rust API instead of this Arrow convenience method. + * + *

{@code name} is parsed as a SQL identifier, so an unquoted name is lowercased ({@code IDs} + * registers the table {@code ids}) while a quoted one keeps its case ({@code "IDs"} registers the + * table {@code IDs}). If two names resolve to the same table, the later registration replaces the + * earlier one (deduplicated when the query is built). A name that resolves to the query's own + * {@link #tableName(String)} would hide the dataset, so {@link #intoBatchRecords()} rejects it + * with an {@link IllegalArgumentException}. + * + *

Because the registered stream is consumed on the first {@link #intoBatchRecords()} call, a + * query with registered relations is single-use; calling {@link #intoBatchRecords()} a second + * time throws. + * + * @throws IllegalArgumentException if {@code name} is null or blank, or {@code stream} is null + * @throws IllegalStateException if the query was already run via {@link #intoBatchRecords()} + */ + public SqlQuery registerArrow(String name, ArrowArrayStream stream) { + if (consumed) { + throw new IllegalStateException( + "registerArrow cannot be called after intoBatchRecords(); build a new query and " + + "re-register relations."); + } + if (name == null || name.trim().isEmpty()) { + throw new IllegalArgumentException( + "registerArrow: table name must be non-empty, got: " + name); + } + if (stream == null) { + throw new IllegalArgumentException( + "registerArrow: stream must not be null (table name: " + name + ")"); + } + this.extraTableNames.add(name); + this.extraStreamAddresses.add(stream.memoryAddress()); + return this; + } + public SqlQuery withRowId(boolean withRowId) { this.withRowId = withRowId; return this; @@ -51,10 +101,29 @@ public SqlQuery withRowAddr(boolean withAddr) { } public ArrowReader intoBatchRecords() throws IOException { + if (consumed) { + throw new IllegalStateException( + "intoBatchRecords() was already called on this SqlQuery; a query with registered " + + "Arrow relations is single-use because the registered streams are consumed. " + + "Build a new query and re-register them."); + } try (LockManager.ReadLock readLock = dataset.acquireReadLock(); ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { + // A query with registered streams is one-shot: the native call below consumes them. Mark it + // consumed just before that call, so a failed output-stream allocation does not brick a + // retry. + if (!extraTableNames.isEmpty()) { + consumed = true; + } intoBatchRecords( - dataset, sql, Optional.ofNullable(table), withRowId, withRowAddr, s.memoryAddress()); + dataset, + sql, + Optional.ofNullable(table), + withRowId, + withRowAddr, + s.memoryAddress(), + extraTableNames, + extraStreamAddresses); return Data.importArrayStream(dataset.allocator(), s); } } @@ -65,7 +134,9 @@ private static native void intoBatchRecords( Optional 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..9ab457b6a47 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,152 @@ public void testToRecordBatches() throws IOException { reader.getVectorSchemaRoot().getSchema().toString()); reader.close(); } + + @Test + public void testRegisterArrow() throws Exception { + // An in-memory Arrow relation (column `id`) to semi-join against the dataset; 100 is absent. + Schema idSchema = + new Schema(Collections.singletonList(Field.nullable("id", new ArrowType.Int(32, true)))); + try (VectorSchemaRoot ids = VectorSchemaRoot.create(idSchema, allocator)) { + ids.allocateNew(); + IntVector idVector = (IntVector) ids.getVector("id"); + int[] wanted = {1, 5, 39, 100}; + for (int i = 0; i < wanted.length; i++) { + idVector.setSafe(i, wanted[i]); + } + ids.setRowCount(wanted.length); + + // The native side consumes the stream; we allocate and close it (try-with-resources). + try (ArrowArrayStream stream = toStream(ids)) { + ArrowReader reader = + dataset + .sql( + "select id from " + + NAME + + " where id in (select id from filter_ids) order by id") + .tableName(NAME) + .registerArrow("filter_ids", stream) + .intoBatchRecords(); + List got = new ArrayList<>(); + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + for (int i = 0; i < root.getRowCount(); i++) { + got.add((Integer) root.getVector(0).getObject(i)); + } + } + reader.close(); + // 100 is not in the dataset; the rest are the intersection, in order. + Assertions.assertEquals(Arrays.asList(1, 5, 39), got); + } + } + } + + @Test + public void testRegisterArrowMultiple() throws Exception { + // Register two relations in one query and join both; only ids present in all three survive. + try (VectorSchemaRoot a = idTable(1, 2, 3, 10); + VectorSchemaRoot b = idTable(2, 3, 4, 10); + ArrowArrayStream sa = toStream(a); + ArrowArrayStream sb = toStream(b)) { + ArrowReader reader = + dataset + .sql( + "select id from " + + NAME + + " where id in (select id from a) and id in (select id from b) order by id") + .tableName(NAME) + .registerArrow("a", sa) + .registerArrow("b", sb) + .intoBatchRecords(); + List got = new ArrayList<>(); + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + for (int i = 0; i < root.getRowCount(); i++) { + got.add((Integer) root.getVector(0).getObject(i)); + } + } + reader.close(); + // a = {1,2,3,10}, b = {2,3,4,10}, dataset ids 0..39, so the intersection is {2,3,10}. + Assertions.assertEquals(Arrays.asList(2, 3, 10), got); + } + } + + @Test + public void testRegisterArrowValidationAndReuse() throws Exception { + // registerArrow rejects invalid inputs at the boundary. + try (VectorSchemaRoot ids = idTable(1, 2); + ArrowArrayStream stream = toStream(ids)) { + SqlQuery q = dataset.sql("select id from " + NAME).tableName(NAME); + Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow("", stream)); + Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow(null, stream)); + Assertions.assertThrows(IllegalArgumentException.class, () -> q.registerArrow("ids", null)); + } + + // A query with registered relations is single-use: the stream is consumed on the first call, + // so a second intoBatchRecords() throws rather than handing JNI a dead stream. + try (VectorSchemaRoot ids = idTable(1, 2); + ArrowArrayStream stream = toStream(ids)) { + SqlQuery q = + dataset + .sql("select id from " + NAME + " where id in (select id from ids)") + .tableName(NAME) + .registerArrow("ids", stream); + q.intoBatchRecords().close(); + Assertions.assertThrows(IllegalStateException.class, q::intoBatchRecords); + + // Registering another relation after the query is consumed is also rejected. + try (VectorSchemaRoot more = idTable(3); + ArrowArrayStream moreStream = toStream(more)) { + Assertions.assertThrows( + IllegalStateException.class, () -> q.registerArrow("more", moreStream)); + } + } + + // A relation whose name resolves to the query's own table name would hide the dataset, so the + // native build rejects it (the Rust InvalidInput surfaces as IllegalArgumentException). Assert + // on the message too: DataFusion's own "table already exists" error maps to the same exception. + try (VectorSchemaRoot ids = idTable(1, 2); + ArrowArrayStream stream = toStream(ids)) { + SqlQuery q = + dataset.sql("select id from " + NAME).tableName(NAME).registerArrow(NAME, stream); + IllegalArgumentException error = + Assertions.assertThrows(IllegalArgumentException.class, q::intoBatchRecords); + Assertions.assertTrue( + error.getMessage().contains("resolves to") + && error.getMessage().contains("the query's table name '" + NAME + "'"), + "unexpected message: " + error.getMessage()); + } + } + + /** Build a single-column (`id`: Int32) in-memory relation. */ + private VectorSchemaRoot idTable(int... ids) { + Schema schema = + new Schema(Collections.singletonList(Field.nullable("id", new ArrowType.Int(32, true)))); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + root.allocateNew(); + IntVector idVector = (IntVector) root.getVector("id"); + for (int i = 0; i < ids.length; i++) { + idVector.setSafe(i, ids[i]); + } + root.setRowCount(ids.length); + return root; + } + + /** + * Serialize a single-batch root to a self-contained Arrow C-Data stream (mirrors + * MergeInsertTest). + */ + private ArrowArrayStream toStream(VectorSchemaRoot root) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray()), allocator); + ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); + Data.exportArrayStream(allocator, reader, stream); + return stream; + } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 7c063a907e2..dc2881294be 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6322,6 +6322,52 @@ def blob_handling( self._builder = self._builder.blob_handling(blob_handling) 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. + The input is fully materialized in memory when this method is called, + so ``data`` must be bounded to available memory. + + ``name`` is parsed as a SQL identifier, so an unquoted name is lowercased + (``IDs`` registers the table ``ids``) while a quoted one keeps its case + (``"IDs"``). If two names resolve to the same table, the later + registration wins. A name that resolves to the query's own + :meth:`table_name` would hide the dataset and is rejected. + + Registered names are validated when the query is executed, not when + :meth:`build` returns, so a rejected name raises from + :meth:`SqlQuery.to_batch_records` (as ``ValueError``) or + :meth:`SqlQuery.to_stream_reader` (as ``OSError``). + + 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 batch_size(self, batch_size: int) -> "SqlQueryBuilder": """ Set the maximum number of rows produced by each query batch. diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index d41ba47b1a1..958a37c9fdf 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -7472,6 +7472,63 @@ 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 + + +@pytest.mark.parametrize("name", ["t", "T", '"t"'], ids=["exact", "case", "quoted"]) +def test_dataset_sql_register_arrow_table_name_collision(tmp_path: Path, name: str): + table = pa.table({"id": [1, 2, 3]}) + ds = lance.write_dataset(table, tmp_path / "test") + + # A registered name that resolves to the query's own table name would hide the + # dataset. The check lives in the Rust builder, which runs when the query is + # executed rather than when build() returns, so build() itself succeeds. + builder = ( + ds.sql("SELECT id FROM t") + .table_name("t") + .register_arrow(name, pa.table({"id": [4]})) + ) + query = builder.build() + + # to_batch_records surfaces Lance errors as ValueError, to_stream_reader as OSError. + with pytest.raises(ValueError, match="resolves to"): + query.to_batch_records() + with pytest.raises(OSError, match="resolves to"): + builder.build().to_stream_reader() + + +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_dataset_sql_batch_size_rows(tmp_path: Path): table = pa.table({"id": range(50)}) ds = lance.write_dataset(table, tmp_path / "test_sql_batch_size_rows") diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 7f81aad6652..d4a0da7b47f 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -4478,6 +4478,38 @@ 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. + /// The input is fully materialized in memory before registration, so it + /// must be bounded to available memory. + /// + /// `name` is parsed as a SQL identifier, so an unquoted name is lowercased + /// (`IDs` registers the table `ids`) while a quoted one keeps its case + /// (`"IDs"`). If two names resolve to the same table, the later registration + /// wins. A name that resolves to the query's own `table_name` would hide the + /// dataset and is rejected. + /// + /// Registered names are validated when the query is executed, not when + /// `build()` returns, so a rejected name raises from `to_batch_records()` + /// (as `ValueError`) or `to_stream_reader()` (as `OSError`). + #[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 }) + } + #[pyo3(signature = (batch_size))] fn batch_size(&self, batch_size: usize) -> Self { Self { diff --git a/rust/lance/src/dataset/sql.rs b/rust/lance/src/dataset/sql.rs index b309029ba39..ddfd2f8f957 100644 --- a/rust/lance/src/dataset/sql.rs +++ b/rust/lance/src/dataset/sql.rs @@ -6,8 +6,9 @@ use crate::datafusion::LanceTableProvider; use crate::dataset::scanner::validate_batch_size; use crate::dataset::utils::SchemaAdapter; use arrow_array::RecordBatch; -use datafusion::common::DataFusionError; +use datafusion::common::{DataFusionError, ResolvedTableReference, TableReference}; use datafusion::dataframe::DataFrame; +use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::SendableRecordBatchStream; use datafusion::logical_expr::{Expr as LogicalExpr, LogicalPlan}; use datafusion::physical_plan::execute_stream; @@ -42,6 +43,11 @@ pub struct SqlQueryBuilder { /// Override how blob columns are materialized for this query. pub(crate) blob_handling: Option, + /// 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)>, /// Override the maximum number of rows in each scan batch. pub(crate) batch_size: Option, @@ -58,11 +64,116 @@ impl SqlQueryBuilder { with_row_id: false, with_row_addr: false, blob_handling: None, + extra_tables: Vec::new(), batch_size: None, batch_size_bytes: None, } } + /// 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. + /// + /// `name` is parsed as a SQL identifier, so an unquoted name is lowercased (`IDs` registers the table + /// `ids`) while a quoted one keeps its case (`"IDs"` registers the table `IDs`). May be called more than + /// once; if two names resolve to the same table, the later registration wins. Registering a name that + /// resolves to the query's own [`Self::table_name`] would hide the dataset, so [`Self::build`] rejects it + /// with an invalid-input error. + /// + /// # 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. @@ -135,6 +246,18 @@ impl SqlQueryBuilder { }; let row_id = self.with_row_id; let row_addr = self.with_row_addr; + // DataFusion parses a table name as a SQL identifier, lowercasing the unquoted parts, and + // resolves it to a catalog.schema.table slot. Resolve names the same way here so the + // duplicate detection below keys on the identifier DataFusion will actually register. + let config = ctx.copied_config(); + let catalog_options = &config.options().catalog; + let resolve = |name: &str| { + TableReference::from(name).resolve( + &catalog_options.default_catalog, + &catalog_options.default_schema, + ) + }; + let dataset_table = resolve(&self.table_name); let mut provider = LanceTableProvider::new(self.dataset.clone(), row_id, row_addr); if let Some(blob_handling) = self.blob_handling { provider = provider.with_blob_handling(blob_handling); @@ -145,7 +268,31 @@ impl SqlQueryBuilder { if let Some(batch_size_bytes) = self.batch_size_bytes { provider = provider.with_batch_size_bytes(batch_size_bytes); } - ctx.register_table(self.table_name, Arc::new(provider))?; + ctx.register_table(self.table_name.as_str(), Arc::new(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 name.trim().is_empty() { + return Err(lance_core::Error::invalid_input(format!( + "registered table name must be non-empty, got '{name}'" + ))); + } + let table_ref = resolve(&name); + // An extra table on the dataset's own slot would hide the dataset from the query, so + // reject it rather than letting either one silently win. + if table_ref == dataset_table { + return Err(lance_core::Error::invalid_input(format!( + "registered table name '{name}' resolves to '{table_ref}', the same table as \ + the query's table name '{}'; register it under a different name", + self.table_name + ))); + } + if registered.insert(table_ref) { + ctx.register_table(name, provider)?; + } + } register_functions(&ctx); let state = ctx.state(); let dialect = state.config_options().sql_parser.dialect; @@ -158,8 +305,13 @@ impl SqlQueryBuilder { .statement_to_plan(statement) .await .map_err(planning_error)?; - let plan = if safe_to_inject_system_columns(&plan, &columns) - && project_system_columns(&mut projected, &columns) + let plan = if safe_to_inject_system_columns( + &plan, + &columns, + &dataset_table, + &catalog_options.default_catalog, + &catalog_options.default_schema, + ) && project_system_columns(&mut projected, &columns) { // Fall back to the original plan when the rewritten statement // fails to plan (e.g. another expression aliased to a system @@ -186,9 +338,22 @@ impl SqlQueryBuilder { /// 2. Name lineage: no intermediate projection redefines an enabled system /// column name (e.g. `SELECT (_rowid + 1) AS _rowid` in a subquery), so /// the injected identifiers can only bind to the real scan columns. -fn safe_to_inject_system_columns(plan: &LogicalPlan, columns: &[(bool, &str)]) -> bool { +/// 3. Table identity: the scan targets the query's resolved dataset table, not +/// an additional table whose ordinary columns may use system-like names. +fn safe_to_inject_system_columns( + plan: &LogicalPlan, + columns: &[(bool, &str)], + dataset_table: &ResolvedTableReference, + default_catalog: &str, + default_schema: &str, +) -> bool { match plan { - LogicalPlan::TableScan(_) => true, + LogicalPlan::TableScan(scan) => { + scan.table_name + .clone() + .resolve(default_catalog, default_schema) + == *dataset_table + } LogicalPlan::Projection(projection) => { let shadows_system_column = projection .schema @@ -207,15 +372,27 @@ fn safe_to_inject_system_columns(plan: &LogicalPlan, columns: &[(bool, &str)]) - } !matches!(expr, LogicalExpr::Column(column) if &column.name == field.name()) }); - !shadows_system_column && safe_to_inject_system_columns(&projection.input, columns) + !shadows_system_column + && safe_to_inject_system_columns( + &projection.input, + columns, + dataset_table, + default_catalog, + default_schema, + ) } LogicalPlan::Filter(_) | LogicalPlan::Sort(_) | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) => plan - .inputs() - .iter() - .all(|input| safe_to_inject_system_columns(input, columns)), + | LogicalPlan::SubqueryAlias(_) => plan.inputs().iter().all(|input| { + safe_to_inject_system_columns( + input, + columns, + dataset_table, + default_catalog, + default_schema, + ) + }), _ => false, } } @@ -370,10 +547,11 @@ mod tests { use all_asserts::assert_true; use arrow_array::cast::AsArray; use arrow_array::types::{Int32Type, Int64Type, UInt64Type}; - use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray}; + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray, UInt64Array}; use arrow_schema::Schema as ArrowSchema; use arrow_schema::{DataType, Field}; use datafusion::common::DataFusionError; + use datafusion::datasource::{MemTable, TableProvider}; use lance_arrow::json::ARROW_JSON_EXT_NAME; use lance_arrow::{ARROW_EXT_NAME_KEY, SchemaExt}; use lance_core::datatypes::BlobHandling; @@ -981,6 +1159,314 @@ mod tests { 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}" + ); + } + + /// DataFusion parses a registered name as a SQL identifier and resolves it to a + /// catalog.schema.table slot, so two spellings of the same table are one registration and the + /// later one wins. + #[rstest] + #[case::later_lowercase("IDs", "ids")] + #[case::later_mixed_case("ids", "IDs")] + #[case::later_quoted("ids", "\"ids\"")] + #[case::later_schema_qualified("ids", "public.ids")] + #[case::earlier_fully_qualified("datafusion.public.ids", "ids")] + #[case::both_qualified("public.IDs", "datafusion.public.ids")] + #[tokio::test] + async fn test_sql_register_duplicate_name_resolution( + #[case] first: &str, + #[case] second: &str, + ) { + let ds = ids_dataset("memory://").await; + + let results = ds + .sql("SELECT id FROM ids ORDER BY id") + .table_name("t") + .register_arrow(first, vec![ids_batch(vec![1, 2])]) + .unwrap() + .register_arrow(second, vec![ids_batch(vec![7, 8, 9])]) + .unwrap() + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + pretty_assertions::assert_eq!(collect_ids(&results), vec![7, 8, 9]); + } + + #[tokio::test] + async fn test_sql_register_quoted_name_keeps_case() { + let ds = ids_dataset("memory://").await; + + // A quoted identifier is not lowercased, so `"IDs"` and `ids` are two different tables and + // both stay registered. + let results = ds + .sql(r#"SELECT id FROM "IDs" UNION ALL SELECT id FROM ids ORDER BY id"#) + .table_name("t") + .register_arrow(r#""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(&results), vec![1, 2, 7, 8, 9]); + } + + /// A registered name that resolves to the query's own table name would shadow the dataset, so + /// it is rejected instead of reaching DataFusion's opaque "table already exists" error. + #[rstest] + #[case::exact("t")] + #[case::different_case("T")] + #[case::quoted("\"t\"")] + #[tokio::test] + async fn test_sql_register_collides_with_table_name(#[case] name: &str) { + let ds = ids_dataset("memory://").await; + + let err = ds + .sql("SELECT id FROM t") + .table_name("t") + .register_arrow(name, vec![ids_batch(vec![1, 2])]) + .unwrap() + .build() + .await + .err() + .expect("expected an error for a name colliding with the query's table name"); + assert!( + matches!(&err, lance_core::Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + let message = err.to_string(); + assert!( + message.contains(name) && message.contains("table name 't'"), + "message should name the registered table and the query table: {message}" + ); + } + + /// System-column injection must only target the registered dataset provider. + /// A query whose top-level scan is an extra provider's ordinary `_rowid` / + /// `_rowaddr` columns must not gain them as dataset system identifiers. + #[tokio::test] + async fn test_sql_register_arrow_extra_only_not_injected_as_system_columns() { + let ds = ids_dataset("memory://test_sql_register_arrow_extra_system_columns").await; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new(ROW_ID, DataType::UInt64, false), + Field::new(ROW_ADDR, DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![7, 8])), + Arc::new(UInt64Array::from(vec![70, 80])), + Arc::new(UInt64Array::from(vec![700, 800])), + ], + ) + .unwrap(); + let extra: Arc = + Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap()); + + // The top-level scan comes from the extra provider, so requesting a + // dataset row id/address must not append the extra provider's ordinary + // system-like columns. + let results = ds + .sql("SELECT id FROM extra ORDER BY id") + .table_name("t") + .register_table("extra", Arc::clone(&extra)) + .with_row_id(true) + .with_row_addr(true) + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + let names: Vec = results[0] + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + pretty_assertions::assert_eq!(names, vec!["id".to_string()]); + pretty_assertions::assert_eq!(collect_ids(&results), vec![7, 8]); + + // Ordinary extra-table columns with system-like names remain directly + // selectable. + let results = ds + .sql("SELECT _rowid, _rowaddr FROM extra ORDER BY _rowid") + .table_name("t") + .register_table("extra", extra) + .with_row_id(true) + .with_row_addr(true) + .build() + .await + .unwrap() + .into_batch_records() + .await + .unwrap(); + let batch = &results[0]; + pretty_assertions::assert_eq!( + batch[ROW_ID].as_primitive::().values().to_vec(), + vec![70, 80] + ); + pretty_assertions::assert_eq!( + batch[ROW_ADDR] + .as_primitive::() + .values() + .to_vec(), + vec![700, 800] + ); + } + /// A malformed query is user error and must surface as [`Error::InvalidInput`] /// so that bindings and servers can map it to a client error rather than a /// generic failure.