feat(sql): register additional tables for Dataset.sql queries#7839
feat(sql): register additional tables for Dataset.sql queries#7839sbrunk wants to merge 5 commits into
Conversation
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<dyn TableProvider>) and a register_arrow(name, Vec<RecordBatch>) 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.
github-actions
Bot
added
A-python
📝 WalkthroughWalkthroughSQL query builders now accept in-memory Arrow relations. Rust registers them with DataFusion, Python exposes table and reader registration, and Java imports Arrow C streams through JNI. Tests cover joins, multiple relations, empty inputs, validation, and single-use behavior. ChangesArrow relation registration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant SqlQuery
participant JNI
participant ArrowArrayStreamReader
participant SqlQueryBuilder
participant DataFusion
SqlQuery->>JNI: pass extra relation names and stream addresses
JNI->>ArrowArrayStreamReader: read Arrow C stream batches
ArrowArrayStreamReader->>SqlQueryBuilder: register Arrow batches
SqlQueryBuilder->>DataFusion: register relation and execute SQL
DataFusion->>SqlQuery: return query batches
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (4)
rust/lance/src/dataset/sql.rs-458-460 (1)
458-460: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert the invalid-input error contract.
Line 460 accepts any failure. Assert both the invalid-input variant and message content so an unrelated regression cannot satisfy this test. As per coding guidelines, “Assert on both the error variant and the message content in tests; do not check only
is_err().”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/sql.rs` around lines 458 - 460, Strengthen the test around register_arrow in the SQL test by matching the returned error’s invalid-input variant and verifying its message contains the expected schema-related text, rather than asserting only err.is_err().Source: Coding guidelines
rust/lance/src/dataset/sql.rs-425-456 (1)
425-456: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winTest duplicate registration semantics.
register_tabledocuments that the later duplicate name wins, but no test locks down repeated registration. Add a chain with the same name twice and assert the later relation is queried. As per coding guidelines, “Every bugfix and feature must have corresponding tests.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/sql.rs` around lines 425 - 456, Add coverage in test_sql_register_arrow for duplicate registration by chaining register_arrow with the same relation name twice and different data, then query that relation and assert the results come from the later registration. Keep the test focused on verifying that repeated registration replaces the earlier relation.Source: Coding guidelines
python/python/tests/test_dataset.py-5384-5401 (1)
5384-5401: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCover the advertised reader input path.
This only exercises
pa.Table; addpa.RecordBatchReadercoverage and parameterize the shared success case. Also assert the empty-readerValueErrorcontract, since that error is translated by the PyO3 binding. As per coding guidelines, “Use@pytest.mark.parametrizefor Python tests that differ only by inputs” and “Every bugfix and feature must have corresponding tests.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/tests/test_dataset.py` around lines 5384 - 5401, The test_dataset_sql_register_arrow success case should cover both pa.Table and pa.RecordBatchReader inputs. Parameterize the shared test using these two input forms, and add a separate assertion that registering an empty pa.RecordBatchReader raises the expected ValueError translated by the PyO3 binding.Source: Coding guidelines
java/src/main/java/org/lance/SqlQuery.java-52-55 (1)
52-55: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winValidate Arrow registration inputs and avoid reusing consumed streams.
registerArrow()forwardsnameandstreamunchecked: a nullstreamNPEs onmemoryAddress(), and blank names should be rejected at the boundary.intoBatchRecords()consumes the registered C stream, but the raw address stays cached on the object, so a second call can reuse a dead stream. Clear the registration state after use or make the query one-shot.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/java/org/lance/SqlQuery.java` around lines 52 - 55, Update SqlQuery.registerArrow to reject null streams and blank names before storing registration state. In intoBatchRecords, clear extraTableName and extraStreamAddress after consuming the registered Arrow stream, ensuring subsequent calls cannot reuse the dead C stream; preserve normal behavior for valid first use.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@java/src/main/java/org/lance/SqlQuery.java`:
- Around line 68-78: The intoBatchRecords() flow must reject reuse of the
consumed stream registered by registerArrow(). Track whether extraStreamAddress
has been consumed, mark it consumed after the first JNI invocation, and fail
fast on later calls until registerArrow() provides a new stream; preserve normal
behavior for the initial use.
In `@rust/lance/src/dataset/sql.rs`:
- Around line 53-64: Add concise, runnable Rust documentation examples for
DatasetSql::register_table and DatasetSql::register_arrow, using their actual
signatures and demonstrating SQL through a join or subquery. In
python/src/dataset.rs at lines 3957-3963, add a synchronized Python-facing
example showing relation registration and subsequent SQL use; update both public
API docs without changing implementation behavior.
- Around line 64-72: The SQL registration flow must avoid materializing reader
streams into memory. In rust/lance/src/dataset/sql.rs lines 64-72, retain
register_arrow as the batch convenience API and add a reader/provider-backed
registration path; in python/src/dataset.rs lines 3965-3973, pass the converted
reader to that path instead of collecting batches; and in
java/lance-jni/src/sql.rs lines 111-120, retain the imported C-stream reader in
the registered provider rather than accumulating batches.
- Around line 418-423: Update the collect_ids function to access each
RecordBatch’s test output field by the named “id” accessor (b["id"]) instead of
positional column(0), while preserving the existing Int32 extraction and
collection behavior.
---
Other comments:
In `@java/src/main/java/org/lance/SqlQuery.java`:
- Around line 52-55: Update SqlQuery.registerArrow to reject null streams and
blank names before storing registration state. In intoBatchRecords, clear
extraTableName and extraStreamAddress after consuming the registered Arrow
stream, ensuring subsequent calls cannot reuse the dead C stream; preserve
normal behavior for valid first use.
In `@python/python/tests/test_dataset.py`:
- Around line 5384-5401: The test_dataset_sql_register_arrow success case should
cover both pa.Table and pa.RecordBatchReader inputs. Parameterize the shared
test using these two input forms, and add a separate assertion that registering
an empty pa.RecordBatchReader raises the expected ValueError translated by the
PyO3 binding.
In `@rust/lance/src/dataset/sql.rs`:
- Around line 458-460: Strengthen the test around register_arrow in the SQL test
by matching the returned error’s invalid-input variant and verifying its message
contains the expected schema-related text, rather than asserting only
err.is_err().
- Around line 425-456: Add coverage in test_sql_register_arrow for duplicate
registration by chaining register_arrow with the same relation name twice and
different data, then query that relation and assert the results come from the
later registration. Keep the test focused on verifying that repeated
registration replaces the earlier relation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 8889e12e-6fef-45cc-9a1a-6ea5f292b178
📒 Files selected for processing (7)
java/lance-jni/src/sql.rsjava/src/main/java/org/lance/SqlQuery.javajava/src/test/java/org/lance/SqlQueryTest.javapython/python/lance/dataset.pypython/python/tests/test_dataset.pypython/src/dataset.rsrust/lance/src/dataset/sql.rs
| /// 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<dyn TableProvider>) -> 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<RecordBatch>) -> lance_core::Result<Self> { |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add runnable examples for the new public registration APIs.
The Rust and PyO3-visible APIs describe the behavior but provide no synchronized usage example.
rust/lance/src/dataset/sql.rs#L53-L64: add concise Rust examples forregister_tableandregister_arrow, including a join or subquery.python/src/dataset.rs#L3957-L3963: add a Python-facing example showing relation registration and SQL use.
As per coding guidelines, “Document all public APIs with examples and links to relevant structs and methods; keep examples synchronized with actual signatures.”
📍 Affects 2 files
rust/lance/src/dataset/sql.rs#L53-L64(this comment)python/src/dataset.rs#L3957-L3963
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/sql.rs` around lines 53 - 64, Add concise, runnable
Rust documentation examples for DatasetSql::register_table and
DatasetSql::register_arrow, using their actual signatures and demonstrating SQL
through a join or subquery. In python/src/dataset.rs at lines 3957-3963, add a
synchronized Python-facing example showing relation registration and subsequent
SQL use; update both public API docs without changing implementation behavior.
Source: Coding guidelines
| pub fn register_arrow(self, name: &str, batches: Vec<RecordBatch>) -> lance_core::Result<Self> { | ||
| 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])?))) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Do not materialize arbitrary Arrow readers before registration.
Python and Java accept reader/stream inputs, but the core API forces both bindings to collect every batch into a Vec before DataFusion can execute. Large or unbounded readers can exhaust memory.
rust/lance/src/dataset/sql.rs#L64-L72: retain the batch convenience API, but add a reader/provider-backed registration path.python/src/dataset.rs#L3965-L3973: forward the converted reader to that path instead of collecting it.java/lance-jni/src/sql.rs#L111-L120: retain the imported C-stream reader in the registered provider rather than accumulating batches.
As per coding guidelines, “Avoid collecting streams of RecordBatch into memory.”
📍 Affects 3 files
rust/lance/src/dataset/sql.rs#L64-L72(this comment)python/src/dataset.rs#L3965-L3973java/lance-jni/src/sql.rs#L111-L120
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/sql.rs` around lines 64 - 72, The SQL registration
flow must avoid materializing reader streams into memory. In
rust/lance/src/dataset/sql.rs lines 64-72, retain register_arrow as the batch
convenience API and add a reader/provider-backed registration path; in
python/src/dataset.rs lines 3965-3973, pass the converted reader to that path
instead of collecting batches; and in java/lance-jni/src/sql.rs lines 111-120,
retain the imported C-stream reader in the registered provider rather than
accumulating batches.
Source: Coding guidelines
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.
- 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.
There was a problem hiding this comment.
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/dataset/sql.rs (1)
141-146: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPreserve full input context in registration validation errors.
rust/lance/src/dataset/sql.rs#L141-L146: includenameandbatches.len()in theInvalidInputmessage.java/src/main/java/org/lance/SqlQuery.java#L61-L66: include the parameter names and rejected values in bothIllegalArgumentExceptionmessages.As per coding guidelines, “Include full error context, including variable names, values, sizes, and types.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/sql.rs` around lines 141 - 146, The registration validation errors lack full input context. In rust/lance/src/dataset/sql.rs lines 141-146, update the register_arrow validation message to include the relation name and batches.len(); in java/src/main/java/org/lance/SqlQuery.java lines 61-66, update both IllegalArgumentException messages to include the relevant parameter names and rejected values.Source: Coding guidelines
🟡 Other comments (1)
java/src/main/java/org/lance/SqlQuery.java-53-54 (1)
53-54: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDo not attribute last-registration-wins behavior to DataFusion.
This behavior is implemented by
SqlQueryBuilder::builddeduplicating registrations in reverse order; the Rust comment states duplicate DataFusion registration would otherwise error.- * the later registration replaces the earlier one at query time (DataFusion register semantics). + * the later registration replaces the earlier one at query time.As per coding guidelines, “Ensure doc comments match actual semantics.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/java/org/lance/SqlQuery.java` around lines 53 - 54, Update the JavaDoc for SqlQuery registration to remove the attribution of last-registration-wins behavior to DataFusion. Describe duplicate-name handling as the behavior implemented by SqlQueryBuilder::build through reverse-order deduplication, while preserving the guidance that multiple tables may be registered.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance/src/dataset/sql.rs`:
- Around line 141-146: The registration validation errors lack full input
context. In rust/lance/src/dataset/sql.rs lines 141-146, update the
register_arrow validation message to include the relation name and
batches.len(); in java/src/main/java/org/lance/SqlQuery.java lines 61-66, update
both IllegalArgumentException messages to include the relevant parameter names
and rejected values.
---
Other comments:
In `@java/src/main/java/org/lance/SqlQuery.java`:
- Around line 53-54: Update the JavaDoc for SqlQuery registration to remove the
attribution of last-registration-wins behavior to DataFusion. Describe
duplicate-name handling as the behavior implemented by SqlQueryBuilder::build
through reverse-order deduplication, while preserving the guidance that multiple
tables may be registered.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: d78597b2-beec-48f5-af2c-6c91b4837917
📒 Files selected for processing (6)
java/lance-jni/src/sql.rsjava/src/main/java/org/lance/SqlQuery.javajava/src/test/java/org/lance/SqlQueryTest.javapython/python/lance/dataset.pypython/python/tests/test_dataset.pyrust/lance/src/dataset/sql.rs
- 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).
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
rust/lance/src/dataset/sql.rs (2)
93-96: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReject invalid relation names at the API boundary.
register_tableaccepts empty or whitespace-only names and defers failure untilbuild(), whileregister_arrowinherits that behavior. Validate names before storing them or duringbuild(), returning a descriptiveInvalidInputcontaining the offending name.As per coding guidelines, “Validate inputs at API boundaries and reject invalid values with descriptive errors.”
Also applies to: 140-148
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/sql.rs` around lines 93 - 96, Update register_table to reject empty or whitespace-only relation names at the API boundary before pushing into extra_tables, returning a descriptive InvalidInput error that includes the offending name. Ensure register_arrow receives the same validation behavior, either by reusing the shared validation or validating during build without storing invalid entries.Source: Coding guidelines
562-568: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winBorrow
errin the variant assertion.
matches!(err, ...)moves the non-Copyerror, so the format argument and subsequenterr.to_string()use a moved value. Match by reference instead.Proposed fix
- matches!(err, lance_core::Error::InvalidInput { .. }), + matches!(&err, lance_core::Error::InvalidInput { .. }),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/sql.rs` around lines 562 - 568, Update the variant assertion in the affected test to match against a reference to err rather than moving the non-Copy error. Preserve the existing InvalidInput validation and subsequent err formatting and to_string checks.java/src/main/java/org/lance/SqlQuery.java (2)
89-94: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSet
consumedafter output-stream allocation.The flag is set before
ArrowArrayStream.allocateNew(...); if allocation fails, JNI was never called and the input streams remain unconsumed, but retries are permanently blocked. Move the assignment inside thetryblock immediately before the native invocation.Proposed fix
- if (!extraTableNames.isEmpty()) { - consumed = true; - } try (ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { + if (!extraTableNames.isEmpty()) { + consumed = true; + } intoBatchRecords(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/java/org/lance/SqlQuery.java` around lines 89 - 94, Move the consumed assignment for registered streams from before ArrowArrayStream allocation into the try block immediately before the native invocation, preserving the extraTableNames guard. This ensures failed output-stream allocation does not block retries while still marking streams consumed before JNI execution.
61-70: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReject registration after the query is consumed.
registerArrowcan append a new stream afterconsumedis alreadytrue. The method succeeds, but the nextintoBatchRecords()call will always reject it, leaving an unexecutable registration and a raw stream address to clean up. Fail fast before mutating either list.Proposed fix
public SqlQuery registerArrow(String name, ArrowArrayStream stream) { + if (consumed) { + throw new IllegalStateException("registerArrow cannot be called after intoBatchRecords()"); + } if (name == null || name.trim().isEmpty()) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/java/org/lance/SqlQuery.java` around lines 61 - 70, Update registerArrow to reject calls when the query’s consumed state is already true, before modifying extraTableNames or extraStreamAddresses. Throw the established argument/state exception with a clear message, while preserving the existing name and stream validation for unconsumed queries.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@java/src/main/java/org/lance/SqlQuery.java`:
- Around line 89-94: Move the consumed assignment for registered streams from
before ArrowArrayStream allocation into the try block immediately before the
native invocation, preserving the extraTableNames guard. This ensures failed
output-stream allocation does not block retries while still marking streams
consumed before JNI execution.
- Around line 61-70: Update registerArrow to reject calls when the query’s
consumed state is already true, before modifying extraTableNames or
extraStreamAddresses. Throw the established argument/state exception with a
clear message, while preserving the existing name and stream validation for
unconsumed queries.
In `@rust/lance/src/dataset/sql.rs`:
- Around line 93-96: Update register_table to reject empty or whitespace-only
relation names at the API boundary before pushing into extra_tables, returning a
descriptive InvalidInput error that includes the offending name. Ensure
register_arrow receives the same validation behavior, either by reusing the
shared validation or validating during build without storing invalid entries.
- Around line 562-568: Update the variant assertion in the affected test to
match against a reference to err rather than moving the non-Copy error. Preserve
the existing InvalidInput validation and subsequent err formatting and to_string
checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 6295ee18-e1b0-47c6-8144-0a28ad62459b
📒 Files selected for processing (2)
java/src/main/java/org/lance/SqlQuery.javarust/lance/src/dataset/sql.rs
- 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).
1 participant
Summary
Adds a way to register extra in-memory / provider-backed tables in the DataFusion
SessionContextthat backsDataset.sql(...), so a query can join a caller-supplied relation instead of encoding it as a large literalIN (...)list.New API:
SqlQueryBuilder::register_table(name: &str, provider: Arc<dyn TableProvider>) -> Selfand a convenienceSqlQueryBuilder::register_arrow(name: &str, batches: Vec<RecordBatch>) -> Result<Self>.SqlQuery.registerArrow(String name, ArrowArrayStream stream).SqlQueryBuilder.register_arrow(name, data)(accepts a pyarrowTableorRecordBatchReader).Example (Rust): semi-join a caller-supplied id set.
Motivation
Dataset.sqltoday can only reference the dataset itself. To scope a query by a set of ids (or to join another table) the only options are:IN (...)list, which is expensive to parse/plan for large sets, orSqlQueryBuilderentirely and hand-build aSessionContext, registering aLanceTableProviderplus the other table (this is what the "join two datasets" example indocs/src/integrations/datafusion.mdcurrently requires).Option 2 is Rust-only. The Java and Python bindings have no way to add a table to a
Dataset.sqlquery at all. This change exposes the capability through the builder so all three bindings can do it, and so the common case (join an in-memory Arrow relation) is a one-liner.Design decisions (and alternatives considered)
Table type is
Arc<dyn TableProvider>, notVec<RecordBatch>or a reader.SqlQueryBuilderis#[derive(Clone, Debug)]and the Python binding clones it on every builder call, so any stored field must beClone + Debug. That rules outimpl RecordBatchReader/SendableRecordBatchStream(neither isClone).Arc<dyn TableProvider>isClone + Debug, is the most general option (callers can register aMemTable, another dataset'sLanceTableProvider, a view, etc.), and matches DataFusion's ownSessionContext::register_table(_, Arc<dyn TableProvider>).register_arrowis a thin, fallible convenience.register_tablestays infallible like the other builder methods (table_name,with_row_id, ...).register_arrowwrapsVec<RecordBatch>in aMemTable, which is the fallible part, so it returnsResult<Self>and keeps that fallibility at the boundary rather than inbuild().Naming.
register_tablefollows the existing in-tree precedent (FtsQueryUDTFBuilder::register_table) and DataFusion'sregister_table, rather thanwith_table.with_*in this builder is reserved for scalar toggles.Empty relations.
register_arrowrequires at least one batch (it derives the schema from the first) and returns an error otherwise, since an emptyVec<RecordBatch>has no schema. If you need to register an empty relation, build aMemTablewith a known schema and useregister_table; the query then returns zero rows rather than failing with "table not found". (An earlier approach storedVec<RecordBatch>and skipped empty inputs inbuild(), which silently produced a "table not found" error at query time; moving toTableProviderremoves that footgun because a provider always carries a schema.)name: &strmatches the other builder methods.impl Into<TableReference>would allow schema-qualified names but would be the only method here not taking&str; happy to switch if preferred.Duplicate names / multiple calls.
register_tablemay be called more than once. A later duplicate name wins, deduplicated inbuild()(registering the same name into the DataFusion context twice would otherwise error, since DataFusion'sregister_tablerejects a duplicate rather than replacing it).In-memory vs streaming.
register_arrowbuilds aMemTable, so its input is materialized. That is by design for an in-memory relation the caller already holds; for a large or streaming source, pass a customTableProvidertoregister_tablerather than collecting it. AMemTableis used here because aTableProvidermust be re-scannable (self-joins, multi-partition planning), which a one-shot reader cannot satisfy.Cross-language notes
ArrowArrayStreamimported through the JNI (ArrowArrayStreamReader::from_raw), the same patternMergeInsertuses.Ownership: the native call consumes (takes ownership of) the C stream; the caller owns the Java
ArrowArrayStreamhandle and closes it afterwards (e.g. try-with-resources), asMergeInsertTestdoes.registerArrowmay be called multiple times to register multiple tables, matching Rust and Python.registerArrowrejects a null/blank name or a null stream. A query with registered relations is single-use: the streams are consumed on the firstintoBatchRecords(), so a secondintoBatchRecords(), or aregisterArrowafter it, throws.register_arrowaccepts a pyarrowTable/RecordBatchReadervia the existingconvert_readerhelper. A generalregister_tablewas not added because Python has no RustTableProviderto pass; the existingFFILanceTableProvider+ the externaldatafusionpackage already cover provider-level joins (seedocs/src/integrations/datafusion.md). A Pythonregister_tableaccepting an object implementing__datafusion_table_provider__could be a follow-up.