From 951ae51ef603cba4af4582182ba9c54af9b68eab Mon Sep 17 00:00:00 2001 From: Dan Rammer Date: Tue, 20 Oct 2026 21:40:20 -0500 Subject: [PATCH 01/33] docs: add realtime engine (RTE) getting-started guide (#4548) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a user-facing realtime engine (RTE) page at `docs/web/tables/realtime-engine.mdx`, listed under **Management** after Branches. It is adapted from the internal WAL field guide (lancedb/sophon#6213), with the deployment and operations material left out. The page covers: - what RTE is and how it works (memtable → SSTables → compaction; reads merge the base table with the fresh tier) - getting started: `set_unenforced_primary_key`, indexes, `enable_rte` / `ALTER TABLE ... SET RTE`, `merge_insert` upserts and queries, partial updates, deletes, and `checkpoint_rte` / `disable_rte` - performance: write throughput and shard scaling, and read latency with and without RTE - a note on operations not yet supported (time travel, branching, `count_rows`, row IDs, `TABLESAMPLE`) Also aligns related docs with current RTE naming and behavior: - `sql/ddl/table.mdx`: `SET MEMWAL` / `UNSET MEMWAL` → `SET RTE` / `UNSET RTE` - `maintenance/background-jobs.mdx`: job names are `memwal_compaction` and `memwal_catchup_` - `feature-engineering/materialized-views.mdx`: MemWAL → realtime engine - Python `AsyncTable.enable_rte` / `RTEWriteSpec.with_maintained_indexes` docstrings: by default, indexes created after enabling RTE are maintained too The examples are written directly in the page rather than generated from `docs/web-tests`, because RTE needs an Enterprise server to run. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: XY Zhan --- docs/web/docs.json | 1 + .../materialized-views.mdx | 2 +- docs/web/maintenance/background-jobs.mdx | 6 +- docs/web/sql/ddl/table.mdx | 14 +- docs/web/tables/realtime-engine.mdx | 229 ++++++++++++++++++ python/python/lancedb/_lancedb.pyi | 6 +- python/python/lancedb/table.py | 12 +- 7 files changed, 251 insertions(+), 19 deletions(-) create mode 100644 docs/web/tables/realtime-engine.mdx diff --git a/docs/web/docs.json b/docs/web/docs.json index a8fb44007..9c9468a3b 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -327,6 +327,7 @@ "management/namespace", "management/table", "tables/branching", + "tables/realtime-engine", "management/materialized-view", "management/view" ] diff --git a/docs/web/feature-engineering/materialized-views.mdx b/docs/web/feature-engineering/materialized-views.mdx index ffe4fc959..1552c3638 100644 --- a/docs/web/feature-engineering/materialized-views.mdx +++ b/docs/web/feature-engineering/materialized-views.mdx @@ -52,7 +52,7 @@ In Python SDK 0.40.0, create the source with above does. Stable row IDs let the view track source rows across updates and compaction. Set this option when creating the table; it cannot be added later. -A source with a MemWAL write spec or retained, uncompacted write-ahead-log +A source with the [realtime engine](/tables/realtime-engine) enabled or retained, uncompacted write-ahead-log rows is not accepted by the local materialized-view implementation. ## Refresh behavior {#incremental} diff --git a/docs/web/maintenance/background-jobs.mdx b/docs/web/maintenance/background-jobs.mdx index dcfa31321..1a710fe79 100644 --- a/docs/web/maintenance/background-jobs.mdx +++ b/docs/web/maintenance/background-jobs.mdx @@ -27,13 +27,13 @@ also start cache-warming jobs through the admin API. ## Write-ahead log maintenance {#wal} -For a table with a [write-ahead log](/sql/ddl/table#properties) enabled, two jobs keep the +For a table with the [realtime engine](/tables/realtime-engine) enabled, two jobs keep the log from growing without bound: | Job type | What it does | | --- | --- | -| `sstable_compaction` | Folds a table shard's SSTables into the base table by primary key. Distinct from `compact`, which merges base-table fragments. The two are unrelated. | -| `wal_index_catchup` | Extends a base-table index over the rows `sstable_compaction` just folded in, and records how far the index has caught up. Nothing else advances that record. | +| `memwal_compaction` | Folds a table shard's SSTables into the base table by primary key. Distinct from `compact`, which merges base-table fragments. The two are unrelated. | +| `memwal_catchup_` | Extends a base-table index over the rows `memwal_compaction` just folded in, and records how far the index has caught up. Nothing else advances that record. The suffix names the index type, for example `memwal_catchup_btree`, `memwal_catchup_fts`, or `memwal_catchup_ivfpq`. | ## Index maintenance {#index} diff --git a/docs/web/sql/ddl/table.mdx b/docs/web/sql/ddl/table.mdx index ecb522fe6..792568220 100644 --- a/docs/web/sql/ddl/table.mdx +++ b/docs/web/sql/ddl/table.mdx @@ -102,18 +102,20 @@ duplicates impossible. It tells the planner what to assume. The column must be declared `NOT NULL`. ```sql -ALTER TABLE t SET MEMWAL -ALTER TABLE t SET MEMWAL WITH (sharding = 'bucket', column = 'id', buckets = 4) -ALTER TABLE t SET MEMWAL WITH (maintained_indexes = '') -ALTER TABLE t UNSET MEMWAL +ALTER TABLE t SET RTE +ALTER TABLE t SET RTE WITH (sharding = 'bucket', column = 'id', buckets = 4) +ALTER TABLE t SET RTE WITH (maintained_indexes = '') +ALTER TABLE t UNSET RTE ``` -A write-ahead log makes small frequent writes cheap by batching them out of the +The [realtime engine](/tables/realtime-engine) puts a write-ahead log in front +of the table, making small frequent writes cheap by batching them out of the commit path. `sharding` and `buckets` spread the log so concurrent writers do not contend on one place; `bucket` sharding names its `column`, which must be the table's unenforced primary key, and inserts into the table then need a scalar index on that key. `maintained_indexes` chooses which indexes are kept -current against WAL rows. A table with a computed column cannot take a +current against WAL rows; omitted, every supported index is kept current, +including ones created later. A table with a computed column cannot take a write-ahead log. See [Background jobs](/maintenance/background-jobs#wal) for how the log gets folded back into the table. diff --git a/docs/web/tables/realtime-engine.mdx b/docs/web/tables/realtime-engine.mdx new file mode 100644 index 000000000..88738791c --- /dev/null +++ b/docs/web/tables/realtime-engine.mdx @@ -0,0 +1,229 @@ +--- +title: "Realtime engine (RTE)" +sidebarTitle: "Realtime Engine" +description: "Absorb small, frequent writes cheaply and make them queryable within milliseconds." +icon: "bolt" +keywords: ["rte", "realtime engine", "wal", "memwal", "lsm", "write-ahead log", "streaming writes", "upsert"] +--- + +Enterprise + +## What is the realtime engine? {#what-is-rte} + +Every write to a LanceDB table is normally a commit to the underlying Lance +dataset. That is ideal for bulk loads, but a stream of small upserts becomes a +stream of small commits, and each write is only visible once its commit lands. + +The realtime engine (RTE) puts an LSM-style write tier in front of the table. +It turns many small writes into a few large commits, and makes each write +**searchable within milliseconds** of being acknowledged. It is enabled per +table, and nothing about how you query that table changes. + +## How it works {#how-it-works} + +- **Writes.** A `merge_insert` or `delete` lands in an in-memory **memtable** + and is acknowledged once it is durable in the write-ahead log on object + storage. +- **Flush and compaction.** Memtables freeze into **SSTables**, which background + jobs compact into the base table, catching its indexes up as they go. See + [Background jobs](/maintenance/background-jobs#wal). +- **Reads.** Every query runs against the base table and the fresh tier + (memtables and SSTables) in parallel. Fresh rows supersede base rows with the + same primary key, and the two are merged into one result. +- **Indexes stay fresh.** Vector, full-text, BTree, bitmap, and most other + scalar indexes are maintained on the memtable, so new rows are searchable + through the index as soon as they are written. + +## Getting started {#getting-started} + +### 1. Create a table with a primary key {#primary-key} + +RTE deduplicates writes by primary key. Create the table with a non-nullable +key column, then declare it with `set_unenforced_primary_key`. The key is a +single `int32`, `int64`, string, or binary column, and cannot be changed once +set. + +```python +from datetime import timedelta + +import lancedb +import numpy as np +import pyarrow as pa +from lancedb import RTEWriteSpec +from lancedb.index import FTS, BTree, IvfPq + +db = lancedb.connect("db://my_db", api_key="...", host_override="https://...") + +n, dim = 1024, 128 +schema = pa.schema([ + pa.field("id", pa.int64(), nullable=False), + pa.field("vector", pa.list_(pa.float32(), dim)), + pa.field("text", pa.utf8()), +]) +initial_data = pa.table({ + "id": pa.array(range(n), pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(np.random.rand(n * dim).astype("float32")), dim), + "text": pa.array([f"doc {i} alpha bravo" for i in range(n)]), +}, schema=schema) +table = db.create_table("docs", initial_data, schema=schema) +table.set_unenforced_primary_key("id") +``` + +Or in SQL: + +```sql +ALTER TABLE docs SET UNENFORCED PRIMARY KEY (id) +``` + +### 2. Create indexes {#indexes} + +By default, RTE maintains every supported index on the table, including +indexes created after it is enabled. A scalar index on the primary key is +required for partial updates and deletes. + +```python +table.create_index("vector", config=IvfPq()) +table.create_index("text", config=FTS()) +table.create_index("id", config=BTree()) +table.wait_for_index(["vector_idx", "text_idx", "id_idx"], timeout=timedelta(minutes=5)) +``` + +### 3. Enable RTE {#enable} + +```python +table.enable_rte(RTEWriteSpec.unsharded()) +``` + +Or in SQL: + +```sql +ALTER TABLE docs SET RTE +``` + +### 4. Write and query {#write-and-query} + + +On an RTE table, an upsert through `merge_insert` is the only way to write +rows. Use the form below exactly: `when_matched_update_all()` with no filter, +plus `when_not_matched_insert_all()`. Other `merge_insert` forms, such as +conditional updates or deleting rows missing from the source, are rejected. + + +Writes go through `merge_insert` on the primary key, and are visible to queries +immediately: + +```python +new_rows = pa.table({ + "id": pa.array([2000, 2001], pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(np.random.rand(2 * dim).astype("float32")), dim), + "text": pa.array(["fresh charlie delta", "fresh echo foxtrot"]), +}, schema=schema) +query_vector = np.random.rand(dim).astype("float32") + +table.merge_insert("id") \ + .when_matched_update_all() \ + .when_not_matched_insert_all() \ + .execute(new_rows) + +table.search(query_vector).limit(10).to_list() +table.search("charlie").limit(10).to_list() +``` + +**Partial updates.** Send the primary key plus only the columns you want to +change; omitted columns keep their current values. This requires a scalar +index on the primary key. + +```python +patch = pa.table( + {"id": [0, 1], "text": ["doc 0 revised", "doc 1 revised"]}, + schema=pa.schema([schema.field("id"), schema.field("text")]), +) +table.merge_insert("id").when_matched_update_all().when_not_matched_insert_all().execute(patch) +``` + +**Deletes.** The predicate must bind the primary key to literal values. Ranges, +`OR`, and filters on other columns are rejected. Deletes require a scalar index +on the primary key. + +```python +table.delete("id IN (2000, 2001)") +``` + +### Disabling RTE {#disable} + +Stop writes and wait for in-flight writes to finish, then flush pending writes +into the base table and disable. `checkpoint_rte()` only drains writes that +existed when it started, so anything written during the checkpoint can be left +behind. `disable_rte()` is rejected while the write-ahead log still holds rows, +so run `checkpoint_rte()` first. + +```python +table.checkpoint_rte() +table.disable_rte() +``` + +## Performance {#performance} + +### Writes {#write-performance} + +RTE acknowledges a write as soon as it is durable in the write-ahead log, so +ingest is no longer limited by how fast the table can commit. + +**Compared with writing directly to the table.** On tables ranging from +hundreds of thousands to tens of millions of rows, RTE delivered 4–12× the +write throughput and 24–52× lower p99 write latency, depending on the mix of +inserts, updates, and deletes. + +**Scaling with shards.** An RTE table can be split into +[shards](/sql/ddl/table#properties) that write independently, and a single +table's ingest grows with its shard count. In our testing on AWS, with every +index on the table (vector, full-text, and BTree) kept current, a single table +reached these peak rates: + +| Shards | Peak ingest | +| --: | --: | +| 1 | 90 MB/s | +| 2 | 170 MB/s | +| 4 | 340 MB/s | +| 8 | 625 MB/s | +| 16 | 690 MB/s | +| 32 | 885 MB/s | + +Adding shards scales a table's ingest nearly linearly. + +### Reads {#read-performance} + +Reads on an RTE table stay close to those on a fully compacted table, even +with the fresh tier at its fullest before compaction. + +The table below compares a 5M-row table with a 128-dimension IVF-PQ index plus +BTree and FTS indexes, with and without RTE. The RTE table holds one +unflushed memtable and three SSTables, each from a 100k-row write wave, half of +which are updates to existing rows. Latencies are p50 in milliseconds. + +| Query | Without RTE | With RTE | +| --- | --: | --: | +| Primary-key lookup | 62 | 63 | +| Vector search | 68 | 71 | +| Vector search with prefilter | 90 | 91 | +| Full-text term | 63 | 64 | +| Full-text phrase | 68 | 82 | +| Hybrid search (RRF) | 68 | 83 | +| SQL point lookup | 189 | 187 | + +Most queries are within a few milliseconds of the base table. Hybrid and phrase +queries add roughly 15 ms to merge the two tiers. + + +Support for all table operations on RTE tables is coming soon. Time travel, +branching, `count_rows`, row IDs (`with_row_id`), and `TABLESAMPLE` are not yet +available. + + +## Next steps {#next-steps} + +- [Merge insert](/ingestion/merge-insert) +- [Background jobs](/maintenance/background-jobs#wal) +- [`ALTER TABLE ... SET RTE`](/sql/ddl/table#properties) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 988482a97..5a613b210 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -850,9 +850,9 @@ class RTEWriteSpec: @staticmethod def unsharded() -> "RTEWriteSpec": ... def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "RTEWriteSpec": - """Set which indexes the RTE keeps up to date. None resolves every - index on the table at install, failing if one cannot be maintained; - a list is verbatim, empty means none.""" + """Set which indexes the RTE keeps up to date. None is every index on + the table, now and later, skipping any it cannot maintain; a list is + verbatim, empty means none.""" ... def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "RTEWriteSpec": """Return a copy of this spec recording the given default diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 48d5582d1..60404b73c 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5564,12 +5564,12 @@ class AsyncTable: via [`set_unenforced_primary_key`]; bucket sharding additionally requires it to be the single column being bucketed. - By default the RTE maintains every index on the table, resolved - here — a snapshot, so an index created afterwards needs the spec unset - and set again. This fails if one cannot be maintained; name the set - with ``with_maintained_indexes`` to install anyway. That pins an exact - set (a still-building index is rejected, not omitted); ``[]`` maintains - none. + By default the RTE maintains every index on the table, re-read as the + table changes, so an index created later is maintained too; one of a + kind the RTE cannot mirror is skipped. Name the set with + ``with_maintained_indexes`` to pin it instead: each name must already + exist and be maintainable (a still-building index is rejected, not + omitted); ``[]`` maintains none. Parameters ---------- From bedb6ab3b1d9f1feeafc380a3453ee9161f09e7d Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Tue, 20 Oct 2026 21:47:11 -0700 Subject: [PATCH 02/33] docs(python): rework PyTorch page around StreamingDataset (#4606) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Rework `docs/web/training/torch.mdx` for people who already know the PyTorch `DataLoader`, built around `lancedb.streaming.StreamingDataset`. - The quickstart trains a state-to-action policy on the real `lance-format/lerobot-xvla-soft-fold` dataset loaded via `hf://`, using a zero-copy Arrow-to-tensor transform. - Add a short table mapping PyTorch idioms (`IterableDataset`, `DistributedSampler`, `collate_fn`, `random_split`, `StatefulDataLoader`) to their LanceDB equivalents. - Add an Optimizations section: Transform (including an `ExtractFrames` example using blob v2 and torchcodec), Distributed Training, and the existing Output Formats content moved underneath. ## Notes - The quickstart depends on #4602 (permutation building over tables with schema metadata), which is merged. - `ExtractFrames` takes a URI and opens the table lazily because a local `LanceTable` cannot yet be pickled into DataLoader workers. - The Permutation-based "Selecting columns" and "Using multiple DataLoader workers" sections further down are unchanged from main. ## Test plan - [x] Ran the quickstart, transform, `ExtractFrames` (main process and forkserver workers) and distributed snippets locally. - [x] Rendered the page with `mint dev` and checked for console errors. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5.5 --- docs/web/training/torch.mdx | 163 +++++++++++++++++++++++++++++------- 1 file changed, 135 insertions(+), 28 deletions(-) diff --git a/docs/web/training/torch.mdx b/docs/web/training/torch.mdx index 53d33e39c..64bf9f5fc 100644 --- a/docs/web/training/torch.mdx +++ b/docs/web/training/torch.mdx @@ -1,7 +1,7 @@ --- title: "PyTorch Integration" sidebarTitle: "PyTorch" -description: Learn how to use LanceDB with PyTorch for training and inference. +description: Use a LanceDB table as a PyTorch IterableDataset, mapped onto the `torch.utils.data` concepts you already know. icon: fire availability: oss: available @@ -11,44 +11,151 @@ availability: tables in the Python SDK. --- -Use LanceDB tables as input to a PyTorch `DataLoader`. Select the columns you -need, choose a batch format, and read data in parallel for training or inference. +`StreamingDataset` is a LanceDB table exposed as a PyTorch +[`IterableDataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset), +so it drops into a stock [`DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader). +It handles per-epoch shuffling, sharding across ranks and workers, prefetching, and exact resume, so you +don't write those yourself. The shuffle is fixed when the dataset is constructed, so each epoch needs a new +`StreamingDataset` built with that epoch's number. ## Quickstart {#quickstart} -The `Table` class in LanceDB implements a contract for a PyTorch -[Dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset). - This means you can simply use a LanceDB table in a PyTorch dataloader directly. +```py Python icon=Python +import lancedb, pyarrow as pa, torch +from torch import nn +from lancedb.streaming import StreamingDataset -```py Python icon=Python -import lancedb -import torch -import pyarrow as pa -from lancedb.util import tbl_to_tensor +def to_samples(rb: pa.RecordBatch): + f = lambda c: torch.from_dlpack(rb[c].flatten()).view(len(rb), -1) + rows = zip(f("observation_state"), f("action")) + return [{"state": s, "action": a} for s, a in rows] -mem_db = lancedb.connect("memory://") -table = mem_db.create_table("test_table", pa.table({"a": range(1000)})) +uri = "hf://datasets/lance-format/lerobot-xvla-soft-fold/data" +db = lancedb.connect(uri, storage_options={"hf_enable_resolve_cache": "true"}) +table = db.open_table("frames") +train = StreamingDataset( + table, columns=["observation_state", "action"], + transform=to_samples) -# Any LanceDB table can be used as a PyTorch Dataset -dataloader = torch.utils.data.DataLoader( - table, batch_size=1024, shuffle=True, collate_fn=tbl_to_tensor -) - -for batch in dataloader: - print(batch) +model = nn.Sequential(nn.Linear(96, 256), nn.ReLU(), nn.Linear(256, 14)) +opt = torch.optim.AdamW(model.parameters(), lr=1e-3) +for batch in torch.utils.data.DataLoader(train, batch_size=256): + loss = nn.MSELoss()(model(batch["state"]), batch["action"]) + opt.zero_grad() + loss.backward() + opt.step() ``` -Although the `Table` class in LanceDB implements the `torch.utils.data.Dataset` interface, you may find that using -a table [Permutation](/training/data-loading) is more flexible. +`StreamingDataset` is a standard PyTorch `IterableDataset`, so it works with everything built for one: a stock +`DataLoader`, its default `collate_fn`, and worker processes. The transform returns plain tensors, so the +model, loss and optimizer are unmodified PyTorch. Here it streams the +[`lerobot-xvla-soft-fold`](https://huggingface.co/datasets/lance-format/lerobot-xvla-soft-fold) dataset +straight from the Hugging Face Hub to train a policy that predicts each `action` from the robot's +`observation_state`. -```py Python icon=Python -from lancedb.permutation import Permutation +## PyTorch concepts in LanceDB {#concept-map} -permutation = Permutation.identity(table) -dataloader = torch.utils.data.DataLoader(permutation) +| In PyTorch you would write | In LanceDB you use | +| --- | --- | +| An `IterableDataset` that reads files in `__iter__` | `StreamingDataset(table)` | +| `DistributedSampler`, or sharding by `get_worker_info()` in `__iter__` | `rank=` and `world_size=`. Each rank and each DataLoader worker reads its own disjoint splits | +| Decoding in `__getitem__` or `collate_fn` | `transform=`, which receives a whole Arrow `RecordBatch` and runs in a thread pool | +| `random_split`, `Subset`, or reading only the fields a model needs | `filter=""` and `columns=[...]`. Unselected columns are never read from storage | +| `StatefulDataLoader` | `state_dict()` and `load_state_dict()`, plus `StreamingDataLoader` for workers | + +The dataset does not copy your data. At construction it computes a shuffled assignment of row IDs to +`num_splits` fixed splits, then reads rows from the table as it iterates. It reads the table version +that exists at construction, so rows appended later do not appear mid-epoch. + +## Optimizations {#optimizations} + +The quickstart works as written. These sections cover what to tune as training scales up: the transform that +builds each sample, running across many GPUs, and the output format of a `Table` or `Permutation`. + +### Transform {#transform} + +A `transform` turns a batch of Arrow rows into the per-sample items your `DataLoader` collates. It receives a +whole `pyarrow.RecordBatch` (`read_batch_size` rows, 64 by default) rather than one sample at a time, so +vectorized code is cheap, and it must return exactly one item per row. The quickstart's transform is the +whole pattern: columns become tensors and each row becomes a dict, which the default `collate_fn` stacks into +a batch. Without a transform, rows arrive as plain Python dicts. + +Heavier per-sample work belongs in the same function, such as decoding video. Store each clip in a blob v2 +column (`lancedb.blob("video")`). A scan returns small descriptors instead of the video bytes, so reading +other columns never pulls video in. `fetch_blob_files` opens the clips for a batch of rows as file handles, +and torchcodec's `VideoDecoder` decodes just the frames you ask for. `ExtractFrames` samples 8 evenly spaced +frames from each clip: + +```py Python icon=Python +import functools +from torchcodec.decoders import VideoDecoder + +class ExtractFrames: + def __init__(self, uri: str, name: str, num_frames: int = 8) -> None: + self.uri, self.name, self.num_frames = uri, name, num_frames + + @functools.cached_property + def table(self) -> lancedb.table.Table: + return lancedb.connect(self.uri).open_table(self.name) + + def __call__(self, rb: pa.RecordBatch) -> list[dict[str, torch.Tensor]]: + samples = [] + for blob in self.table.fetch_blob_files("video", rb["_rowid"]): + decoder = VideoDecoder(blob) + idx = torch.linspace(0, len(decoder) - 1, self.num_frames).long() + samples.append({"frames": decoder.get_frames_at(idx.tolist()).data}) + return samples + +# `clips` is the table "clips" in the database at `uri`, with a blob v2 column "video" +train = StreamingDataset(clips, columns=["_rowid"], transform=ExtractFrames(uri, "clips")) ``` -## Output Formats {#output-formats} +Transforms run in a thread pool, one thread per CPU by default (`transform_parallelism`), overlapping with +storage reads and your training step. NumPy, PyArrow and PyTorch release the GIL, so vectorized +work runs in parallel. A pure-Python loop does not. If a profile shows the transform is the bottleneck, move +it into [DataLoader workers](#using-multiple-dataloader-workers). + +Real data has bad rows, such as a truncated frame. Raise inside the transform and set +`on_transform_error="skip"` (or `"warn"` to log each failure). LanceDB re-runs the failing batch one row at a +time, drops only the rows that raise, and counts them in `dataset.rows_skipped`. When the bad rows match a SQL +predicate, prefer `filter=`: it runs before the splits are built, so every split keeps the same length. + +### Distributed Training {#distributed-training} + +`StreamingDataset` takes the place of `DistributedSampler`. Pass the rank and world size, and each rank reads +its own disjoint part of the table, together covering it once per epoch: + +```py Python icon=Python +import torch.distributed as dist + +dist.init_process_group("nccl") +train = StreamingDataset( + table, columns=["observation_state", "action"], transform=to_samples, + num_splits=48, shuffle_seed=42, epoch=epoch, + rank=dist.get_rank(), world_size=dist.get_world_size()) +``` + +Everything after the dataset is standard distributed PyTorch. Wrap the model in `DistributedDataParallel` +and pass `batch_size` to the `DataLoader` per GPU, so the global batch is `batch_size * world_size`. Nothing +in the data path changes. + +Three arguments matter in distributed runs: + +- **`num_splits`** is how many fixed streams the table is divided into. It must be divisible by `world_size`, + and by `world_size * num_workers` when you use DataLoader workers. Choose a number with many divisors, such + as 48, and runs on 8, 4 or 2 GPUs see the same data at every step, provided the global batch size is a multiple of + `num_splits`. +- **`shuffle_seed` and `epoch`** must be identical on every rank. Pass a fixed `shuffle_seed` and increase + `epoch` each epoch. `shuffle_seed=None` draws a different random seed in each process, so ranks would read + overlapping rows and miss others. +- **`rank` and `world_size`** come from your launcher, for example `torchrun`. + +To resume after a failure, save `train.state_dict()` next to the model. It records progress per split, so the +resumed run continues from the same row even on a different number of GPUs. See +[Parallelism & Elasticity](/training/data-loading#parallelism-&-elasticity) and +[Checkpointing](/training/data-loading#checkpointing) for the batch-size rules and the multi-worker case. + +### Output Formats {#output-formats} By default, a `Table` data loader will emit Arrow data. `collate_fn` is PyTorch's batching hook: PyTorch calls it to turn the fetched items into one batch. PyTorch's default collate function only knows how to combine tensors, NumPy @@ -68,7 +175,7 @@ the Arrow data in different ways. The `arrow` and `polars` formats will always `pandas`, and `torch_col` formats will also avoid data copies in most cases. The `python`, `python_col`, and `torch` formats will all require at least one full copy of the data and are the slowest options. -### Using the torch column format with a torch data loader {#using-the-torch-column-format-with-a-torch-data-loader} +#### Using the torch column format with a torch data loader {#using-the-torch-column-format-with-a-torch-data-loader} The `torch_col` format is the most efficient way to convert from Arrow to a `torch.Tensor`. It will convert the entire Arrow batch to a _column-major_ `torch.Tensor`. In other words, given C columns and R rows, the resulting From b533fe431a26ee794c0cb229d729947c82cbc81e Mon Sep 17 00:00:00 2001 From: IOCRYPOSTO Date: Wed, 21 Oct 2026 11:49:28 +0530 Subject: [PATCH 03/33] docs: fix TypeScript file paths in AGENTS.md (#4612) Lines 105-106 of AGENTS.md tell agents to add TypeScript methods in `nodejs/src/table.ts` and `nodejs/src/native_table.ts`. Neither file exists; the abstract `Table` class and the concrete `LocalTable` class are both in `nodejs/lancedb/table.ts`. An agent following these steps searches for files that are not there. This updates both paths. Docs-only, no code change. Suggested by LanceDB support (Daisuke) after I reported it by email. --- AGENTS.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f6d01db03..24465a28f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -102,8 +102,8 @@ TypeScript bindings changes: 1. Add napi-rs method binding on `Table` in `nodejs/src/table.rs`. 2. Run `pnpm build` to generate TypeScript definitions. -3. Add typescript method on abstract class `Table` in `nodejs/src/table.ts`. -4. Add concrete method on `LocalTable` class in `nodejs/src/native_table.ts`. +3. Add typescript method on abstract class `Table` in `nodejs/lancedb/table.ts`. +4. Add concrete method on `LocalTable` class in `nodejs/lancedb/table.ts`. * Note: despite the name, this class is also used for remote tables. 5. Add test in `nodejs/__test__/table.test.ts`. 6. Run `pnpm run docs` to generate TypeScript documentation. From b00061fb28fc09208fc9e8eed12391734de7554f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 21 Oct 2026 16:12:12 +0800 Subject: [PATCH 04/33] chore: update lance dependency to v14.0.0-beta.22 (#4614) Pin the Rust workspace and Java lance-core dependency to [Lance v14.0.0-beta.22](https://github.com/lance-format/lance/releases/tag/v14.0.0-beta.22), including managed data file parts and descriptor reuse from lance-format/lance#10176. This enables the next LanceDB preview release to provide the tagged dependency needed by lancedb/sophon#8898. No compatibility code changes are required. Update rustls-webpki to 0.103.16 and webpki-roots to 1.0.10 in the lockfile to resolve the two TLS advisories blocking dependency CI (RUSTSEC-2026-0382 and RUSTSEC-2026-0381), without adding advisory exceptions. The Java dependency availability blocker is resolved: both the beta.22 POM and JAR return HTTP 200 from maven.lance.org. On `ea043bf93`, the review reproducer (`cd java && ./mvnw -U -q test -pl lancedb-core -Dspotless.skip=true`) passes all 37 tests, and [Java CI attempt 2](https://github.com/lancedb/lancedb/actions/runs/39556336205/attempts/2) passes. --- Cargo.lock | 109 ++++++++++++++++++++++++++------------------------- Cargo.toml | 28 ++++++------- java/pom.xml | 2 +- 3 files changed, 70 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 61aee21be..8722b0e49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3584,8 +3584,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4441,7 +4441,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.7", + "webpki-roots 1.0.10", ] [[package]] @@ -5059,8 +5059,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arc-swap", "arrow", @@ -5110,6 +5110,7 @@ dependencies = [ "lance-tokenizer", "log", "moka", + "num-traits", "object_store 0.14.2", "permutation", "pin-project", @@ -5134,8 +5135,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-buffer", @@ -5157,7 +5158,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-buffer", @@ -5171,7 +5172,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-schema", @@ -5180,8 +5181,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrayref", "crunchy", @@ -5191,8 +5192,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-buffer", @@ -5229,8 +5230,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow", "arrow-array", @@ -5261,8 +5262,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow", "arrow-array", @@ -5279,8 +5280,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "proc-macro2", "quote", @@ -5289,8 +5290,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-arith", "arrow-array", @@ -5323,8 +5324,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-arith", "arrow-array", @@ -5356,8 +5357,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "datafusion", "geo-traits", @@ -5371,8 +5372,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arc-swap", "arrow", @@ -5443,8 +5444,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-schema", @@ -5467,8 +5468,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow", "arrow-array", @@ -5511,8 +5512,8 @@ dependencies = [ [[package]] name = "lance-jsonb" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "byteorder", "ethnum", @@ -5529,8 +5530,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-schema", @@ -5544,8 +5545,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow", "async-trait", @@ -5559,8 +5560,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arc-swap", "arrow", @@ -5586,7 +5587,7 @@ dependencies = [ "object_store 0.14.2", "quick-xml 0.40.1", "rand 0.9.5", - "reqwest 0.12.28", + "reqwest 0.13.4", "roaring", "serde", "serde_json", @@ -5614,8 +5615,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-buffer", @@ -5629,8 +5630,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow", "arrow-array", @@ -5670,8 +5671,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "arrow-array", "arrow-schema", @@ -5684,8 +5685,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "14.0.0-beta.21" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.21#8c233c2e6aa6948524f352e5e4cd52bd1364f6ff" +version = "14.0.0-beta.22" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" dependencies = [ "frostem", "icu_segmenter", @@ -8909,7 +8910,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots 1.0.7", + "webpki-roots 1.0.10", ] [[package]] @@ -9239,9 +9240,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "aws-lc-rs", "ring", @@ -11156,14 +11157,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.7", + "webpki-roots 1.0.10", ] [[package]] name = "webpki-roots" -version = "1.0.7" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52f5ee44c96cf55f1b349600768e3ece3a8f26010c05265ab73f945bb1a2eb9d" +checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9" dependencies = [ "rustls-pki-types", ] diff --git a/Cargo.toml b/Cargo.toml index e302d2891..1307de8f9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.21", default-features = false, "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.21", default-features = false, "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.21", default-features = false, "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.21", "tag" = "v14.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/java/pom.xml b/java/pom.xml index 4e4a82b63..eb3208627 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 14.0.0-beta.21 + 14.0.0-beta.22 2.18.11 false 2.30.0 From 46ccbe4e2c375911b22a517e04dfb4785849d11d Mon Sep 17 00:00:00 2001 From: Lance Release Date: Wed, 21 Oct 2026 08:20:35 +0000 Subject: [PATCH 05/33] =?UTF-8?q?Bump=20version:=200.41.0-beta.10=20?= =?UTF-8?q?=E2=86=92=200.41.0-beta.11?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.toml | 2 +- Cargo.lock | 6 +++--- docs/src/java/java.md | 2 +- java/lancedb-core/pom.xml | 2 +- java/pom.xml | 2 +- nodejs/Cargo.toml | 2 +- nodejs/npm/darwin-arm64/package.json | 2 +- nodejs/npm/linux-arm64-gnu/package.json | 2 +- nodejs/npm/linux-arm64-musl/package.json | 2 +- nodejs/npm/linux-x64-gnu/package.json | 2 +- nodejs/npm/linux-x64-musl/package.json | 2 +- nodejs/npm/win32-arm64-msvc/package.json | 2 +- nodejs/npm/win32-x64-msvc/package.json | 2 +- nodejs/package.json | 2 +- python/Cargo.toml | 2 +- rust/lancedb/Cargo.toml | 2 +- 16 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 413d11f82..52cb1eb2f 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.10" +current_version = "0.41.0-beta.11" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index 8722b0e49..cf0e79f3b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5699,7 +5699,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" dependencies = [ "ahash", "anyhow", @@ -5797,7 +5797,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" dependencies = [ "arrow-array", "arrow-buffer", @@ -5822,7 +5822,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" dependencies = [ "arc-swap", "arrow", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index b6cbd8586..d6f3e7ef7 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.10 + 0.41.0-beta.11 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index e64f65852..6dd92bbac 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.10 + 0.41.0-beta.11 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index eb3208627..12a4272d4 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.10 + 0.41.0-beta.11 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 79a594f14..51dd5a4c2 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 403b1d2c2..1b1794d77 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index ef2d4ca40..12e4eda00 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index 5b247be2c..e19f0efce 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index 0c36b7213..09c39080d 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index 274877c83..bed7c26c7 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 67f0176f5..c5a64dcb3 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 5f3cf6032..4cdaaeffd 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index df083600a..1128a5299 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.10", + "version": "0.41.0-beta.11", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 615b7c5aa..4c221ef11 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 4649e560d..4ec9a70b1 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.10" +version = "0.41.0-beta.11" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true From 22b15d4df00e25bf1226e3bf3b82bf701ca29b2b Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Wed, 21 Oct 2026 18:45:57 +0800 Subject: [PATCH 06/33] chore: update lance dependency to v14.0.0-beta.23 (#4616) Update the Rust workspace and Java lance-core from beta.22 to [Lance v14.0.0-beta.23](https://github.com/lance-format/lance/releases/tag/v14.0.0-beta.23), including the supplied-sample IVF training API from lance-format/lance#10188. Non-Lance lockfile entries match current main; no compatibility code changes are required. Validation on `f591f5d5`: full-workspace Clippy (`--tests --all-features -- -D warnings`), Rust formatting, and Java checkstyle passed locally. The beta.23 [POM](https://maven.lance.org/org/lance/lance-core/14.0.0-beta.23/lance-core-14.0.0-beta.23.pom) and [JAR](https://maven.lance.org/org/lance/lance-core/14.0.0-beta.23/lance-core-14.0.0-beta.23.jar) are now published and both return HTTP 200. The review reproducer (`cd java && ./mvnw -U -q test -pl lancedb-core -Dspotless.skip=true`) passes all 37 tests with zero failures or skips; JUnit classpaths confirm beta.23 was loaded. [Java CI attempt 2](https://github.com/lancedb/lancedb/actions/runs/39568468567/attempts/2) also passes, resolving the Java dependency publication finding. --- Cargo.lock | 92 ++++++++++++++++++++++++++-------------------------- Cargo.toml | 28 ++++++++-------- java/pom.xml | 2 +- 3 files changed, 61 insertions(+), 61 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cf0e79f3b..9fbde0847 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3584,8 +3584,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5059,8 +5059,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arc-swap", "arrow", @@ -5135,8 +5135,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-buffer", @@ -5158,7 +5158,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-buffer", @@ -5172,7 +5172,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-schema", @@ -5181,8 +5181,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrayref", "crunchy", @@ -5192,8 +5192,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-buffer", @@ -5230,8 +5230,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow", "arrow-array", @@ -5262,8 +5262,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow", "arrow-array", @@ -5280,8 +5280,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "proc-macro2", "quote", @@ -5290,8 +5290,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-arith", "arrow-array", @@ -5324,8 +5324,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-arith", "arrow-array", @@ -5357,8 +5357,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "datafusion", "geo-traits", @@ -5372,8 +5372,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arc-swap", "arrow", @@ -5444,8 +5444,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-schema", @@ -5468,8 +5468,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow", "arrow-array", @@ -5512,8 +5512,8 @@ dependencies = [ [[package]] name = "lance-jsonb" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "byteorder", "ethnum", @@ -5530,8 +5530,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-schema", @@ -5545,8 +5545,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow", "async-trait", @@ -5560,8 +5560,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arc-swap", "arrow", @@ -5615,8 +5615,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-buffer", @@ -5630,8 +5630,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow", "arrow-array", @@ -5671,8 +5671,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "arrow-array", "arrow-schema", @@ -5685,8 +5685,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "14.0.0-beta.22" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.22#72c5b645600bde5d2e0e0b629ed986e2018c62bf" +version = "14.0.0-beta.23" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" dependencies = [ "frostem", "icu_segmenter", diff --git a/Cargo.toml b/Cargo.toml index 1307de8f9..49f7ea6a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.22", default-features = false, "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.22", "tag" = "v14.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/java/pom.xml b/java/pom.xml index 12a4272d4..c3e3dd771 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 14.0.0-beta.22 + 14.0.0-beta.23 2.18.11 false 2.30.0 From 1480a9ff5768b8c5016db6e1e4c44bb617c343f1 Mon Sep 17 00:00:00 2001 From: Lance Release Date: Wed, 21 Oct 2026 10:48:07 +0000 Subject: [PATCH 07/33] =?UTF-8?q?Bump=20version:=200.41.0-beta.11=20?= =?UTF-8?q?=E2=86=92=200.41.0-beta.12?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.toml | 2 +- Cargo.lock | 6 +++--- docs/src/java/java.md | 2 +- java/lancedb-core/pom.xml | 2 +- java/pom.xml | 2 +- nodejs/Cargo.toml | 2 +- nodejs/npm/darwin-arm64/package.json | 2 +- nodejs/npm/linux-arm64-gnu/package.json | 2 +- nodejs/npm/linux-arm64-musl/package.json | 2 +- nodejs/npm/linux-x64-gnu/package.json | 2 +- nodejs/npm/linux-x64-musl/package.json | 2 +- nodejs/npm/win32-arm64-msvc/package.json | 2 +- nodejs/npm/win32-x64-msvc/package.json | 2 +- nodejs/package.json | 2 +- python/Cargo.toml | 2 +- rust/lancedb/Cargo.toml | 2 +- 16 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 52cb1eb2f..f2c9e44e1 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.11" +current_version = "0.41.0-beta.12" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index 9fbde0847..1d5d05bdf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5699,7 +5699,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" dependencies = [ "ahash", "anyhow", @@ -5797,7 +5797,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" dependencies = [ "arrow-array", "arrow-buffer", @@ -5822,7 +5822,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" dependencies = [ "arc-swap", "arrow", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index d6f3e7ef7..5c1874e6f 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.11 + 0.41.0-beta.12 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 6dd92bbac..ad5adbc92 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.11 + 0.41.0-beta.12 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index c3e3dd771..db0ee33d7 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.11 + 0.41.0-beta.12 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 51dd5a4c2..85ee5df04 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 1b1794d77..080340224 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 12e4eda00..a2b5b8bf6 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index e19f0efce..96bda03c7 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index 09c39080d..edd596044 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index bed7c26c7..953c46250 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index c5a64dcb3..236d356be 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 4cdaaeffd..205f0ee11 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 1128a5299..023c4e400 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.11", + "version": "0.41.0-beta.12", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 4c221ef11..e6a0e457a 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 4ec9a70b1..5e6dd315f 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.11" +version = "0.41.0-beta.12" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true From 0fbc0981ddd7577d3bc24dc6e0d31ddfe24b3455 Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Wed, 21 Oct 2026 14:26:47 -0700 Subject: [PATCH 08/33] feat: document all OAuthFlowType entries in authentication.mdx (#4565) Previously, authentication.mdx only documented a subset of the existing OAuthFlowType entries. We should document all of them. --- docs/web/enterprise/authentication.mdx | 312 ++++++++++++++++++++++--- 1 file changed, 282 insertions(+), 30 deletions(-) diff --git a/docs/web/enterprise/authentication.mdx b/docs/web/enterprise/authentication.mdx index 3a570d1d9..9074c641c 100644 --- a/docs/web/enterprise/authentication.mdx +++ b/docs/web/enterprise/authentication.mdx @@ -3,7 +3,7 @@ title: "Authentication" sidebarTitle: "Authentication" description: "Authentication modes for LanceDB Enterprise: API keys and OAuth credentials." icon: key -keywords: ["authentication", "oauth", "api key", "azure managed identity", "client credentials"] +keywords: ["authentication", "oauth", "api key", "authorization code", "device code", "azure managed identity", "client credentials"] availability: oss: unavailable enterprise: available @@ -13,9 +13,9 @@ availability: [Architecture](/architecture). --- -A deployment authenticates every request before -[authorization](/governance/authorization) decides what it may do, against a -`db://` remote table, in one of two ways: +A deployment authenticates each request over a `db://` connection before +[authorization](/governance/authorization) decides what it may do. +Choose one of these methods: - **API keys**: a long-lived shared secret passed as `api_key` on connect. Works in every SDK. - **OAuth 2.0**: short-lived bearer tokens obtained from your identity provider, refreshed automatically by the client. @@ -28,55 +28,91 @@ identities so no secret material lives on the client. Pass the API key from your Enterprise tenant on `connect`. This works with both the synchronous and asynchronous Python clients, as well as TypeScript -and Rust. +and Rust. Load the key from an environment variable or secret manager. -```python + +```python Python icon="python" +import os import lancedb db = lancedb.connect( uri="db://your-database-uri", - api_key="your-api-key", + api_key=os.environ["LANCEDB_API_KEY"], region="us-east-1", host_override="https://your-enterprise-endpoint.com", ) ``` -An API key belongs in the environment rather than in source. See -[Database](/management/database#deployment) for the connection options. +```typescript TypeScript icon="square-js" +import * as lancedb from "@lancedb/lancedb"; + +const apiKey = process.env.LANCEDB_API_KEY; +if (!apiKey) throw new Error("LANCEDB_API_KEY is required"); + +const db = await lancedb.connect("db://your-database-uri", { + apiKey, + region: "us-east-1", + hostOverride: "https://your-enterprise-endpoint.com", +}); +``` + +```rust Rust icon="rust" +let api_key = std::env::var("LANCEDB_API_KEY")?; + +let db = lancedb::connect("db://your-database-uri") + .api_key(&api_key) + .region("us-east-1") + .host_override("https://your-enterprise-endpoint.com") + .execute() + .await?; +``` + + +For Rust, enable the `remote` feature on the `lancedb` dependency to use +API key or OAuth authentication. + +See [Database](/management/database#deployment) for the connection options. ## OAuth {#oauth} -The async Python client and the TypeScript client can obtain bearer tokens +The async Python, TypeScript, and Rust clients can obtain bearer tokens from an OIDC issuer and attach them to every request. Token acquisition, caching, and refresh are handled inside the client. Your application code only provides the configuration. A deployment accepts these tokens only when its administrator has configured -it with your identity provider. Ask them for the issuer URL, client ID and -scopes to use. +it with your identity provider. Ask them for the issuer URL, client ID, scopes, +and any resource or audience required by the provider. For Authorization Code, +register the loopback redirect URI with the provider. For Device Code, enable +the provider's device authorization grant for your application. In Python, OAuth is supported through `lancedb.connect_async`. The synchronous `connect` entry point continues to use API key authentication for `db://` URIs. In TypeScript, `lancedb.connect` accepts `oauthConfig` directly. +In Rust, use `.oauth_config(config)` on the connection builder. ### Supported flows {#supported-flows} -`OAuthFlowType` selects how the client acquires tokens: +`OAuthFlowType` (Python and TypeScript) or `OAuthFlow` (Rust) selects how the +client acquires tokens: -| Flow | Python value | TypeScript value | When to use | -| --- | --- | --- | --- | -| Client Credentials | `OAuthFlowType.CLIENT_CREDENTIALS` | `OAuthFlowType.ClientCredentials` | Service-to-service / machine-to-machine. Requires a client ID and client secret registered with your identity provider. | -| Azure Managed Identity | `OAuthFlowType.AZURE_MANAGED_IDENTITY` | `OAuthFlowType.AzureManagedIdentity` | Workloads running on Azure compute (VMs, AKS, App Service, Container Apps). Tokens are fetched from the Azure IMDS endpoint, so no client secret is stored on the client. | +| Flow | Python value | TypeScript value | Rust value | When to use | +| --- | --- | --- | --- | --- | +| Client Credentials | `OAuthFlowType.CLIENT_CREDENTIALS` | `OAuthFlowType.ClientCredentials` | `OAuthFlow::ClientCredentials` | Service-to-service / machine-to-machine. Requires a client ID and client secret registered with your identity provider. | +| Authorization Code | `OAuthFlowType.AUTHORIZATION_CODE` | `OAuthFlowType.AuthorizationCode` | `OAuthFlow::AuthorizationCode(AuthorizationCodeOptions::new())` | Interactive user sign-in through a browser and local callback server. Uses S256 PKCE by default. | +| Device Code | `OAuthFlowType.DEVICE_CODE` | `OAuthFlowType.DeviceCode` | `OAuthFlow::DeviceCode` | CLI or headless sign-in. The client displays a verification URL and user code for the user to enter on another device. | +| Azure Managed Identity | `OAuthFlowType.AZURE_MANAGED_IDENTITY` | `OAuthFlowType.AzureManagedIdentity` | `OAuthFlow::AzureManagedIdentity { client_id: None }` | Azure workloads with access to the IMDS endpoint. Tokens are fetched from IMDS, so no client secret is stored on the client. | ### Configure OAuth {#configure-oauth} Build an OAuth config and pass it on connect. Use `oauth_config` in Python and -`oauthConfig` in TypeScript. +Rust, and `oauthConfig` in TypeScript. ```python Python icon="python" +import os import lancedb from lancedb.remote import OAuthConfig, OAuthFlowType @@ -84,7 +120,7 @@ from lancedb.remote import OAuthConfig, OAuthFlowType oauth_config = OAuthConfig( issuer_url="https://login.microsoftonline.com/{tenant}/v2.0", client_id="your-application-client-id", - client_secret="your-application-client-secret", + client_secret=os.environ["LANCEDB_OAUTH_CLIENT_SECRET"], scopes=["api://lancedb-api/.default"], flow=OAuthFlowType.CLIENT_CREDENTIALS, ) @@ -101,11 +137,14 @@ db = await lancedb.connect_async( import * as lancedb from "@lancedb/lancedb"; import { OAuthConfig, OAuthFlowType } from "@lancedb/lancedb"; +const clientSecret = process.env.LANCEDB_OAUTH_CLIENT_SECRET; +if (!clientSecret) throw new Error("LANCEDB_OAUTH_CLIENT_SECRET is required"); + // Client Credentials (service-to-service) const oauthConfig: OAuthConfig = { issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0", clientId: "your-application-client-id", - clientSecret: "your-application-client-secret", + clientSecret, scopes: ["api://lancedb-api/.default"], flow: OAuthFlowType.ClientCredentials, }; @@ -116,9 +155,34 @@ const db = await lancedb.connect("db://your-database-uri", { oauthConfig, }); ``` + +```rust Rust icon="rust" +use lancedb::remote::{OAuthConfig, OAuthFlow}; + +// Client Credentials (service-to-service) +let oauth_config = OAuthConfig { + issuer_url: "https://login.microsoftonline.com/{tenant}/v2.0".into(), + client_id: "your-application-client-id".into(), + client_secret: Some(std::env::var("LANCEDB_OAUTH_CLIENT_SECRET")?), + scopes: vec!["api://lancedb-api/.default".into()], + flow: OAuthFlow::ClientCredentials, + client_auth_method: None, + resource: None, + audience: None, + refresh_buffer_secs: None, + token_cache: None, +}; + +let db = lancedb::connect("db://your-database-uri") + .region("us-east-1") + .host_override("https://your-enterprise-endpoint.com") + .oauth_config(oauth_config) + .execute() + .await?; +``` -For workloads running on Azure compute, use a managed identity instead of a +For Azure workloads with access to IMDS, use a managed identity instead of a client secret: @@ -130,8 +194,8 @@ oauth_config = OAuthConfig( client_id="your-application-client-id", scopes=["api://lancedb-api/.default"], flow=OAuthFlowType.AZURE_MANAGED_IDENTITY, - # Optional: required only for user-assigned managed identities. - managed_identity_client_id="your-user-assigned-identity-client-id", + # Set only for a user-assigned managed identity: + # managed_identity_client_id="your-user-assigned-identity-client-id", ) ``` @@ -143,31 +207,219 @@ const oauthConfig: OAuthConfig = { clientId: "your-application-client-id", scopes: ["api://lancedb-api/.default"], flow: OAuthFlowType.AzureManagedIdentity, - // Optional: required only for user-assigned managed identities. - managedIdentityClientId: "your-user-assigned-identity-client-id", + // Set only for a user-assigned managed identity: + // managedIdentityClientId: "your-user-assigned-identity-client-id", +}; +``` + +```rust Rust icon="rust" +use lancedb::remote::{OAuthConfig, OAuthFlow}; + +let oauth_config = OAuthConfig { + issuer_url: "https://login.microsoftonline.com/{tenant}/v2.0".into(), + client_id: "your-application-client-id".into(), + client_secret: None, + scopes: vec!["api://lancedb-api/.default".into()], + flow: OAuthFlow::AzureManagedIdentity { + // Set to Some("your-user-assigned-identity-client-id".into()) + // only for a user-assigned managed identity. + client_id: None, + }, + client_auth_method: None, + resource: None, + audience: None, + refresh_buffer_secs: None, + token_cache: None, }; ``` +`managed_identity_client_id` / `managedIdentityClientId` selects a user-assigned +identity. Omit it for a system-assigned identity. The required `client_id` / +`clientId` field does not select the Azure managed identity. +In Rust, set the `client_id` inside `OAuthFlow::AzureManagedIdentity` to +`Some(...)` for a user-assigned identity, or `None` for a system-assigned +identity. The top-level `OAuthConfig::client_id` is the application client ID. + +For interactive sign-in, configure Authorization Code or Device Code. The +authorization URL or device verification instructions are written to standard +error. Authorization Code also tries to open a browser and listens for the +callback at `http://127.0.0.1:8400/callback` by default. Register that redirect +URI with your identity provider. Public clients can omit `client_secret`; +request a refresh token from your provider (often by including +`offline_access`) if you want to reuse the session. + + +```python Python icon="python" +import lancedb +from lancedb.remote import OAuthConfig, OAuthFlowType, TokenCacheOptions + +# Use DEVICE_CODE for a CLI without a local browser or callback server. +oauth_config = OAuthConfig( + issuer_url="https://your-issuer.example.com", + client_id="your-public-client-id", + scopes=["openid", "offline_access", "lancedb.read"], + flow=OAuthFlowType.AUTHORIZATION_CODE, + token_cache=TokenCacheOptions(), +) + +db = await lancedb.connect_async( + uri="db://your-database-uri", + region="us-east-1", + host_override="https://your-enterprise-endpoint.com", + oauth_config=oauth_config, +) +``` + +```typescript TypeScript icon="square-js" +import * as lancedb from "@lancedb/lancedb"; +import { OAuthFlowType, type OAuthConfig } from "@lancedb/lancedb"; + +// Use OAuthFlowType.DeviceCode for a CLI without a local browser or callback server. +const oauthConfig: OAuthConfig = { + issuerUrl: "https://your-issuer.example.com", + clientId: "your-public-client-id", + scopes: ["openid", "offline_access", "lancedb.read"], + flow: OAuthFlowType.AuthorizationCode, + tokenCache: {}, +}; + +const db = await lancedb.connect("db://your-database-uri", { + region: "us-east-1", + hostOverride: "https://your-enterprise-endpoint.com", + oauthConfig, +}); +``` + +```rust Rust icon="rust" +use lancedb::remote::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow, TokenCacheOptions}; + +// Use OAuthFlow::DeviceCode for a CLI without a local browser or callback server. +let oauth_config = OAuthConfig { + issuer_url: "https://your-issuer.example.com".into(), + client_id: "your-public-client-id".into(), + client_secret: None, + scopes: vec![ + "openid".into(), + "offline_access".into(), + "lancedb.read".into(), + ], + flow: OAuthFlow::AuthorizationCode(AuthorizationCodeOptions::new()), + client_auth_method: None, + resource: None, + audience: None, + refresh_buffer_secs: None, + token_cache: Some(TokenCacheOptions::new()), +}; + +let db = lancedb::connect("db://your-database-uri") + .region("us-east-1") + .host_override("https://your-enterprise-endpoint.com") + .oauth_config(oauth_config.clone()) + .execute() + .await?; +``` + + +With `DEVICE_CODE` / `DeviceCode` / `OAuthFlow::DeviceCode`, the same +configuration and connection calls apply; only the `flow` value changes. + +### Reuse an interactive session {#reuse-an-interactive-session} + +By default, tokens stay in process memory. Set `token_cache` / `tokenCache` +for Authorization Code or Device Code to keep refresh tokens across process +restarts. Only refresh tokens are persisted; access tokens and client secrets +are not written to the cache. Set `cache_dir` / `cacheDir` to choose the +directory, or leave it unset to use the platform cache directory. Your +identity provider must issue a refresh token for a later process to sign in +without prompting again. + +Use `OAuthSession` with the same config to log in before connecting, inspect +whether a refresh token is cached, or remove the local cached credential: + + +```python Python icon="python" +from lancedb.remote import OAuthSession + +session = OAuthSession(oauth_config) +status = await session.login() +print(status.refreshable) +status = await session.status() # Does not contact the provider. +print(status.refreshable) +logout = await session.logout() +print(logout.removed) +``` + +```typescript TypeScript icon="square-js" +import { OAuthSession } from "@lancedb/lancedb"; + +const session = new OAuthSession(oauthConfig); +const status = await session.login(); +console.log(status.refreshable); +const current = await session.status(); // Does not contact the provider. +console.log(current.refreshable); +const logout = await session.logout(); +console.log(logout.removed); +``` + +```rust Rust icon="rust" +use lancedb::remote::OAuthSession; + +let session = OAuthSession::new(oauth_config)?; +let status = session.login().await?; +println!("{}", status.refreshable); +let status = session.status().await?; // Does not contact the provider. +println!("{}", status.refreshable); +let logout = session.logout().await?; +println!("{}", logout.removed); +``` + + +`login()` runs the interactive flow even when a session is cached. `logout()` +removes the local cache entry; it does not revoke the provider's token or end +a browser sign-in. `OAuthSession` requires a persistent cache and supports +only Authorization Code and Device Code. A cache is unused for Client +Credentials and rejected for Azure Managed Identity. + ### Configuration reference {#configuration-reference} -The same configuration is available in both SDKs. Python uses `snake_case` -field names; TypeScript uses `camelCase`. +The same configuration is available in all three SDKs. Python and Rust use +`snake_case` field names; TypeScript uses `camelCase`. In Rust, set optional +`OAuthConfig` fields to `Some(value)` or `None`. Configure `redirect_uri`, +`callback_port`, and `use_pkce` through the builder methods on +`AuthorizationCodeOptions`, and select a managed identity through +`OAuthFlow::AzureManagedIdentity { client_id: ... }`. | Python field / TypeScript field | Required | Description | | --- | --- | --- | | `issuer_url` / `issuerUrl` | Yes | OIDC issuer URL or OAuth authority URL. For Azure, use `https://login.microsoftonline.com/{tenant_id}/v2.0`. | | `client_id` / `clientId` | Yes | Application / client ID registered with your identity provider. | | `scopes` | Yes | List of OAuth scopes to request. For Azure managed identity, provide exactly one scope or resource, for example `["api://your-app-id/.default"]`. | -| `flow` | No | Selects the OAuth flow. Defaults to Client Credentials. | -| `client_secret` / `clientSecret` | Conditional | Required for Client Credentials. Redacted from the config's `repr` (Python) and `Debug` output (TypeScript native binding) so it does not leak into logs. | +| `flow` | Yes in Rust; otherwise no | Selects the OAuth flow. Defaults to Client Credentials in Python and TypeScript; Rust requires an explicit `OAuthFlow`. | +| `client_secret` / `clientSecret` | Conditional | Required for Client Credentials; only required for Authorization Code and Device Code when the provider registers a confidential client. Redacted from the config's `repr` (Python) and `Debug` output (Rust and TypeScript native binding). | +| `client_auth_method` / `clientAuthMethod` | No | Token endpoint authentication: `ClientAuthMethod.NONE` / `ClientAuthMethod.None` for a public client, `ClientAuthMethod.CLIENT_SECRET_BASIC` / `ClientAuthMethod.ClientSecretBasic`, or `ClientAuthMethod.CLIENT_SECRET_POST` / `ClientAuthMethod.ClientSecretPost`. Defaults to Basic when a client secret is set and no client authentication otherwise. The secret methods require `client_secret` / `clientSecret`. Omit for Azure Managed Identity. | +| `redirect_uri` / `redirectUri` | No | Authorization Code loopback callback URI. Defaults to `http://127.0.0.1:{callback_port}/callback`; register it with the provider. | +| `callback_port` / `callbackPort` | No | Authorization Code callback server port (default 8400). If `redirect_uri` / `redirectUri` has an explicit port, the two ports must match. | +| `use_pkce` / `usePkce` | No | Enable S256 PKCE for Authorization Code (default `true`). | | `managed_identity_client_id` / `managedIdentityClientId` | No | Client ID of a user-assigned managed identity. Only used with Azure Managed Identity; omit for system-assigned identities. | +| `resource` | No | RFC 8707 resource indicator, an absolute URI without a fragment. Sent to authorization and token endpoints, including refresh requests. Not supported for Azure Managed Identity. | +| `audience` | No | Provider-specific audience sent to authorization and token endpoints, including refresh requests. Not supported for Azure Managed Identity. | | `refresh_buffer_secs` / `refreshBufferSecs` | No | How many seconds before token expiry to proactively refresh. Defaults to 300. Keep this well below the token TTL: setting it greater than or equal to the TTL forces a refresh on every request. | +| `token_cache` / `tokenCache` | No | Opt in to a persistent refresh-token cache for Authorization Code or Device Code. Use `TokenCacheOptions()` in Python, `{}` in TypeScript, or `Some(TokenCacheOptions::new())` in Rust. Defaults to an in-memory session. | + +`TokenCacheOptions` also accepts `cache_dir` / `cacheDir` to select a cache +directory and `lock_timeout_secs` / `lockTimeoutSecs` to set the wait for a +cross-process refresh lock (default 30 seconds). +In Rust, use the `TokenCacheOptions::cache_dir(...)` and +`TokenCacheOptions::lock_timeout_secs(...)` builder methods. Token endpoint +authentication uses `Some(ClientAuthMethod::None)`, +`Some(ClientAuthMethod::ClientSecretBasic)`, or +`Some(ClientAuthMethod::ClientSecretPost)`; `None` selects the default. Treat the client secret like any other production credential. Load it from a secret manager or environment variable rather than committing it to source -control. Both SDKs deliberately omit the secret from their debug/repr output +control. All three SDKs deliberately omit the secret from their debug/repr output so accidental log lines do not expose the value. From d97df0c82b3f7b6525bb7b108ddc154f7f16a2ca Mon Sep 17 00:00:00 2001 From: Esteban Gutierrez Date: Wed, 21 Oct 2026 17:11:36 -0500 Subject: [PATCH 09/33] fix: update trust-store pins for the eMudhra CA compromise (#4624) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cargo.lock: webpki-root-certs 1.0.10; webpki-roots 1.0.10 and rustls-webpki 0.103.16 are already current. docs/web-tests/rs/Cargo.lock: rustls-webpki 0.103.16. python/uv.lock: certifi 2026.10.19, the only entry that changes. The root crate drops the root CA whose intermediates were compromised (Mozilla bug 2083449), and rustls-webpki 0.103.16 rejects chains that carry the compromised eMudhra intermediate keys. certifi 2026.10.19 is the release the certifi advisory names as fixed: it drops the same root and keeps the issuer's other roots, so HTTPS connections that chain to them keep working. Relative to the previous 2026.6.17 pin it also adds Telia EC and RSA TLS Root CA v3 and SECOM TLS ECC and RSA Root CA 2024 and drops Atos TrustedRoot 2011. The webpki-root-certs entry is written by hand with the crates.io checksum. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Esteban Gutierrez Co-authored-by: Claude Opus 5.1 --- Cargo.lock | 4 ++-- docs/web-tests/rs/Cargo.lock | 4 ++-- python/uv.lock | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1d5d05bdf..3b4086203 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11144,9 +11144,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.7" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f31141ce3fc3e300ae89b78c0dd67f9708061d1d2eda54b8209346fd6be9a92c" +checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" dependencies = [ "rustls-pki-types", ] diff --git a/docs/web-tests/rs/Cargo.lock b/docs/web-tests/rs/Cargo.lock index 447bae7f6..6e288950e 100644 --- a/docs/web-tests/rs/Cargo.lock +++ b/docs/web-tests/rs/Cargo.lock @@ -5561,9 +5561,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "ring", "rustls-pki-types", diff --git a/python/uv.lock b/python/uv.lock index ace8b98b9..e808f5f9e 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -443,11 +443,11 @@ wheels = [ [[package]] name = "certifi" -version = "2026.6.17" +version = "2026.10.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c9/c7/424b75da314c1045981bd9777432fad05a9e0c69daa4ed7e308bbaffe405/certifi-2026.6.17.tar.gz", hash = "sha256:024c88eeec92ca068db80f02b8b07c9cef7b9fe261d1d535abfd5abd6f6af432", size = 134594, upload-time = "2026-06-17T10:31:07.894Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a6/bd/1436c50cc46699de0230fbaacc1d1de7bfec413176099949e9d5112e88f2/certifi-2026.10.19.tar.gz", hash = "sha256:035c1e85572b141bf1c3a3eff0e33a8eb74917522f2559075ff8d2fea780a1ae", size = 137598, upload-time = "2026-10-20T02:24:16.604Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/2f/c5464532e965badff2f4c4c1a3a83f5697f0d7c407ed0cda44aaa99bb451/certifi-2026.6.17-py3-none-any.whl", hash = "sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db", size = 133289, upload-time = "2026-06-17T10:31:06.348Z" }, + { url = "https://files.pythonhosted.org/packages/c4/51/29046e48f693c54e5b7412216d447a1d8362c9653148d98c02cdf65fa733/certifi-2026.10.19-py3-none-any.whl", hash = "sha256:42d35340cc5b9ba78e818ddbff0c5b1c2b779de448a340fd1b989f5fa8670e1b", size = 136464, upload-time = "2026-10-20T02:24:15.072Z" }, ] [[package]] From a4544c54722c4b29c88f0190bf58f1ddd7d50a66 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Wed, 21 Oct 2026 20:27:11 -0700 Subject: [PATCH 10/33] feat: configure automatic refresh when declaring computed columns (#4607) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `auto_refresh_enabled` and `auto_refresh_interval_seconds` to the Rust builder and Python sync/async add-column APIs for computed columns and registered Function outputs. Python users can configure refresh when declaring a column: ```python table.add_columns( computed={"total": "price * quantity"}, auto_refresh_enabled=True, auto_refresh_interval_seconds=300, ) ``` Omit the interval to let the server determine when to refresh: ```python table.add_columns(computed={"discounted_total": "total * 0.9"}) ``` The async API accepts the same options: ```python await async_table.add_columns( computed={"total": "price * quantity"}, auto_refresh_enabled=True, auto_refresh_interval_seconds=300, ) ``` For an existing registered `FunctionApplication`, the options also apply to Function columns. Disabling automatic refresh still permits explicit refresh: ```python table.add_columns( {"features": function_application}, auto_refresh_enabled=False, auto_refresh_interval_seconds=300, ) table.refresh_column("features") ``` Options apply to every output declared by a call. Automatic refresh is enabled by default; when the interval is omitted, the server determines when to refresh. An explicit interval sets a minimum delay after a successful refresh. The interval does not force recomputation of unchanged inputs. Invalid intervals and options on one-time transforms are rejected before declaration. Native declarations persist `computed_column.auto_refresh_enabled` and an explicitly supplied `computed_column.auto_refresh_interval_seconds` in each output field’s metadata in the same commit as the declaration. An omitted interval is neither stored nor sent in remote requests. Remote declarations forward the settings to Cloud/Enterprise servers. Local tables store these settings; automatic execution requires a supporting server with scheduling enabled. --- .../feature-engineering/refresh-column.mdx | 48 +++++++ docs/web/sql/ddl/table.mdx | 45 +++++- python/python/lancedb/_lancedb.pyi | 11 +- python/python/lancedb/remote/table.py | 11 +- python/python/lancedb/table.py | 73 +++++++++- .../tests/test_first_class_function_slice1.py | 23 +++- python/python/tests/test_table.py | 50 +++++++ python/src/table.rs | 29 +++- rust/lancedb/src/remote/table.rs | 86 ++++++++---- rust/lancedb/src/table.rs | 10 +- rust/lancedb/src/table/add_columns.rs | 128 +++++++++++++++++- rust/lancedb/src/table/computed_columns.rs | 100 ++++++++++++++ rust/lancedb/src/table/schema_evolution.rs | 6 + 13 files changed, 573 insertions(+), 47 deletions(-) diff --git a/docs/web/feature-engineering/refresh-column.mdx b/docs/web/feature-engineering/refresh-column.mdx index f4550ed24..ce36daa32 100644 --- a/docs/web/feature-engineering/refresh-column.mdx +++ b/docs/web/feature-engineering/refresh-column.mdx @@ -57,6 +57,54 @@ That is what makes the call idempotent and safe to retry. However, unlike a spreadsheet formula, nothing is recomputed until you refresh, so a derived column is stale between an update and the next refresh. +{/* +Draft for automatic computed-column refresh. Keep the released manual-refresh +instructions above until the server feature is available. + +## Automatic refresh {#automatic} + +Set scheduling options when declaring a computed column: + +```python +table.add_columns( + computed={"doubled": "x * 2"}, + auto_refresh_enabled=True, + auto_refresh_interval_seconds=300, +) +``` + +Both options apply to every computed output declared in that call, including +registered Function outputs. Automatic refresh is enabled by default. When +the interval is omitted, the server determines when to refresh. Local tables +persist only an explicitly configured interval; automatic execution requires a +Cloud or Enterprise server with its separate `column_refresh.enabled` switch +turned on. That server switch is disabled in the base configuration. + +The interval is a positive minimum delay after the last successful refresh +of that output. A first refresh is immediately eligible. The scheduler still +checks input freshness and waits for computed ancestors to settle; the timer +does not make unchanged inputs require recomputation. Each Function output's +interval is independent, so refreshing a sibling does not postpone it. + +Pacing uses retained live Job Registry history. If completed jobs are archived +before a configured interval expires, the scheduler can lose the last-success +timestamp and refresh earlier. Leave archival disabled, or retain live history +for at least the longest configured column interval. Archival is disabled in +the base configuration; its default retention threshold is seven days. + +Change a stored setting through SQL: + +```sql +ALTER TABLE my_table ALTER COLUMN doubled SET auto_refresh_enabled = false; +ALTER TABLE my_table ALTER COLUMN doubled SET auto_refresh_interval_seconds = 600; +ALTER TABLE my_table ALTER COLUMN doubled SET auto_refresh_enabled = true; +``` + +Disabling automatic refresh prevents new automatic jobs. Inflight jobs may +finish, and manual refresh remains available. Re-enabling wakes maintenance +to consider accumulated changes under the interval and dependency rules. +*/} + ## Refreshing part of a table {#partial} Enterprise In SQL a refresh can be scoped to a diff --git a/docs/web/sql/ddl/table.mdx b/docs/web/sql/ddl/table.mdx index 792568220..11e7a3900 100644 --- a/docs/web/sql/ddl/table.mdx +++ b/docs/web/sql/ddl/table.mdx @@ -51,6 +51,43 @@ The column is committed empty and the expression is stored, not evaluated: declaring one on a billion-row table costs what declaring one on an empty table costs. [`REFRESH COLUMN`](/sql/dml/refresh) fills it. +{/* +Draft for automatic computed-column refresh. Enable this section when the +server feature is released. + +### Automatic refresh options {#automatic-refresh-options} + +```sql +ALTER TABLE t ADD COLUMN doubled AS value * 2 +WITH (auto_refresh_enabled = true, auto_refresh_interval_seconds = 300); + +ALTER TABLE t ALTER COLUMN doubled SET auto_refresh_enabled = false; +ALTER TABLE t ALTER COLUMN doubled SET auto_refresh_interval_seconds = 600; +ALTER TABLE t ALTER COLUMN doubled SET auto_refresh_enabled = true; +``` + +Automatic refresh is enabled by default. An explicitly configured interval is +stored in each computed output's field metadata; when omitted, the server +determines when to refresh. The interval must be a positive integer +and is minimum pacing after a successful refresh of that output; unchanged +inputs do not force work. A new column is immediately eligible for its first +refresh. The server's separate automatic-refresh switch must also be enabled. + +These options also work beside a Function binding's `version` and `branch`: + +```sql +ALTER TABLE t ADD COLUMN (label, score) AS classify(text) +WITH (version = 'fv_1', auto_refresh_enabled = true, + auto_refresh_interval_seconds = 300); +ALTER TABLE t ALTER COLUMN score SET auto_refresh_enabled = false; +``` + +Declaration settings initialize every output in the group; later SET commands +change one output. Disabling automatic refresh prevents new submissions while +allowing inflight jobs to finish and explicit REFRESH COLUMN to run. Column +SUSPEND/RESUME syntax is not supported. +*/} + {/* Draft for the Functions release. @@ -81,12 +118,14 @@ ALTER TABLE t ADD COLUMN (label, score) AS classify(text) WITH (version = 'fv_1' ```sql ALTER TABLE t ADD COLUMN c AS f(x) FILTER (WHERE keep) WITH (version = 'fv_1') -ALTER TABLE t ADD COLUMN c AS f(x) WITH (num_workers = '2') +REFRESH COLUMN c ON t WITH (num_workers = 2) ``` -`WITH` on a plain computed column accepts only `version`. Resource options such -as `num_workers` apply to function columns, where there are workers to size. +Declaration `WITH` accepts `auto_refresh_enabled` and +`auto_refresh_interval_seconds`. Function bindings also accept `version` and +`branch`. Worker resource options such as `num_workers` belong on +`REFRESH COLUMN ... WITH (...)`. */} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 5a613b210..89c69e721 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -501,10 +501,17 @@ class Table: async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ... async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ... async def add_computed_columns( - self, columns: list[tuple[str, str]] + self, + columns: list[tuple[str, str]], + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[int] = None, ) -> AddColumnsResult: ... async def add_function_columns( - self, application_json: str, output_name: Optional[str] + self, + application_json: str, + output_name: Optional[str], + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[int] = None, ) -> AddColumnsResult: ... async def refresh_column(self, column: str) -> RefreshColumnResult: ... async def refresh_column_async(self, column: str) -> Job: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index eec958799..96b988911 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1001,8 +1001,17 @@ class RemoteTable(Table): | None = None, *, computed: Dict[str, str] | None = None, + auto_refresh_enabled: bool | None = None, + auto_refresh_interval_seconds: int | None = None, ) -> AddColumnsResult: - return LOOP.run(self._table.add_columns(transforms, computed=computed)) + return LOOP.run( + self._table.add_columns( + transforms, + computed=computed, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, + ) + ) def refresh_column(self, column: str): return LOOP.run(self._table.refresh_column(column)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 60404b73c..d5872e6f0 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2561,6 +2561,8 @@ class Table(ABC): | None = None, *, computed: Dict[str, str] | None = None, + auto_refresh_enabled: bool | None = None, + auto_refresh_interval_seconds: int | None = None, ): """ Add new columns with defined values. @@ -2605,6 +2607,15 @@ class Table(ABC): [`refresh_column_async`][lancedb.table.Table.refresh_column_async]. Cannot be combined with ``transforms``. + auto_refresh_enabled: bool, optional + Enable automatic refresh for each declared computed output. Defaults + to True. Local tables store this option; scheduling requires a server. + auto_refresh_interval_seconds: int, optional + Positive minimum interval after a successful refresh, in seconds. + When omitted, the server determines when to refresh. Applies to every + computed output in this call. + Scheduling options require computed or Function columns. + Returns ------- AddColumnsResult @@ -4803,8 +4814,17 @@ class LanceTable(Table): | None = None, *, computed: Dict[str, str] | None = None, + auto_refresh_enabled: bool | None = None, + auto_refresh_interval_seconds: int | None = None, ) -> AddColumnsResult: - return LOOP.run(self._table.add_columns(transforms, computed=computed)) + return LOOP.run( + self._table.add_columns( + transforms, + computed=computed, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, + ) + ) def refresh_column(self, column: str) -> "RefreshColumnResult": """Fill a computed column's unfilled rows and recompute those whose @@ -6852,6 +6872,8 @@ class AsyncTable: | None = None, *, computed: dict[str, str] | None = None, + auto_refresh_enabled: bool | None = None, + auto_refresh_interval_seconds: int | None = None, ) -> AddColumnsResult: """ Add new columns with defined values. @@ -6891,12 +6913,46 @@ class AsyncTable: On LanceDB Cloud and Enterprise the expression is planned by the server. Cannot be combined with ``transforms``. + auto_refresh_enabled: bool, optional + Enable automatic refresh for each computed output declared by this + call. Defaults to True. Local tables store the setting; automatic + scheduling requires Cloud or Enterprise. + auto_refresh_interval_seconds: int, optional + Positive minimum seconds after a successful refresh before another + automatic refresh is eligible. When omitted, the server determines when + to refresh. Unchanged inputs do + not force recomputation. Applies to every declared computed output. + These options cannot be used with one-time transforms or ordinary + all-null columns. + + Examples + -------- + >>> await table.add_columns( # doctest: +SKIP + ... computed={"total": "price * quantity"}, + ... auto_refresh_enabled=True, + ... auto_refresh_interval_seconds=300, + ... ) + Returns ------- AddColumnsResult version: the new version number of the table after adding columns. """ + if auto_refresh_enabled is not None and type(auto_refresh_enabled) is not bool: + raise ValueError("auto_refresh_enabled must be a boolean") + if auto_refresh_interval_seconds is not None and ( + type(auto_refresh_interval_seconds) is not int + or not 0 < auto_refresh_interval_seconds <= (2**64 - 1) // 1000 + ): + raise ValueError( + "auto_refresh_interval_seconds must be a positive integer " + "no greater than u64::MAX / 1000" + ) + options_supplied = ( + auto_refresh_enabled is not None + or auto_refresh_interval_seconds is not None + ) function_application = None function_output_name = None if isinstance(transforms, FunctionApplication): @@ -6920,7 +6976,10 @@ class AsyncTable: ) function_application._ensure_declarable() return await self._inner.add_function_columns( - function_application.to_canonical_json(), function_output_name + function_application.to_canonical_json(), + function_output_name, + auto_refresh_enabled, + auto_refresh_interval_seconds, ) if isinstance(transforms, pa.Field): @@ -6934,7 +6993,15 @@ class AsyncTable: raise ValueError( "add_columns cannot take both transforms and computed columns" ) - return await self._inner.add_computed_columns(list(computed.items())) + return await self._inner.add_computed_columns( + list(computed.items()), + auto_refresh_enabled, + auto_refresh_interval_seconds, + ) + if options_supplied: + raise ValueError( + "automatic refresh options require computed or Function columns" + ) if transforms is None: raise ValueError("add_columns requires transforms or computed columns") if isinstance(transforms, pa.Schema): diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py index b08dadf36..06730a5f1 100644 --- a/python/python/tests/test_first_class_function_slice1.py +++ b/python/python/tests/test_first_class_function_slice1.py @@ -398,8 +398,21 @@ class _FunctionDeclarationInner: def __init__(self): self.calls = [] - async def add_function_columns(self, application_json, output_name): - self.calls.append((json.loads(application_json), output_name)) + async def add_function_columns( + self, + application_json, + output_name, + auto_refresh_enabled, + auto_refresh_interval_seconds, + ): + self.calls.append( + ( + json.loads(application_json), + output_name, + auto_refresh_enabled, + auto_refresh_interval_seconds, + ) + ) return "declared" @@ -420,14 +433,18 @@ async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomica ) assert result == "declared" assert inner.calls[-1][1] == "features" + assert inner.calls[-1][2:] == (None, None) bare = application._copy(update={"columns": {}}).rename( columns={"normalized_text": "search_text"} ) - result = await table.add_columns(bare) + result = await table.add_columns( + bare, auto_refresh_enabled=False, auto_refresh_interval_seconds=300 + ) assert result == "declared" assert inner.calls[-1][1] is None assert inner.calls[-1][0]["columns"] == {"normalized_text": "search_text"} + assert inner.calls[-1][2:] == (False, 300) @pytest.mark.asyncio diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7ec2003b7..755e94d7f 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4869,3 +4869,53 @@ async def test_refresh_column_async_job_async_table(tmp_path): assert result.published_version == 4 assert await job.status() == "finished" assert (await table.to_arrow())["tripled"].to_pylist() == [9] + + +@pytest.mark.parametrize("interval", [None, 300]) +def test_computed_auto_refresh_options(tmp_path, interval): + db = lancedb.connect(tmp_path) + table = db.create_table("refresh_options", [{"x": 1}]) + before = table.version + table.add_columns( + computed={"doubled": "x * 2", "tripled": "x * 3"}, + auto_refresh_enabled=False, + auto_refresh_interval_seconds=interval, + ) + assert table.version == before + 1 + for column in ("doubled", "tripled"): + metadata = table.schema.field(column).metadata + assert metadata[b"computed_column.auto_refresh_enabled"] == b"false" + if interval is None: + assert b"computed_column.auto_refresh_interval_seconds" not in metadata + else: + assert metadata[b"computed_column.auto_refresh_interval_seconds"] == b"300" + table.refresh_column("doubled") + assert table.to_arrow()["doubled"].to_pylist() == [2] + + +@pytest.mark.parametrize("value", [0, -1, True, 1.5, "300", 2**64]) +def test_computed_auto_refresh_rejects_invalid_interval(tmp_path, value): + table = lancedb.connect(tmp_path).create_table("invalid_options", [{"x": 1}]) + with pytest.raises(ValueError, match="auto_refresh_interval_seconds"): + table.add_columns( + computed={"doubled": "x * 2"}, auto_refresh_interval_seconds=value + ) + assert table.schema.names == ["x"] + + +def test_auto_refresh_options_require_computed_columns(tmp_path): + table = lancedb.connect(tmp_path).create_table("invalid_kind", [{"x": 1}]) + with pytest.raises(ValueError, match="require computed"): + table.add_columns({"doubled": "x * 2"}, auto_refresh_enabled=True) + + +@pytest.mark.asyncio +async def test_async_computed_auto_refresh_uses_server_interval(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("server_interval", [{"x": 1}]) + before = await table.version() + await table.add_columns(computed={"doubled": "x * 2"}) + assert await table.version() == before + 1 + metadata = (await table.schema()).field("doubled").metadata + assert metadata[b"computed_column.auto_refresh_enabled"] == b"true" + assert b"computed_column.auto_refresh_interval_seconds" not in metadata diff --git a/python/src/table.rs b/python/src/table.rs index 6cea168ec..c1d23d151 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use lancedb::table::computed_columns::ComputedColumnOptions; use std::{ collections::HashMap, sync::{ @@ -1803,35 +1804,59 @@ impl Table { }) } + #[pyo3(signature = (columns, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] pub fn add_computed_columns( self_: PyRef<'_, Self>, columns: Vec<(String, String)>, + auto_refresh_enabled: Option, + auto_refresh_interval_seconds: Option, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); + let options = ComputedColumnOptions { + auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true), + auto_refresh_interval_seconds, + }; + options.validate().infer_error()?; future_into_py(self_.py(), async move { let mut builder = inner.add_columns(); for (name, expression) in columns { builder = builder.computed(name, expression); } - let result = builder.execute().await.infer_error()?; + let result = builder + .computed_options(options) + .execute() + .await + .infer_error()?; Ok(AddColumnsResult::from(result)) }) } + #[pyo3(signature = (application_json, output_name, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] pub fn add_function_columns( self_: PyRef<'_, Self>, application_json: String, output_name: Option, + auto_refresh_enabled: Option, + auto_refresh_interval_seconds: Option, ) -> PyResult> { let application = lancedb::function::FunctionApplication::from_json(&application_json).infer_error()?; let inner = self_.inner_ref()?.clone(); + let options = ComputedColumnOptions { + auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true), + auto_refresh_interval_seconds, + }; + options.validate().infer_error()?; future_into_py(self_.py(), async move { let builder = match output_name { Some(name) => inner.add_columns().function_as(name, application), None => inner.add_columns().function(application), }; - let result = builder.execute().await.infer_error()?; + let result = builder + .computed_options(options) + .execute() + .await + .infer_error()?; Ok(AddColumnsResult::from(result)) }) } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 3325f6943..1c569f6e5 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3497,7 +3497,11 @@ impl BaseTable for RemoteTable { Ok(result) } - async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { + async fn add_computed_columns( + &self, + columns: &[(String, String)], + options: crate::table::computed_columns::ComputedColumnOptions, + ) -> Result { self.check_mutable().await?; crate::table::computed_columns::ensure_not_function_bound( self.schema().await?.as_ref(), @@ -3516,7 +3520,10 @@ impl BaseTable for RemoteTable { }, ) .collect::>(); - let mut body = serde_json::json!({ "new_columns": entries }); + let mut body = serde_json::json!({ "new_columns": entries, "auto_refresh_enabled": options.auto_refresh_enabled }); + if let Some(interval) = options.auto_refresh_interval_seconds { + body["auto_refresh_interval_seconds"] = serde_json::json!(interval); + } self.apply_branch_body(&mut body); let request = self .client @@ -3550,6 +3557,7 @@ impl BaseTable for RemoteTable { &self, application: &crate::function::FunctionApplication, output_name: Option<&str>, + options: crate::table::computed_columns::ComputedColumnOptions, ) -> Result { self.check_mutable().await?; let schema = self.schema().await?; @@ -3570,6 +3578,7 @@ impl BaseTable for RemoteTable { .collect::>(); let mut body = serde_json::json!({ "new_columns": new_columns, + "auto_refresh_enabled": options.auto_refresh_enabled, "function": { "application": plan.application, "binding_metadata_version": plan.binding_metadata_version, @@ -3579,6 +3588,9 @@ impl BaseTable for RemoteTable { "outputs": plan.outputs, }, }); + if let Some(interval) = options.auto_refresh_interval_seconds { + body["auto_refresh_interval_seconds"] = serde_json::json!(interval); + } self.apply_branch_body(&mut body); let request = self .client @@ -9263,31 +9275,44 @@ mod tests { /// client never types the expression itself. #[tokio::test] async fn test_add_computed_columns_sends_the_expression() { - let table = Table::new_with_handler("my_table", |request| match request.url().path() { - "/v1/table/my_table/describe/" => simple_describe_response(), - "/v1/table/my_table/add_columns/" => { - assert_eq!(request.method(), "POST"); - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!( - value["new_columns"], - serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) - ); - http::Response::builder() - .status(200) - .body(r#"{"version": 7}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - }); + for interval in [None, Some(300)] { + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" => { + assert_eq!(request.method(), "POST"); + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + value["new_columns"], + serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) + ); + assert_eq!(value["auto_refresh_enabled"], false); + if let Some(interval) = interval { + assert_eq!(value["auto_refresh_interval_seconds"], interval); + } else { + assert!(value.get("auto_refresh_interval_seconds").is_none()); + } + http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); - let result = table - .add_columns() - .computed("doubled", "x * 2") - .execute() - .await - .unwrap(); - assert_eq!(result.version, 7); + let result = table + .add_columns() + .computed("doubled", "x * 2") + .computed_options(crate::table::computed_columns::ComputedColumnOptions { + auto_refresh_enabled: false, + auto_refresh_interval_seconds: interval, + }) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 7); + } } #[tokio::test] @@ -9305,10 +9330,11 @@ mod tests { request.body().unwrap().as_bytes().unwrap(), ) .unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( + let mut expected: serde_json::Value = serde_json::from_str(include_str!( "../../tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json" )) .unwrap(); + expected["auto_refresh_enabled"] = serde_json::json!(true); assert_eq!(actual, expected); http::Response::builder() .status(200) @@ -9504,10 +9530,11 @@ mod tests { request.body().unwrap().as_bytes().unwrap(), ) .unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( + let mut expected: serde_json::Value = serde_json::from_str(include_str!( "../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" )) .unwrap(); + expected["auto_refresh_enabled"] = serde_json::json!(true); assert_eq!(actual, expected); http::Response::builder() .status(200) @@ -9550,10 +9577,11 @@ mod tests { "/v1/table/my_table/add_columns/" => { let actual: serde_json::Value = serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - let expected: serde_json::Value = serde_json::from_str(include_str!( + let mut expected: serde_json::Value = serde_json::from_str(include_str!( "../../tests/fixtures/first_class_functions/v1/remote_multi_output_declaration_request.json" )) .unwrap(); + expected["auto_refresh_enabled"] = serde_json::json!(true); assert_eq!(actual, expected); http::Response::builder() .status(200) diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index d99926ae1..61a710dcd 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -814,6 +814,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn add_computed_columns( &self, _columns: &[(String, String)], + _options: computed_columns::ComputedColumnOptions, ) -> Result { Err(Error::NotSupported { message: "computed columns are not supported on this table type".into(), @@ -824,6 +825,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { &self, _application: &crate::function::FunctionApplication, _output_name: Option<&str>, + _options: computed_columns::ComputedColumnOptions, ) -> Result { Err(Error::NotSupported { message: "Function columns are supported only on LanceDB Cloud and Enterprise".into(), @@ -3768,8 +3770,12 @@ impl BaseTable for NativeTable { Ok(result) } - async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { - let result = schema_evolution::execute_declare(self, columns).await?; + async fn add_computed_columns( + &self, + columns: &[(String, String)], + options: computed_columns::ComputedColumnOptions, + ) -> Result { + let result = schema_evolution::execute_declare(self, columns, options).await?; self.bump_freshness(); Ok(result) } diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index 635a99bea..4f9cbb772 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use lance::dataset::NewColumnTransform; use super::BaseTable; +use super::computed_columns::ComputedColumnOptions; use super::schema_evolution::AddColumnsResult; use crate::function::FunctionApplication; use crate::{Error, Result}; @@ -19,6 +20,7 @@ pub struct AddColumnsBuilder { computed: Vec<(String, String)>, function: Option<(FunctionApplication, Option)>, read_columns: Option>, + options: Option, } impl std::fmt::Debug for AddColumnsBuilder { @@ -41,6 +43,7 @@ impl AddColumnsBuilder { computed: Vec::new(), function: None, read_columns: None, + options: None, } } @@ -139,6 +142,26 @@ impl AddColumnsBuilder { self } + /// Configure automatic refresh for every output declared in this call. + /// Settings are stored with the declaration. They require computed or + /// Function columns; local tables persist them without starting a scheduler. + /// + /// ``` + /// # use lancedb::{Table, table::computed_columns::ComputedColumnOptions}; + /// # async fn declare(table: &Table) -> lancedb::Result<()> { + /// table.add_columns().computed("doubled", "x * 2") + /// .computed_options(ComputedColumnOptions { + /// auto_refresh_enabled: true, + /// auto_refresh_interval_seconds: Some(300), + /// }).execute().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn computed_options(mut self, options: ComputedColumnOptions) -> Self { + self.options = Some(options); + self + } + /// Add the columns. pub async fn execute(self) -> Result { let Self { @@ -147,7 +170,18 @@ impl AddColumnsBuilder { computed, function, read_columns, + options, } = self; + if let Some(options) = options { + options.validate()?; + if computed.is_empty() && function.is_none() { + return Err(Error::InvalidInput { + message: "automatic refresh options require computed or Function columns" + .into(), + }); + } + } + let options = options.unwrap_or_default(); let declaration_count = usize::from(!computed.is_empty()) + usize::from(function.is_some()); if transform.is_some() && declaration_count != 0 || declaration_count > 1 { @@ -179,7 +213,7 @@ impl AddColumnsBuilder { .into(), }); } - parent.add_computed_columns(&computed).await + parent.add_computed_columns(&computed, options).await } (None, true, Some((application, output_name))) => { if read_columns.is_some() { @@ -189,7 +223,7 @@ impl AddColumnsBuilder { }); } parent - .add_function_columns(&application, output_name.as_deref()) + .add_function_columns(&application, output_name.as_deref(), options) .await } _ => unreachable!("mixed add_columns modes were rejected above"), @@ -201,6 +235,9 @@ impl AddColumnsBuilder { mod tests { use std::sync::Arc; + use super::ComputedColumnOptions; + use crate::table::computed_columns::AUTO_REFRESH_INTERVAL_META_KEY; + use arrow_array::{Int32Array, RecordBatch, record_batch}; use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BatchUDF, NewColumnTransform}; @@ -214,6 +251,93 @@ mod tests { conn.create_table(name, batch).execute().await.unwrap() } + #[tokio::test] + async fn computed_options_commit_with_declaration() { + let table = table_with_two_columns("options").await; + let version = table.version().await.unwrap(); + let options = ComputedColumnOptions { + auto_refresh_enabled: false, + auto_refresh_interval_seconds: Some(300), + }; + let result = table + .add_columns() + .computed("doubled", "x * 2") + .computed_options(options) + .execute() + .await + .unwrap(); + assert_eq!(result.version, version + 1); + let schema = table.schema().await.unwrap(); + assert_eq!( + ComputedColumnOptions::from_metadata( + schema.field_with_name("doubled").unwrap().metadata() + ) + .unwrap(), + options + ); + for interval in [0, u64::MAX] { + let invalid = ComputedColumnOptions { + auto_refresh_interval_seconds: Some(interval), + ..options + }; + assert!( + table + .add_columns() + .computed("bad", "x") + .computed_options(invalid) + .execute() + .await + .is_err() + ); + } + assert!( + table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "once".into(), + "x".into() + )])) + .computed_options(options) + .execute() + .await + .is_err() + ); + assert!( + table + .schema() + .await + .unwrap() + .field_with_name("bad") + .is_err() + ); + } + + #[tokio::test] + async fn omitted_refresh_interval_stays_unset() { + let table = table_with_two_columns("server_interval").await; + let version = table.version().await.unwrap(); + let result = table + .add_columns() + .computed("doubled", "x * 2") + .execute() + .await + .unwrap(); + assert_eq!(result.version, version + 1); + let schema = table.schema().await.unwrap(); + let metadata = schema.field_with_name("doubled").unwrap().metadata(); + assert!(!metadata.contains_key(AUTO_REFRESH_INTERVAL_META_KEY)); + assert_eq!( + ComputedColumnOptions::from_metadata(metadata).unwrap(), + ComputedColumnOptions::default() + ); + assert!( + serde_json::to_value(ComputedColumnOptions::default()) + .unwrap() + .get("auto_refresh_interval_seconds") + .is_none() + ); + } + #[tokio::test] async fn test_requires_a_transform() { let table = table_with_two_columns("no_transform").await; diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 9a3922b2b..e86b350ee 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -55,6 +55,106 @@ pub const EXPRESSION_META_KEY: &str = "computed_column.expression"; /// Field metadata key holding the column's inputs, as a JSON array of names. pub const INPUTS_META_KEY: &str = "computed_column.inputs"; +/// Whether the server may automatically refresh this output. +pub const AUTO_REFRESH_ENABLED_META_KEY: &str = "computed_column.auto_refresh_enabled"; +/// Minimum interval between successful automatic refreshes, in seconds. +pub const AUTO_REFRESH_INTERVAL_META_KEY: &str = "computed_column.auto_refresh_interval_seconds"; + +/// Scheduling settings persisted on each computed output. Local tables store +/// these settings; automatic scheduling requires a Cloud or Enterprise server. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ComputedColumnOptions { + /// Permit automatic refresh. Manual refresh is always available. + pub auto_refresh_enabled: bool, + /// Minimum seconds after a successful refresh before another is eligible. + /// When omitted, the server determines when to refresh. + #[serde(skip_serializing_if = "Option::is_none")] + pub auto_refresh_interval_seconds: Option, +} + +impl Default for ComputedColumnOptions { + fn default() -> Self { + Self { + auto_refresh_enabled: true, + auto_refresh_interval_seconds: None, + } + } +} + +impl ComputedColumnOptions { + /// Reject intervals that cannot be represented by the server scheduler. + pub fn validate(&self) -> Result<()> { + if self + .auto_refresh_interval_seconds + .is_some_and(|interval| interval == 0 || interval.checked_mul(1000).is_none()) + { + return Err(Error::InvalidInput { + message: "auto_refresh_interval_seconds must be a positive integer no greater than u64::MAX / 1000".into(), + }); + } + Ok(()) + } + + /// Read scheduling settings, using defaults for omitted properties. + pub fn from_metadata(metadata: &HashMap) -> Result { + let mut options = Self::default(); + if let Some(value) = metadata.get(AUTO_REFRESH_ENABLED_META_KEY) { + options.auto_refresh_enabled = value.parse().map_err(|_| Error::InvalidInput { + message: "auto_refresh_enabled must be true or false".into(), + })?; + } + if let Some(value) = metadata.get(AUTO_REFRESH_INTERVAL_META_KEY) { + options.auto_refresh_interval_seconds = + Some(value.parse().map_err(|_| Error::InvalidInput { + message: "auto_refresh_interval_seconds must be a positive integer".into(), + })?); + } + options.validate()?; + Ok(options) + } + + /// Attach validated settings to newly declared computed output fields. + pub fn apply_to_schema(&self, schema: &ArrowSchema) -> Result { + self.validate()?; + let fields = schema + .fields() + .iter() + .map(|field| { + if computed_column_from_field(field).is_none() + || field + .metadata() + .get(FUNCTION_OUTPUT_ORDINAL_META_KEY) + .is_some_and(|ordinal| { + ordinal == &FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL.to_string() + }) + { + return field.clone(); + } + let mut metadata = field.metadata().clone(); + metadata.extend(self.metadata()); + Arc::new(field.as_ref().clone().with_metadata(metadata)) + }) + .collect::>(); + Ok(ArrowSchema::new_with_metadata( + fields, + schema.metadata().clone(), + )) + } + + /// Field metadata values shared by native and server declaration paths. + pub fn metadata(&self) -> HashMap { + let mut metadata = HashMap::from([( + AUTO_REFRESH_ENABLED_META_KEY.into(), + self.auto_refresh_enabled.to_string(), + )]); + if let Some(interval) = self.auto_refresh_interval_seconds { + metadata.insert(AUTO_REFRESH_INTERVAL_META_KEY.into(), interval.to_string()); + } + metadata + } +} + /// Field metadata key holding the Function binding identity. pub const FUNCTION_BINDING_ID_META_KEY: &str = "computed_column.function.binding_id"; diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index c541a9a43..1ad532d0a 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -13,6 +13,7 @@ use lance::dataset::transaction::{Operation, Transaction, UpdateMap, UpdateMapEn use lance::dataset::{ColumnAlteration, CommitBuilder, Dataset, NewColumnTransform}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::sync::Arc; use super::computed_columns; use super::{BaseTable, NativeTable}; @@ -128,6 +129,7 @@ pub(crate) async fn execute_add_columns( pub(crate) async fn execute_declare( table: &NativeTable, columns: &[(String, String)], + options: computed_columns::ComputedColumnOptions, ) -> Result { use lance::dataset::mem_wal::DatasetMemWalExt; @@ -157,6 +159,10 @@ pub(crate) async fn execute_declare( }); } let transform = computed_columns::declare(table.schema().await?, columns)?; + let NewColumnTransform::AllNulls(schema) = transform else { + unreachable!("declarations are all-null") + }; + let transform = NewColumnTransform::AllNulls(Arc::new(options.apply_to_schema(&schema)?)); commit_add_columns(table, transform, None).await } From 1564ec10c8534a828aa38b2fb05c36b464844945 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Wed, 21 Oct 2026 23:18:30 -0700 Subject: [PATCH 11/33] feat: add Bootstrap split strategy for sampling with replacement (#4597) ## Context `StreamingDataset`'s shipped split strategies (`split_random`, `split_sequential`, `split_hash`, `split_calculated`) all share one trait: every row in the base table is assigned to exactly one split, exactly once. They only differ in *how* that one-to-one assignment is chosen -- none of them can express "give me a resampled stream of N rows, allowing repeats." That's a distinct, common need in ML training: **the bootstrap** -- the statistical technique (Efron, 1979) of estimating a dataset's properties by resampling it *with replacement*. It's the same sense scikit-learn's bagging classifiers mean by `bootstrap=True`. `SplitStrategy::Bootstrap` adds exactly this to the data loader: sample rows uniformly with replacement from the (optionally filtered) base table, independently per split. ## Main use cases - **Bagging / ensemble training** (e.g. random forests): each split trains on an independent bootstrap resample of the training set. - **Oversampling a rare class**: filter to the minority class, then sample it up to whatever count balances it against the majority, instead of exhausting its small set once per epoch. - **Fixed-size training steps independent of dataset size**: want exactly N batches per epoch regardless of whether the table has 50 rows or 50 million? Set `num_samples` directly instead of it being implicitly `len(dataset)`. - **Weighted sampling without a weight column**: filter to "eligible" rows (e.g. a windowed-sampling workload filtering to valid window-start positions), then bootstrap the filtered set -- an item with more eligible rows is naturally drawn proportionally more often, with no explicit per-row weight needed. ## Example (an actual run, not a hypothetical one) A 5-row base table (`row_id` 0..4), sampled via `Bootstrap { num_samples: 8, sizes: Fixed(2), seed: Some(42) }`: ``` row_id split_id 0 0 2 0 4 0 4 0 <- row 4 repeated within split 0 4 1 <- row 4 also drawn in split 1 (independent draw) 4 1 1 1 2 1 ``` 8 output rows from a 5-row table -- not 5. Row 4 drawn 4 times total, row 3 never drawn. That's with-replacement sampling actually behaving like with-replacement sampling, not an even round-robin dressed up with a new name. ## What changed Adds `SplitStrategy::Bootstrap` (uniform sampling with replacement, optionally filtered, independently seeded per split -- a given split's content doesn't depend on how many other splits there are, the same elastic-determinism property documented for the loader's other strategies). No new `PermutationBuilder` Rust method was needed: `with_split_strategy` already accepts any `SplitStrategy` value generically. Python binding: `PermutationBuilder.bootstrap_sample(num_samples, *, ratios=None, counts=None, fixed=None, seed=None, split_names=None)`, mirroring `split_random`'s existing shape. Also adds a generic `split_strategy` parameter to `StreamingDataset`'s constructor (a callable taking/returning a `PermutationBuilder`), since it previously could only reach `split_random`/`split_sequential` even though `split_hash`/`split_calculated` already shipped. This unblocks `Bootstrap` and any future strategy without further constructor changes. --- python/python/lancedb/_lancedb.pyi | 9 + python/python/lancedb/permutation.py | 44 +++ python/python/lancedb/streaming.py | 150 ++++++- .../python/tests/test_elastic_dataloader.py | 227 +++++++++++ python/python/tests/test_permutation.py | 70 ++++ python/src/permutation.rs | 44 +++ .../src/dataloader/permutation/builder.rs | 6 +- .../src/dataloader/permutation/split.rs | 365 +++++++++++++++++- 8 files changed, 909 insertions(+), 6 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 89c69e721..96bc7752b 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -947,6 +947,15 @@ class AsyncPermutationBuilder: counts: Optional[List[int]] = None, fixed: Optional[int] = None, ) -> "AsyncPermutationBuilder": ... + def bootstrap_sample( + self, + num_samples: int, + *, + ratios: Optional[List[float]] = None, + counts: Optional[List[int]] = None, + fixed: Optional[int] = None, + seed: Optional[int] = None, + ) -> "AsyncPermutationBuilder": ... def split_calculated(self, calculation: str) -> "AsyncPermutationBuilder": ... def shuffle( self, seed: Optional[int], clump_size: Optional[int] diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 103ec0aa4..355ebf798 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -169,6 +169,50 @@ class PermutationBuilder: ) return self + def bootstrap_sample( + self, + num_samples: int, + *, + ratios: Optional[list[float]] = None, + counts: Optional[list[int]] = None, + fixed: Optional[int] = None, + seed: Optional[int] = None, + split_names: Optional[list[str]] = None, + ) -> "PermutationBuilder": + """ + Sample rows uniformly **with replacement**, independently per split. + + Unlike the other split strategies, the number of output rows is + ``num_samples``, not the number of (optionally filtered) input rows: the + same row can appear zero, one, or many times, and ``num_samples`` may be + larger or smaller than the input row count. Useful for bootstrap + resampling, oversampling rare classes, or any training workload that + wants a fixed-size stream of random draws rather than a single + deterministic pass over every row. Combine with :meth:`filter` to + restrict the eligible population. + + One of ratios, counts, or fixed must be provided, with the same meaning + as in :meth:`split_random`, except they divide ``num_samples`` across + splits rather than the input row count: if fixed is 4, each split gets + ``num_samples // 4`` samples. + + Each split draws its own sample using a seed derived independently from + ``(seed, split_index)``, so a split's content does not change if the samples are + later redistributed across a different number of splits. + + The optional split_names can be provided to name the splits. If not provided, + the splits can only be referenced by their index. + """ + self._async.bootstrap_sample( + num_samples, + ratios=ratios, + counts=counts, + fixed=fixed, + seed=seed, + split_names=split_names, + ) + return self + def split_calculated( self, calculation: str, split_names: Optional[list[str]] = None ) -> "PermutationBuilder": diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 91ca6ce28..4a4ffd1f5 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -19,6 +19,7 @@ above. """ import ctypes +import hashlib import heapq import logging import os @@ -39,6 +40,7 @@ from torch.utils.data import DataLoader, IterableDataset, get_worker_info from .permutation import ( Permutation, + PermutationBuilder, Transforms, permutation_builder, _drop_base_version, @@ -244,6 +246,20 @@ class StreamingDataset(IterableDataset): default) rows are shuffled using ``shuffle_seed`` and ``epoch``. When ``False`` rows are divided into splits sequentially in storage order, which can be useful for deterministic debugging or evaluation. + Ignored when ``split_strategy`` is provided. + split_strategy: + Escape hatch for any ``PermutationBuilder`` split strategy this class + doesn't otherwise expose a dedicated parameter for -- ``split_hash``, + ``split_calculated``, ``bootstrap_sample`` (uniform sampling with + replacement), or any strategy added later. A callable that receives a + fresh ``PermutationBuilder`` for ``table`` and returns a configured one; + it is called and ``.execute()``'d directly, taking over permutation + construction entirely. When provided, ``shuffle``, ``shuffle_seed``, + ``epoch``, and ``shuffle_clump_size`` are ignored -- the callable is + solely responsible for how rows are split and ordered. For example, + ``split_strategy=lambda b: b.bootstrap_sample(10_000, fixed=num_splits, + seed=42)`` draws 10,000 samples with replacement, divided evenly across + ``num_splits`` splits. shuffle_seed: Base seed for the random permutation. Combined with ``epoch`` so each epoch produces a different ordering. Pass ``None`` to generate @@ -393,6 +409,9 @@ class StreamingDataset(IterableDataset): *, num_splits: Optional[int] = None, shuffle: bool = True, + split_strategy: Optional[ + Callable[[PermutationBuilder], PermutationBuilder] + ] = None, shuffle_seed: Optional[int] = 0, epoch: int = 0, rank: int = 0, @@ -583,13 +602,70 @@ class StreamingDataset(IterableDataset): builder = permutation_builder(table) if filter is not None: builder = builder.filter(filter) - if shuffle: + if split_strategy is not None: + # Escape hatch: lets any PermutationBuilder split strategy drive this + # dataset (split_hash, split_calculated, bootstrap_sample, or any future + # strategy) without StreamingDataset needing a bespoke parameter for each + # one. shuffle/shuffle_seed/epoch/shuffle_clump_size are ignored in this + # path -- the caller's split_strategy is solely responsible for ordering. + self._perm_table = split_strategy(builder).execute() + # _rank_splits below assumes split ids are exactly 0..num_splits-1 with no + # gaps (every other path guarantees this by always building with + # fixed=num_splits) -- a split_strategy callable has no such guarantee, and + # a mismatch would otherwise silently drop rows (ids >= num_splits) or leave + # some ranks reading an empty split (missing ids < num_splits). + split_value_counts = ( + self._perm_table.to_arrow().column("split_id").value_counts() + ) + counts_by_split = dict( + zip( + split_value_counts.field("values").to_pylist(), + split_value_counts.field("counts").to_pylist(), + ) + ) + expected_splits = set(range(num_splits)) + if set(counts_by_split) != expected_splits: + raise ValueError( + f"split_strategy produced split ids {sorted(counts_by_split)}, " + f"but num_splits={num_splits} requires exactly " + f"{{0, ..., {num_splits - 1}}}. Pass num_splits matching the " + "strategy's own split count (e.g. `fixed=num_splits` on " + "bootstrap_sample/split_random/etc.)." + ) + # The round-robin iteration below (and in _iter_owned) assumes every + # split has the same row count -- it stops as soon as any one owned + # split is exhausted (see its own comment: "all splits exhaust + # simultaneously"). Every other path guarantees this by always building + # with fixed=num_splits; a split_strategy using ratios/counts/split_hash + # weights has no such guarantee, and unequal sizes would otherwise + # silently drop rows from the larger splits -- or, with world_size > 1 + # where each rank owns different splits, end ranks after different row + # counts and desync or hang a distributed run's end-of-epoch collectives. + if len(set(counts_by_split.values())) > 1: + raise ValueError( + f"split_strategy produced unequal split sizes " + f"{counts_by_split} (split id -> row count). Every split must " + "have the same row count. Use `fixed=num_splits` rather than " + "`ratios`/`counts`/weighted `split_hash`, which can produce " + "unequal sizes, or split unevenly-sized groups into separate " + "StreamingDatasets instead." + ) + elif shuffle: perm_seed = shuffle_seed + epoch * _EPOCH_PRIME self._perm_table = builder.split_random( fixed=num_splits, seed=perm_seed, clump_size=shuffle_clump_size ).execute() else: self._perm_table = builder.split_sequential(fixed=num_splits).execute() + # Only the split_strategy path needs a permutation fingerprint in the + # checkpoint (see `_checkpoint_snapshot`): split_random/split_sequential are + # already fully pinned by shuffle_seed/num_splits, and a fingerprint over raw + # row ids would spuriously reject an otherwise-valid resume after those row + # ids get renumbered by compaction (row ids are not stable by default). + self._uses_split_strategy = split_strategy is not None + # Lazily computed and cached by `_perm_fingerprint` -- `_perm_table` never + # changes after construction, so this is correct to compute once. + self._perm_fingerprint_cache: Optional[str] = None if self._blocks_per_epoch == "auto": self._blocks_per_epoch = self._estimate_blocks_per_epoch() @@ -1406,11 +1482,51 @@ class StreamingDataset(IterableDataset): ) return state + @property + def _perm_fingerprint(self) -> str: + """Cheap identity check for the built permutation, computed once and cached + (``_perm_table`` never changes after construction). + + Only meaningful -- and only recorded in a checkpoint -- for the + ``split_strategy`` path (see ``_checkpoint_snapshot``): ``num_splits``/ + ``shuffle_seed`` alone cannot detect a `load_state_dict` resuming saved + per-split progress against a *different* permutation than the one that + produced it, since a ``split_strategy`` callable ignores ``shuffle_seed`` + entirely -- e.g. a ``bootstrap_sample`` with no seed (or any changed + callable) can rebuild a completely different permutation on the next run + while ``shuffle_seed`` keeps its same default value and the existing check + passes anyway. ``split_random``/``split_sequential`` are already fully + pinned by ``shuffle_seed``/``num_splits`` and must *not* additionally be + fingerprinted: this hashes physical row ids, which are not stable across + compaction by default, so checking it on those paths would reject an + otherwise-valid resume after the table gets compacted between save and + load even though the seeded permutation itself is unchanged. + """ + if self._perm_fingerprint_cache is None: + arrow_tbl = self._perm_table.to_arrow() + hasher = hashlib.sha256() + hasher.update(str(arrow_tbl.num_rows).encode()) + for name in ("row_id", "split_id"): + column = arrow_tbl.column(name).combine_chunks() + hasher.update(column.to_numpy(zero_copy_only=False).tobytes()) + self._perm_fingerprint_cache = hasher.hexdigest() + return self._perm_fingerprint_cache + def _checkpoint_snapshot(self) -> dict: + # Only recorded for split_strategy: split_random/split_sequential are + # already fully pinned by shuffle_seed/num_splits below, and row ids are not + # stable across compaction by default, so fingerprinting them unconditionally + # would reject an otherwise-valid resume after the table got compacted. + perm_fingerprint = ( + {"perm_fingerprint": self._perm_fingerprint} + if (self._uses_split_strategy) + else {} + ) if self._pack_sequences is not None: return { "shuffle_seed": self._shuffle_seed, "num_splits": self._num_splits, + **perm_fingerprint, "epoch": self._epoch, "pack_sequences": self._pack_sequences, "eos_id": self._eos_id, @@ -1431,6 +1547,7 @@ class StreamingDataset(IterableDataset): return { "shuffle_seed": self._shuffle_seed, "num_splits": self._num_splits, + **perm_fingerprint, "epoch": self._epoch, "samples_consumed_per_split": samples, "positions_consumed_per_split": positions, @@ -1560,8 +1677,16 @@ class StreamingDataset(IterableDataset): Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ from the checkpoint, since a different split structure or shuffle order - makes mid-epoch resumption meaningless. Packed checkpoints - pin ``pack_sequences``, ``eos_id``, ``pad_id``, + makes mid-epoch resumption meaningless. Also raises if the checkpoint + was built with ``split_strategy`` and records a permutation fingerprint + that does not match the current dataset's permutation -- this catches a + changed or non-deterministically-reseeded ``split_strategy`` that + ``num_splits``/``shuffle_seed`` alone cannot detect, since ``shuffle_seed`` + is unused on that path. Checkpoints from ``split_random``/ + ``split_sequential`` never carry this fingerprint (row ids are not stable + across compaction, so it would reject an otherwise-valid resume), and + continue to rely on ``shuffle_seed``/``num_splits`` alone as before. + Packed checkpoints pin ``pack_sequences``, ``eos_id``, ``pad_id``, ``blocks_per_epoch``, and ``epoch``. """ if state["num_splits"] != self._num_splits: @@ -1574,6 +1699,19 @@ class StreamingDataset(IterableDataset): f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, " f"current dataset has {self._shuffle_seed}" ) + checkpoint_fingerprint = state.get("perm_fingerprint") + if ( + checkpoint_fingerprint is not None + and checkpoint_fingerprint != self._perm_fingerprint + ): + raise ValueError( + "permutation mismatch: this checkpoint was saved against a " + "different permutation (e.g. a split_strategy that resamples " + "differently between runs, such as bootstrap_sample with no " + "seed, or a changed callable) than the current dataset built. " + "Resuming would apply saved per-split positions to the wrong " + "rows." + ) self._consumer_checkpoint_requires_uniform = False if "pack_buffers" in state or self._pack_sequences is not None: @@ -1675,6 +1813,10 @@ class StreamingDataset(IterableDataset): first = states[0] packed = "pack_buffers" in first config_keys = ["shuffle_seed", "num_splits", "epoch"] + if "perm_fingerprint" in first: + # Absent only in checkpoints saved before this key existed; skip rather + # than KeyError so old checkpoints keep merging. + config_keys.append("perm_fingerprint") if packed: config_keys.extend( ["pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"] @@ -1684,6 +1826,8 @@ class StreamingDataset(IterableDataset): if ("pack_buffers" in state) != packed: raise ValueError("cannot merge packed and unpacked state dicts") for key in config_keys: + if key not in state: + continue if state[key] != first[key]: raise ValueError( f"{key} mismatch across state dicts: " diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 9d18752f1..5eb3f095c 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -3297,3 +3297,230 @@ def test_streaming_dataset_over_remote_table(): assert len(server.scans) == 1, "the permutation is built with one row-id scan" assert server.takes, "rows must be fetched with row-id takes" assert_server_safe_row_id_requests(server) + + +# --------------------------------------------------------------------------- +# split_strategy escape hatch +# --------------------------------------------------------------------------- + + +def test_split_strategy_bootstrap_sample(lance_table): + """split_strategy lets StreamingDataset use bootstrap_sample directly, producing + more rows than the base table has (sampling with replacement).""" + num_samples = 4 * NUM_ROWS + ds = StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + num_samples, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS, + ) + ids = [row["id"] for row in ds] + + assert len(ids) == num_samples + # 4x oversampling a small table must produce at least one repeat. + assert len(set(ids)) < len(ids) + assert set(ids).issubset(set(range(NUM_ROWS))) + + +def test_split_strategy_matches_direct_permutation_builder_use(lance_table): + """StreamingDataset(split_strategy=...), read single-process (no DataLoader + workers, so this one process owns every split and round-robins across all of + them), must yield exactly what calling the same strategy directly on a + PermutationBuilder would read, interleaved one row per split per cycle. + + Reads each split via __getitems__ (explicit indices), not plain iteration: + Permutation's sequential __iter__ returns empty for a permutation with repeated + row_ids (confirmed directly -- len() is correct, __getitems__ returns the right + rows, but `for row in perm` yields nothing). That's a pre-existing limitation of + Permutation's iterator, unrelated to Bootstrap sampling's correctness: it doesn't + affect StreamingDataset or this benchmark, since both only ever read via + explicit-index batched fetches (take_offsets/__getitems__), never the plain + iterator -- flagged separately, not fixed here (out of scope for this change). + """ + from lancedb.permutation import Permutation, permutation_builder + + def strategy(builder): + return builder.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=1) + + direct_perm_tbl = strategy(permutation_builder(lance_table)).execute() + expected_by_split = [] + for split in range(NUM_SPLITS): + perm = Permutation.from_tables(lance_table, direct_perm_tbl, split=split) + rows = perm.__getitems__(list(range(len(perm)))) + expected_by_split.append([row["id"] for row in rows]) + samples_per_split = len(expected_by_split[0]) + expected = [ + expected_by_split[split][cycle] + for cycle in range(samples_per_split) + for split in range(NUM_SPLITS) + ] + + ds = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + actual = [row["id"] for row in ds] + assert actual == expected + + +def test_split_strategy_is_reproducible_across_instances(lance_table): + def strategy(builder): + return builder.bootstrap_sample(3 * NUM_ROWS, fixed=NUM_SPLITS, seed=99) + + ds1 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + ds2 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS) + + assert [row["id"] for row in ds1] == [row["id"] for row in ds2] + + +def test_split_strategy_num_splits_mismatch_raises(lance_table): + """StreamingDataset(split_strategy=...) must reject a strategy whose split + count doesn't match num_splits. `_rank_splits` assumes split ids are exactly + 0..num_splits-1 with no gaps; a mismatch would otherwise silently drop rows + (split ids >= num_splits) or leave some ranks reading an empty split (missing + ids < num_splits).""" + with pytest.raises(ValueError, match="num_splits"): + StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS + 1, # the strategy above always builds NUM_SPLITS + ) + + +def test_split_strategy_unequal_split_sizes_raises(lance_table): + """StreamingDataset(split_strategy=...) must reject a strategy whose splits + all exist (ids 0..num_splits-1) but aren't the same size. The round-robin + iteration (see _iter_owned) assumes every owned split has the same row count + and stops as soon as any one is exhausted -- `ratios`/`counts` (unlike + `fixed`) can produce unequal split sizes, which would otherwise silently + drop rows from the larger splits.""" + with pytest.raises(ValueError, match="unequal split sizes"): + StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 100, ratios=[0.7, 0.3], seed=SHUFFLE_SEED + ), + num_splits=2, + ) + + +def test_split_strategy_permutation_mismatch_on_resume_raises(lance_table): + """load_state_dict must reject resuming against a different permutation than + the one that produced the checkpoint. num_splits/shuffle_seed alone can't catch + this for split_strategy, since that path ignores shuffle_seed entirely: a + bootstrap_sample reseeded differently resamples completely differently while + shuffle_seed keeps whatever value the caller passed either way, so the + pre-existing checks would pass despite the mismatch.""" + + def strategy(seed): + return lambda b: b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=seed) + + ds1 = StreamingDataset( + lance_table, + split_strategy=strategy(1), + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + it = iter(ds1) + for _ in range(NUM_SPLITS): # one full cycle: every split advances uniformly + next(it) + checkpoint = ds1.state_dict() + + ds2 = StreamingDataset( + lance_table, + split_strategy=strategy(2), # different seed -> a different permutation + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, # unchanged -- irrelevant on this path + ) + with pytest.raises(ValueError, match="permutation mismatch"): + ds2.load_state_dict(checkpoint) + + +def test_split_strategy_resume_with_matching_permutation_succeeds(lance_table): + """The common case: resuming against the same split_strategy (same seed) must + work exactly like the non-split_strategy checkpoint/resume path, continuing + from exactly the same position with no gaps or repeats.""" + + def strategy(b): + return b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=7) + + ds1 = StreamingDataset( + lance_table, + split_strategy=strategy, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + it = iter(ds1) + for _ in range(NUM_SPLITS): + next(it) + checkpoint = ds1.state_dict() + remaining_original = [row["id"] for row in it] + + ds2 = StreamingDataset( + lance_table, + split_strategy=strategy, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + ) + ds2.load_state_dict(checkpoint) # must not raise + remaining_resumed = [row["id"] for row in ds2] + + assert remaining_original == remaining_resumed + + +def test_default_checkpoint_omits_perm_fingerprint(lance_table): + """split_random/split_sequential checkpoints must not carry a + perm_fingerprint -- see test_default_checkpoint_resumes_after_compaction for + why checking one on this path would be wrong, not just unnecessary.""" + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + assert "perm_fingerprint" not in ds.state_dict() + + +def test_split_strategy_checkpoint_includes_perm_fingerprint(lance_table): + ds = StreamingDataset( + lance_table, + split_strategy=lambda b: b.bootstrap_sample( + 2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED + ), + num_splits=NUM_SPLITS, + ) + assert "perm_fingerprint" in ds.state_dict() + + +def test_default_checkpoint_resumes_after_compaction(lance_table): + """split_random/split_sequential checkpoints must resume normally even after + the table is compacted between save and load. Lance row ids are not stable + across compaction by default (confirmed empirically: a 2-fragment table's + row ids differ before and after `optimize()`), so fingerprinting raw row ids + on this path -- instead of relying on shuffle_seed/num_splits alone, as it + does -- would reject this resume even though the seeded permutation itself + is unchanged.""" + lance_table.add(pa.table({"id": [NUM_ROWS]})) # a second fragment to compact + + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it = iter(ds) + for _ in range(NUM_SPLITS): + next(it) + checkpoint = ds.state_dict() + assert "perm_fingerprint" not in checkpoint + + before_compaction = set( + lance_table.to_lance() + .to_table(columns=[], with_row_id=True)["_rowid"] + .to_pylist() + ) + lance_table.optimize() + after_compaction = set( + lance_table.to_lance() + .to_table(columns=[], with_row_id=True)["_rowid"] + .to_pylist() + ) + assert before_compaction != after_compaction, ( + "test setup didn't actually renumber row ids -- fix the fixture, not " + "the assertion below" + ) + + ds_resumed = StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED + ) + ds_resumed.load_state_dict(checkpoint) # must not raise diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 5e8bb1d76..ed1894340 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -161,6 +161,76 @@ def test_split_random_with_seed(mem_db): assert data1["split_id"] == data2["split_id"] +def test_bootstrap_sample_fixed(mem_db): + """Output row count is num_samples, not the input row count, and splits divide + it evenly.""" + tbl = mem_db.create_table("test_table", pa.table({"x": range(20)})) + permutation_tbl = ( + permutation_builder(tbl).bootstrap_sample(1000, fixed=4, seed=42).execute() + ) + + assert permutation_tbl.count_rows() == 1000 + + data = permutation_tbl.search(None).to_arrow().to_pydict() + split_ids = data["split_id"] + assert set(split_ids) == {0, 1, 2, 3} + for split_id in range(4): + assert split_ids.count(split_id) == 250 + + # Every sampled row_id must still be one of the 20 real rows. + assert set(data["row_id"]).issubset(set(range(20))) + + +def test_bootstrap_sample_draws_with_replacement(mem_db): + """Sampling more rows than exist must produce repeats.""" + tbl = mem_db.create_table("test_table", pa.table({"x": range(5)})) + permutation_tbl = ( + permutation_builder(tbl).bootstrap_sample(500, fixed=1, seed=1).execute() + ) + + data = permutation_tbl.search(None).to_arrow().to_pydict() + assert len(data["row_id"]) == 500 + # Only 5 distinct rows exist, so with 500 draws every one of them must repeat. + for row_id in range(5): + assert data["row_id"].count(row_id) > 1 + + +def test_bootstrap_sample_with_seed_is_reproducible(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(50)})) + + perm1 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute() + perm2 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute() + + data1 = perm1.search(None).to_arrow().to_pydict() + data2 = perm2.search(None).to_arrow().to_pydict() + assert data1["row_id"] == data2["row_id"] + assert data1["split_id"] == data2["split_id"] + + +def test_bootstrap_sample_respects_filter(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(20)})) + permutation_tbl = ( + permutation_builder(tbl) + .filter("x >= 15") + .bootstrap_sample(100, fixed=1, seed=7) + .execute() + ) + + data = permutation_tbl.search(None).to_arrow().to_pydict() + assert len(data["row_id"]) == 100 + assert set(data["row_id"]).issubset(set(range(15, 20))) + + +def test_bootstrap_sample_requires_exactly_one_size_arg(mem_db): + tbl = mem_db.create_table("test_table", pa.table({"x": range(10)})) + with pytest.raises(ValueError): + permutation_builder(tbl).bootstrap_sample(10).execute() + with pytest.raises(ValueError): + permutation_builder(tbl).bootstrap_sample( + 10, ratios=[0.5, 0.5], fixed=2 + ).execute() + + def test_split_hash(mem_db): """Test hash-based splitting.""" tbl = mem_db.create_table( diff --git a/python/src/permutation.rs b/python/src/permutation.rs index 24ca2b5a3..4e07d439e 100644 --- a/python/src/permutation.rs +++ b/python/src/permutation.rs @@ -178,6 +178,50 @@ impl PyAsyncPermutationBuilder { }) } + #[pyo3(signature = (num_samples, *, ratios=None, counts=None, fixed=None, seed=None, split_names=None))] + pub fn bootstrap_sample( + slf: PyRefMut<'_, Self>, + num_samples: u64, + ratios: Option>, + counts: Option>, + fixed: Option, + seed: Option, + split_names: Option>, + ) -> PyResult { + // Check that exactly one split type is provided + let split_args_count = [ratios.is_some(), counts.is_some(), fixed.is_some()] + .iter() + .filter(|&&x| x) + .count(); + + if split_args_count != 1 { + return Err(pyo3::exceptions::PyValueError::new_err( + "Exactly one of 'ratios', 'counts', or 'fixed' must be provided", + )); + } + + let sizes = if let Some(ratios) = ratios { + SplitSizes::Percentages(ratios) + } else if let Some(counts) = counts { + SplitSizes::Counts(counts) + } else if let Some(fixed) = fixed { + SplitSizes::Fixed(fixed) + } else { + unreachable!("One of the split arguments must be provided"); + }; + + slf.modify(|builder| { + builder.with_split_strategy( + SplitStrategy::Bootstrap { + num_samples, + sizes, + seed, + }, + split_names, + ) + }) + } + pub fn split_calculated( slf: PyRefMut<'_, Self>, calculation: String, diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index 32e7063a3..15f8c3cd3 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/rust/lancedb/src/dataloader/permutation/builder.rs @@ -316,6 +316,10 @@ impl PermutationBuilder { } else { self.sort_by_column(rows, ROW_ID).await? }; + // Bootstrap's output row count (num_samples) differs from its input row count + // (num_rows); every other strategy's output matches its input. The shuffler + // must be told the former, not assume it always equals the latter. + let output_rows = splitter.output_row_count(num_rows); let split_data = splitter.apply(rows, num_rows).await?; // Shuffle data if requested @@ -328,7 +332,7 @@ impl PermutationBuilder { temp_dir: self.config.temp_dir.clone(), max_rows_per_file: 10 * 1024 * 1024, }); - shuffler.shuffle(split_data, num_rows).await? + shuffler.shuffle(split_data, output_rows).await? } }; diff --git a/rust/lancedb/src/dataloader/permutation/split.rs b/rust/lancedb/src/dataloader/permutation/split.rs index cfa8e5235..b4e5deac1 100644 --- a/rust/lancedb/src/dataloader/permutation/split.rs +++ b/rust/lancedb/src/dataloader/permutation/split.rs @@ -11,6 +11,7 @@ use arrow_schema::{DataType, Field, Schema}; use datafusion_common::hash_utils::{RandomState, create_hashes}; use futures::{StreamExt, TryStreamExt}; use lance_arrow::SchemaExt; +use rand::Rng; use lance_core::ROW_ID; @@ -19,7 +20,7 @@ use crate::{ arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}, dataloader::{ permutation::shuffle::{Shuffler, ShufflerConfig}, - permutation::util::TemporaryDirectory, + permutation::util::{TemporaryDirectory, non_crypto_rng}, }, query::{Query, QueryBase, Select}, }; @@ -76,6 +77,59 @@ pub enum SplitStrategy { /// /// If this strategy is used then the counts/percentages in the SplitSizes are ignored. Calculated { calculation: String }, + /// Rows are sampled uniformly **with replacement** from the base table (combine with + /// [`PermutationBuilder::with_filter`] to restrict the eligible population), independently + /// per split. Named after "the bootstrap" -- the statistical technique of estimating a + /// dataset's properties by resampling it with replacement -- the same sense `bootstrap=True` + /// has in scikit-learn's bagging classifiers. + /// + /// Every other strategy partitions the table: each row goes to exactly one split, exactly + /// once. `Bootstrap` does not partition -- it resamples. The number of output rows is + /// `num_samples`, not the number of input rows: the same row can appear zero, one, or many + /// times, and `num_samples` may be larger or smaller than the (filtered) input row count. + /// Typical uses: + /// + /// - **Bagging / ensemble training** (e.g. random forests): each split is an independent + /// bootstrap resample of the training set. + /// - **Oversampling a rare class**: filter to the minority class, then sample it up to + /// whatever count balances it against the majority, instead of exhausting its small set + /// once per epoch. + /// - **Fixed-size training steps independent of dataset size**: `num_samples` is set + /// directly rather than being implicitly `len(dataset)`. + /// - **Weighted sampling without a weight column**: filter to "eligible" rows (e.g. a + /// windowed-sampling workload filtering to valid window-start positions), then bootstrap + /// the filtered set -- an item with more eligible rows is naturally drawn proportionally + /// more often, with no explicit per-row weight needed. + /// + /// `sizes` divides `num_samples` across splits (reusing [`SplitSizes`], so `Fixed(n)` means + /// "n splits, each getting `num_samples / n` samples", the same convention `Random` and + /// `Sequential` use). Each split draws its own sample using a seed derived independently + /// from `(seed, split_index)`, so a given split's content does not change if `num_samples` + /// is redistributed across a different number of splits -- the same elastic-determinism + /// property documented for the data loader's other strategies. + /// + /// Example: a 5-row base table (`row_id` 0..4), `Bootstrap { num_samples: 8, sizes: + /// Fixed(2), seed: Some(42) }` produced (an actual run, not a hypothetical one): + /// + /// ```text + /// row_id split_id + /// 0 0 + /// 2 0 + /// 4 0 + /// 4 0 <- row 4 repeated within split 0 + /// 4 1 <- row 4 also drawn in split 1 (independent draw) + /// 4 1 + /// 1 1 + /// 2 1 + /// ``` + /// + /// 8 output rows from a 5-row table, row 4 drawn 4 times total, row 3 never drawn -- + /// exactly the with-replacement behavior the name promises, not an even round-robin. + Bootstrap { + num_samples: u64, + sizes: SplitSizes, + seed: Option, + }, } impl SplitStrategy { @@ -107,6 +161,25 @@ impl SplitStrategy { } Self::Sequential { sizes } => sizes.validate(num_rows), Self::Calculated { .. } => Ok(()), + Self::Bootstrap { + num_samples, sizes, .. + } => { + if num_rows == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy requires at least one eligible row to \ + sample from (check the filter, if one is set)" + .to_string(), + }); + } + if *num_samples == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy requires num_samples > 0".to_string(), + }); + } + // Validated against num_samples, not num_rows: sampling is with replacement, + // so num_samples may exceed the number of eligible rows. + sizes.validate(*num_samples) + } } } } @@ -314,6 +387,95 @@ impl Splitter { ))) } + /// Materializes the full (filtered) row id stream -- random access is required since + /// with-replacement sampling can draw any position, including ones already drawn -- then + /// independently samples each split's share with its own derived RNG seed. + async fn apply_bootstrap( + &self, + source: SendableRecordBatchStream, + num_samples: u64, + sizes: &SplitSizes, + seed: Option, + ) -> Result { + let schema = source.schema(); + let row_id_field = schema.field(0).clone(); + let batches: Vec = source.try_collect().await?; + let row_id_arrays = batches + .iter() + .map(|b| b.column(0).as_ref()) + .collect::>(); + let row_ids = arrow::compute::concat(&row_id_arrays)?; + // Sample against however many rows this scan actually returned, not the + // separately computed `num_rows` count checked by `validate` -- that count and + // this scan are two different queries, not pinned to one snapshot, so a + // concurrent append or delete between them could otherwise push + // `random_range` out of `row_ids`'s real bounds. + let num_eligible_rows = row_ids.len() as u64; + if num_eligible_rows == 0 { + return Err(Error::InvalidInput { + message: "Bootstrap strategy found no eligible rows to sample from in \ + this scan (the table or its filter may have changed since \ + validation)" + .to_string(), + }); + } + + let split_sample_counts = sizes.to_counts(num_samples); + + // `try_reserve_exact` (fallible) rather than `with_capacity` (infallible -- + // aborts the whole process via `handle_alloc_error` on failure, which no + // caller can catch) so an unreasonably large `num_samples` becomes a normal, + // catchable error instead of a crash. This bounds the failure mode, not peak + // memory itself: the output is still one fully materialized RecordBatch, not + // streamed in chunks. + let too_large = || Error::InvalidInput { + message: format!( + "Bootstrap strategy's num_samples ({num_samples}) is too large to \ + allocate in memory" + ), + }; + let mut position_indices = Vec::::new(); + position_indices + .try_reserve_exact(num_samples as usize) + .map_err(|_| too_large())?; + let mut split_ids = Vec::::new(); + split_ids + .try_reserve_exact(num_samples as usize) + .map_err(|_| too_large())?; + for (split_idx, &count) in split_sample_counts.iter().enumerate() { + // Derive an independent seed per split (rather than one RNG shared across all + // splits) so a split's sample doesn't depend on how many other splits there are. + let split_seed = seed.map(|s| { + s.wrapping_add( + (split_idx as u64) + .wrapping_mul(0x9E37_79B9_7F4A_7C15) + .wrapping_add(1), + ) + }); + let mut rng = non_crypto_rng(&split_seed); + for _ in 0..count { + position_indices.push(rng.random_range(0..num_eligible_rows)); + split_ids.push(split_idx as u64); + } + } + + let indices = UInt64Array::from(position_indices); + let sampled_row_ids = arrow::compute::take(row_ids.as_ref(), &indices, None)?; + let split_id_array: Arc = Arc::new(UInt64Array::from(split_ids)); + + let new_schema = Arc::new(Schema::new(vec![ + row_id_field, + Field::new(SPLIT_ID_COLUMN, DataType::UInt64, false), + ])); + let result_batch = + RecordBatch::try_new(new_schema.clone(), vec![sampled_row_ids, split_id_array])?; + + Ok(Box::pin(SimpleRecordBatchStream::new( + futures::stream::once(std::future::ready(Ok(result_batch))), + new_schema, + ))) + } + pub async fn apply( &self, source: SendableRecordBatchStream, @@ -357,6 +519,14 @@ impl Splitter { self.apply_hash(source, split_weights, *discard_weight) .await } + SplitStrategy::Bootstrap { + num_samples, + sizes, + seed, + } => { + self.apply_bootstrap(source, *num_samples, sizes, *seed) + .await + } } } @@ -377,7 +547,13 @@ impl Splitter { pub fn orders_by_split_id(&self) -> bool { match &self.strategy { - SplitStrategy::Hash { .. } | SplitStrategy::Calculated { .. } => true, + // Bootstrap emits samples grouped split-by-split already (see apply_bootstrap), + // rather than relying on the shared pipeline's sort-by-split-id step: within-split + // draw order is meaningful to callers (e.g. a warmup/measured split within one + // split's stream), and sort stability for ties isn't documented behavior to lean on. + SplitStrategy::Hash { .. } + | SplitStrategy::Calculated { .. } + | SplitStrategy::Bootstrap { .. } => true, SplitStrategy::NoSplit | SplitStrategy::Sequential { .. } // It may be strange but for random we shuffle and then assign splits so the result is @@ -385,6 +561,26 @@ impl Splitter { | SplitStrategy::Random { .. } => false, } } + + /// The row count `apply` produces from `num_rows` input rows. Equal to `num_rows` + /// for every strategy except `Bootstrap`, whose output size is `num_samples` + /// regardless of how many (filtered) input rows it draws from. Callers that need + /// to validate or pre-size `apply`'s output (e.g. the shuffler) must use this + /// instead of assuming the input and output row counts match. + /// + /// For `Bootstrap`, this is `sizes.to_counts(num_samples)`'s sum, not + /// `num_samples` itself: `to_counts` can emit less than its input (e.g. + /// `Fixed(n)` when `n` does not divide `num_samples` evenly, or explicit + /// `counts`/`ratios` that do not use every sample) and `apply_bootstrap` only + /// ever emits exactly what `to_counts` returns. + pub fn output_row_count(&self, num_rows: u64) -> u64 { + match &self.strategy { + SplitStrategy::Bootstrap { + num_samples, sizes, .. + } => sizes.to_counts(*num_samples).iter().sum(), + _ => num_rows, + } + } } /// Split configuration - either percentages or absolute counts @@ -807,4 +1003,169 @@ mod tests { assert_eq!(num_in_split_0, 12); // 24% assert_eq!(num_in_split_1, 30); // 60% } + + #[tokio::test] + async fn test_bootstrap_split_schema_and_counts() { + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(4), + seed: Some(42), + }, + ); + + let split_batches = splitter + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = split_batches[0].schema(); + let split_batch = concat_batches(&schema, &split_batches).unwrap(); + + // Schema stays exactly [row_id, split_id], same as every other strategy. + let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(field_names, vec!["id", SPLIT_ID_COLUMN]); + + // Output row count is num_samples, not num_rows (50) -- sampling is with replacement. + assert_eq!(split_batch.num_rows(), 1000); + + let split_ids = split_batch.column(1).as_primitive::().values(); + for split_id in [0u64, 1, 2, 3] { + assert_eq!(split_ids.iter().filter(|v| **v == split_id).count(), 250); + } + + // With 1000 draws from only 50 distinct rows, some row id must repeat. + let row_ids = split_batch.column(0).as_primitive::().values(); + let mut seen = std::collections::HashSet::new(); + assert!(!row_ids.iter().all(|v| seen.insert(*v))); + } + + #[tokio::test] + async fn test_bootstrap_split_seed_is_reproducible() { + let make_splitter = || { + Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 200, + sizes: SplitSizes::Fixed(2), + seed: Some(7), + }, + ) + }; + + let first = make_splitter() + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let second = make_splitter() + .apply(test_data(), 50) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let schema = first[0].schema(); + let first = concat_batches(&schema, &first).unwrap(); + let second = concat_batches(&schema, &second).unwrap(); + assert_eq!(first.column(0), second.column(0)); + assert_eq!(first.column(1), second.column(1)); + } + + #[tokio::test] + async fn test_bootstrap_split_respects_filter() { + // Bootstrap is driven through a filter the same way every other strategy is -- the + // base table rows reaching Splitter::apply are already filtered by the time they get + // here, so this exercises the same num_rows=0 guard a fully-filtered-out population + // would hit. + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 10, + sizes: SplitSizes::Fixed(1), + seed: Some(1), + }, + ); + let Err(err) = splitter.apply(test_data(), 0).await else { + panic!("expected Bootstrap over zero eligible rows to fail"); + }; + assert!(err.to_string().contains("at least one eligible row")); + } + + fn empty_data() -> SendableRecordBatchStream { + lance_datagen::gen_batch() + .with_seed(Seed::from(42)) + .col(ID_COLUMN, lance_datagen::array::step::()) + .into_ldb_stream(RowCount::from(0), BatchCount::from(1)) + } + + #[test] + fn test_output_row_count() { + // Bootstrap's output row count is independent of the input size -- every + // other strategy's output matches its input. The shuffler relies on this + // (see builder.rs) instead of assuming input and output row counts match. + let bootstrap = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(4), + seed: Some(42), + }, + ); + assert_eq!(bootstrap.output_row_count(50), 1000); + + // Not num_samples itself: `Fixed(3)` doesn't divide 1000 evenly, so + // `to_counts` emits 333 * 3 = 999, one short of num_samples -- + // `apply_bootstrap` only ever emits exactly what `to_counts` returns. + let bootstrap_uneven = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 1000, + sizes: SplitSizes::Fixed(3), + seed: Some(42), + }, + ); + assert_eq!(bootstrap_uneven.output_row_count(50), 999); + + let sequential = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Sequential { + sizes: SplitSizes::Fixed(2), + }, + ); + assert_eq!(sequential.output_row_count(50), 50); + } + + #[tokio::test] + async fn test_bootstrap_rejects_empty_scan_cleanly() { + // The eligibility count `validate` checks and this scan are two separate + // queries against the table, not pinned to one snapshot -- a concurrent + // delete could shrink the real scan to zero rows after `validate` already + // passed on a nonzero count. Calling `apply_bootstrap` directly (bypassing + // `apply`'s `validate`) reproduces exactly that: a scan that returns zero + // rows regardless of what count was validated upstream. It must return a + // clean `Err`, not panic on an out-of-range `take` or an empty-range + // `random_range`. + let splitter = Splitter::new( + TemporaryDirectory::None, + SplitStrategy::Bootstrap { + num_samples: 10, + sizes: SplitSizes::Fixed(1), + seed: Some(1), + }, + ); + let Err(err) = splitter + .apply_bootstrap(empty_data(), 10, &SplitSizes::Fixed(1), Some(1)) + .await + else { + panic!("expected Bootstrap over an empty scan to fail cleanly, not panic"); + }; + assert!(err.to_string().contains("no eligible rows")); + } } From f9680806c982d30a85fd70c0542a8dc2b20ee5d9 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Wed, 21 Oct 2026 23:47:09 -0700 Subject: [PATCH 12/33] feat(python): expose materialized view auto-refresh creation options (#4621) Expose optional `auto_refresh_enabled` and `auto_refresh_interval_seconds` when creating materialized views from Python and Rust. Omitted intervals remain unset: the server determines when to refresh. ```python view = db.create_materialized_view( "recent_orders", "orders", auto_refresh_enabled=True, auto_refresh_interval_seconds=300, ) ``` ## Architecture and workflow - Python synchronous, asynchronous, job-returning and namespace wrappers forward both optional values through PyO3 to the Rust creation builder. Both default to `None`; no client interval is selected. - Remote creation calls `POST /v1/materialized_view/{id}/create` with the defining query and only the supplied refresh properties. A `202 Accepted` response must contain a job ID. Existing job tracking waits for the server's ordinary SQL CREATE job before synchronous creation opens the view. - The server validates and authorizes the CREATE, then uses the existing SQL worker to populate the view and persist its refresh settings. The Lance agent subsequently uses those settings to submit REFRESH SQL jobs. - Explicit positive finite fractional seconds are preserved. An interval alone does not enable refresh. Embedded connections retain manual refresh, accept explicit disabled state, and reject requests for automatic scheduling rather than ignore them. Includes API documentation and coverage for omission, false, fractional/invalid intervals, namespace forwarding, server job IDs and invalid acceptance responses. Requires the server REST-create endpoint to be deployed before SDK release. Existing remote `with_no_data=True` remains unsupported. --- docs/src/python/python.md | 16 ++ python/python/lancedb/_lancedb.pyi | 4 + python/python/lancedb/db.py | 27 ++++ python/python/lancedb/namespace.py | 16 ++ python/python/lancedb/remote/db.py | 8 + .../python/tests/test_materialized_views.py | 105 ++++++++++++- python/src/connection.rs | 22 ++- rust/lancedb/src/materialized_view.rs | 44 +++++- rust/lancedb/src/remote/db.rs | 140 ++++++++++++++---- 9 files changed, 349 insertions(+), 33 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index c8a82851f..f4c697497 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -377,6 +377,22 @@ listing a storage directory. ## Materialized Views (Synchronous) +On Cloud and Enterprise, creation accepts `auto_refresh_enabled` and +`auto_refresh_interval_seconds`: + +```python +view = db.create_materialized_view( + "recent_orders", + "orders", + auto_refresh_enabled=True, + auto_refresh_interval_seconds=300, +) +``` + +Omit `auto_refresh_interval_seconds` to let the server determine when to refresh. +These options also apply to asynchronous creation and creation jobs. Embedded +connections support manual refresh and reject automatic scheduling requests. + ::: lancedb.materialized_view.MaterializedView ::: lancedb.materialized_view.MaterializedViewDefinition diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96bc7752b..96eef064c 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -270,6 +270,8 @@ class Connection(object): filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Table: ... async def create_materialized_view_async( self, @@ -279,6 +281,8 @@ class Connection(object): filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Job: ... def list_materialized_views( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 47d24b762..5def89b01 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -607,6 +607,8 @@ class DBConnection(EnforceOverrides): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> MaterializedView: """Define a materialized view named ``name`` over the table ``source``. @@ -631,6 +633,13 @@ class DBConnection(EnforceOverrides): Cap the view at this many rows, in materialization order. with_no_data: bool, default False Skip the initial refresh and leave the backing table empty. + auto_refresh_enabled: bool, optional + Enable or disable automatic refresh on Cloud and Enterprise. + When omitted, the server determines whether automatic refresh is enabled. + Embedded databases do not support automatic refresh. + auto_refresh_interval_seconds: float, optional + Positive, finite refresh interval in seconds on Cloud and Enterprise. + When omitted, the server determines when to refresh. Returns ------- @@ -649,6 +658,8 @@ class DBConnection(EnforceOverrides): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Job[None]: """Submit materialized-view creation and return its job. @@ -1706,6 +1717,8 @@ class LanceDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> MaterializedView: """Define a materialized view named ``name`` over the table ``source``. See @@ -1734,6 +1747,8 @@ class LanceDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ) return MaterializedView(self.open_table(name)) @@ -1748,6 +1763,8 @@ class LanceDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Job[None]: job = LOOP.run( self._conn.create_materialized_view_async( @@ -1757,6 +1774,8 @@ class LanceDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ) return Job(job) @@ -2767,6 +2786,8 @@ class AsyncConnection(object): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> AsyncMaterializedView: """Define a materialized view named ``name`` over the table ``source``. See @@ -2782,6 +2803,8 @@ class AsyncConnection(object): filter=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, **native_options, ) return AsyncMaterializedView(AsyncTable(inner)) @@ -2795,6 +2818,8 @@ class AsyncConnection(object): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> AsyncJob[None]: """Submit materialized-view creation and return its job. @@ -2810,6 +2835,8 @@ class AsyncConnection(object): filter=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, **native_options, ) return AsyncJob(inner) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 26c62925f..f8675ebb5 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -665,6 +665,8 @@ class LanceNamespaceDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> "MaterializedView": """Define a materialized view over a table in the root namespace. See @@ -680,6 +682,8 @@ class LanceNamespaceDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ).name ) @@ -695,6 +699,8 @@ class LanceNamespaceDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Job[None]: job = LOOP.run( self._inner.create_materialized_view_async( @@ -704,6 +710,8 @@ class LanceNamespaceDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ) return Job(job) @@ -1288,6 +1296,8 @@ class AsyncLanceNamespaceDBConnection: where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> "AsyncMaterializedView": """Define a materialized view over a table in the root namespace.""" view = await self._inner.create_materialized_view( @@ -1297,6 +1307,8 @@ class AsyncLanceNamespaceDBConnection: where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) # Reopen through the namespace so the view's table carries the # namespace client and pushdown configuration a bare inner table lacks. @@ -1311,6 +1323,8 @@ class AsyncLanceNamespaceDBConnection: where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> AsyncJob[None]: """Submit materialized-view creation and return its job.""" return await self._inner.create_materialized_view_async( @@ -1320,6 +1334,8 @@ class AsyncLanceNamespaceDBConnection: where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) async def open_materialized_view(self, name: str) -> "AsyncMaterializedView": diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index ac6b8b99d..89c7f8a78 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -723,6 +723,8 @@ class RemoteDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> MaterializedView: from .table import RemoteTable @@ -734,6 +736,8 @@ class RemoteDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ) return MaterializedView( @@ -755,6 +759,8 @@ class RemoteDBConnection(DBConnection): where: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, + auto_refresh_enabled: Optional[bool] = None, + auto_refresh_interval_seconds: Optional[float] = None, ) -> Job[None]: job = LOOP.run( self._conn.create_materialized_view_async( @@ -764,6 +770,8 @@ class RemoteDBConnection(DBConnection): where=where, limit=limit, with_no_data=with_no_data, + auto_refresh_enabled=auto_refresh_enabled, + auto_refresh_interval_seconds=auto_refresh_interval_seconds, ) ) return Job(job) diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index bd1b2a0bb..46972a15c 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -129,8 +129,21 @@ def test_sync_remote_create_uses_public_async_connection(): where=None, limit=None, with_no_data=False, + auto_refresh_enabled=None, + auto_refresh_interval_seconds=None, ): - calls.append((name, source, select, where, limit, with_no_data)) + calls.append( + ( + name, + source, + select, + where, + limit, + with_no_data, + auto_refresh_enabled, + auto_refresh_interval_seconds, + ) + ) return StubAsyncMaterializedView() async def drop_materialized_view(self, name, *, namespace_path=None): @@ -148,9 +161,11 @@ def test_sync_remote_create_uses_public_async_connection(): where="age >= 18", limit=10, with_no_data=True, + auto_refresh_enabled=True, + auto_refresh_interval_seconds=0.5, ) assert view.name == "adults" - assert calls == [("adults", "people", ["name"], "age >= 18", 10, True)] + assert calls == [("adults", "people", ["name"], "age >= 18", 10, True, True, 0.5)] db.drop_materialized_view("adults", namespace_path=["analytics"]) assert calls[-1] == ("drop", "adults", ["analytics"]) @@ -550,3 +565,89 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): assert sorted([first, *resumed]) == ["a_view", "z_view"] assert resumed.page_token() is None assert list(resumed) == [] + + +@pytest.mark.parametrize( + "options", + [ + {}, + {"auto_refresh_enabled": True}, + {"auto_refresh_enabled": False, "auto_refresh_interval_seconds": 0.5}, + ], +) +def test_remote_create_forwards_refresh_options_without_interval_default(options): + with mock_remote_materialized_view_create() as (host, requests): + db = lancedb.connect("db://dev", api_key="fake", host_override=host) + job = db.create_materialized_view_async("adults", "people", **options) + assert job.id == "mv-create-123" + assert requests == [ + ( + "/v1/materialized_view/adults/create", + { + "query": 'SELECT * FROM "people"', + **options, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_async_remote_create_keeps_interval_unspecified(): + with mock_remote_materialized_view_create() as (host, requests): + db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host) + job = await db.create_materialized_view_async( + "adults", + "people", + auto_refresh_enabled=True, + ) + assert job.id == "mv-create-123" + assert requests[0][1] == { + "query": 'SELECT * FROM "people"', + "auto_refresh_enabled": True, + } + + +@pytest.mark.parametrize( + "options", + [ + {"auto_refresh_enabled": True}, + {"auto_refresh_interval_seconds": 300}, + ], +) +def test_embedded_create_refuses_automatic_refresh_before_creating_view( + tmp_path, options +): + db = make_db(tmp_path) + with pytest.raises(NotImplementedError, match="Cloud or Enterprise"): + db.create_materialized_view("adults", "people", **options) + assert "adults" not in db.table_names() + + +@pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")]) +def test_remote_create_rejects_invalid_interval_without_request(seconds): + with mock_remote_materialized_view_create() as (host, requests): + db = lancedb.connect("db://dev", api_key="fake", host_override=host) + with pytest.raises(ValueError, match="positive, finite"): + db.create_materialized_view_async( + "adults", "people", auto_refresh_interval_seconds=seconds + ) + assert requests == [] + + +def test_namespace_forwards_automatic_refresh_request(tmp_path): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + with pytest.raises(NotImplementedError, match="Cloud or Enterprise"): + db.create_materialized_view("mv", "source", auto_refresh_enabled=True) + + +@pytest.mark.asyncio +async def test_async_namespace_forwards_automatic_refresh_request(tmp_path): + db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)}) + with pytest.raises(NotImplementedError, match="Cloud or Enterprise"): + await db.create_materialized_view("mv", "source", auto_refresh_enabled=True) + + +def test_embedded_explicitly_disabled_refresh_is_supported(tmp_path): + db = make_db(tmp_path) + view = db.create_materialized_view("mv", "people", auto_refresh_enabled=False) + assert view.table.count_rows() == 3 diff --git a/python/src/connection.rs b/python/src/connection.rs index 9ac081b9d..7e20f5f99 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -412,7 +412,8 @@ impl Connection { }) } - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] pub fn create_materialized_view( self_: PyRef<'_, Self>, name: String, @@ -421,10 +422,18 @@ impl Connection { filter: Option, limit: Option, with_no_data: bool, + auto_refresh_enabled: Option, + auto_refresh_interval_seconds: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { let mut builder = inner.create_materialized_view(name, source); + if let Some(enabled) = auto_refresh_enabled { + builder = builder.auto_refresh_enabled(enabled); + } + if let Some(seconds) = auto_refresh_interval_seconds { + builder = builder.auto_refresh_interval_seconds(seconds); + } if let Some(projections) = projections { builder = builder.select(projections); } @@ -440,7 +449,8 @@ impl Connection { }) } - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] pub fn create_materialized_view_async( self_: PyRef<'_, Self>, name: String, @@ -449,10 +459,18 @@ impl Connection { filter: Option, limit: Option, with_no_data: bool, + auto_refresh_enabled: Option, + auto_refresh_interval_seconds: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { let mut builder = inner.create_materialized_view(name, source); + if let Some(enabled) = auto_refresh_enabled { + builder = builder.auto_refresh_enabled(enabled); + } + if let Some(seconds) = auto_refresh_interval_seconds { + builder = builder.auto_refresh_interval_seconds(seconds); + } if let Some(projections) = projections { builder = builder.select(projections); } diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 95e3baac7..ecdd1783b 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -280,7 +280,7 @@ pub struct MaterializedViewInfo { /// The backend-independent request used to create a remote materialized view. #[doc(hidden)] -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct CreateMaterializedViewRequest { /// Name of the new view. pub name: String, @@ -290,6 +290,10 @@ pub struct CreateMaterializedViewRequest { pub query: String, /// Whether to skip the initial population job. pub with_no_data: bool, + /// Whether the server should refresh the view automatically. + pub auto_refresh_enabled: Option, + /// Refresh interval in seconds. When absent, the server determines when to refresh. + pub auto_refresh_interval_seconds: Option, } /// Prefix of the internal columns holding source columns a computed column @@ -1570,6 +1574,8 @@ pub struct CreateMaterializedViewBuilder { filter: Option, limit: Option, with_no_data: bool, + auto_refresh_enabled: Option, + auto_refresh_interval_seconds: Option, } fn source_sql_identifier( @@ -1603,6 +1609,8 @@ impl CreateMaterializedViewBuilder { filter: None, limit: None, with_no_data: false, + auto_refresh_enabled: None, + auto_refresh_interval_seconds: None, } } @@ -1651,6 +1659,32 @@ impl CreateMaterializedViewBuilder { self } + /// Enable or disable automatic refresh on Cloud and Enterprise. + /// When omitted, the server determines whether automatic refresh is enabled. + /// Embedded databases do not support automatic refresh. + pub fn auto_refresh_enabled(mut self, enabled: bool) -> Self { + self.auto_refresh_enabled = Some(enabled); + self + } + + /// Set a positive, finite refresh interval in seconds on Cloud and Enterprise. + /// When omitted, the server determines when to refresh. + /// + /// ```no_run + /// # #![recursion_limit = "256"] + /// # async fn example(db: &lancedb::Connection) -> lancedb::Result<()> { + /// db.create_materialized_view("recent_orders", "orders") + /// .auto_refresh_enabled(true) + /// .auto_refresh_interval_seconds(300.0) + /// .execute().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn auto_refresh_interval_seconds(mut self, seconds: f64) -> Self { + self.auto_refresh_interval_seconds = Some(seconds); + self + } + fn query(&self) -> Result { fn quote(name: &str) -> String { format!("\"{}\"", name.replace('"', "\"\"")) @@ -1693,6 +1727,8 @@ impl CreateMaterializedViewBuilder { namespace_path: self.namespace.clone(), query: self.query()?, with_no_data: self.with_no_data, + auto_refresh_enabled: self.auto_refresh_enabled, + auto_refresh_interval_seconds: self.auto_refresh_interval_seconds, }) .await; } @@ -1719,6 +1755,12 @@ impl CreateMaterializedViewBuilder { } async fn execute_native(self) -> Result { + if self.auto_refresh_enabled == Some(true) || self.auto_refresh_interval_seconds.is_some() { + return Err(Error::NotSupported { + message: "automatic materialized-view refresh requires Cloud or Enterprise" + .to_string(), + }); + } let source = self .connection .open_table(&self.source) diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index a0f4aa953..4a8cbf32e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -51,10 +51,6 @@ use super::table::RemoteTable; use super::util::parse_server_version; use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; -fn quote_sql_identifier(identifier: &str) -> String { - format!("\"{}\"", identifier.replace('"', "\"\"")) -} - // Request structure for the remote clone table API #[derive(serde::Serialize)] struct RemoteCloneTableRequest { @@ -67,6 +63,15 @@ struct RemoteCloneTableRequest { is_shallow: Option, } +#[derive(serde::Serialize)] +struct RemoteCreateMaterializedViewRequest { + query: String, + #[serde(skip_serializing_if = "Option::is_none")] + auto_refresh_enabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + auto_refresh_interval_seconds: Option, +} + // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. @@ -877,16 +882,43 @@ impl Database for RemoteDatabase { .to_string(), }); } - let client = self.sql_client.clone().ok_or_else(|| Error::NotSupported { - message: "SQL is unavailable for this remote database client".to_string(), + if let Some(seconds) = request.auto_refresh_interval_seconds + && (!seconds.is_finite() || seconds <= 0.0) + { + return Err(Error::InvalidInput { + message: "auto_refresh_interval_seconds must be positive, finite seconds".into(), + }); + } + let identifier = build_table_identifier(&request.name, &request.namespace_path)?; + let body = RemoteCreateMaterializedViewRequest { + query: request.query, + auto_refresh_enabled: request.auto_refresh_enabled, + auto_refresh_interval_seconds: request.auto_refresh_interval_seconds, + }; + let request = self + .client + .post(&format!("/v1/materialized_view/{identifier}/create")) + .json(&body); + let (request_id, response) = self.client.send(request).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + if status != StatusCode::ACCEPTED { + return Err(Error::Http { + source: "materialized-view creation must return 202 Accepted".into(), + request_id, + status_code: Some(status), + }); + } + let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { + source: "materialized-view create response did not contain a valid job_id".into(), + request_id, + status_code: Some(status), })?; - let namespace = request.namespace_path; - let statement = format!( - "CREATE MATERIALIZED VIEW {} AS {}", - quote_sql_identifier(&request.name), - request.query - ); - Ok(client.submit_as_job(statement, namespace, || async { Ok(()) })) + Ok(Job::new(Box::new(RemoteJob::new( + self.client.clone(), + job_id, + )))) } async fn drop_materialized_view_async( @@ -1964,21 +1996,73 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_requires_sql_client() { - let db = super::RemoteDatabase::new_mock(|request| -> http::Response { - panic!("unexpected REST request: {}", request.url().path()) - }); - let error = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10" - .into(), - with_no_data: false, - }) - .await - .unwrap_err(); - assert!(error.to_string().contains("SQL is unavailable")); + async fn test_create_materialized_view_uses_rest_and_preserves_optional_settings() { + for (enabled, interval) in [(None, None), (Some(true), None), (Some(false), Some(0.5))] { + let db = super::RemoteDatabase::new_mock(move |request| { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.url().path(), + "/v1/materialized_view/analytics$adults/create" + ); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["query"], "SELECT * FROM people"); + assert_eq!( + body.get("auto_refresh_enabled"), + enabled.map(serde_json::Value::from).as_ref() + ); + assert_eq!( + body.get("auto_refresh_interval_seconds"), + interval.map(serde_json::Value::from).as_ref() + ); + http::Response::builder() + .status(202) + .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) + .unwrap() + }); + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["analytics".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: enabled, + auto_refresh_interval_seconds: interval, + }) + .await + .unwrap(); + assert_eq!(job.id(), Some("j1-mv-create")); + } + } + + #[tokio::test] + async fn test_create_materialized_view_rejects_invalid_acceptance() { + for (status, body) in [ + (200, r#"{"job_id":"j1"}"#), + (202, "{}"), + (202, r#"{"job_id":""}"#), + ] { + let db = super::RemoteDatabase::new_mock(move |_| { + http::Response::builder().status(status).body(body).unwrap() + }); + let error = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec![], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: None, + auto_refresh_interval_seconds: None, + }) + .await + .err() + .unwrap(); + assert!(error.to_string().contains(if status == 200 { + "202 Accepted" + } else { + "job_id" + })); + } } #[tokio::test] From dd4a527727e3409b2a329225d802be45ca48ee86 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 22 Oct 2026 16:13:27 +0800 Subject: [PATCH 13/33] fix(index): preserve registered scalar plugin type names (#4627) Native `list_indices()` drops registered scalar plugin types absent from LanceDB's fixed enum, and `index_stats()` reports them as `Unknown`. Preserve their registered display names through `IndexType::Custom(String)`. Resolve these names through Lance's shared `IndexDetails` registry entry point. Registered types such as MinHashLSH become visible; system and unregistered types remain excluded. Existing builtin aliases retain their canonical names. Rust consumers with exhaustive `IndexType` matches need to handle the added variant. Manual validation with [Lance #10224](https://github.com/lance-format/lance/pull/10224) at `fcc070f88a`: a downstream statistics-only plugin appears as `Custom("Sketch")` in native listing and statistics. Its `load_index` deliberately returns an error, and repeated calls do not construct additional plugin instances, verifying metadata-only lookup through the shared registry. --- rust/lancedb/src/index.rs | 30 +++++++++ rust/lancedb/src/table.rs | 10 +-- .../lancedb/tests/registered_scalar_plugin.rs | 66 +++++++++++++++++++ 3 files changed, 98 insertions(+), 8 deletions(-) create mode 100644 rust/lancedb/tests/registered_scalar_plugin.rs diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index 82f8ad845..712a15fa0 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors use chrono::{DateTime, Utc}; +use lance::index::scalar::IndexDetails; use scalar::FtsIndexBuilder; use serde::Deserialize; use serde_with::skip_serializing_none; @@ -382,10 +383,38 @@ pub enum IndexType { // FTS #[serde(alias = "INVERTED", alias = "Inverted")] FTS, + /// A registered scalar index plugin's display name, preserved verbatim. + /// + /// For example, `Custom("Sketch".into()).to_string()` returns `"Sketch"`. + Custom(String), /// Catch-all for index types not recognized by this version of LanceDB. Unknown, } +impl IndexType { + pub(crate) fn from_description( + description: &dyn lance_index::traits::IndexDescription, + ) -> Self { + let index_type = description.index_type().parse().unwrap_or(Self::Unknown); + if index_type != Self::Unknown { + return index_type; + } + let Some(metadata) = description.segments().first() else { + return Self::Unknown; + }; + if lance_index::is_system_index(metadata) { + return Self::Unknown; + } + let Some(details) = metadata.index_details.as_ref() else { + return Self::Unknown; + }; + IndexDetails(details.clone()) + .get_plugin() + .map(|plugin| Self::Custom(plugin.name().to_string())) + .unwrap_or(Self::Unknown) + } +} + impl std::fmt::Display for IndexType { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { match self { @@ -405,6 +434,7 @@ impl std::fmt::Display for IndexType { Self::BloomFilter => write!(f, "BLOOM_FILTER"), Self::RTree => write!(f, "RTREE"), Self::FTS => write!(f, "FTS"), + Self::Custom(name) => f.write_str(name), Self::Unknown => write!(f, "UNKNOWN"), } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 61a710dcd..39e6f1911 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3821,10 +3821,7 @@ impl BaseTable for NativeTable { let mut indices: Vec = descriptions .iter() .filter_map(|idx_desc| { - let index_type: crate::index::IndexType = idx_desc - .index_type() - .parse() - .unwrap_or(crate::index::IndexType::Unknown); + let index_type = crate::index::IndexType::from_description(idx_desc.as_ref()); if index_type == crate::index::IndexType::Unknown { // Internal or future index types that this version doesn't recognize // (e.g. Lance's internal FragReuseIndex) are silently excluded from @@ -3960,10 +3957,7 @@ impl BaseTable for NativeTable { return Ok(None); }; - let index_type: crate::index::IndexType = description - .index_type() - .parse() - .unwrap_or(crate::index::IndexType::Unknown); + let index_type = crate::index::IndexType::from_description(description.as_ref()); let is_vector = matches!( index_type, diff --git a/rust/lancedb/tests/registered_scalar_plugin.rs b/rust/lancedb/tests/registered_scalar_plugin.rs new file mode 100644 index 000000000..2f5620cc8 --- /dev/null +++ b/rust/lancedb/tests/registered_scalar_plugin.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use arrow_array::{RecordBatchIterator, record_batch}; +use lance::{Dataset, index::DatasetIndexExt}; +use lance_index::scalar::ScalarIndexParams; +use lancedb::index::IndexType; + +// MinHashLSH is a registered scalar plugin without a LanceDB IndexType variant. +// It exercises the same metadata-only fallback as a downstream plugin. +#[tokio::test] +async fn list_registered_scalar_plugin_with_its_type_name() { + let directory = tempfile::tempdir().unwrap(); + let batch = record_batch!(( + "value", + Utf8, + ["alpha beta", "beta gamma", "gamma delta", "delta alpha"] + )) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch.clone())], batch.schema()), + directory.path().join("test.lance").to_str().unwrap(), + None, + ) + .await + .unwrap(); + dataset + .create_index_builder( + &["value"], + lance_index::IndexType::Scalar, + &ScalarIndexParams::new("minhashlsh".into()), + ) + .name("profile".into()) + .train(false) + .await + .unwrap(); + + let connection = lancedb::connect(directory.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = connection.open_table("test").execute().await.unwrap(); + let indices = table.list_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].name, "profile"); + assert_eq!( + indices[0].index_type, + IndexType::Custom("MinHashLsh".into()) + ); + assert_eq!(indices[0].index_type.to_string(), "MinHashLsh"); + assert_eq!(indices[0].columns, ["value"]); + assert!( + indices[0] + .type_url + .as_deref() + .unwrap() + .ends_with("MinHashLshIndexDetails") + ); + let statistics = table.index_stats("profile").await.unwrap().unwrap(); + assert_eq!( + statistics.index_type, + IndexType::Custom("MinHashLsh".into()) + ); + assert_eq!(statistics.num_indexed_rows, 0); + assert_eq!(statistics.num_unindexed_rows, 4); +} From a5079b0fc6c47100f8f5030b8dd99e40c9d576de Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 22 Oct 2026 19:41:40 +0800 Subject: [PATCH 14/33] chore: update lance dependency to v15.0.0-beta.1 (#4628) --- .github/workflows/npm-publish.yml | 32 +++---- Cargo.lock | 152 +++++++++++++++++++++--------- Cargo.toml | 28 +++--- docs/web/indexing/fts-index.mdx | 8 +- java/pom.xml | 2 +- rust/lancedb/src/index/scalar.rs | 11 +++ 6 files changed, 154 insertions(+), 79 deletions(-) diff --git a/.github/workflows/npm-publish.yml b/.github/workflows/npm-publish.yml index c3724de8d..a8dd3993e 100644 --- a/.github/workflows/npm-publish.yml +++ b/.github/workflows/npm-publish.yml @@ -71,9 +71,9 @@ jobs: # https://github.com/napi-rs/napi-rs/blob/main/debian.Dockerfile docker: ghcr.io/napi-rs/napi-rs/nodejs-rust:lts-debian pre_build: |- - set -e && - apt-get update && - apt-get install -y protobuf-compiler pkg-config && + set -e + apt-get update + apt-get install -y protobuf-compiler pkg-config # The base image (manylinux2014-cross) sets TARGET_CC to the old # GCC 4.8 cross-compiler. aws-lc-sys checks TARGET_CC before CC, # so it picks up GCC even though the napi-rs image sets CC=clang. @@ -84,10 +84,10 @@ jobs: host: ubuntu-2404-8x-x64 features: fp16kernels pre_build: |- - set -e && - sudo apt-get update && - sudo apt-get install -y protobuf-compiler pkg-config && - rustup target add x86_64-unknown-linux-musl && + set -e + sudo apt-get update + sudo apt-get install -y protobuf-compiler pkg-config + rustup target add x86_64-unknown-linux-musl export EXTRA_ARGS="-x" - target: aarch64-unknown-linux-gnu host: ubuntu-2404-8x-x64 @@ -103,13 +103,13 @@ jobs: # linker process 7.0 -> 4.0 GiB (lancedb/sophon#7313). linker: /tmp/aarch64-lld-clang pre_build: |- - set -e && - apt-get update && - apt-get install -y protobuf-compiler pkg-config && - export TARGET_CC=clang TARGET_CXX=clang++ && + set -e + apt-get update + apt-get install -y protobuf-compiler pkg-config + export TARGET_CC=clang TARGET_CXX=clang++ # The manylinux2014 sysroot has glibc 2.17 headers which lack # AT_HWCAP2 (added in Linux 3.17). Define it for aws-lc-sys. - export CFLAGS="$CFLAGS -DAT_HWCAP2=26" && + export CFLAGS="$CFLAGS -DAT_HWCAP2=26" rustup target add aarch64-unknown-linux-gnu # Not `&&`-chained: in dash, errexit does not fire for a # non-final command in an `&&` list, so failures were ignored. @@ -136,10 +136,10 @@ jobs: lto: thin codegen_units: 16 pre_build: |- - set -e && - sudo apt-get update && - sudo apt-get install -y protobuf-compiler && - rustup target add aarch64-unknown-linux-musl && + set -e + sudo apt-get update + sudo apt-get install -y protobuf-compiler + rustup target add aarch64-unknown-linux-musl export EXTRA_ARGS="-x" name: build - ${{ matrix.settings.target }} runs-on: ${{ matrix.settings.host }} diff --git a/Cargo.lock b/Cargo.lock index 3b4086203..c72df37e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3584,8 +3584,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5059,8 +5059,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arc-swap", "arrow", @@ -5135,8 +5135,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5158,7 +5158,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5172,7 +5172,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-schema", @@ -5181,8 +5181,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrayref", "crunchy", @@ -5192,8 +5192,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5230,8 +5230,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow", "arrow-array", @@ -5262,8 +5262,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow", "arrow-array", @@ -5280,8 +5280,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "proc-macro2", "quote", @@ -5290,8 +5290,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-arith", "arrow-array", @@ -5324,8 +5324,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-arith", "arrow-array", @@ -5357,8 +5357,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "datafusion", "geo-traits", @@ -5372,8 +5372,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arc-swap", "arrow", @@ -5387,6 +5387,7 @@ dependencies = [ "async-recursion", "async-trait", "bitvec", + "blake3", "bytes", "crossbeam-queue", "datafusion", @@ -5421,6 +5422,7 @@ dependencies = [ "lance-tokenizer", "libsais-rs", "log", + "nalgebra", "ndarray", "num-traits", "object_store 0.14.2", @@ -5444,8 +5446,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-schema", @@ -5468,8 +5470,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow", "arrow-array", @@ -5512,8 +5514,8 @@ dependencies = [ [[package]] name = "lance-jsonb" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "byteorder", "ethnum", @@ -5530,8 +5532,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-schema", @@ -5545,8 +5547,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow", "async-trait", @@ -5560,8 +5562,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arc-swap", "arrow", @@ -5615,8 +5617,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5630,8 +5632,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow", "arrow-array", @@ -5671,8 +5673,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-schema", @@ -5685,8 +5687,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "14.0.0-beta.23" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.23#ea81833500fe251568854dfd889524bc4ebcbaea" +version = "15.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "frostem", "icu_segmenter", @@ -6411,6 +6413,21 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "nalgebra" +version = "0.33.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d43ddcacf343185dfd6de2ee786d9e8b1c2301622afab66b6c73baf9882abfd" +dependencies = [ + "approx", + "matrixmultiply", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + [[package]] name = "napi" version = "3.12.5" @@ -6650,6 +6667,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -9268,6 +9296,15 @@ version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d" +[[package]] +name = "safe_arch" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +dependencies = [ + "bytemuck", +] + [[package]] name = "safetensors" version = "0.7.0" @@ -9676,6 +9713,19 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simba" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c99284beb21666094ba2b75bbceda012e610f5479dfcc2d6e2426f53197ffd95" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "paste", + "wide", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -11182,6 +11232,16 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wide" +version = "0.7.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index 49f7ea6a6..3ac246149 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.23", default-features = false, "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.23", "tag" = "v14.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/docs/web/indexing/fts-index.mdx b/docs/web/indexing/fts-index.mdx index 832b67315..4660e8b01 100644 --- a/docs/web/indexing/fts-index.mdx +++ b/docs/web/indexing/fts-index.mdx @@ -99,7 +99,7 @@ await async_table.create_index("payload.text", config=FTS(with_position=True)) | `stem` | bool | `True` | Apply stemming (`running` → `run`) | | `remove_stop_words` | bool | `True` | Drop common stop words | | `ascii_folding` | bool | `True` | Normalize accented characters | -| `custom_stop_words` | list[str] | `None` | Extra stop words to drop in addition to the language defaults. Requires `remove_stop_words=True`. | +| `custom_stop_words` | list[str] | `None` | Explicit stop-word list that replaces the built-in list. Requires `remove_stop_words=True`. | | `ngram_min_length` | int | `3` | Minimum n-gram length. Applies only when `base_tokenizer="ngram"`. | | `ngram_max_length` | int | `3` | Maximum n-gram length. Applies only when `base_tokenizer="ngram"`. | | `prefix_only` | bool | `False` | Index only prefix n-grams rather than all substrings. Applies only when `base_tokenizer="ngram"`. | @@ -123,7 +123,11 @@ await async_table.create_index("payload.text", config=FTS(with_position=True)) Model-backed tokenizers such as `jieba/default`, `lindera/ipadic`, and `lindera/ko-dic` require tokenizer model files in Lance's language model home. Lance looks under the default platform data directory for `lance/language_models`, or you can set `LANCE_LANGUAGE_MODEL_HOME` to point to another model root. For example, `jieba/default` is resolved under `/jieba/default/...`. For a table in LanceDB Enterprise, the files must be present where the deployment builds and queries the index, so ask its administrator whether they are installed. -`language` is used by token filters, not by the base tokenizer. Stemming supports Arabic, Danish, Dutch, English, Finnish, French, German, Greek, Hungarian, Italian, Norwegian, Portuguese, Romanian, Russian, Spanish, Swedish, Tamil, and Turkish. Built-in stop-word removal supports Danish, Dutch, English, Finnish, French, German, Hungarian, Italian, Norwegian, Portuguese, Russian, Spanish, and Swedish. For other stemming languages, set `remove_stop_words=False` or pass `custom_stop_words`. +`language` is used by token filters, not by the base tokenizer. Stemming and language-specific stop-word removal support Arabic, Danish, Dutch, English, Finnish, French, German, Greek, Hungarian, Italian, Norwegian, Portuguese, Romanian, Russian, Spanish, Swedish, Tamil, and Turkish. The ICU tokenizer uses a combined multilingual stop-word list. An explicit `custom_stop_words` list replaces the built-in list for either tokenizer. + + +Lance `15.0.0-beta.1` revises the built-in stop-word lists, including the multilingual list used by ICU. Existing index metadata records tokenizer settings, but does not freeze these built-in lists. Rebuild affected FTS indexes with the same configuration before serving queries with the updated SDK, so stored terms and query tokens use the same lists. Indexes with stop-word removal disabled or an explicit `custom_stop_words` list are unaffected by this change. + ### Posting block size {#posting-block-size} diff --git a/java/pom.xml b/java/pom.xml index db0ee33d7..5cc7e365e 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 14.0.0-beta.23 + 15.0.0-beta.1 2.18.11 false 2.30.0 diff --git a/rust/lancedb/src/index/scalar.rs b/rust/lancedb/src/index/scalar.rs index 5bfc2b4e6..5d2256381 100644 --- a/rust/lancedb/src/index/scalar.rs +++ b/rust/lancedb/src/index/scalar.rs @@ -65,6 +65,17 @@ pub struct FmIndexBuilder {} /// A ZoneMap index stores min/max summaries for ranges of rows and can /// accelerate range predicates by pruning zones that cannot match. /// +/// ZoneMap indexes use per-fragment write seeds by default for strings, binary +/// data, and supported types wider than eight bytes. Creating a seed-enabled +/// index activates a table capability immediately, before any version 1 seed +/// is written. A mutation retaining a legacy seed-enabled index also activates +/// it. Older clients cannot read or write the resulting snapshots, so upgrade +/// all readers and writers before creating such an index or mutating a table +/// that already has one. The capability +/// remains set in subsequent snapshots, including after dropping the index or +/// restoring an older version; dropping the index does not restore access for +/// older clients. Opening a table alone does not activate the capability. +/// /// ``` /// use lancedb::{ /// index::{scalar::ZoneMapIndexBuilder, Index}, From 9f0f6400bc3490962542fff9e3f0ce26da1182a1 Mon Sep 17 00:00:00 2001 From: Lance Release Date: Thu, 22 Oct 2026 11:53:10 +0000 Subject: [PATCH 15/33] =?UTF-8?q?Bump=20version:=200.41.0-beta.12=20?= =?UTF-8?q?=E2=86=92=200.41.0-beta.13?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.toml | 2 +- Cargo.lock | 6 +++--- docs/src/java/java.md | 2 +- java/lancedb-core/pom.xml | 2 +- java/pom.xml | 2 +- nodejs/Cargo.toml | 2 +- nodejs/npm/darwin-arm64/package.json | 2 +- nodejs/npm/linux-arm64-gnu/package.json | 2 +- nodejs/npm/linux-arm64-musl/package.json | 2 +- nodejs/npm/linux-x64-gnu/package.json | 2 +- nodejs/npm/linux-x64-musl/package.json | 2 +- nodejs/npm/win32-arm64-msvc/package.json | 2 +- nodejs/npm/win32-x64-msvc/package.json | 2 +- nodejs/package.json | 2 +- python/Cargo.toml | 2 +- rust/lancedb/Cargo.toml | 2 +- 16 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index f2c9e44e1..0909579cb 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.12" +current_version = "0.41.0-beta.13" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index c72df37e3..7f36dc088 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5701,7 +5701,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" dependencies = [ "ahash", "anyhow", @@ -5799,7 +5799,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" dependencies = [ "arrow-array", "arrow-buffer", @@ -5824,7 +5824,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" dependencies = [ "arc-swap", "arrow", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 5c1874e6f..0c887f66f 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.12 + 0.41.0-beta.13 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index ad5adbc92..87c5d4cb8 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.12 + 0.41.0-beta.13 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 5cc7e365e..1f3e0efb0 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.12 + 0.41.0-beta.13 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 85ee5df04..1b9d47e46 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 080340224..bd4b95463 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index a2b5b8bf6..6e4e57684 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index 96bda03c7..f77c50ee2 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index edd596044..b671b7ebf 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index 953c46250..b75738ff1 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 236d356be..7158748f9 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 205f0ee11..83d0fb6f2 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 023c4e400..9bade7200 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.12", + "version": "0.41.0-beta.13", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index e6a0e457a..ad1790a99 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 5e6dd315f..486eb1e8e 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.12" +version = "0.41.0-beta.13" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true From ba7b7f15c2a576062a480e6863abfa092b6e0e58 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:07:18 +0800 Subject: [PATCH 16/33] fix(node): reject musl binaries with unresolved AVX-512 symbol (#3791) ## Summary - inspect the assembled x64 musl addon for the unresolved AVX-512 distance-table symbol before npm publication - fail the publish lifecycle with an actionable error if the broken symbol is present - add regression coverage using the dynamic-symbol output from the affected 0.27.x artifact ## Root cause LanceDB 0.27.x could enable the Rust AVX-512 distance-table callsite when the separate C kernel build did not emit its object, leaving `sum_4bit_dist_table_32bytes_batch_avx512` undefined in the musl addon. The current Lance dependency includes the upstream split-kernel cfg fix, and the published 0.33.0 musl artifact no longer has the undefined symbol, but the LanceDB publish path did not verify that property before releasing an addon. This adds a release guard around the corrected implementation so the same invalid binary cannot be published again. ## Validation - `pnpm lint` - `pnpm run build` - `pnpm test --runInBand` (680 passed, 5 skipped) - `pnpm run docs` - checked published artifacts: 0.27.1 is rejected for the unresolved symbol; 0.33.0 is accepted Fixes #3200 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- nodejs/.npmignore | 1 + nodejs/__test__/native_symbols.test.ts | 47 ++++++++++++++++ nodejs/package.json | 3 +- nodejs/scripts/check-native-symbols.js | 77 ++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 nodejs/__test__/native_symbols.test.ts create mode 100644 nodejs/scripts/check-native-symbols.js diff --git a/nodejs/.npmignore b/nodejs/.npmignore index 41f6643a6..128c0d36d 100644 --- a/nodejs/.npmignore +++ b/nodejs/.npmignore @@ -15,6 +15,7 @@ renovate.json src lancedb examples +scripts nodejs-artifacts Cargo.toml biome.json diff --git a/nodejs/__test__/native_symbols.test.ts b/nodejs/__test__/native_symbols.test.ts new file mode 100644 index 000000000..5445867f6 --- /dev/null +++ b/nodejs/__test__/native_symbols.test.ts @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const { + checkNativeBinary, + findForbiddenUndefinedSymbols, +} = require("../scripts/check-native-symbols.js"); + +test("detects the unresolved AVX-512 symbol from broken musl binaries", () => { + const output = [ + " U napi_create_function", + " U sum_4bit_dist_table_32bytes_batch_avx512", + " U strlen", + ].join("\n"); + + expect(findForbiddenUndefinedSymbols(output)).toEqual([ + "sum_4bit_dist_table_32bytes_batch_avx512", + ]); +}); + +test("accepts native binaries without unresolved internal Lance symbols", () => { + const runNm = jest.fn(() => ({ + status: 0, + stdout: + " U napi_create_function\n U strlen\n", + stderr: "", + })); + + expect(() => checkNativeBinary("lancedb.node", runNm)).not.toThrow(); + expect(runNm).toHaveBeenCalledWith( + "nm", + ["-D", "--undefined-only", "lancedb.node"], + { encoding: "utf8" }, + ); +}); + +test("rejects native binaries with the unresolved AVX-512 symbol", () => { + const runNm = jest.fn(() => ({ + status: 0, + stdout: " U sum_4bit_dist_table_32bytes_batch_avx512\n", + stderr: "", + })); + + expect(() => checkNativeBinary("lancedb.node", runNm)).toThrow( + "lancedb.node contains unresolved internal Lance symbols: sum_4bit_dist_table_32bytes_batch_avx512", + ); +}); diff --git a/nodejs/package.json b/nodejs/package.json index 9bade7200..21f903235 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -90,7 +90,8 @@ "postdocs": "node typedoc_post_process.js", "lint": "biome check . && biome format .", "lint-fix": "biome check --write . && biome format --write .", - "prepublishOnly": "napi prepublish -t npm", + "check:native-symbols": "node scripts/check-native-symbols.js", + "prepublishOnly": "pnpm check:native-symbols && napi prepublish -t npm", "test": "jest --verbose", "integration": "S3_TEST=1 pnpm test", "universal": "napi universalize", diff --git a/nodejs/scripts/check-native-symbols.js b/nodejs/scripts/check-native-symbols.js new file mode 100644 index 000000000..2f6dbf9b8 --- /dev/null +++ b/nodejs/scripts/check-native-symbols.js @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const { spawnSync } = require("node:child_process"); +const { existsSync } = require("node:fs"); +const path = require("node:path"); + +const FORBIDDEN_UNDEFINED_SYMBOLS = new Set([ + "sum_4bit_dist_table_32bytes_batch_avx512", +]); + +function findForbiddenUndefinedSymbols(output) { + const found = new Set(); + + for (const line of output.split(/\r?\n/)) { + const columns = line.trim().split(/\s+/); + const symbol = columns.at(-1)?.split("@")[0]; + if (symbol && FORBIDDEN_UNDEFINED_SYMBOLS.has(symbol)) { + found.add(symbol); + } + } + + return [...found].sort(); +} + +function checkNativeBinary(binaryPath, runNm = spawnSync) { + const result = runNm("nm", ["-D", "--undefined-only", binaryPath], { + encoding: "utf8", + }); + + if (result.error) { + throw new Error(`Unable to inspect ${binaryPath}: ${result.error.message}`); + } + if (result.status !== 0) { + throw new Error( + `Unable to inspect ${binaryPath}: nm exited with status ${result.status}\n${result.stderr}`, + ); + } + + const forbidden = findForbiddenUndefinedSymbols(result.stdout); + if (forbidden.length > 0) { + throw new Error( + `${binaryPath} contains unresolved internal Lance symbols: ${forbidden.join( + ", ", + )}`, + ); + } +} + +function main() { + const binaryPath = path.resolve( + __dirname, + "..", + "npm", + "linux-x64-musl", + "lancedb.linux-x64-musl.node", + ); + + if (!existsSync(binaryPath)) { + throw new Error( + `Missing ${binaryPath}; assemble the native artifacts before publishing`, + ); + } + + checkNativeBinary(binaryPath); +} + +if (require.main === module) { + try { + main(); + } catch (error) { + console.error(error instanceof Error ? error.message : error); + process.exitCode = 1; + } +} + +module.exports = { checkNativeBinary, findForbiddenUndefinedSymbols }; From a98667749d923512c1f26b65167da5aea0c6669c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:01 +0800 Subject: [PATCH 17/33] fix(python): verify Lance-created IVF_RQ compatibility (#3793) ## Summary - add regression coverage for an IVF_RQ index created through the Lance Python API and consumed through LanceDB - verify a fresh LanceDB handle can list and query the externally-created index ## Root cause Lance 4.0.0-beta.5 could serialize `rotate_mat_position` as null, while the Lance 2.0.0 reader used by LanceDB 0.29.2 required a `u32`. Upstream Lance fixed the incompatible IVF_RQ metadata handling and index versioning, and the current LanceDB dependency includes that repair. LanceDB was missing coverage for the reported mixed-API workflow, allowing the compatibility regression to go undetected. ## Validation - `cd python && uv run --extra tests pytest python/tests/test_index.py -q` (24 passed) - `uv run --project python --extra dev ruff format .` - `uv run --project python --extra dev ruff check .` - `uv run --project python --extra dev ruff format --check python/python/tests/test_index.py` - `uv run --project python --extra dev ruff check python/python/tests/test_index.py` Fixes #3092 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/tests/test_index.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index fe6ebe87a..0252a09ad 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -430,6 +430,28 @@ async def test_create_ivfrq_index(some_table: AsyncTable): assert indices[0].name == "vector_idx" +@pytest.mark.asyncio +async def test_read_ivfrq_index_created_with_lance(db_async, some_table: AsyncTable): + pytest.importorskip("lance") + dataset = await some_table.to_lance() + dataset.create_index( + "vector", + index_type="IVF_RQ", + metric="cosine", + num_partitions=1, + ) + + # Open a fresh LanceDB handle so the index is discovered from disk. + table = await db_async.open_table("some_table") + indices = await table.list_indices() + assert len(indices) == 1 + assert indices[0].index_type == "IvfRq" + assert indices[0].columns == ["vector"] + + results = await table.query().nearest_to(list(range(DIM))).limit(8).to_arrow() + assert len(results) == 8 + + @pytest.mark.asyncio async def test_create_hnswpq_index(some_table: AsyncTable): await some_table.create_index("vector", config=HnswPq(num_partitions=10)) From 102518ebda39513aa0bf337be96706df5ad558a0 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:16 +0800 Subject: [PATCH 18/33] fix(python): preserve nested fields in filtered scans (#3816) ## Summary - add a Python regression for filtered scans over deeply nested list/struct values - verify the public query path returns every nested sibling field without an Arrow schema mismatch ## Root cause Older Lance scanner code split filter projections into partial nested arrays and did not recursively merge List children, so a sibling field could be omitted from the returned array while remaining in the declared schema. That merge was fixed upstream in lance-format/lance#4838, and the Lance revision currently pinned by LanceDB contains that fix. This test makes the repair durable at the LanceDB Python API boundary. ## Validation - cd python && uv run --extra tests pytest python/tests/test_query.py::test_filter_preserves_deeply_nested_struct_fields -q - cd python && uv run --extra tests pytest python/tests/test_query.py -q - uv run --project python --extra tests --extra dev ruff check . - uv run --project python --extra tests --extra dev ruff format --check python/python/tests/test_query.py Fixes #2217 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/tests/test_query.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index c8d0ad159..58ea7deb5 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -833,6 +833,28 @@ def test_with_row_id(table: lancedb.table.Table): assert rs["_rowid"].to_pylist() == [0, 1] +def test_filter_preserves_deeply_nested_struct_fields(tmp_db): + # Regression test for https://github.com/lancedb/lancedb/issues/2217. + data = [ + { + "item": "test", + "sub_obj": { + "second_sub_obj": [ + { + "third_sub_obj": {"int_val": 1}, + "str_val": "test", + } + ] + }, + } + ] + table = tmp_db.create_table("deeply_nested_struct", data=data) + + result = table.search().where("item = 'test'").limit(1).to_arrow() + + assert result.to_pylist() == data + + def test_blob_v2_query_omits_auto_row_id(tmp_db): table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid") From 40af9f350e785d762cc887e029718fc9904bbe0b Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:32 +0800 Subject: [PATCH 19/33] fix(python): avoid deadlock for query-backed iterators (#3820) ## Summary - move synchronous iterator preparation off LanceDB's singleton background event loop - keep the loop available when an input iterator executes a synchronous table query - cover query-backed iterators for both `Table.add` and `create_table` ## Root cause The async implementations prepared input data on the singleton background loop and peeked at the first item of iterators there. If that iterator called a synchronous LanceDB query, `LOOP.run()` scheduled the query on the same loop and blocked waiting for it, so neither operation could progress. ## Validation - `ruff format` on the changed Python files - `ruff check .` - `cd python && uv run --frozen --extra tests pytest python/tests/test_table.py -q` (138 passed) - the broader non-slow suite was attempted but hit a native segmentation fault in the external DataFusion `register_table_provider` path; an isolated control using the pre-change inline preparation behavior reproduced the same unrelated crash Fixes #2107 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/db.py | 13 +++++++++++-- python/python/lancedb/table.py | 16 ++++++++++++---- python/python/tests/test_table.py | 24 ++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 6 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 5def89b01..e1e5d270a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -5,6 +5,7 @@ from __future__ import annotations from abc import abstractmethod +import asyncio from datetime import timedelta from pathlib import Path import sys @@ -2652,8 +2653,16 @@ class AsyncConnection(object): if fill_value is None: fill_value = 0.0 - data, schema = sanitize_create_table( - data, schema, metadata, on_bad_vectors, fill_value + # Input preparation may advance a user-provided iterator. Keep that work + # off the background event loop so an iterator can use the synchronous + # LanceDB API without blocking the loop that API needs to make progress. + data, schema = await asyncio.to_thread( + sanitize_create_table, + data, + schema, + metadata, + on_bad_vectors, + fill_value, ) validate_schema(schema) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index d5872e6f0..006ef8cb9 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6213,8 +6213,12 @@ class AsyncTable: if mode == "overwrite": # For overwrite, apply the same preprocessing as create_table # so vector columns are inferred as FixedSizeList. - data, _ = sanitize_create_table( - data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value + data, _ = await asyncio.to_thread( + sanitize_create_table, + data, + None, + on_bad_vectors=on_bad_vectors, + fill_value=fill_value, ) if blob_v2_column_paths(schema): # Plain bytes cannot be inferred as blob v2. Keep matching blob @@ -6224,7 +6228,8 @@ class AsyncTable: elif on_bad_vectors != "error" or ( schema.metadata is not None and b"embedding_functions" in schema.metadata ): - data = _sanitize_data( + data = await asyncio.to_thread( + _sanitize_data, data, schema, metadata=schema.metadata, @@ -6233,7 +6238,10 @@ class AsyncTable: allow_subschema=True, ) _register_optional_converters() - data = to_scannable(data) + # Converting an iterator peeks at its first item. A synchronous query in + # that iterator schedules work on LOOP, so peeking on LOOP's own thread + # would deadlock waiting for itself. + data = await asyncio.to_thread(to_scannable, data) progress, owns = _normalize_progress(progress) try: return await self._inner.add( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 755e94d7f..095f0b911 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -518,6 +518,30 @@ def test_add(mem_db: DBConnection): _add(table, schema) +def test_add_from_iterator_that_queries_table(mem_db: DBConnection): + source = mem_db.create_table("source", data=pa.table({"id": range(16)})) + target = mem_db.create_table("target", schema=source.schema) + + def batches(): + for _ in range(5): + yield source.search().limit(10).to_arrow() + + target.add(batches()) + + assert target.count_rows() == 50 + + +def test_create_table_from_iterator_that_queries_table(mem_db: DBConnection): + source = mem_db.create_table("source", data=pa.table({"id": range(16)})) + + def batches(): + yield source.search().limit(10).to_arrow() + + target = mem_db.create_table("target", data=batches()) + + assert target.count_rows() == 10 + + def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection): """Regression test for https://github.com/lancedb/lancedb/issues/2512.""" schema = pa.schema([pa.field("x", pa.int64())]) From a1230cf9a490c1768cbec2709edf772543d27331 Mon Sep 17 00:00:00 2001 From: Li Zinuo Date: Fri, 23 Oct 2026 01:09:46 +0800 Subject: [PATCH 20/33] fix: preserve unbounded plain queries in namespace pushdown (#3927) ## What Preserve unbounded plain-query semantics in the Rust Namespace `QueryTable` adapter. Plain queries without an explicit limit are now encoded with `i32::MAX`. Explicit limits remain unchanged, while vector and FTS queries continue to use the default top-k of 10. ## Root cause The Rust Namespace adapter converted a missing limit using: ```rust q.limit.unwrap_or(10) ``` This turned an unbounded plain query into an explicit `k=10`, silently truncating the results returned by Namespace pushdown. ## Tests Added regression coverage for unbounded and explicitly limited plain queries, while verifying that vector and FTS queries retain their default limit. ```shell cargo fmt --all -- --check cargo test --quiet --features remote -p lancedb \ test_convert_to_namespace_query --lib ``` close #3926. --- rust/lancedb/src/table/query.rs | 56 ++++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index cf70c6626..f625a4aff 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -671,6 +671,12 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { }); } + let k = match q.limit { + Some(limit) => limit as i32, + None if q.full_text_search.is_some() => DEFAULT_TOP_K as i32, + None => i32::MAX, + }; + let filter = q.filter.as_ref().map(filter_to_sql).transpose()?; let columns = match &q.select { @@ -723,7 +729,7 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { Ok(NsQueryTableRequest { id: None, // Will be set by caller vector, - k: q.limit.unwrap_or(10) as i32, + k, filter, columns, prefilter: Some(q.prefilter), @@ -1073,6 +1079,54 @@ mod tests { assert!(ns_request.vector.single_vector.as_ref().unwrap().is_empty()); } + #[test] + fn test_convert_to_namespace_query_plain_query_without_limit() { + let any_query = AnyQuery::Query(QueryRequest::default()); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, i32::MAX); + } + + #[test] + fn test_convert_to_namespace_query_plain_query_with_explicit_limit() { + let any_query = AnyQuery::Query(QueryRequest { + limit: Some(10), + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + + #[test] + fn test_convert_to_namespace_query_vector_without_limit() { + let query_vector = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])); + let any_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![query_vector as ArrayRef], + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + + #[test] + fn test_convert_to_namespace_query_fts_without_limit() { + let any_query = AnyQuery::Query(QueryRequest { + full_text_search: Some(lance_index::scalar::FullTextSearchQuery::new( + "hello world".to_string(), + )), + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + #[tokio::test] async fn test_execute_query_local_routing() { use crate::connect; From 548a1903216075756ae7e6e6eecb671a78f38dd8 Mon Sep 17 00:00:00 2001 From: primorLee Date: Fri, 23 Oct 2026 01:10:00 +0800 Subject: [PATCH 21/33] fix(python): refresh zero-lifetime OAuth tokens (#3971) ## Summary - treat an OAuth `expires_in` value of zero as an explicit immediate expiration - refresh zero-lifetime tokens instead of caching them indefinitely - add a regression test covering two consecutive header requests ## Tests - `pytest python/tests/test_header_provider.py -q` (16 passed) - `ruff format --check python/lancedb/remote/header.py python/tests/test_header_provider.py` - `ruff check python/lancedb/remote/header.py python/tests/test_header_provider.py` --- python/python/lancedb/remote/header.py | 2 +- python/python/tests/test_header_provider.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/remote/header.py b/python/python/lancedb/remote/header.py index 206b7a7cb..5f5bd12d8 100644 --- a/python/python/lancedb/remote/header.py +++ b/python/python/lancedb/remote/header.py @@ -141,7 +141,7 @@ class OAuthProvider(HeaderProvider): # Set expiration if provided expires_in = token_data.get("expires_in") - if expires_in: + if expires_in is not None: self._token_expires_at = time.time() + expires_in else: # Token doesn't expire or expiration unknown diff --git a/python/python/tests/test_header_provider.py b/python/python/tests/test_header_provider.py index 187b0a50a..e50ff3270 100644 --- a/python/python/tests/test_header_provider.py +++ b/python/python/tests/test_header_provider.py @@ -114,6 +114,24 @@ class TestOAuthProvider: "x-lancedb-credential-type": "oidc", } + def test_zero_lifetime_token_is_refreshed(self): + """A token that expires immediately must not be cached indefinitely.""" + call_count = 0 + + def fetcher(): + nonlocal call_count + call_count += 1 + return { + "access_token": f"token{call_count}", + "expires_in": 0, + } + + provider = OAuthProvider(fetcher, refresh_buffer_seconds=0) + + assert provider.get_headers() == {"Authorization": "Bearer token1"} + assert provider.get_headers() == {"Authorization": "Bearer token2"} + assert call_count == 2 + def test_missing_access_token(self): """Test error handling when access_token is missing.""" From 6d92307fc34b81adc0c040da367c6f9b64f0aeba Mon Sep 17 00:00:00 2001 From: Zhewen Tan <1090179959@qq.com> Date: Fri, 23 Oct 2026 01:10:15 +0800 Subject: [PATCH 22/33] fix(python): accept tuple and Arrow vector queries (#3972) ## Summary Sync `Table.search()` declares tuple, `pyarrow.Array`, and `pyarrow.ChunkedArray` query vectors as supported, and the async query path already accepts them. The sync query builder only classified lists and NumPy arrays as vectors, however, so these inputs were routed through embedding/FTS resolution and then rejected. Arrow and tuple inputs also did not contribute a query dimension when selecting among multiple vector columns. This change: - recognizes the supported tuple and Arrow containers as sync vector queries for both `auto` and explicit `vector` query types; - converts them to the existing internal `float32` NumPy representation at the sync builder boundary; - infers dimensions from tuple, primitive Arrow, and list-valued Arrow queries so schemas with differently sized vector columns select the matching column. The regression test uses two vector columns with dimensions 2 and 3. It verifies tuple, Arrow array, and chunked Arrow queries through both sync query modes and confirms the dimension-2 column is selected. ## Test plan - `pytest python/tests/test_table.py::test_sync_search_accepts_supported_vector_containers -q` (6 passed) - `pytest python/tests/test_util.py python/tests/test_query.py python/tests/test_table.py -q -m "not slow and not s3_test"` (308 passed, 1 skipped, 1 deselected) - `pytest python/tests -q -m "not slow and not s3_test" --ignore=python/tests/docs` (981 passed, 25 skipped, 65 deselected) - `ruff format --check python` - `ruff check python` - `cargo fmt --all --check` Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 9 +++++--- python/python/lancedb/util.py | 11 +++++++++- python/python/tests/test_table.py | 36 +++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index c427820e8..53f56a2c0 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -82,6 +82,7 @@ if TYPE_CHECKING: T = TypeVar("T", bound="LanceModel") AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"] +_SYNC_VECTOR_QUERY_TYPES = (list, tuple, np.ndarray, pa.Array, pa.ChunkedArray) # Number of rows a hybrid query returns when no limit was set on it. This # mirrors the default the Rust query builder applies to its sub-queries. @@ -1024,10 +1025,12 @@ class LanceQueryBuilder(ABC): fast_search=fast_search, ) - if isinstance(query, list): + if isinstance(query, (list, tuple)): query = np.array(query, dtype=np.float32) elif isinstance(query, np.ndarray): query = query.astype(np.float32) + elif isinstance(query, (pa.Array, pa.ChunkedArray)): + query = np.array(query.to_pylist(), dtype=np.float32) else: raise TypeError(f"Unsupported query type: {type(query)}") @@ -1049,7 +1052,7 @@ class LanceQueryBuilder(ABC): query = cls._query_to_vector(table, query, vector_column_name) return query, query_type elif query_type == "auto": - if isinstance(query, (list, np.ndarray)): + if isinstance(query, _SYNC_VECTOR_QUERY_TYPES): return query, "vector" else: conf = table.embedding_functions.get(vector_column_name) @@ -1065,7 +1068,7 @@ class LanceQueryBuilder(ABC): @classmethod def _query_to_vector(cls, table, query, vector_column_name): - if isinstance(query, (list, np.ndarray)): + if isinstance(query, _SYNC_VECTOR_QUERY_TYPES): return query conf = table.embedding_functions.get(vector_column_name) if conf is not None: diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index d4533887b..3babc3872 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -309,7 +309,16 @@ def _query_vector_dim(query: Optional[Any]) -> Optional[int]: if query.ndim == 0: return None return query.shape[-1] - if isinstance(query, list) and query: + if isinstance(query, (pa.Array, pa.ChunkedArray)): + if len(query) == 0: + return None + if pa.types.is_fixed_size_list(query.type): + return query.type.list_size + if pa.types.is_list(query.type) or pa.types.is_large_list(query.type): + first = query[0].as_py() + return len(first) if first is not None else None + return len(query) + if isinstance(query, (list, tuple)) and query: first = query[0] if isinstance(first, (list, tuple, np.ndarray)): return len(first) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 095f0b911..7a4f5981c 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3697,6 +3697,42 @@ def test_multiple_vector_columns(mem_db: DBConnection): assert result1["text"][0] != result2["text"][0] +@pytest.mark.parametrize( + "query", + [ + (0.0, 0.0), + pa.array([0.0, 0.0]), + pa.chunked_array([[0.0], [0.0]]), + ], +) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +def test_sync_search_accepts_supported_vector_containers( + mem_db: DBConnection, query, query_type +): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("small", pa.list_(pa.float32(), 2)), + pa.field("large", pa.list_(pa.float32(), 3)), + ] + ) + table = mem_db.create_table( + "supported_vector_containers", + data=pa.table( + { + "id": [1, 2], + "small": [[0.0, 0.0], [1.0, 1.0]], + "large": [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]], + }, + schema=schema, + ), + ) + + result = table.search(query, query_type=query_type).limit(1).to_arrow() + + assert result["id"].to_pylist() == [1] + + def test_create_scalar_index(mem_db: DBConnection): vec_array = pa.array( [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], pa.list_(pa.float32(), 2) From ae97d43f4bc3c1c01950c3a0d8369f21ceba014d Mon Sep 17 00:00:00 2001 From: Vinayak Vathare <137673662+VathareVinayak@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:40:30 +0530 Subject: [PATCH 23/33] fix(python): preserve valid embeddings when an OpenAI batch is rejected (#4005) Retries individual inputs one at a time when the OpenAI API rejects a batch with BadRequestError, instead of discarding the whole batch. Valid inputs now still get embeddings; only the actually-bad input resolves to None, and the log names the exact failing input. Fixes #1677 --- python/python/lancedb/embeddings/openai.py | 25 ++++++++---- python/python/tests/test_embeddings.py | 45 ++++++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py index 1edfaa2ca..d4ecbdf3f 100644 --- a/python/python/lancedb/embeddings/openai.py +++ b/python/python/lancedb/embeddings/openai.py @@ -101,24 +101,35 @@ class OpenAIEmbeddings(TextEmbeddingFunction): valid_indices.append(idx) # TODO retry, rate limit, token limit - try: + def embed_batch(texts_batch, indices_batch): kwargs = { - "input": valid_texts, + "input": texts_batch, "model": self.name, } if self.name != "text-embedding-ada-002": kwargs["dimensions"] = self.dim rs = self._openai_client.embeddings.create(**kwargs) - valid_embeddings = { - idx: v.embedding for v, idx in zip(rs.data, valid_indices) - } + return {idx: v.embedding for v, idx in zip(rs.data, indices_batch)} + + def embed_one_or_log(idx, text): + # Retry a single rejected input + try: + return embed_batch([text], [idx]) + except openai.BadRequestError: + logging.exception("Bad request for input at index %d: %r", idx, text) + return {} + + try: + valid_embeddings = embed_batch(valid_texts, valid_indices) except openai.AuthenticationError: logging.error("Authentication failed: Invalid API key provided") raise except openai.BadRequestError: - logging.exception("Bad request: %s", texts) - return [None] * len(texts) + # Retry each input individually + valid_embeddings = {} + for idx, text in zip(valid_indices, valid_texts): + valid_embeddings.update(embed_one_or_log(idx, text)) except Exception: logging.exception("OpenAI embeddings error") raise diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 9434033e1..c42172f77 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -596,6 +596,51 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +def test_openai_bad_request_partial_batch_recovery(): + # Regression test for issue #1677 + + # Fake exception types, avoids requiring the optional openai/httpx deps + class FakeBadRequestError(Exception): + pass + + class FakeAuthenticationError(Exception): + pass + + fake_openai_module = MagicMock() + fake_openai_module.BadRequestError = FakeBadRequestError + fake_openai_module.AuthenticationError = FakeAuthenticationError + + registry = get_registry() + model = registry.get("openai").create(name="text-embedding-ada-002") + + def fake_create(input, model=None, **kwargs): + # Full batch always rejected + if len(input) > 1: + raise FakeBadRequestError("bad request") + if input[0] == "bad text": + raise FakeBadRequestError("bad request") + embedding = MagicMock() + embedding.embedding = [0.1, 0.2, 0.3] + result = MagicMock() + result.data = [embedding] + return result + + mock_client = MagicMock() + mock_client.embeddings.create.side_effect = fake_create + + with ( + patch( + "lancedb.embeddings.openai.attempt_import_or_raise", + return_value=fake_openai_module, + ), + patch.object(type(model), "_openai_client", new_callable=lambda: mock_client), + ): + embeddings = model.generate_embeddings(["good text", "bad text"]) + + assert embeddings[0] == [0.1, 0.2, 0.3] + assert embeddings[1] is None + + def test_ollama_embeddings_pickle(): """OllamaEmbeddings must pickle even after the cached client is created.""" registry = get_registry() From 0c3a1c92c862d1f4d251d55eacafabe225277722 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Thu, 22 Oct 2026 22:40:45 +0530 Subject: [PATCH 24/33] fix(python): don't crash inferring schema for an all-null vector column (#4034) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #4033 `create_table(data=...)` without an explicit schema raised `IndexError: index out of bounds` when a `vector`/`embedding`-named list column was all-null. `_modal_list_size` infers the dimension with `pc.mode(pc.list_value_length(arr))[0]`, and pyarrow's `mode` skips nulls, so an all-null column produces a zero-length result and the `[0]` subscript blows up. The sibling `_infer_vector_dim` already guards the identical `pc.mode(...)` with a `len == 0` check and returns `None`; `_modal_list_size` just didn't. This makes `_modal_list_size` return `Optional[int]` and return `None` when every value is null, and the call site then leaves the column as a variable-length list instead of inventing a fixed size — the dimension is genuinely unknowable when there are no non-null rows. Added a regression test that creates a table from an all-null `list` vector column and asserts it succeeds, the column stays a `list`, and the rows are queryable. It fails on `main` with the `IndexError` and passes with this change. --- python/python/lancedb/table.py | 17 ++++++++++++++--- python/python/tests/test_table.py | 24 ++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 006ef8cb9..ff866d4db 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5335,6 +5335,12 @@ def _infer_target_schema( # Use the most common length of the list as the dimensions dim = _modal_list_size(peeked.column(i)) + if dim is None: + # Every list entry is null (e.g. embeddings not computed yet), + # so the dimension is unknowable. Leave the column as a + # variable-length list instead of guessing a fixed size. + continue + # Determine target type based on value type if pa.types.is_floating(field.type.value_type): target_type = pa.list_(pa.float32(), dim) @@ -5374,9 +5380,14 @@ def _infer_target_schema( return schema, reader -def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> int: - # Use the most common length of the list as the dimensions - return pc.mode(pc.list_value_length(arr))[0].as_py()["mode"] +def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> Optional[int]: + # Use the most common length of the list as the dimensions. pyarrow's + # `mode` skips nulls, so an all-null column yields an empty result; return + # None in that case so the caller can keep a variable-length list. + modes = pc.mode(pc.list_value_length(arr)) + if len(modes) == 0: + return None + return modes[0].as_py()["mode"] def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]: diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7a4f5981c..827f2ef58 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -280,6 +280,30 @@ def test_create_table_infers_large_int_vectors(mem_db: DBConnection): assert vector_column.to_pylist() == [[0.0, 300.0]] +def test_create_table_all_null_vector_column(mem_db: DBConnection): + # A vector column can legitimately arrive with every value null, e.g. when + # the embeddings have not been computed yet. The dimension is unknowable in + # that case, so schema inference must not crash and should leave the column + # as a variable-length list rather than a fixed-size one. + data = pa.table( + { + "id": pa.array([1, 2]), + "vector": pa.array([None, None], type=pa.list_(pa.float32())), + } + ) + + table = mem_db.create_table("all_null_vector", data=data) + + vector_field = table.schema.field("vector") + assert pa.types.is_list(vector_field.type) + assert vector_field.type.value_type == pa.float32() + + result = table.to_arrow() + assert result.num_rows == 2 + assert result.column("vector").to_pylist() == [None, None] + assert table.count_rows() == 2 + + @pytest.mark.asyncio async def test_create_table_async_infers_large_int_vectors( mem_db_async: AsyncConnection, From 79e256959fa44f6f335e4b0e6409a6c21b6dcde3 Mon Sep 17 00:00:00 2001 From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:54:57 +0530 Subject: [PATCH 25/33] fix(python): accept pathlib.Path in get_uri_scheme and get_uri_location (#4132) This PR converts the uri parameter to string before calling urlparse in get_uri_scheme and get_uri_location, allowing pathlib.Path objects to be passed seamlessly without raising a TypeError. --------- Co-authored-by: Xuanwo --- python/python/lancedb/util.py | 15 ++++++++------- python/python/tests/test_util.py | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index 3babc3872..f1a652924 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -34,20 +34,20 @@ def safe_import_adlfs(): adlfs = safe_import_adlfs() -def get_uri_scheme(uri: str) -> str: +def get_uri_scheme(uri: Union[str, pathlib.Path]) -> str: """ Get the scheme of a URI. If the URI does not have a scheme, assume it is a file URI. Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: The scheme of the URI. """ - parsed = urlparse(uri) + parsed = urlparse(str(uri)) scheme = parsed.scheme if not scheme: scheme = "file" @@ -61,25 +61,26 @@ def get_uri_scheme(uri: str) -> str: return scheme -def get_uri_location(uri: str) -> str: +def get_uri_location(uri: Union[str, pathlib.Path]) -> str: """ Get the location of a URI. If the parameter is not a url, assumes it is just a path Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: Location part of the URL, without scheme """ - parsed = urlparse(uri) + uri_str = str(uri) + parsed = urlparse(uri_str) if len(parsed.scheme) == 1: # Windows drive names are parsed as the scheme # e.g. "c:\path" -> ParseResult(scheme="c", netloc="", path="/path", ...) # So we add special handling here for schemes that are a single character - return uri + return uri_str if not parsed.netloc: return parsed.path diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 741214794..4671109b8 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -30,6 +30,7 @@ import pytest import lancedb from lancedb.util import ( flatten_columns, + get_uri_location, get_uri_scheme, inf_vector_column_query, join_uri, @@ -86,6 +87,22 @@ def test_normalize_uri(): assert parsed_scheme == expected_scheme +def test_get_uri_scheme_and_location_pathlib(): + # Test pathlib.Path instances + abs_path = pathlib.Path("/tmp/data") + rel_path = pathlib.Path("relative/data") + + assert get_uri_scheme(abs_path) == "file" + assert get_uri_scheme(rel_path) == "file" + + assert get_uri_location(abs_path) == str(abs_path) + assert get_uri_location(rel_path) == str(rel_path) + + # String URIs still behave identically + assert get_uri_scheme("s3://bucket/data") == "s3" + assert get_uri_location("s3://bucket/data") == "bucket/data" + + def test_join_uri_remote(): schemes = ["s3", "az", "gs"] for scheme in schemes: From d4b2290babbc8ab83a4287551b0145fd7ed05f44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=BF=97=E8=B0=A6?= <89645338+simpleqt@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:13 +0800 Subject: [PATCH 26/33] fix(python): correct the voyageai reranker env var docstring (#4201) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `VoyageAIReranker`'s `api_key` docstring said it falls back to `OPENAI_API_KEY`, but the class actually reads `VOYAGE_API_KEY` (and raises a "VOYAGE_API_KEY not set" error without it) — the docstring was copy-pasted from the OpenAI reranker. Docstring-only fix. --- python/python/lancedb/rerankers/voyageai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/python/lancedb/rerankers/voyageai.py b/python/python/lancedb/rerankers/voyageai.py index 0f41ba927..d3dbd59ac 100644 --- a/python/python/lancedb/rerankers/voyageai.py +++ b/python/python/lancedb/rerankers/voyageai.py @@ -32,7 +32,7 @@ class VoyageAIReranker(Reranker): return_score : str, default "relevance" options are "relevance" or "all". Only "relevance" is supported for now. api_key : str, default None - The API key to use. If None, will use the OPENAI_API_KEY environment variable. + The API key to use. If None, will use the VOYAGE_API_KEY environment variable. truncation : Optional[bool], default None """ From 9a590913767d31a0ba5ff524c7884052cec2451f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=BF=97=E8=B0=A6?= <89645338+simpleqt@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:29 +0800 Subject: [PATCH 27/33] fix(python): add missing space in the hybrid search error message (#4202) `LanceHybridQueryBuilder._validate_query`'s ValueError concatenated two adjacent string literals without a separator, rendering "...explicitly for hybrid search.**But** not both.". Added the trailing space. String-only change. --- python/python/lancedb/query.py | 2 +- python/python/tests/test_hybrid_query.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 53f56a2c0..fe8c33861 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2441,7 +2441,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): if query is not None and (vector is not None or text is not None): raise ValueError( "You can either provide a string query in search() method " - "or set `vector()` and `text()` explicitly for hybrid search." + "or set `vector()` and `text()` explicitly for hybrid search. " "But not both." ) diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 8d670a136..19050093c 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -536,3 +536,15 @@ def test_normalize_scores(): assert pc.equal(result, expected), ( f"Expected {expected} but got {result} for invert={invert}" ) + + +def test_hybrid_query_rejects_query_with_vector(sync_table: Table): + # Pins the full message: the space between "hybrid search." and "But" is + # load-bearing (see the _validate_query literal concatenation). + with pytest.raises(ValueError) as excinfo: + sync_table.search("dog", query_type="hybrid").vector([0.5, -0.5]).to_arrow() + assert str(excinfo.value) == ( + "You can either provide a string query in search() method " + "or set `vector()` and `text()` explicitly for hybrid search. " + "But not both." + ) From 82bfe7d38e38285424e843f2491b9680c59c30d2 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:25:43 +0100 Subject: [PATCH 28/33] fix: table_names should not list loose .lance files as tables (#4203) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `table_names()` enumerates with `read_dir`, which returns files as well as directories, filtered only on the `.lance` extension — so a loose `foo.lance` *file* at the database prefix is reported as a table. `list_tables()` already applies the correct rule ("Only child directories can be tables"), and `test_listing_ignores_non_table_children` documents it, so the two listing APIs currently disagree. The practical damage is in `drop_all_tables()`, which is built on `table_names`: with a loose `.lance` file present it deletes some real tables, fails on the file with `os error 20 (NotADirectory)`, and then fails the same way on every retry — the database can never be emptied through the API, and the partial drop is not recoverable. `table_names` now lists via `list_with_delimiter` and filters `common_prefixes` through the same `table_name` helper `list_tables` uses; the sort/`start_after`/`limit` handling is unchanged. `drop_tables` additionally attempts every name before returning, still reporting the first error unchanged (kept rather than an aggregate error so the mapped exception types in the bindings stay exactly as they are). Extended the non-table-children test to cover `table_names`, plus tests for the drop behavior — all three fail without the fix. Full `cargo test -p lancedb --lib --features remote` 1320 passed on the repo-pinned 1.97 toolchain; clippy/fmt clean. --- rust/lancedb/src/database/listing.rs | 120 +++++++++++++++++++++------ 1 file changed, 94 insertions(+), 26 deletions(-) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 71e9f0ef1..700c855dd 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -750,26 +750,39 @@ impl ListingDatabase { uri.push_str(&format!("?{}", query_string)); } let commit_handler = commit_handler_from_url(&uri, &Some(object_store_params)).await?; + // One name that cannot be dropped must not strand the names after it: a caller + // dropping every table would otherwise be left with a partly dropped database that + // the same call cannot finish. Every name is attempted and the first failure is + // still what the call reports, so a caller matching on the error sees what it saw + // before. + let mut first_error = None; for name in names { let dir_name = format!("{}.{}", name, LANCE_EXTENSION); let full_path = self.base_path.clone().join(dir_name.clone()); - commit_handler.delete(&full_path).await?; + let dropped: Result<()> = async { + commit_handler.delete(&full_path).await?; - self.object_store - .remove_dir_all(full_path.clone()) - .await - .map_err(|err| match err { - // this error is not lance::Error::DatasetNotFound, as the method - // `remove_dir_all` may be used to remove something not be a dataset - lance::Error::NotFound { .. } => Error::TableNotFound { - name: name.clone(), - source: Box::new(err), - }, - _ => Error::from(err), - })?; + self.object_store + .remove_dir_all(full_path.clone()) + .await + .map_err(|err| match err { + // this error is not lance::Error::DatasetNotFound, as the method + // `remove_dir_all` may be used to remove something not be a dataset + lance::Error::NotFound { .. } => Error::TableNotFound { + name: name.clone(), + source: Box::new(err), + }, + _ => Error::from(err), + }) + } + .await; + + if let Err(err) = dropped { + first_error.get_or_insert(err); + } } - Ok(()) + first_error.map_or(Ok(()), Err) } /// Inherit storage options from the connection into the target map @@ -930,20 +943,18 @@ impl Database for ListingDatabase { if !request.namespace_path.is_empty() { return self.namespace_database().table_names(request).await; } - let mut f = self + // Only child directories can be tables, the same rule `list_tables` lists by: a + // loose object named like a table is not one, and a name handed back for it is a + // name the caller can neither open nor drop. + let dir_suffix = format!(".{LANCE_EXTENSION}"); + let listing = self .object_store - .read_dir(self.base_path.clone()) - .await? + .list_with_delimiter(Some(&self.base_path)) + .await?; + let mut f = listing + .common_prefixes .iter() - .map(Path::new) - .filter(|path| { - let is_lance = path - .extension() - .and_then(|e| e.to_str()) - .map(|e| e == LANCE_EXTENSION); - is_lance.unwrap_or(false) - }) - .filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from))) + .filter_map(|location| table_name(location, &dir_suffix)) .collect::>(); f.sort(); if let Some(start_after) = request.start_after { @@ -1697,6 +1708,63 @@ mod tests { .unwrap(); assert_eq!(page.tables, vec!["real"]); + + // The deprecated name listing answers the same question and so follows the same + // rule: a name it hands back has to be one a caller can open and drop. + #[allow(deprecated)] + let names = db.table_names(TableNamesRequest::default()).await.unwrap(); + assert_eq!(names, vec!["real"]); + } + + /// Dropping the database drops its tables and leaves everything else alone. A loose + /// `*.lance` object under the database prefix is not a table, so it is neither dropped + /// nor allowed to strand the tables that are. + #[tokio::test] + async fn test_drop_all_tables_leaves_loose_lance_files() { + let (tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + // Sorts between the two tables, so a drop that stops at the first failure would + // leave `b` behind. + let loose = tempdir.path().join("aaa-loose.lance"); + std::fs::write(&loose, b"not a table").unwrap(); + + db.drop_all_tables(&[]).await.unwrap(); + + #[allow(deprecated)] + let names = db.table_names(TableNamesRequest::default()).await.unwrap(); + assert!(names.is_empty(), "a table survived the drop: {names:?}"); + assert!(!tempdir.path().join("a.lance").exists()); + assert!(!tempdir.path().join("b.lance").exists()); + assert!( + loose.exists(), + "a loose object was dropped as if it were a table" + ); + } + + /// A name that cannot be dropped does not strand the names after it: every name is + /// attempted, and the failure is still what the call reports. + #[tokio::test] + async fn test_drop_tables_attempts_every_name() { + let (tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + + let result = db + .drop_tables(vec![ + "a".to_string(), + "missing".to_string(), + "b".to_string(), + ]) + .await; + + assert!( + matches!(result, Err(Error::TableNotFound { .. })), + "dropping a table that is not there did not report it: {result:?}" + ); + assert!(!tempdir.path().join("a.lance").exists()); + assert!( + !tempdir.path().join("b.lance").exists(), + "the table after the failing name was left behind" + ); } #[tokio::test] From d37a6b76a465a1fc62372ffe51d4cf4bb5233633 Mon Sep 17 00:00:00 2001 From: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:25:58 -0400 Subject: [PATCH 29/33] fix: gate remote job error construction behind the remote feature (#4233) ## Summary Fixes #4142: `lancedb` failed to compile without the `remote` feature. `TerminalResult::decode` in `rust/lancedb/src/job.rs` constructed the `Error::Http` variant unconditionally, but that variant is declared behind `#[cfg(feature = "remote")]` in `rust/lancedb/src/error.rs`. Since the crate's default feature set is empty, a plain `lancedb = "0.38"` dependency failed with: ``` error[E0599]: no variant named `Http` found for enum `error::Error` --> rust/lancedb/src/job.rs:102:40 ``` The module itself is not remote-only (it carries local job support), so instead of gating `pub mod job`, both `Error::Http` constructions now go through a small `remote_job_error` helper that resolves to `Error::Http` with the `remote` feature and falls back to `Error::Runtime` without it. A `request_id` is only ever `Some` for remote jobs, so the fallback is unreachable in practice without the feature. ## Testing - Reproduced the reported failure on current `main`: `cargo check -p lancedb --no-default-features` failed with the exact E0599 errors from the issue before the change, and passes after it. - `cargo check -p lancedb --features remote` passes with no errors and no new warnings. - Added two regression tests in `rust/lancedb/src/job.rs` (`decode_missing_result_uses_feature_appropriate_error`, `decode_unparsable_result_uses_feature_appropriate_error`) asserting the error variant matches the feature configuration. The test code compiles under the fixed configuration (`cargo check --tests` passes), and the branching logic was executed standalone against both feature cfgs (3/3 green with and without `remote`). The full `cargo test` binary link could not complete on the small build VM available here (the linker stalled), so CI coverage of the new tests is appreciated. - `cargo fmt --all -- --check` is clean. ## Disclosure AI assistance was used while preparing this change (code search, editing, and verification commands were agent-driven); the diagnosis, fix design, and all test results above were verified by actually running the builds. --- rust/lancedb/src/job.rs | 80 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 9 deletions(-) diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 711a73033..6865e59ee 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -101,21 +101,19 @@ impl TerminalResult { fn decode(self) -> Result { let value = self.value.ok_or_else(|| match &self.request_id { - Some(request_id) => Error::Http { - source: "successful typed job response did not contain a result".into(), - request_id: request_id.clone(), - status_code: None, - }, + Some(request_id) => remote_job_error( + request_id.clone(), + "successful typed job response did not contain a result".to_string(), + ), None => Error::Runtime { message: "successful typed job did not contain a result".to_string(), }, })?; serde_json::from_value(value).map_err(|error| match self.request_id { - Some(request_id) => Error::Http { - source: format!("failed to parse typed job result: {error}").into(), + Some(request_id) => remote_job_error( request_id, - status_code: None, - }, + format!("failed to parse typed job result: {error}"), + ), None => Error::Runtime { message: format!("failed to parse typed job result: {error}"), }, @@ -123,6 +121,28 @@ impl TerminalResult { } } +/// Build the error for a failed remote typed-job result. +/// +/// The `Error::Http` variant is gated behind the `remote` feature, so without +/// it we fall back to `Error::Runtime`. A `request_id` is only ever `Some` +/// for remote jobs, making the fallback unreachable in practice without the +/// feature. +#[cfg(feature = "remote")] +fn remote_job_error(request_id: String, source: String) -> Error { + Error::Http { + source: source.into(), + request_id, + status_code: None, + } +} + +#[cfg(not(feature = "remote"))] +fn remote_job_error(request_id: String, source: String) -> Error { + Error::Runtime { + message: format!("{source} (request_id={request_id})"), + } +} + type ResultDecoder = Arc Result + Send + Sync>; enum JobInner { @@ -634,4 +654,46 @@ mod tests { assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. }))); assert_eq!(job.status().await.unwrap(), "cancelled"); } + + #[test] + fn decode_missing_result_uses_feature_appropriate_error() { + // A `request_id` is only `Some` for remote jobs. The error for a + // missing result must use `Error::Http` with the `remote` feature and + // fall back to `Error::Runtime` without it (where `Error::Http` + // does not exist). See https://github.com/lancedb/lancedb/issues/4142. + let result = TerminalResult { + value: None, + request_id: Some("req-1".to_string()), + }; + let err = result.decode::().unwrap_err(); + #[cfg(feature = "remote")] + assert!( + matches!(err, Error::Http { .. }), + "expected Error::Http, got {err:?}" + ); + #[cfg(not(feature = "remote"))] + assert!( + matches!(err, Error::Runtime { .. }), + "expected Error::Runtime, got {err:?}" + ); + } + + #[test] + fn decode_unparsable_result_uses_feature_appropriate_error() { + let result = TerminalResult { + value: Some(Value::String("not a u64".to_string())), + request_id: Some("req-1".to_string()), + }; + let err = result.decode::().unwrap_err(); + #[cfg(feature = "remote")] + assert!( + matches!(err, Error::Http { .. }), + "expected Error::Http, got {err:?}" + ); + #[cfg(not(feature = "remote"))] + assert!( + matches!(err, Error::Runtime { .. }), + "expected Error::Runtime, got {err:?}" + ); + } } From 66b5b5bbdcb816b547927fb082f3f96925951e5e Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:26:13 +0800 Subject: [PATCH 30/33] fix(python): clarify unavailable CUDA index acceleration (#4249) ## Summary - preflight CUDA availability and initialization before accelerated index training - direct unavailable or misconfigured CUDA environments to the explicit CPU path - cover both legacy and unified Python index APIs ## Root cause LanceDB delegated accelerator requests directly to PyLance. When PyTorch lacked CUDA support or CUDA initialization failed, users received a low-level assertion or runtime error with no indication that omitting `accelerator` or setting `accelerator=None` creates the index on CPU. The repair validates CUDA before training and preserves explicit accelerator semantics without silently changing the requested backend. ## Scope This PR addresses only the separately reproduced CUDA availability and initialization error. It does not close #1766; the reported query failure after successful CUDA index creation remains open pending reproducible evidence. ## Validation - `uv run --extra tests pytest python/tests/test_table.py -k create_index -q` (10 passed) - `ruff format .` - `ruff check .` - CPU-only integration check confirms the actionable fallback message Related to #1766 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/table.py | 35 +++++++++++++++++++++-- python/python/tests/test_table.py | 47 ++++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index ff866d4db..35ec7a466 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -119,6 +119,33 @@ def _should_push_down_query_table( return namespace_client is not None and "QueryTable" in pushdown_operations +_CUDA_UNAVAILABLE_MESSAGE = ( + "CUDA acceleration is unavailable in this Python environment. Install a " + "CUDA-enabled PyTorch build and verify that the CUDA device is accessible, or " + "omit `accelerator` (or set `accelerator=None`) to create the index on CPU." +) + + +def _validate_index_accelerator(accelerator: Any) -> None: + if not isinstance(accelerator, str) or accelerator.split(":", 1)[0] != "cuda": + return + + try: + import torch + + cuda_available = torch.cuda.is_available() + except (ImportError, RuntimeError, AssertionError) as err: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err + + if not cuda_available: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) + + try: + torch.cuda.init() + except (RuntimeError, AssertionError) as err: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err + + def _polars_predicate_pushdown_barrier(frame: Any) -> Any: """Return a Polars frame unchanged while blocking predicate pushdown.""" return frame @@ -3600,7 +3627,9 @@ class LanceTable(Table): # Handle accelerator through pylance if accelerator is not None: - self.to_lance().create_index( + dataset = self.to_lance() + _validate_index_accelerator(accelerator) + dataset.create_index( column=column, index_type=index_type, metric=metric, @@ -3625,6 +3654,8 @@ class LanceTable(Table): acc = getattr(config, "accelerator", None) if acc is not None: # Dispatch to pylance for GPU acceleration + dataset = self.to_lance() + _validate_index_accelerator(acc) index_type_map = { "IvfFlat": "IVF_FLAT", "IvfSq": "IVF_SQ", @@ -3636,7 +3667,7 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") - self.to_lance().create_index( + dataset.create_index( column=column, index_type=lance_index_type, metric=getattr(config, "distance_type", "l2"), diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 827f2ef58..b241adb0a 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -15,7 +15,7 @@ from datetime import date, datetime, timedelta from decimal import Decimal, localcontext from time import sleep from typing import List -from unittest.mock import patch +from unittest.mock import MagicMock, patch import lancedb from lancedb import table as table_module @@ -1773,6 +1773,51 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection): ) +@pytest.mark.parametrize("api", ["legacy", "unified"]) +@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"]) +def test_create_index_cuda_failure_suggests_cpu( + api, cuda_failure, mem_db: DBConnection +): + table = mem_db.create_table( + f"test_cuda_unavailable_{api}", + data=[{"vector": [3.1, 4.1]}, {"vector": [5.9, 26.5]}], + ) + torch = MagicMock() + torch.cuda.is_available.return_value = cuda_failure == "init_error" + if cuda_failure == "init_error": + torch.cuda.init.side_effect = RuntimeError("CUDA driver initialization failed") + + with ( + patch.dict(sys.modules, {"torch": torch}), + patch.object(table, "to_lance") as to_lance, + ): + if api == "legacy": + with ( + pytest.warns(DeprecationWarning, match="create_index"), + pytest.raises(RuntimeError, match="accelerator=None.*CPU"), + ): + table.create_index( + metric="dot", + num_partitions=1, + num_sub_vectors=1, + vector_column_name="vector", + accelerator="cuda", + ) + else: + with pytest.raises(RuntimeError, match="accelerator=None.*CPU"): + table.create_index( + "vector", + config=IvfPq( + distance_type="dot", + num_partitions=1, + num_sub_vectors=1, + accelerator="cuda", + ), + ) + + to_lance.return_value.create_index.assert_not_called() + + @patch("lancedb.table.AsyncTable.create_index") def test_create_index_name_and_train_parameters( mock_create_index, mem_db: DBConnection From 624defd109c161505e84b26e66f43a35b71ac16c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:26:28 +0800 Subject: [PATCH 31/33] fix: preserve sliced child validity during struct merge (#4261) ## Root cause The pinned Lance Arrow crate rebuilt a sliced struct child with its original Arrow offset but supplied a validity bitmap containing only the visible rows. Arrow requires the bitmap to cover the offset as well, so filtered vector searches involving nullable nested fields could panic with `null_bit_buffer size too small`. ## Fix - Carry an offset-aligned validity bitmap when propagating parent nulls to a sliced child. - Vendor the pinned `lance-arrow` crate until a Lance release includes the fix; include it in Maturin source distributions. - Add unit regressions for sliced children and parent bitmaps, plus a record-batch merge regression. ## Validation - The issue's 8,600-row Python reproducer completes its filtered flat dot search and returns 8,448 rows without a panic. - The targeted Rust merge regression and two focused Python query tests pass. - Rust formatting, Ruff formatting and lint, the vendored crate tests, and the workspace Cargo checks passed. The Maturin source archive includes the patched crate. Fixes #4260 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- Cargo.lock | 2 +- Cargo.toml | 7 +- python/Cargo.toml | 2 + python/vendor/lance-arrow/Cargo.toml | 33 + python/vendor/lance-arrow/README.lancedb.md | 7 + python/vendor/lance-arrow/README.md | 6 + python/vendor/lance-arrow/src/bfloat16.rs | 406 +++ python/vendor/lance-arrow/src/deepcopy.rs | 356 +++ python/vendor/lance-arrow/src/floats.rs | 292 ++ python/vendor/lance-arrow/src/ipc.rs | 621 ++++ python/vendor/lance-arrow/src/json.rs | 1538 +++++++++ python/vendor/lance-arrow/src/lib.rs | 3102 +++++++++++++++++++ python/vendor/lance-arrow/src/list.rs | 177 ++ python/vendor/lance-arrow/src/memory.rs | 92 + python/vendor/lance-arrow/src/scalar.rs | 302 ++ python/vendor/lance-arrow/src/schema.rs | 208 ++ python/vendor/lance-arrow/src/stream.rs | 554 ++++ python/vendor/lance-arrow/src/struct.rs | 227 ++ rust/lancedb/tests/nullable_struct_merge.rs | 64 + 19 files changed, 7994 insertions(+), 2 deletions(-) create mode 100644 python/vendor/lance-arrow/Cargo.toml create mode 100644 python/vendor/lance-arrow/README.lancedb.md create mode 100644 python/vendor/lance-arrow/README.md create mode 100644 python/vendor/lance-arrow/src/bfloat16.rs create mode 100644 python/vendor/lance-arrow/src/deepcopy.rs create mode 100644 python/vendor/lance-arrow/src/floats.rs create mode 100644 python/vendor/lance-arrow/src/ipc.rs create mode 100644 python/vendor/lance-arrow/src/json.rs create mode 100644 python/vendor/lance-arrow/src/lib.rs create mode 100644 python/vendor/lance-arrow/src/list.rs create mode 100644 python/vendor/lance-arrow/src/memory.rs create mode 100644 python/vendor/lance-arrow/src/scalar.rs create mode 100644 python/vendor/lance-arrow/src/schema.rs create mode 100644 python/vendor/lance-arrow/src/stream.rs create mode 100644 python/vendor/lance-arrow/src/struct.rs create mode 100644 rust/lancedb/tests/nullable_struct_merge.rs diff --git a/Cargo.lock b/Cargo.lock index 7f36dc088..34ee05cb1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5136,7 +5136,6 @@ dependencies = [ [[package]] name = "lance-arrow" version = "15.0.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5834,6 +5833,7 @@ dependencies = [ "datafusion-common", "env_logger", "futures", + "lance-arrow", "lance-core", "lance-io", "lance-namespace", diff --git a/Cargo.toml b/Cargo.toml index 3ac246149..870ef5d07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ lance-datagen = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" lance-file = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-io = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-index = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-jsonb = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" } lance-linalg = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-namespace = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-namespace-impls = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } @@ -26,7 +27,7 @@ lance-table = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = lance-testing = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-datafusion = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-encoding = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { version = "=15.0.0-beta.1", path = "python/vendor/lance-arrow" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow @@ -79,6 +80,10 @@ tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] } uuid = { version = "1.7.0", features = ["v4", "v7"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } +[patch."https://github.com/lance-format/lance.git"] +# Carry the offset-safe validity fix until it is available in a Lance release. +lance-arrow = { path = "python/vendor/lance-arrow" } + [patch.crates-io] # Lance v14.0.0-beta.18 requires multipart write modes from this upstream revision. object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" } diff --git a/python/Cargo.toml b/python/Cargo.toml index ad1790a99..076299a7d 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -22,6 +22,8 @@ bytes.workspace = true lancedb.workspace = true datafusion-common.workspace = true lance-core.workspace = true +# The workspace path dependency includes the patched crate in Maturin sources. +lance-arrow.workspace = true lance-namespace.workspace = true lance-namespace-impls.workspace = true lance-io.workspace = true diff --git a/python/vendor/lance-arrow/Cargo.toml b/python/vendor/lance-arrow/Cargo.toml new file mode 100644 index 000000000..9cb156f2e --- /dev/null +++ b/python/vendor/lance-arrow/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "lance-arrow" +version = "15.0.0-beta.1" +edition = "2024" +authors = ["Lance Devs "] +license = "Apache-2.0" +repository = "https://github.com/lance-format/lance" +readme = "README.md" +description = "Arrow Extension for Lance" +keywords = ["data-format", "data-science", "machine-learning", "apache-arrow", "data-analytics"] +categories = ["database-implementations"] + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +arrow-array.workspace = true +arrow-buffer.workspace = true +arrow-data.workspace = true +arrow-ipc = { workspace = true, features = ["zstd"] } +arrow-ord.workspace = true +arrow-schema.workspace = true +arrow-select.workspace = true +bytemuck = { version = "1", default-features = false, features = ["extern_crate_alloc"] } +bytes.workspace = true +futures.workspace = true +half = { workspace = true, features = ["std", "bytemuck"] } +lance-jsonb.workspace = true +num-traits.workspace = true +rand = { workspace = true, features = ["small_rng"] } + +[target.'cfg(target_arch = "wasm32")'.dependencies] +getrandom = { version = "0.2", features = ["js"] } diff --git a/python/vendor/lance-arrow/README.lancedb.md b/python/vendor/lance-arrow/README.lancedb.md new file mode 100644 index 000000000..5ed714e08 --- /dev/null +++ b/python/vendor/lance-arrow/README.lancedb.md @@ -0,0 +1,7 @@ +This directory contains `rust/lance-arrow` from Lance tag `v15.0.0-beta.1` +(commit `46a7cfe1fba5efa1f90deea4c7b14d414b991409`), licensed under Apache-2.0. + +LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps +until a Lance release includes the fix. The patch is selected in the root +`Cargo.toml`; the rest of Lance stays on the pinned release. An explicit unsafe +block in `src/bfloat16.rs` lets this copy pass LanceDB's Rust 2024 warning gate. diff --git a/python/vendor/lance-arrow/README.md b/python/vendor/lance-arrow/README.md new file mode 100644 index 000000000..e0f129a3a --- /dev/null +++ b/python/vendor/lance-arrow/README.md @@ -0,0 +1,6 @@ +# lance-arrow + +`lance-arrow` is an internal sub-crate, containing [Apache-Arrow](https://github.com/apache/arrow-rs) +extensions used by [Lance](https://github.com/lance-format/lance). + +**Important Note**: This crate is **not intended for external usage**. diff --git a/python/vendor/lance-arrow/src/bfloat16.rs b/python/vendor/lance-arrow/src/bfloat16.rs new file mode 100644 index 000000000..1a60e5ece --- /dev/null +++ b/python/vendor/lance-arrow/src/bfloat16.rs @@ -0,0 +1,406 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! bfloat16 support for Apache Arrow. + +use std::fmt::Formatter; +use std::slice; + +use arrow_array::{Array, FixedSizeBinaryArray, builder::BooleanBufferBuilder}; +use arrow_buffer::{Buffer, MutableBuffer}; +use arrow_data::ArrayData; +use arrow_schema::{ArrowError, DataType, Field as ArrowField}; +use half::bf16; + +use crate::{ARROW_EXT_NAME_KEY, FloatArray}; + +/// The name of the bfloat16 extension in Arrow metadata +pub const BFLOAT16_EXT_NAME: &str = "lance.bfloat16"; + +/// Check whether the given field is a bfloat16 field +/// +/// A field is a bfloat16 field if it has a data type of `FixedSizeBinary(2)` and the metadata +/// contains the bfloat16 extension name. +pub fn is_bfloat16_field(field: &ArrowField) -> bool { + field.data_type() == &DataType::FixedSizeBinary(2) + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == BFLOAT16_EXT_NAME) + .unwrap_or_default() +} + +/// The bfloat16 data type +/// +/// This implements the [`ArrowFloatType`](crate::floats::ArrowFloatType) trait for bfloat16 values. +#[derive(Debug)] +pub struct BFloat16Type {} + +/// An array of bfloat16 values +/// +/// Note that bfloat16 is not the same thing as fp16 which is supported natively by arrow-rs. +#[derive(Clone)] +pub struct BFloat16Array { + inner: FixedSizeBinaryArray, +} + +impl std::fmt::Debug for BFloat16Array { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "BFloat16Array\n[\n")?; + from_arrow::print_long_array(&self.inner, f, |array, i, f| { + if array.is_null(i) { + write!(f, "null") + } else { + let binary_values = array.value(i); + let value = + bf16::from_bits(u16::from_le_bytes([binary_values[0], binary_values[1]])); + write!(f, "{:?}", value) + } + })?; + write!(f, "]") + } +} + +impl BFloat16Array { + pub fn from_iter_values(iter: impl IntoIterator) -> Self { + let values: Vec = iter.into_iter().collect(); + values.into() + } + + pub fn len(&self) -> usize { + self.inner.len() + } + + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + pub fn is_null(&self, i: usize) -> bool { + self.inner.is_null(i) + } + + pub fn null_count(&self) -> usize { + self.inner.null_count() + } + + pub fn iter(&self) -> BFloat16Iter<'_> { + BFloat16Iter { + array: self, + index: 0, + } + } + + pub fn value(&self, i: usize) -> bf16 { + assert!( + i < self.len(), + "Trying to access an element at index {} from a BFloat16Array of length {}", + i, + self.len() + ); + // Safety: + // `i < self.len() + unsafe { self.value_unchecked(i) } + } + + /// # Safety + /// Caller must ensure that `i < self.len()` + pub unsafe fn value_unchecked(&self, i: usize) -> bf16 { + // SAFETY: the caller guarantees that `i` is within this array's bounds. + let binary_value = unsafe { self.inner.value_unchecked(i) }; + bf16::from_bits(u16::from_le_bytes([binary_value[0], binary_value[1]])) + } + + pub fn into_inner(self) -> FixedSizeBinaryArray { + self.inner + } +} + +impl FromIterator> for BFloat16Array { + fn from_iter>>(iter: I) -> Self { + let mut buffer = MutableBuffer::new(10); + // No null buffer builder :( + let mut nulls = BooleanBufferBuilder::new(10); + let mut len = 0; + + for maybe_value in iter { + if let Some(value) = maybe_value { + let bytes = value.to_le_bytes(); + buffer.extend(bytes); + } else { + buffer.extend([0u8, 0u8]); + } + nulls.append(maybe_value.is_some()); + len += 1; + } + + let null_buffer = nulls.finish(); + let num_valid = null_buffer.count_set_bits(); + let null_buffer = if num_valid == len { + None + } else { + Some(null_buffer.into_inner()) + }; + + let array_data = ArrayData::builder(DataType::FixedSizeBinary(2)) + .len(len) + .add_buffer(buffer.into()) + .null_bit_buffer(null_buffer); + // SAFETY: the value buffer contains exactly `2 * len` bytes (two bytes + // pushed per iteration of the loop above, including the zero-fill for + // null slots), which matches the `FixedSizeBinary(2)` storage layout. + // The null bit buffer, when present, has `len` bits appended above, so + // its length covers the array's logical range. + let array_data = unsafe { array_data.build_unchecked() }; + Self { + inner: FixedSizeBinaryArray::from(array_data), + } + } +} + +impl FromIterator for BFloat16Array { + fn from_iter>(iter: I) -> Self { + Self::from_iter_values(iter) + } +} + +impl From> for BFloat16Array { + fn from(data: Vec) -> Self { + let len = data.len(); + // Zero-copy: `bf16` is `#[repr(transparent)]` over `u16` and derives + // `bytemuck::Pod`, so `cast_vec` reinterprets the allocation in place — + // no per-element copy or heap alloc. The crate-root `compile_error!` + // pins `target_endian = "little"`, so the resulting bytes match the + // `FixedSizeBinary(2)` on-disk order Lance writes elsewhere. + let raw: Vec = bytemuck::cast_vec(data); + let array_data = ArrayData::builder(DataType::FixedSizeBinary(2)) + .len(len) + .add_buffer(Buffer::from_vec(raw)); + // SAFETY: the value buffer contains exactly `2 * len` bytes — one + // `u16` per element after the layout-compatible cast — matching the + // `FixedSizeBinary(2)` storage layout. No null buffer is attached, so + // every element is logically valid. + let array_data = unsafe { array_data.build_unchecked() }; + Self { + inner: FixedSizeBinaryArray::from(array_data), + } + } +} + +impl TryFrom for BFloat16Array { + type Error = ArrowError; + + fn try_from(value: FixedSizeBinaryArray) -> Result { + if value.value_length() == 2 { + Ok(Self { inner: value }) + } else { + Err(ArrowError::InvalidArgumentError( + "FixedSizeBinaryArray must have a value length of 2".to_string(), + )) + } + } +} + +impl PartialEq for BFloat16Array { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +pub struct BFloat16Iter<'a> { + array: &'a BFloat16Array, + index: usize, +} + +impl<'a> Iterator for BFloat16Iter<'a> { + type Item = Option; + + fn next(&mut self) -> Option { + if self.index >= self.array.len() { + return None; + } + let i = self.index; + self.index += 1; + if self.array.is_null(i) { + Some(None) + } else { + Some(Some(self.array.value(i))) + } + } +} + +/// Methods that are lifted from arrow-rs temporarily until they are made public. +mod from_arrow { + use arrow_array::Array; + + /// Helper function for printing potentially long arrays. + pub(super) fn print_long_array( + array: &A, + f: &mut std::fmt::Formatter, + print_item: F, + ) -> std::fmt::Result + where + A: Array, + F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result, + { + let head = std::cmp::min(10, array.len()); + + for i in 0..head { + if array.is_null(i) { + writeln!(f, " null,")?; + } else { + write!(f, " ")?; + print_item(array, i, f)?; + writeln!(f, ",")?; + } + } + if array.len() > 10 { + if array.len() > 20 { + writeln!(f, " ...{} elements...,", array.len() - 20)?; + } + + let tail = std::cmp::max(head, array.len() - 10); + + for i in tail..array.len() { + if array.is_null(i) { + writeln!(f, " null,")?; + } else { + write!(f, " ")?; + print_item(array, i, f)?; + writeln!(f, ",")?; + } + } + } + Ok(()) + } +} + +impl FloatArray for FixedSizeBinaryArray { + type FloatType = BFloat16Type; + + /// Returns the underlying `bf16` values as a borrowed slice. + /// + /// # Preconditions + /// + /// - `value_length()` must be 2 (the `FixedSizeBinary(2)` storage shape + /// used by [`BFloat16Array`]). Asserted at entry. + /// - The value buffer must be at least 2-byte aligned. Lance's in-tree + /// constructors always satisfy this: value buffers are built either via + /// `MutableBuffer` (aligned to arrow-buffer's `ALIGNMENT` constant, â‰Ĩ32 + /// bytes) or via `Buffer::from_vec::` (aligned to `align_of::()` + /// == 2); both meet `bf16`'s 2-byte requirement. Externally-built + /// `FixedSizeBinaryArray`s arriving via FFI, IPC, or + /// `Buffer::from_custom_allocation` are not required by arrow-rs to be + /// aligned beyond a single byte; passing one to this method violates the + /// precondition. A `debug_assert` below catches such inputs in debug and + /// test builds. + /// + /// # Endianness + /// + /// `lance-arrow` is gated on `target_endian = "little"` at the crate root, + /// so this method always returns values in the same byte order Lance writes + /// (see [`BFloat16Array::value`] and the [`FromIterator`] impls). + fn as_slice(&self) -> &[bf16] { + assert_eq!( + self.value_length(), + 2, + "BFloat16 arrays must use FixedSizeBinary(2) storage" + ); + debug_assert_eq!( + (self.value_data().as_ptr() as usize) % std::mem::align_of::(), + 0, + "BFloat16 value buffer must be at least 2-byte aligned" + ); + // SAFETY: + // - The assert above pins `value_size == 2`, so `value_data().len() / 2` + // equals the array's logical element count. + // `FixedSizeBinaryArray::From` constructs its value buffer + // as `buffers[0].slice_with_length(offset * 2, len * 2)` (arrow-array + // `fixed_size_binary_array.rs`), so `value_data()` already returns + // the offset-adjusted slice. Do not replace `value_data()` with an + // accessor that returns the un-sliced backing buffer. + // - `bf16` is `#[repr(transparent)]` over `u16` (size 2, alignment 2); + // every `u16` bit pattern is a valid `bf16`, so any byte content + // yields a defined value — never UB. + // - Alignment is the caller's responsibility per the precondition + // documented above. The `debug_assert_eq!` immediately preceding this + // block catches violations in debug and test builds only — release + // builds rely on callers honoring the precondition. arrow-rs + // declares `FixedSizeBinary(n)`'s + // `BufferSpec::FixedWidth { alignment: align_of::() == 1 }` + // (arrow-data `data.rs`), so arrow-rs alone does not guarantee + // 2-byte alignment. Lance's in-tree construction paths build value + // buffers via `MutableBuffer` (arrow-buffer `ALIGNMENT` constant, + // â‰Ĩ32 bytes) or `Buffer::from_vec::` (2-byte aligned), both of + // which satisfy `bf16`'s 2-byte requirement. + // - The returned slice borrows from `self`; the underlying ref-counted, + // immutable Arrow buffer cannot be mutated or freed for the slice's + // lifetime. + unsafe { + slice::from_raw_parts( + self.value_data().as_ptr() as *const bf16, + self.value_data().len() / 2, + ) + } + } + + fn from_values(values: Vec) -> Self { + BFloat16Array::from(values).into_inner() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basics() { + let values: Vec = vec![1.0, 2.0, 3.0]; + let values: Vec = values.iter().map(|v| bf16::from_f32(*v)).collect(); + + let array = BFloat16Array::from_iter_values(values.clone()); + let array2 = BFloat16Array::from(values.clone()); + assert_eq!(array, array2); + assert_eq!(array.len(), 3); + + // Pin the raw little-endian bytes emitted by `From>` (rewritten to + // reinterpret the Vec via `bytemuck::cast_vec`), so a layout/byte-order + // regression is caught directly rather than only through Debug formatting. + // bf16 is the high 16 bits of the f32: 1.0->0x3F80, 2.0->0x4000, 3.0->0x4040. + let inner = array2.clone().into_inner(); + let raw_bytes: Vec = (0..inner.len()) + .flat_map(|i| inner.value(i).to_vec()) + .collect(); + assert_eq!(raw_bytes, vec![0x80, 0x3F, 0x00, 0x40, 0x40, 0x40]); + + let expected_fmt = "BFloat16Array\n[\n 1.0,\n 2.0,\n 3.0,\n]"; + assert_eq!(expected_fmt, format!("{:?}", array)); + + for (expected, value) in values.iter().zip(array.iter()) { + assert_eq!(Some(*expected), value); + } + + for (expected, value) in values.as_slice().iter().zip(array2.iter()) { + assert_eq!(Some(*expected), value); + } + + let arrow_array = array.into_inner(); + assert_eq!(arrow_array.as_slice(), values.as_slice()); + } + + #[test] + fn test_nulls() { + let values: Vec> = + vec![Some(bf16::from_f32(1.0)), None, Some(bf16::from_f32(3.0))]; + let array = BFloat16Array::from_iter(values.clone()); + assert_eq!(array.len(), 3); + assert_eq!(array.null_count(), 1); + + let expected_fmt = "BFloat16Array\n[\n 1.0,\n null,\n 3.0,\n]"; + assert_eq!(expected_fmt, format!("{:?}", array)); + + for (expected, value) in values.iter().zip(array.iter()) { + assert_eq!(*expected, value); + } + } +} diff --git a/python/vendor/lance-arrow/src/deepcopy.rs b/python/vendor/lance-arrow/src/deepcopy.rs new file mode 100644 index 000000000..5df4339b5 --- /dev/null +++ b/python/vendor/lance-arrow/src/deepcopy.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_array::{Array, RecordBatch, make_array}; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer}; +use arrow_data::{ArrayData, ArrayDataBuilder, transform::MutableArrayData}; +use arrow_schema::DataType; + +pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer { + Buffer::from(buffer.as_slice()) +} + +pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option { + let nulls = nulls?; + let bit_buffer = deep_copy_buffer(nulls.inner().inner()); + // SAFETY: `null_count` is taken from the source `NullBuffer`, which already + // upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over + // the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice` + // adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances + // the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s + // `as_slice()` view from byte 0) reproduces the exact bit pattern at the same + // bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new` + // panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the + // copy has the same length, so that check still passes. + Some(unsafe { + NullBuffer::new_unchecked( + BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()), + nulls.null_count(), + ) + }) +} + +pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData { + let data_type = data.data_type().clone(); + let len = data.len(); + let nulls = deep_copy_nulls(data.nulls()); + let offset = data.offset(); + let buffers = data + .buffers() + .iter() + .map(deep_copy_buffer) + .collect::>(); + let child_data = data + .child_data() + .iter() + .map(deep_copy_array_data) + .collect::>(); + // SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract — + // `(data_type, len, offset, nulls, buffers, child_data)` must form a valid + // Arrow array. This call reproduces `data` structurally: `data_type`, `len`, + // and `offset` are forwarded unchanged; each buffer is replaced by a byte- + // identical copy of its offset-applied `as_slice()` view (the output buffer + // is `MutableBuffer`-allocated, at least as aligned as the source); `nulls` + // is deep-copied with the same bit offset/length and unset-bit count (see + // `deep_copy_nulls`); `child_data` is recursively cloned with the same + // guarantee. Every value-level invariant the source upheld — UTF-8 validity, + // monotonic offsets, in-bounds dictionary indices, run-end monotonicity, + // struct child-length matching — therefore transfers to the copy. If the + // source `ArrayData` was itself constructed via `new_unchecked` with an + // invalid payload, this function faithfully reproduces that invalidity. + unsafe { + ArrayDataBuilder::new(data_type) + .len(len) + .nulls(nulls) + .offset(offset) + .buffers(buffers) + .child_data(child_data) + .build_unchecked() + } +} + +pub fn deep_copy_array(array: &dyn Array) -> Arc { + let data = array.to_data(); + let data = deep_copy_array_data(&data); + make_array(data) +} + +pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result { + let arrays = batch + .columns() + .iter() + .map(|array| deep_copy_array(array)) + .collect::>(); + RecordBatch::try_new(batch.schema(), arrays) +} + +/// Deep copy array data, extracting only the sliced portion using MutableArrayData +/// This is the most efficient and correct way to copy just the sliced data +pub fn deep_copy_array_data_sliced(data: &ArrayData) -> ArrayData { + // Use MutableArrayData to efficiently copy just the slice + let mut mutable = MutableArrayData::new(vec![data], false, data.len()); + + // Which index this takes depends on the layout, because arrow's extenders + // disagree about who applies `ArrayData::offset()`. The bit-packed one adds + // it to the raw values buffer itself, so an offset-applied index there + // reads from twice the offset and runs off the end of the buffer. Every + // other layout either reads through an already-offset-applied view or + // forwards the index to children carrying their own offsets, and takes the + // absolute index it has always been given. + let start = match data.data_type() { + DataType::Boolean => 0, + _ => data.offset(), + }; + mutable.extend(0, start, start + data.len()); + + // Freeze into immutable ArrayData + mutable.freeze() +} + +/// Deep copy an array, extracting only the sliced portion using MutableArrayData +pub fn deep_copy_array_sliced(array: &dyn Array) -> Arc { + let data = array.to_data(); + let data = deep_copy_array_data_sliced(&data); + make_array(data) +} + +/// Deep copy a RecordBatch, extracting only the sliced portion using MutableArrayData +pub fn deep_copy_batch_sliced(batch: &RecordBatch) -> crate::Result { + let arrays = batch + .columns() + .iter() + .map(|array| deep_copy_array_sliced(array)) + .collect::>(); + RecordBatch::try_new(batch.schema(), arrays) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{Array, BooleanArray, Int32Array, RecordBatch, StringArray}; + use arrow_data::ArrayDataBuilder; + use arrow_schema::{DataType, Field, Schema}; + + #[test] + fn raw_sliced_fixed_size_list_data_keeps_its_child_offset() { + // `ArrayData::slice` records the offset on the parent and leaves the + // child whole -- a shape `FixedSizeListArray::slice` never produces, + // but a legal one this public helper can be handed directly. Arrow's + // extender forwards `start * size` to that unsliced child and adds + // nothing, so the index it gets has to be the absolute one; this is + // here to stop the boolean fix from being generalised over it. + let child = Int32Array::from(vec![10, 11, 20, 21]).to_data(); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let data = ArrayDataBuilder::new(DataType::FixedSizeList(field, 2)) + .len(2) + .add_child_data(child) + .build() + .unwrap(); + let sliced = data.slice(1, 1); + + let copied = super::deep_copy_array_data_sliced(&sliced); + let copied_child = Int32Array::from(copied.child_data()[0].clone()); + assert_eq!(copied_child.values().as_ref(), &[20, 21]); + } + + #[test] + fn raw_sliced_struct_of_fixed_size_list_reaches_the_right_values() { + // Slicing raw struct data pushes the window into the struct's children, + // and a fixed-size-list child takes it as its own parent offset with + // its values left whole. Two levels of offset, and the absolute index + // has to remain right through both of them. + let values = Int32Array::from(vec![10, 11, 20, 21, 30, 31]).to_data(); + let item = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list = ArrayDataBuilder::new(DataType::FixedSizeList(item, 2)) + .len(3) + .add_child_data(values) + .build() + .unwrap(); + let field = Arc::new(Field::new("l", list.data_type().clone(), false)); + let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into())) + .len(3) + .add_child_data(list) + .build() + .unwrap(); + let sliced = data.slice(2, 1); + + let copied = super::deep_copy_array_data_sliced(&sliced); + let copied_values = Int32Array::from(copied.child_data()[0].child_data()[0].clone()); + assert_eq!(copied_values.values().as_ref(), &[30, 31]); + } + + #[test] + fn raw_parent_offset_struct_keeps_the_selected_child_row() { + // A struct whose children are whole and whose window is the parent + // offset: arrow's struct extender forwards the index to those children + // untouched, so it has to be the absolute one. + let field = Arc::new(Field::new("a", DataType::Int32, false)); + let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into())) + .len(1) + .offset(1) + .add_child_data(Int32Array::from(vec![10, 20]).to_data()) + .build() + .unwrap(); + + let copied = super::deep_copy_array_data_sliced(&data); + let copied_child = Int32Array::from(copied.child_data()[0].clone()); + assert_eq!(copied_child.values().as_ref(), &[20]); + } + + #[test] + fn sliced_boolean_deep_copy_reads_from_the_slice() { + // A boolean slice keeps its `ArrayData::offset()` -- the buffer cannot + // advance by a fraction of a byte -- so a copy that adds the offset on + // top of what `MutableArrayData` already applies reads past the end of + // the values buffer and panics. + let array = BooleanArray::from(vec![true, false, true, false, true, false, true, false]); + for (offset, len) in [(0usize, 8usize), (1, 7), (3, 5), (7, 1)] { + let sliced = array.slice(offset, len); + let copied = super::deep_copy_array_sliced(&sliced); + let copied = copied.as_any().downcast_ref::().unwrap(); + let expected: Vec = (0..len).map(|i| sliced.value(i)).collect(); + let actual: Vec = (0..len).map(|i| copied.value(i)).collect(); + assert_eq!(actual, expected, "offset={offset} len={len}"); + } + } + + #[test] + fn sliced_boolean_deep_copy_keeps_its_nulls() { + let array = BooleanArray::from(vec![ + Some(true), + None, + Some(false), + Some(true), + None, + Some(false), + ]); + let sliced = array.slice(1, 4); + let copied = super::deep_copy_array_sliced(&sliced); + let copied = copied.as_any().downcast_ref::().unwrap(); + let expected: Vec> = (0..sliced.len()) + .map(|i| (!sliced.is_null(i)).then(|| sliced.value(i))) + .collect(); + let actual: Vec> = (0..copied.len()) + .map(|i| (!copied.is_null(i)).then(|| copied.value(i))) + .collect(); + assert_eq!(actual, expected); + } + + #[test] + fn test_deep_copy_sliced_array_with_nulls() { + let array = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + ])); + let sliced_array = array.slice(1, 3); + let copied_array = super::deep_copy_array(&sliced_array); + assert_eq!(sliced_array.len(), copied_array.len()); + assert_eq!(sliced_array.nulls(), copied_array.nulls()); + } + + #[test] + fn test_deep_copy_array_data_sliced() { + let array = Int32Array::from((0..1000).collect::>()); + let sliced = array.slice(100, 10); + + let sliced_data = sliced.to_data(); + let copied_data = super::deep_copy_array_data_sliced(&sliced_data); + + assert_eq!(copied_data.len(), 10); + assert_eq!(copied_data.offset(), 0); + + // Verify data correctness + let copied_array = Int32Array::from(copied_data); + for i in 0..10 { + assert_eq!(copied_array.value(i), 100 + i as i32); + } + } + + #[test] + fn test_deep_copy_array_sliced() { + let array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + let sliced = array.slice(1, 3); + + let copied = super::deep_copy_array_sliced(&sliced); + + assert_eq!(copied.len(), 3); + let copied_int = copied.as_any().downcast_ref::().unwrap(); + assert_eq!(copied_int.value(0), 2); + assert_eq!(copied_int.value(1), 3); + assert_eq!(copied_int.value(2), 4); + } + + #[test] + fn test_deep_copy_batch_sliced() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ])); + + let id_array = Arc::new(Int32Array::from((0..100).collect::>())); + let name_array = Arc::new(StringArray::from( + (0..100) + .map(|i| format!("name_{}", i)) + .collect::>(), + )); + + let batch = RecordBatch::try_new( + schema, + vec![id_array as Arc, name_array as Arc], + ) + .unwrap(); + + let sliced = batch.slice(10, 5); + let copied = super::deep_copy_batch_sliced(&sliced).unwrap(); + + assert_eq!(copied.num_rows(), 5); + assert_eq!(copied.num_columns(), 2); + + // Verify data correctness + let id_col = copied + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let name_col = copied + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + for i in 0..5 { + assert_eq!(id_col.value(i), 10 + i as i32); + assert_eq!(name_col.value(i), format!("name_{}", 10 + i)); + } + } + + #[test] + fn test_deep_copy_array_sliced_with_nulls() { + let array = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + ])); + let sliced = array.slice(1, 3); // [None, Some(3), None] + + let copied = super::deep_copy_array_sliced(&sliced); + + assert_eq!(copied.len(), 3); + assert_eq!(copied.null_count(), 2); // Two nulls in the slice + + let copied_int = copied.as_any().downcast_ref::().unwrap(); + assert!(!copied_int.is_valid(0)); // None + assert!(copied_int.is_valid(1)); // Some(3) + assert!(!copied_int.is_valid(2)); // None + assert_eq!(copied_int.value(1), 3); + } +} diff --git a/python/vendor/lance-arrow/src/floats.rs b/python/vendor/lance-arrow/src/floats.rs new file mode 100644 index 000000000..0024fbb02 --- /dev/null +++ b/python/vendor/lance-arrow/src/floats.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Floats Array + +use std::fmt::{Debug, Display}; +use std::iter::Sum; +use std::sync::Arc; +use std::{ + fmt::Formatter, + ops::{AddAssign, DivAssign}, +}; + +use arrow_array::{ + Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array, + types::{Float16Type, Float32Type, Float64Type}, +}; +use arrow_schema::{DataType, Field}; +use half::{bf16, f16}; +use num_traits::{AsPrimitive, Bounded, Float, FromPrimitive}; + +use super::bfloat16::{BFloat16Array, BFloat16Type}; +use crate::Result; +use crate::bfloat16::is_bfloat16_field; + +/// Float data type. +/// +/// This helps differentiate between the different float types, +/// because bf16 is not officially supported [DataType] in arrow-rs. +#[derive(Debug)] +pub enum FloatType { + BFloat16, + Float16, + Float32, + Float64, +} + +impl std::fmt::Display for FloatType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::BFloat16 => write!(f, "bfloat16"), + Self::Float16 => write!(f, "float16"), + Self::Float32 => write!(f, "float32"), + Self::Float64 => write!(f, "float64"), + } + } +} + +/// Try to convert a [DataType] to a [FloatType]. To support bfloat16, always +/// prefer using the `TryFrom<&Field>` implementation. +impl TryFrom<&DataType> for FloatType { + type Error = crate::ArrowError; + + fn try_from(value: &DataType) -> Result { + match *value { + DataType::Float16 => Ok(Self::Float16), + DataType::Float32 => Ok(Self::Float32), + DataType::Float64 => Ok(Self::Float64), + _ => Err(crate::ArrowError::InvalidArgumentError(format!( + "{:?} is not a floating type", + value + ))), + } + } +} + +impl TryFrom<&Field> for FloatType { + type Error = crate::ArrowError; + + fn try_from(field: &Field) -> Result { + match field.data_type() { + DataType::FixedSizeBinary(2) if is_bfloat16_field(field) => Ok(Self::BFloat16), + _ => Self::try_from(field.data_type()), + } + } +} + +/// Trait for float types used in Lance indexes +/// +/// This mimics the utilities provided by [`arrow_array::ArrowPrimitiveType`] +/// but applies to all float types (including bfloat16) +pub trait ArrowFloatType: Debug { + type Native: FromPrimitive + + FloatToArrayType + + AsPrimitive + + Debug + + Display; + + const FLOAT_TYPE: FloatType; + const MIN: Self::Native; + const MAX: Self::Native; + + /// Arrow Float Array Type. + type ArrayType: FloatArray; + + /// Returns empty array of this type. + fn empty_array() -> Self::ArrayType { + >::from_values(Vec::new()) + } +} + +/// Trait to be implemented by native types that have a corresponding [`ArrowFloatType`] +/// implementation. +/// +/// This helps define what operations are supported by native floats and also helps convert +/// from a native type back to the corresponding Arrow float type. +pub trait FloatToArrayType: + Float + + Bounded + + Sum + + AddAssign + + AsPrimitive + + AsPrimitive + + DivAssign + + Send + + Sync + + Copy +{ + /// The corresponding [`ArrowFloatType`] implementation for this native type + type ArrowType: ArrowFloatType; +} + +impl FloatToArrayType for bf16 { + type ArrowType = BFloat16Type; +} + +impl FloatToArrayType for f16 { + type ArrowType = Float16Type; +} + +impl FloatToArrayType for f32 { + type ArrowType = Float32Type; +} + +impl FloatToArrayType for f64 { + type ArrowType = Float64Type; +} + +impl ArrowFloatType for BFloat16Type { + type Native = bf16; + + const FLOAT_TYPE: FloatType = FloatType::BFloat16; + const MIN: Self::Native = bf16::MIN; + const MAX: Self::Native = bf16::MAX; + + type ArrayType = FixedSizeBinaryArray; +} + +impl ArrowFloatType for Float16Type { + type Native = f16; + + const FLOAT_TYPE: FloatType = FloatType::Float16; + const MIN: Self::Native = f16::MIN; + const MAX: Self::Native = f16::MAX; + + type ArrayType = Float16Array; +} + +impl ArrowFloatType for Float32Type { + type Native = f32; + + const FLOAT_TYPE: FloatType = FloatType::Float32; + const MIN: Self::Native = f32::MIN; + const MAX: Self::Native = f32::MAX; + + type ArrayType = Float32Array; +} + +impl ArrowFloatType for Float64Type { + type Native = f64; + + const FLOAT_TYPE: FloatType = FloatType::Float64; + const MIN: Self::Native = f64::MIN; + const MAX: Self::Native = f64::MAX; + + type ArrayType = Float64Array; +} + +/// [FloatArray] is a trait that is implemented by all float type arrays +/// +/// This is similar to [`arrow_array::PrimitiveArray`] but applies to all float types (including bfloat16) +/// and is implemented as a trait and not a struct +pub trait FloatArray: Array + Clone + 'static { + type FloatType: ArrowFloatType; + + /// Returns a reference to the underlying data as a slice. + /// + /// # Panics + /// + /// Implementations may panic if the array's storage shape does not match + /// the expected element layout. In particular, the `bf16` impl panics if + /// `value_length() != 2` (the `FixedSizeBinary(2)` shape required by + /// `BFloat16Array`). + /// + /// # Preconditions + /// + /// Implementations may impose additional invariants on the underlying + /// buffer. The `bf16` impl requires the value buffer to be at least + /// 2-byte aligned — satisfied automatically by every in-tree Lance + /// constructor, but external callers passing externally-built arrays + /// (FFI, IPC, `Buffer::from_custom_allocation`) must ensure alignment. + /// See the impl's docstring for details. + fn as_slice(&self) -> &[T::Native]; + + /// Construct an array from a vector of values. + fn from_values(values: Vec) -> Self; + + /// Construct an array from an iterator of values. + fn from_iter_values(values: impl IntoIterator) -> Self + where + Self: Sized, + { + Self::from_values(values.into_iter().collect()) + } +} + +impl FloatArray for Float16Array { + type FloatType = Float16Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +impl FloatArray for Float32Array { + type FloatType = Float32Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +impl FloatArray for Float64Array { + type FloatType = Float64Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +/// Convert a float32 array to another float array +/// +/// This is used during queries as query vectors are always provided as float32 arrays +/// and need to be converted to the appropriate float type for the index. +pub fn coerce_float_vector(input: &Float32Array, float_type: FloatType) -> Result> { + match float_type { + FloatType::BFloat16 => Ok(Arc::new( + BFloat16Array::from_iter_values(input.values().iter().map(|v| bf16::from_f32(*v))) + .into_inner(), + )), + FloatType::Float16 => Ok(Arc::new(Float16Array::from_iter_values( + input.values().iter().map(|v| f16::from_f32(*v)), + ))), + FloatType::Float32 => Ok(Arc::new(input.clone())), + FloatType::Float64 => Ok(Arc::new(Float64Array::from_iter_values( + input.values().iter().map(|v| *v as f64), + ))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_coerce_float_vector_bfloat16() { + let input = Float32Array::from(vec![1.0f32, 2.0, 3.0]); + let array = coerce_float_vector(&input, FloatType::BFloat16).unwrap(); + + assert_eq!(array.data_type(), &DataType::FixedSizeBinary(2)); + + let fixed = array + .as_any() + .downcast_ref::() + .unwrap(); + let expected: Vec = input.values().iter().map(|v| bf16::from_f32(*v)).collect(); + assert_eq!(fixed.as_slice(), expected.as_slice()); + } +} diff --git a/python/vendor/lance-arrow/src/ipc.rs b/python/vendor/lance-arrow/src/ipc.rs new file mode 100644 index 000000000..099c362b4 --- /dev/null +++ b/python/vendor/lance-arrow/src/ipc.rs @@ -0,0 +1,621 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Zero-copy Arrow IPC stream read/write utilities. +//! +//! Provides helpers for serializing and deserializing [`RecordBatch`]es as +//! self-delimiting Arrow IPC streams using synchronous [`Read`]/[`Write`] I/O. +//! +//! These are designed for embedding IPC streams inside larger binary formats +//! (e.g. a cache entry that contains multiple IPC sections). Each stream is +//! self-delimiting (schema + batches + EOS marker) and can be read back +//! independently. +//! +//! # Zero-copy reads +//! +//! [`read_ipc_stream`] and [`read_ipc_stream_single`] take `&Bytes` and use +//! [`Bytes::slice`] to produce each message buffer. Because `Bytes::slice` +//! increments a reference count rather than copying, the resulting +//! [`Buffer`]s — and the array data decoded from them by [`FileDecoder`] — +//! are all backed by the same allocation as the input. + +use std::io::{Read, Write}; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_buffer::Buffer; +use arrow_ipc::convert::fb_to_schema; +use arrow_ipc::reader::FileDecoder; +use arrow_ipc::root_as_message; +use arrow_ipc::writer::StreamWriter; +use arrow_schema::ArrowError; +use bytes::Bytes; + +// --------------------------------------------------------------------------- +// Length-prefixed byte utilities +// --------------------------------------------------------------------------- + +/// Write `data` prefixed by its length as a little-endian `u64`. +/// +/// Paired with [`read_len_prefixed_bytes`]. +pub fn write_len_prefixed_bytes(writer: &mut dyn Write, data: &[u8]) -> Result<(), ArrowError> { + writer + .write_all(&(data.len() as u64).to_le_bytes()) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + writer + .write_all(data) + .map_err(|e| ArrowError::IoError(e.to_string(), e)) +} + +/// Read a byte slice written by [`write_len_prefixed_bytes`]. +/// +/// Reads an 8-byte little-endian length then exactly that many bytes. +pub fn read_len_prefixed_bytes(reader: &mut dyn Read) -> Result, ArrowError> { + let mut len_buf = [0u8; 8]; + reader + .read_exact(&mut len_buf) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + let len = u64::from_le_bytes(len_buf) as usize; + let mut buf = vec![0u8; len]; + reader + .read_exact(&mut buf) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + Ok(buf) +} + +// --------------------------------------------------------------------------- +// IPC stream utilities +// --------------------------------------------------------------------------- + +// 4-byte continuation marker used by modern Arrow IPC streams. +const IPC_CONTINUATION: [u8; 4] = [0xff; 4]; + +/// Write `batch` as a single-batch Arrow IPC stream to `writer`. +pub fn write_ipc_stream(batch: &RecordBatch, writer: &mut dyn Write) -> Result<(), ArrowError> { + let mut sw = StreamWriter::try_new(&mut *writer, batch.schema_ref())?; + sw.write(batch)?; + sw.finish() +} + +/// Write all batches from `iter` as a single Arrow IPC stream to `writer`. +/// +/// `iter` must yield at least one batch; the schema is inferred from the first +/// batch. Returns `ArrowError::InvalidArgumentError` if the iterator is empty. +/// If you need to write an empty stream (schema only, no rows), construct a +/// `StreamWriter` directly. +pub fn write_ipc_stream_batches(iter: I, writer: &mut dyn Write) -> Result<(), ArrowError> +where + I: IntoIterator, +{ + let mut iter = iter.into_iter(); + let first = iter + .next() + .ok_or_else(|| ArrowError::InvalidArgumentError("no batches to serialize".into()))?; + let mut sw = StreamWriter::try_new(&mut *writer, first.schema_ref())?; + sw.write(&first)?; + for batch in iter { + sw.write(&batch)?; + } + sw.finish() +} + +/// Read one complete Arrow IPC stream message from `data` as a zero-copy [`Buffer`]. +/// +/// Parses the first message starting at byte 0 of `data`. Returns `None` on +/// EOS (size field == 0) or empty input. The returned [`Buffer`] is backed by +/// `data`'s allocation — no bytes are copied. +/// +/// The caller should advance its position by `buf.len()` after each call. +fn read_one_ipc_message(data: &Bytes) -> Result, ArrowError> { + let bytes = data.as_ref(); + + if bytes.is_empty() { + return Ok(None); + } + if bytes.len() < 4 { + return Err(ArrowError::IoError( + "IPC: truncated header".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC header"), + )); + } + + let has_continuation = bytes[..4] == IPC_CONTINUATION; + let (size_bytes, prefix_len): ([u8; 4], usize) = if has_continuation { + if bytes.len() < 8 { + return Err(ArrowError::IoError( + "IPC: truncated header after continuation".into(), + std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "truncated after continuation", + ), + )); + } + (bytes[4..8].try_into().unwrap(), 8) + } else { + (bytes[..4].try_into().unwrap(), 4) + }; + + let meta_size = u32::from_le_bytes(size_bytes) as usize; + if meta_size == 0 { + return Ok(None); // EOS + } + + let meta_end = prefix_len + meta_size; + if bytes.len() < meta_end { + return Err(ArrowError::IoError( + "IPC: truncated metadata".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC metadata"), + )); + } + + let msg = root_as_message(&bytes[prefix_len..meta_end]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + + let total = meta_end + body_len; + if bytes.len() < total { + return Err(ArrowError::IoError( + "IPC: truncated body".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC body"), + )); + } + + // Zero-copy: Bytes::slice shares the backing allocation; Buffer::from + // wraps it without copying. + Ok(Some(Buffer::from(data.slice(0..total)))) +} + +/// Read a length-prefixed byte slice at `offset` in `data`, advancing `offset`. +/// +/// Reads an 8-byte little-endian length, then slices exactly that many bytes +/// from `data`. The returned [`Bytes`] is zero-copy (shares `data`'s allocation). +pub fn read_len_prefixed_bytes_at(data: &Bytes, offset: &mut usize) -> Result { + let bytes = data.as_ref(); + let len_end = offset + .checked_add(8) + .filter(|&e| e <= bytes.len()) + .ok_or_else(|| { + ArrowError::IoError( + "length-prefixed bytes: truncated length field".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated length"), + ) + })?; + let len = u64::from_le_bytes(bytes[*offset..len_end].try_into().unwrap()) as usize; + *offset = len_end; + let data_end = offset + .checked_add(len) + .filter(|&e| e <= bytes.len()) + .ok_or_else(|| { + ArrowError::IoError( + "length-prefixed bytes: truncated data".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated data"), + ) + })?; + let result = data.slice(*offset..data_end); + *offset = data_end; + Ok(result) +} + +/// Read all [`RecordBatch`]es from one Arrow IPC stream starting at `offset`, +/// advancing `offset` past the stream (including the EOS marker). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation. +pub fn read_ipc_stream_at( + data: &Bytes, + offset: &mut usize, +) -> Result, ArrowError> { + let batches = read_ipc_stream(&data.slice(*offset..))?; + + // Recompute how many bytes were consumed by re-parsing message sizes. + // We can't get this from read_ipc_stream directly, so we re-walk the + // message headers (metadata only, no body re-read) to sum up lengths. + let slice = &data.as_ref()[*offset..]; + let mut consumed = 0usize; + loop { + let rem = &slice[consumed..]; + if rem.is_empty() { + break; + } + let has_cont = rem.len() >= 4 && rem[..4] == IPC_CONTINUATION; + let (size_bytes, prefix_len): ([u8; 4], usize) = if has_cont { + if rem.len() < 8 { + break; + } + (rem[4..8].try_into().unwrap(), 8) + } else { + if rem.len() < 4 { + break; + } + (rem[..4].try_into().unwrap(), 4) + }; + let meta_size = u32::from_le_bytes(size_bytes) as usize; + if meta_size == 0 { + // EOS — consume it and stop. + consumed += prefix_len; + break; + } + let meta_end = prefix_len + meta_size; + if rem.len() < meta_end { + break; + } + let msg = root_as_message(&rem[prefix_len..meta_end]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + consumed += meta_end + body_len; + } + *offset += consumed; + + Ok(batches) +} + +/// Read exactly one [`RecordBatch`] from one Arrow IPC stream starting at `offset`, +/// advancing `offset` past the stream (including the EOS marker). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation. +pub fn read_ipc_stream_single_at( + data: &Bytes, + offset: &mut usize, +) -> Result { + let mut batches = read_ipc_stream_at(data, offset)?; + match batches.len() { + 1 => Ok(batches.remove(0)), + n => Err(ArrowError::ParseError(format!( + "expected exactly 1 IPC record batch, got {n}" + ))), + } +} + +/// Extract the prefix length and metadata size from a raw IPC message buffer. +/// +/// Modern IPC streams have an 8-byte prefix `[continuation: 4][size: 4]`. +/// Legacy streams have a 4-byte prefix `[size: 4]`. Returns `(prefix_len, meta_size)`. +fn parse_ipc_message_prefix(buf: &Buffer) -> Result<(usize, usize), ArrowError> { + let has_continuation = buf.len() >= 4 && buf[..4] == IPC_CONTINUATION; + if has_continuation { + if buf.len() < 8 { + return Err(ArrowError::ParseError( + "IPC message buffer too short".into(), + )); + } + let meta_size = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize; + Ok((8, meta_size)) + } else { + if buf.len() < 4 { + return Err(ArrowError::ParseError( + "IPC message buffer too short".into(), + )); + } + let meta_size = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize; + Ok((4, meta_size)) + } +} + +/// Read all [`RecordBatch`]es from one Arrow IPC stream. +/// +/// Zero-copy: each batch's array data buffers are borrowed from the input +/// message buffer(s) and not copied during decoding. +/// +/// Uses [`FileDecoder`] directly (rather than `StreamDecoder`) to avoid a +/// known edge case where `StreamDecoder` does not produce a batch for messages +/// with a zero-length body when the message exactly fills the decode buffer. +pub fn read_ipc_stream(data: &Bytes) -> Result, ArrowError> { + let mut offset = 0usize; + + let schema_buf = read_one_ipc_message(&data.slice(offset..))?.ok_or_else(|| { + ArrowError::ParseError("IPC stream: expected schema message, got EOS".into()) + })?; + offset += schema_buf.len(); + + let (prefix_len, meta_size) = parse_ipc_message_prefix(&schema_buf)?; + let schema_msg = root_as_message(&schema_buf[prefix_len..prefix_len + meta_size]) + .map_err(|e| ArrowError::ParseError(format!("IPC schema parse error: {e}")))?; + let schema = Arc::new(fb_to_schema(schema_msg.header_as_schema().ok_or_else( + || ArrowError::ParseError("IPC stream: first message is not a schema".into()), + )?)); + let mut decoder = FileDecoder::new(schema, schema_msg.version()); + + let mut batches = Vec::new(); + + loop { + let Some(buf) = read_one_ipc_message(&data.slice(offset..))? else { + break; + }; + offset += buf.len(); + + let (prefix_len, meta_size) = parse_ipc_message_prefix(&buf)?; + let msg = root_as_message(&buf[prefix_len..prefix_len + meta_size]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + + // Block offset = 0 since the buffer starts at the message boundary. + // metaDataLength = prefix_len + meta_size (prefix + flatbuf + padding). + let block = arrow_ipc::Block::new(0, (prefix_len + meta_size) as i32, body_len as i64); + + match msg.header_type() { + arrow_ipc::MessageHeader::RecordBatch => { + if let Some(batch) = decoder.read_record_batch(&block, &buf)? { + batches.push(batch); + } + } + arrow_ipc::MessageHeader::DictionaryBatch => { + decoder.read_dictionary(&block, &buf)?; + } + _ => break, + } + } + + Ok(batches) +} + +/// Read exactly one [`RecordBatch`] from one Arrow IPC stream. +pub fn read_ipc_stream_single(data: &Bytes) -> Result { + let mut batches = read_ipc_stream(data)?; + match batches.len() { + 1 => Ok(batches.remove(0)), + n => Err(ArrowError::ParseError(format!( + "expected exactly 1 IPC record batch, got {n}" + ))), + } +} + +// --------------------------------------------------------------------------- +// Aligned IPC sections +// --------------------------------------------------------------------------- + +/// Byte alignment that each IPC section's stream start is padded to. +/// +/// When several IPC streams are concatenated into one larger blob (e.g. a +/// cache entry), a section that starts at an arbitrary offset would leave its +/// array data misaligned. [`FileDecoder`] with `require_alignment = false` +/// then silently copies each buffer into a freshly aligned allocation on +/// every read, defeating zero-copy. Padding each section start to a 64-byte +/// boundary keeps the decoded buffers borrowed directly from the input. +pub const IPC_SECTION_ALIGNMENT: usize = 64; + +/// Number of zero-padding bytes needed to advance `pos` to the next +/// [`IPC_SECTION_ALIGNMENT`] boundary. +fn section_padding(pos: usize) -> usize { + (IPC_SECTION_ALIGNMENT - (pos % IPC_SECTION_ALIGNMENT)) % IPC_SECTION_ALIGNMENT +} + +/// A [`Write`] adapter that counts the bytes written through it. +struct CountingWriter<'a> { + inner: &'a mut dyn Write, + count: usize, +} + +impl Write for CountingWriter<'_> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let n = self.inner.write(buf)?; + self.count += n; + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Write zero padding so the next byte lands on an [`IPC_SECTION_ALIGNMENT`] +/// boundary, advancing `pos` past it. +fn write_section_padding(writer: &mut dyn Write, pos: &mut usize) -> Result<(), ArrowError> { + let pad = section_padding(*pos); + if pad > 0 { + const ZEROS: [u8; IPC_SECTION_ALIGNMENT] = [0u8; IPC_SECTION_ALIGNMENT]; + writer + .write_all(&ZEROS[..pad]) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + *pos += pad; + } + Ok(()) +} + +/// Write `batch` as a 64-byte-aligned single-batch Arrow IPC section. +/// +/// `pos` is the absolute byte offset of `writer` within the enclosing blob. +/// Zero padding is written first so the IPC stream begins on an +/// [`IPC_SECTION_ALIGNMENT`] boundary, then the stream itself. `pos` is +/// advanced past both the padding and the stream so the caller can write +/// further aligned sections. +/// +/// Paired with [`read_ipc_section_at`]. For the decoded buffers to be borrowed +/// zero-copy, the blob must ultimately be read back from a buffer whose base +/// address is at least 64-byte aligned. +pub fn write_ipc_section( + writer: &mut dyn Write, + pos: &mut usize, + batch: &RecordBatch, +) -> Result<(), ArrowError> { + write_section_padding(writer, pos)?; + + let mut counting = CountingWriter { + inner: writer, + count: 0, + }; + write_ipc_stream(batch, &mut counting)?; + *pos += counting.count; + Ok(()) +} + +/// Read a single [`RecordBatch`] from an aligned IPC section at `offset`. +/// +/// Skips the alignment padding written by [`write_ipc_section`], then reads +/// the stream, advancing `offset` past the section (padding + stream + EOS). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base +/// address is at least 64-byte aligned (see [`write_ipc_section`]). +pub fn read_ipc_section_at(data: &Bytes, offset: &mut usize) -> Result { + *offset += section_padding(*offset); + read_ipc_stream_single_at(data, offset) +} + +/// Write `batches` as a single 64-byte-aligned multi-batch Arrow IPC section. +/// +/// Like [`write_ipc_section`] but emits every batch from `iter` into one IPC +/// stream (schema + N batches + EOS). `iter` must yield at least one batch. +/// Paired with [`read_ipc_section_batches_at`]. +pub fn write_ipc_section_batches( + writer: &mut dyn Write, + pos: &mut usize, + iter: I, +) -> Result<(), ArrowError> +where + I: IntoIterator, +{ + write_section_padding(writer, pos)?; + + let mut counting = CountingWriter { + inner: writer, + count: 0, + }; + write_ipc_stream_batches(iter, &mut counting)?; + *pos += counting.count; + Ok(()) +} + +/// Read all [`RecordBatch`]es from an aligned multi-batch IPC section at +/// `offset`, advancing `offset` past the section (padding + stream + EOS). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base +/// address is at least 64-byte aligned (see [`write_ipc_section_batches`]). +pub fn read_ipc_section_batches_at( + data: &Bytes, + offset: &mut usize, +) -> Result, ArrowError> { + *offset += section_padding(*offset); + read_ipc_stream_at(data, offset) +} + +#[cfg(test)] +mod tests { + use arrow_array::{ArrayRef, record_batch}; + + use super::*; + + #[test] + fn test_ipc_roundtrip() { + let batch1 = record_batch!( + ("int", Int32, [1, 2, 3]), + ("str", Utf8, ["foo", "bar", "baz"]) + ) + .unwrap(); + let batch2 = record_batch!(("int", Int32, [4, 5]), ("str", Utf8, ["qux", "quux"])).unwrap(); + let batches = vec![batch1.clone(), batch2.clone()]; + + let mut buf = Vec::new(); + write_ipc_stream_batches(batches, &mut buf).unwrap(); + + let data = Bytes::from(buf); + + let batches = read_ipc_stream(&data).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0], batch1); + assert_eq!(batches[1], batch2); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + let assert_col_zero_copy = |array: &ArrayRef| { + for buffer in array.to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "buffer at {ptr:#x} is not backed by the input Bytes allocation \ + [{data_base:#x}..{data_end:#x})" + ); + } + }; + + for batch in &batches { + assert_eq!(batch.schema(), batch1.schema()); + assert_col_zero_copy(batch.column(0)); + assert_col_zero_copy(batch.column(1)); + } + } + + /// Allocate a [`Bytes`] whose base address is 64-byte aligned, modelling a + /// backend that reads cache entries into an aligned buffer. A plain + /// `Bytes::from(vec)` only guarantees the allocator's alignment for `u8`. + fn aligned_bytes(payload: &[u8]) -> Bytes { + let mut v = vec![0u8; payload.len() + IPC_SECTION_ALIGNMENT]; + let pad = section_padding(v.as_ptr() as usize); + v[pad..pad + payload.len()].copy_from_slice(payload); + Bytes::from(v).slice(pad..pad + payload.len()) + } + + #[test] + fn test_aligned_ipc_sections_are_zero_copy() { + // A LargeBinary column exercises the i64-offset buffer whose 8-byte + // alignment requirement triggers a realigning memcpy when misaligned. + let blocks = arrow_array::LargeBinaryArray::from_vec(vec![&b"hello"[..], b"world"]); + let section_a = RecordBatch::try_from_iter([("a", Arc::new(blocks) as ArrayRef)]).unwrap(); + let section_b = record_batch!(("b", Int64, [10i64, 20, 30, 40, 50])).unwrap(); + + let mut buf = Vec::new(); + // Arbitrary, deliberately non-64-aligned preamble so the first section + // must be padded rather than landing at offset 0 by luck. + buf.extend_from_slice(&[0xABu8; 7]); + let mut pos = buf.len(); + // The first section's stream begins after padding the 7-byte preamble + // up to the next 64-byte boundary. + assert_eq!(7 + section_padding(7), IPC_SECTION_ALIGNMENT); + write_ipc_section(&mut buf, &mut pos, §ion_a).unwrap(); + write_ipc_section(&mut buf, &mut pos, §ion_b).unwrap(); + + let data = aligned_bytes(&buf); + assert_eq!( + section_padding(data.as_ptr() as usize), + 0, + "base not aligned" + ); + + let mut offset = 7; + let read_a = read_ipc_section_at(&data, &mut offset).unwrap(); + let read_b = read_ipc_section_at(&data, &mut offset).unwrap(); + assert_eq!(read_a, section_a); + assert_eq!(read_b, section_b); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + for batch in [&read_a, &read_b] { + for buffer in batch.column(0).to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "section buffer at {ptr:#x} was realigned out of the input \ + [{data_base:#x}..{data_end:#x}) — misaligned section", + ); + } + } + } + + #[test] + fn test_aligned_multi_batch_section_roundtrip_zero_copy() { + // A multi-batch section (e.g. IVF SQ storage chunks) must round-trip + // every batch and decode the first batch's buffers zero-copy. + let b1 = record_batch!(("v", Int64, [1i64, 2, 3])).unwrap(); + let b2 = record_batch!(("v", Int64, [4i64, 5])).unwrap(); + let b3 = record_batch!(("v", Int64, [6i64])).unwrap(); + + let mut buf = vec![0xCDu8; 5]; + let mut pos = buf.len(); + write_ipc_section_batches(&mut buf, &mut pos, [b1.clone(), b2.clone(), b3.clone()]) + .unwrap(); + + let data = aligned_bytes(&buf); + let mut offset = 5; + let read = read_ipc_section_batches_at(&data, &mut offset).unwrap(); + assert_eq!(read, vec![b1, b2, b3]); + assert_eq!(offset, buf.len(), "offset should land at section end"); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + for buffer in read[0].column(0).to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "first batch buffer at {ptr:#x} was realigned out of the input", + ); + } + } +} diff --git a/python/vendor/lance-arrow/src/json.rs b/python/vendor/lance-arrow/src/json.rs new file mode 100644 index 000000000..df908c4cb --- /dev/null +++ b/python/vendor/lance-arrow/src/json.rs @@ -0,0 +1,1538 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! JSON support for Apache Arrow. + +use std::convert::TryFrom; +use std::sync::Arc; + +use arrow_array::builder::{GenericStringBuilder, LargeBinaryBuilder}; +use arrow_array::cast::AsArray; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, GenericStringArray, LargeBinaryArray, LargeListArray, + LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray, +}; +use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema}; +use lance_jsonb::OwnedJsonb; +use lance_jsonb::jsonpath::{JsonPath, Selector}; + +use crate::ARROW_EXT_NAME_KEY; + +/// Arrow extension type name for JSON data (Lance internal) +pub const JSON_EXT_NAME: &str = "lance.json"; + +/// Arrow extension type name for JSON data (Arrow official) +pub const ARROW_JSON_EXT_NAME: &str = "arrow.json"; + +/// Check if a field is a JSON extension field (Lance internal JSONB storage) +pub fn is_json_field(field: &ArrowField) -> bool { + field.data_type() == &DataType::LargeBinary + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == JSON_EXT_NAME) + .unwrap_or_default() +} + +/// Check if a field is an Arrow JSON extension field (PyArrow pa.json() type) +pub fn is_arrow_json_field(field: &ArrowField) -> bool { + // Arrow JSON extension type uses Utf8 or LargeUtf8 as storage type + (field.data_type() == &DataType::Utf8 || field.data_type() == &DataType::LargeUtf8) + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == ARROW_JSON_EXT_NAME) + .unwrap_or_default() +} + +/// Check if a field or any of its descendants is a JSON field +pub fn has_json_fields(field: &ArrowField) -> bool { + if is_json_field(field) { + return true; + } + + match field.data_type() { + DataType::Struct(fields) => fields.iter().any(|f| has_json_fields(f)), + DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => { + has_json_fields(f) + } + DataType::Map(f, _) => has_json_fields(f), + _ => false, + } +} + +/// Check if a field or any of its descendants is an Arrow JSON field +pub fn has_arrow_json_fields(field: &ArrowField) -> bool { + if is_arrow_json_field(field) { + return true; + } + + match field.data_type() { + DataType::Struct(fields) => fields.iter().any(|f| has_arrow_json_fields(f)), + DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => { + has_arrow_json_fields(f) + } + DataType::Map(f, _) => has_arrow_json_fields(f), + _ => false, + } +} + +/// The Arrow representation holding a JSON field's values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JsonEncoding { + /// JSONB bytes: [`JSON_EXT_NAME`] over `LargeBinary`. Lance stores JSON + /// this way. + Jsonb, + /// JSON text: [`ARROW_JSON_EXT_NAME`] over `Utf8` or `LargeUtf8`. Callers + /// write JSON this way, and reads return it this way by default. + Text, +} + +impl JsonEncoding { + /// The encoding of `field`'s values, or `None` when `field` does not hold + /// JSON. + pub fn of_field(field: &ArrowField) -> Option { + if is_json_field(field) { + Some(Self::Jsonb) + } else if is_arrow_json_field(field) { + Some(Self::Text) + } else { + None + } + } +} + +/// The values of a JSON column, readable as JSONB or as text whichever +/// encoding they arrived in. +/// +/// Code that computes on JSON, such as tokenizers and indices, reads values +/// through this type, so its behavior does not depend on whether they come +/// from storage or from a caller. +#[derive(Debug, Clone, Copy)] +pub struct JsonValues<'a> { + encoding: JsonEncoding, + array: &'a ArrayRef, +} + +impl<'a> JsonValues<'a> { + /// Read `array` as the values of `field`, or `None` when `field` does not + /// hold JSON. + pub fn try_new(field: &ArrowField, array: &'a ArrayRef) -> Option { + JsonEncoding::of_field(field).map(|encoding| Self { encoding, array }) + } + + pub fn encoding(&self) -> JsonEncoding { + self.encoding + } + + /// The values as JSONB. Text values are parsed, so invalid JSON fails. + pub fn to_jsonb(&self) -> Result { + match self.encoding { + JsonEncoding::Jsonb => Ok(self.array.as_binary::().clone()), + JsonEncoding::Text => Ok(JsonArray::try_from(self.array.clone())?.into_inner()), + } + } + + /// The values as JSON text: the input array when it already holds text, + /// otherwise `LargeUtf8` decoded from JSONB. + pub fn to_text(&self) -> ArrayRef { + match self.encoding { + JsonEncoding::Jsonb => Arc::new(decode_jsonb_array::(self.array.as_binary())), + JsonEncoding::Text => self.array.clone(), + } + } +} + +/// Create a JSON field with the appropriate extension metadata +pub fn json_field(name: &str, nullable: bool) -> ArrowField { + let mut field = ArrowField::new(name, DataType::LargeBinary, nullable); + let mut metadata = std::collections::HashMap::new(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), JSON_EXT_NAME.to_string()); + field.set_metadata(metadata); + field +} + +/// A specialized array for JSON data stored as JSONB binary format +#[derive(Debug, Clone)] +pub struct JsonArray { + inner: LargeBinaryArray, +} + +impl JsonArray { + /// Create a new JsonArray from an iterator of JSON strings + pub fn try_from_iter(iter: I) -> Result + where + I: IntoIterator>, + S: AsRef, + { + let mut builder = LargeBinaryBuilder::new(); + + for json_str in iter { + match json_str { + Some(s) => { + let encoded = encode_json(s.as_ref()).map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) + })?; + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + + Ok(Self { + inner: builder.finish(), + }) + } + + /// Get the underlying LargeBinaryArray + pub fn into_inner(self) -> LargeBinaryArray { + self.inner + } + + /// Get a reference to the underlying LargeBinaryArray + pub fn inner(&self) -> &LargeBinaryArray { + &self.inner + } + + /// Get the value at index i as decoded JSON string + pub fn value(&self, i: usize) -> Result { + if self.inner.is_null(i) { + return Err(ArrowError::InvalidArgumentError( + "Value is null".to_string(), + )); + } + + let jsonb_bytes = self.inner.value(i); + Ok(decode_json(jsonb_bytes)) + } + + /// Get the value at index i as raw JSONB bytes + pub fn value_bytes(&self, i: usize) -> &[u8] { + self.inner.value(i) + } + + /// Get JSONPath value from the JSON at index i + pub fn json_path(&self, i: usize, path: &str) -> Result, ArrowError> { + if self.inner.is_null(i) { + return Ok(None); + } + + let path = parse_json_path(path)?; + Ok(select_json_path(self.inner.value(i), &path)?.map(|value| value.to_string())) + } + + /// Convert to Arrow string array (JSON as UTF-8) + pub fn to_arrow_json(&self) -> ArrayRef { + Arc::new(decode_jsonb_array::(&self.inner)) + } + + pub fn len(&self) -> usize { + self.inner.len() + } + + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + pub fn is_null(&self, i: usize) -> bool { + self.inner.is_null(i) + } +} + +// TryFrom implementations for string arrays +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array: StringArray) -> Result { + Self::try_from(&array) + } +} + +impl TryFrom<&StringArray> for JsonArray { + type Error = ArrowError; + + fn try_from(array: &StringArray) -> Result { + encode_json_text_array(array) + } +} + +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array: LargeStringArray) -> Result { + Self::try_from(&array) + } +} + +impl TryFrom<&LargeStringArray> for JsonArray { + type Error = ArrowError; + + fn try_from(array: &LargeStringArray) -> Result { + encode_json_text_array(array) + } +} + +fn encode_json_text_array( + array: &GenericStringArray, +) -> Result { + let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(json_str) => { + let encoded = encode_json(json_str).map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) + })?; + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + Ok(JsonArray { + inner: builder.finish(), + }) +} + +fn decode_jsonb_array(array: &LargeBinaryArray) -> GenericStringArray { + let mut builder = + GenericStringBuilder::::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(jsonb_bytes) => builder.append_value(decode_json(jsonb_bytes)), + None => builder.append_null(), + } + } + builder.finish() +} + +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array_ref: ArrayRef) -> Result { + match array_ref.data_type() { + DataType::Utf8 => { + // Downcast is guaranteed to succeed after matching on DataType::Utf8 + let string_array = array_ref + .as_any() + .downcast_ref::() + .expect("DataType::Utf8 array must be StringArray"); + Self::try_from(string_array) + } + DataType::LargeUtf8 => { + // Downcast is guaranteed to succeed after matching on DataType::LargeUtf8 + let large_string_array = array_ref + .as_any() + .downcast_ref::() + .expect("DataType::LargeUtf8 array must be LargeStringArray"); + Self::try_from(large_string_array) + } + dt => Err(ArrowError::InvalidArgumentError(format!( + "Unsupported array type for JSON: {:?}. Expected Utf8 or LargeUtf8", + dt + ))), + } + } +} + +/// Encode JSON string to JSONB format +pub fn encode_json(json_str: &str) -> Result, Box> { + let value = lance_jsonb::parse_value(json_str.as_bytes())?; + Ok(value.to_vec()) +} + +/// Decode JSONB bytes to JSON string +pub fn decode_json(jsonb_bytes: &[u8]) -> String { + let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes); + raw_jsonb.to_string() +} + +/// Parse a JSONPath expression. +pub fn parse_json_path(path: &str) -> Result, ArrowError> { + lance_jsonb::jsonpath::parse_json_path(path.as_bytes()) + .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}"))) +} + +/// Select the value at `path` in a JSONB value. +/// +/// A path matching several values selects them together as one JSON array, +/// so no match is silently dropped. +pub fn select_json_path( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .select_value(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) +} + +/// Whether `path` matches any value in a JSONB value. +pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .exists(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}")) + }) +} + +/// Select every value matched by `path` in a JSONB value. +pub fn select_json_path_values( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .select_values(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) +} + +/// Convert an Arrow JSON field to Lance JSON field (with JSONB storage) +pub fn arrow_json_to_lance_json(field: &ArrowField) -> ArrowField { + if is_arrow_json_field(field) { + return field_with_extension(field, DataType::LargeBinary, JSON_EXT_NAME); + } + + let data_type = match field.data_type() { + DataType::Struct(fields) => { + let fields = fields + .iter() + .map(|field| Arc::new(arrow_json_to_lance_json(field))) + .collect::>(); + DataType::Struct(Fields::from(fields)) + } + DataType::List(item) => DataType::List(Arc::new(arrow_json_to_lance_json(item))), + DataType::LargeList(item) => DataType::LargeList(Arc::new(arrow_json_to_lance_json(item))), + DataType::FixedSizeList(item, size) => { + DataType::FixedSizeList(Arc::new(arrow_json_to_lance_json(item)), *size) + } + DataType::Map(entries, keys_sorted) => { + DataType::Map(Arc::new(arrow_json_to_lance_json(entries)), *keys_sorted) + } + _ => return field.clone(), + }; + + field_with_data_type(field, data_type) +} + +/// Convert a Lance JSON field to Arrow JSON field. +pub fn lance_json_to_arrow_json(field: &ArrowField) -> ArrowField { + if is_json_field(field) { + return field_with_extension(field, DataType::Utf8, ARROW_JSON_EXT_NAME); + } + + let data_type = match field.data_type() { + DataType::Struct(fields) => { + let fields = fields + .iter() + .map(|field| Arc::new(lance_json_to_arrow_json(field))) + .collect::>(); + DataType::Struct(Fields::from(fields)) + } + DataType::List(item) => DataType::List(Arc::new(lance_json_to_arrow_json(item))), + DataType::LargeList(item) => DataType::LargeList(Arc::new(lance_json_to_arrow_json(item))), + DataType::FixedSizeList(item, size) => { + DataType::FixedSizeList(Arc::new(lance_json_to_arrow_json(item)), *size) + } + DataType::Map(entries, keys_sorted) => { + DataType::Map(Arc::new(lance_json_to_arrow_json(entries)), *keys_sorted) + } + _ => return field.clone(), + }; + + field_with_data_type(field, data_type) +} + +fn field_with_data_type(field: &ArrowField, data_type: DataType) -> ArrowField { + ArrowField::new(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()) +} + +fn field_with_extension( + field: &ArrowField, + data_type: DataType, + extension_name: &str, +) -> ArrowField { + let mut metadata = field.metadata().clone(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), extension_name.to_string()); + ArrowField::new(field.name(), data_type, field.is_nullable()).with_metadata(metadata) +} + +fn convert_json_array( + field: &ArrowField, + array: &ArrayRef, + convert_leaf: &F, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> +where + F: Fn(&ArrowField, &ArrayRef) -> Result, ArrowError>, +{ + if let Some((field, array)) = convert_leaf(field, array)? { + return Ok((field, array, true)); + } + + match field.data_type() { + DataType::Struct(fields) => { + let struct_array = array.as_struct(); + let mut new_fields = Vec::with_capacity(fields.len()); + let mut new_columns = Vec::with_capacity(fields.len()); + let mut changed = false; + + for (field, column) in fields.iter().zip(struct_array.columns()) { + let (new_field, new_column, field_changed) = + convert_json_array(field, column, convert_leaf)?; + changed |= field_changed; + new_fields.push(Arc::new(new_field)); + new_columns.push(new_column); + } + + if changed { + let fields = Fields::from(new_fields); + let new_field = field_with_data_type(field, DataType::Struct(fields.clone())); + let new_array = + StructArray::new(fields, new_columns, struct_array.nulls().cloned()); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::List(item) => { + let list_array: &ListArray = array.as_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = + field_with_data_type(field, DataType::List(Arc::new(new_item.clone()))); + let new_array = ListArray::new( + Arc::new(new_item), + list_array.offsets().clone(), + new_values, + list_array.nulls().cloned(), + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::LargeList(item) => { + let list_array: &LargeListArray = array.as_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = + field_with_data_type(field, DataType::LargeList(Arc::new(new_item.clone()))); + let new_array = LargeListArray::new( + Arc::new(new_item), + list_array.offsets().clone(), + new_values, + list_array.nulls().cloned(), + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::FixedSizeList(item, size) => { + let list_array: &FixedSizeListArray = array.as_fixed_size_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = field_with_data_type( + field, + DataType::FixedSizeList(Arc::new(new_item.clone()), *size), + ); + let new_array = FixedSizeListArray::try_new_with_length( + Arc::new(new_item), + *size, + new_values, + list_array.nulls().cloned(), + list_array.len(), + )?; + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::Map(entries, keys_sorted) => { + let map_array = array + .as_any() + .downcast_ref::() + .expect("DataType::Map array must be MapArray"); + let entries_array = Arc::new(map_array.entries().clone()) as ArrayRef; + let (new_entries, new_entries_array, changed) = + convert_json_array(entries, &entries_array, convert_leaf)?; + if changed { + let entries_struct = new_entries_array + .as_any() + .downcast_ref::() + .expect("Map entries must be StructArray") + .clone(); + let new_field = field_with_data_type( + field, + DataType::Map(Arc::new(new_entries.clone()), *keys_sorted), + ); + let new_array = MapArray::new( + Arc::new(new_entries), + map_array.offsets().clone(), + entries_struct, + map_array.nulls().cloned(), + *keys_sorted, + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + _ => Ok((field.clone(), array.clone(), false)), + } +} + +fn convert_arrow_json_array( + field: &ArrowField, + array: &ArrayRef, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> { + convert_json_array(field, array, &|field, array| { + if is_arrow_json_field(field) { + let json_array = JsonArray::try_from(array.clone())?; + Ok(Some(( + arrow_json_to_lance_json(field), + Arc::new(json_array.into_inner()) as ArrayRef, + ))) + } else { + Ok(None) + } + }) +} + +fn convert_lance_json_array( + field: &ArrowField, + array: &ArrayRef, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> { + convert_json_array(field, array, &|field, array| { + if is_json_field(field) { + Ok(Some(( + lance_json_to_arrow_json(field), + Arc::new(decode_jsonb_array::(array.as_binary())) as ArrayRef, + ))) + } else { + Ok(None) + } + }) +} + +/// Convert a RecordBatch with Lance JSON columns (JSONB) back to Arrow JSON format (strings) +pub fn convert_lance_json_to_arrow( + batch: &arrow_array::RecordBatch, +) -> Result { + let schema = batch.schema(); + let mut needs_conversion = false; + let mut new_fields = Vec::with_capacity(schema.fields().len()); + let mut new_columns = Vec::with_capacity(batch.num_columns()); + + for (i, field) in schema.fields().iter().enumerate() { + let column = batch.column(i); + let (new_field, new_column, changed) = convert_lance_json_array(field, column)?; + + needs_conversion |= changed; + new_fields.push(new_field); + new_columns.push(new_column); + } + + if needs_conversion { + let new_schema = Arc::new(Schema::new_with_metadata( + new_fields, + schema.metadata().clone(), + )); + RecordBatch::try_new(new_schema, new_columns) + } else { + // No conversion needed, return original batch + Ok(batch.clone()) + } +} + +/// Convert a RecordBatch with Arrow JSON columns to Lance JSON format (JSONB) +pub fn convert_json_columns( + batch: &arrow_array::RecordBatch, +) -> Result { + let schema = batch.schema(); + let mut needs_conversion = false; + let mut new_fields = Vec::with_capacity(schema.fields().len()); + let mut new_columns = Vec::with_capacity(batch.num_columns()); + + for (i, field) in schema.fields().iter().enumerate() { + let column = batch.column(i); + let (new_field, new_column, changed) = convert_arrow_json_array(field, column)?; + + needs_conversion |= changed; + new_fields.push(new_field); + new_columns.push(new_column); + } + + if needs_conversion { + let new_schema = Arc::new(Schema::new_with_metadata( + new_fields, + schema.metadata().clone(), + )); + RecordBatch::try_new(new_schema, new_columns) + } else { + // No conversion needed, return original batch + Ok(batch.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_json_field_creation() { + let field = json_field("data", true); + assert_eq!(field.name(), "data"); + assert_eq!(field.data_type(), &DataType::LargeBinary); + assert!(field.is_nullable()); + assert!(is_json_field(&field)); + } + + #[test] + fn test_json_array_from_strings() { + let json_strings = vec![ + Some(r#"{"name": "Alice", "age": 30}"#), + None, + Some(r#"{"name": "Bob", "age": 25}"#), + ]; + + let array = JsonArray::try_from_iter(json_strings).unwrap(); + assert_eq!(array.len(), 3); + assert!(!array.is_null(0)); + assert!(array.is_null(1)); + assert!(!array.is_null(2)); + + let decoded = array.value(0).unwrap(); + assert!(decoded.contains("Alice")); + } + + #[test] + fn test_json_array_from_string_array() { + let string_array = StringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + None, + ]); + + let json_array = JsonArray::try_from(string_array).unwrap(); + assert_eq!(json_array.len(), 3); + assert!(!json_array.is_null(0)); + assert!(!json_array.is_null(1)); + assert!(json_array.is_null(2)); + } + + #[test] + fn test_json_path_extraction() { + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"user": {"name": "Alice", "age": 30}}"#), + Some(r#"{"user": {"name": "Bob"}}"#), + ]) + .unwrap(); + + let name = json_array.json_path(0, "$.user.name").unwrap(); + assert_eq!(name, Some("\"Alice\"".to_string())); + + let age = json_array.json_path(1, "$.user.age").unwrap(); + assert_eq!(age, None); + + // Several matches are selected together, as the SQL JSON functions do. + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"a": [{"b": 1}, {"b": 2}]}"#)]).unwrap(); + let bs = json_array.json_path(0, "$.a[*].b").unwrap(); + assert_eq!(bs, Some("[1,2]".to_string())); + } + + #[test] + fn test_json_values_read_either_encoding() { + let docs = vec![Some(r#"{"a":1}"#), None]; + let text: ArrayRef = Arc::new(StringArray::from(docs.clone())); + let jsonb: ArrayRef = Arc::new(JsonArray::try_from_iter(docs).unwrap().into_inner()); + let mut text_field = ArrowField::new("j", DataType::Utf8, true); + text_field.set_metadata(std::collections::HashMap::from([( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + )])); + let jsonb_field = json_field("j", true); + + let from_text = JsonValues::try_new(&text_field, &text).unwrap(); + let from_jsonb = JsonValues::try_new(&jsonb_field, &jsonb).unwrap(); + assert_eq!(from_text.encoding(), JsonEncoding::Text); + assert_eq!(from_jsonb.encoding(), JsonEncoding::Jsonb); + + assert_eq!(&from_text.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_jsonb.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_text.to_text(), &text); + let decoded = from_jsonb.to_text(); + let decoded = decoded.as_string::(); + assert_eq!(decoded.value(0), r#"{"a":1}"#); + assert!(decoded.is_null(1)); + + let plain = ArrowField::new("j", DataType::Utf8, true); + assert!(JsonValues::try_new(&plain, &text).is_none()); + } + + #[test] + fn test_convert_json_columns() { + // Create a batch with Arrow JSON column + let json_strings = vec![Some(r#"{"name": "Alice"}"#), Some(r#"{"name": "Bob"}"#)]; + let json_arr = StringArray::from(json_strings); + + // Create field with arrow.json extension + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(json_arr) as ArrayRef]).unwrap(); + + // Convert the batch + let converted = convert_json_columns(&batch).unwrap(); + + // Check the converted schema + assert_eq!(converted.num_columns(), 1); + let converted_schema = converted.schema(); + let converted_field = converted_schema.field(0); + assert_eq!(converted_field.data_type(), &DataType::LargeBinary); + assert_eq!( + converted_field.metadata().get(ARROW_EXT_NAME_KEY), + Some(&JSON_EXT_NAME.to_string()) + ); + + // Check the data was converted + let converted_column = converted.column(0); + assert_eq!(converted_column.data_type(), &DataType::LargeBinary); + assert_eq!(converted_column.len(), 2); + + // Verify the data is valid JSONB + let binary_array = converted_column + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..binary_array.len() { + let jsonb_bytes = binary_array.value(i); + let decoded = decode_json(jsonb_bytes); + assert!(decoded.contains("name")); + } + } + + #[test] + fn test_convert_nested_json_columns() { + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, false)); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let extra_field = + Arc::new(ArrowField::new("extra", DataType::Utf8, true).with_metadata(metadata)); + let item_fields = Fields::from(vec![uri_field, extra_field]); + + let values = StructArray::new( + item_fields.clone(), + vec![ + Arc::new(StringArray::from(vec![Some("a.jpg"), Some("b.jpg")])) as ArrayRef, + Arc::new(StringArray::from(vec![ + Some(r#"{"codec":"h264"}"#), + None::<&str>, + ])) as ArrayRef, + ], + None, + ); + let item = Arc::new(ArrowField::new("item", DataType::Struct(item_fields), true)); + let media = ListArray::new( + item, + OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2])), + Arc::new(values), + None, + ); + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "media", + media.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(media) as ArrayRef]).unwrap(); + + assert!(has_arrow_json_fields(batch.schema().field(0))); + + let converted = convert_json_columns(&batch).unwrap(); + let converted_schema = converted.schema(); + let DataType::List(item) = converted_schema.field(0).data_type() else { + panic!("expected list field"); + }; + let DataType::Struct(fields) = item.data_type() else { + panic!("expected struct item"); + }; + assert!(is_json_field(&fields[1])); + + let list_array: &ListArray = converted.column(0).as_list(); + let values = list_array.values().as_struct(); + let extra = values + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(decode_json(extra.value(0)).contains("h264")); + assert!(extra.is_null(1)); + + let logical = convert_lance_json_to_arrow(&converted).unwrap(); + let logical_schema = logical.schema(); + let DataType::List(item) = logical_schema.field(0).data_type() else { + panic!("expected list field"); + }; + let DataType::Struct(fields) = item.data_type() else { + panic!("expected struct item"); + }; + assert!(is_arrow_json_field(&fields[1])); + + let list_array: &ListArray = logical.column(0).as_list(); + let values = list_array.values().as_struct(); + let extra = values + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(extra.value(0).contains("h264")); + assert!(extra.is_null(1)); + } + + #[test] + fn test_convert_fixed_size_list_zero_json_preserves_length() { + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let item = Arc::new(ArrowField::new("item", DataType::Utf8, true).with_metadata(metadata)); + let values = Arc::new(StringArray::from(Vec::>::new())) as ArrayRef; + let lists = FixedSizeListArray::try_new_with_length(item, 0, values, None, 3).unwrap(); + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "lists", + lists.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(lists) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_rows(), 3); + assert_eq!(converted.column(0).len(), 3); + + let converted_schema = converted.schema(); + let DataType::FixedSizeList(item, size) = converted_schema.field(0).data_type() else { + panic!("expected fixed size list field"); + }; + assert_eq!(*size, 0); + assert!(is_json_field(item)); + + let logical = convert_lance_json_to_arrow(&converted).unwrap(); + assert_eq!(logical.num_rows(), 3); + assert_eq!(logical.column(0).len(), 3); + + let logical_schema = logical.schema(); + let DataType::FixedSizeList(item, size) = logical_schema.field(0).data_type() else { + panic!("expected fixed size list field"); + }; + assert_eq!(*size, 0); + assert!(is_arrow_json_field(item)); + } + + #[test] + fn test_has_json_fields() { + // Test direct JSON field + let json_f = json_field("data", true); + assert!(has_json_fields(&json_f)); + + // Test non-JSON field + let non_json = ArrowField::new("data", DataType::Utf8, true); + assert!(!has_json_fields(&non_json)); + + // Test struct containing JSON field + let struct_field = ArrowField::new( + "struct", + DataType::Struct(vec![json_field("nested_json", true)].into()), + true, + ); + assert!(has_json_fields(&struct_field)); + + // Test struct without JSON field + let struct_no_json = ArrowField::new( + "struct", + DataType::Struct(vec![ArrowField::new("text", DataType::Utf8, true)].into()), + true, + ); + assert!(!has_json_fields(&struct_no_json)); + + // Test List containing JSON field + let list_field = ArrowField::new( + "list", + DataType::List(Arc::new(json_field("item", true))), + true, + ); + assert!(has_json_fields(&list_field)); + + // Test LargeList containing JSON field + let large_list_field = ArrowField::new( + "large_list", + DataType::LargeList(Arc::new(json_field("item", true))), + true, + ); + assert!(has_json_fields(&large_list_field)); + + // Test FixedSizeList containing JSON field + let fixed_list_field = ArrowField::new( + "fixed_list", + DataType::FixedSizeList(Arc::new(json_field("item", true)), 3), + true, + ); + assert!(has_json_fields(&fixed_list_field)); + + // Test Map containing JSON field + let map_field = ArrowField::new( + "map", + DataType::Map( + Arc::new(ArrowField::new( + "entries", + DataType::Struct( + vec![ + ArrowField::new("key", DataType::Utf8, false), + json_field("value", true), + ] + .into(), + ), + false, + )), + false, + ), + true, + ); + assert!(has_json_fields(&map_field)); + } + + #[test] + fn test_json_array_inner() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + let inner = json_array.inner(); + assert_eq!(inner.len(), 1); + } + + #[test] + fn test_json_array_value_null_error() { + let json_array = JsonArray::try_from_iter(vec![None::<&str>]).unwrap(); + let result = json_array.value(0); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("null")); + } + + #[test] + fn test_json_array_value_bytes() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + let bytes = json_array.value_bytes(0); + assert!(!bytes.is_empty()); + } + + #[test] + fn test_json_path_with_null() { + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"user": {"name": "Alice"}}"#), None::<&str>]) + .unwrap(); + + let result = json_array.json_path(1, "$.user.name").unwrap(); + assert_eq!(result, None); + } + + #[test] + fn test_to_arrow_json() { + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + None::<&str>, + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + + let arrow_json = json_array.to_arrow_json(); + assert_eq!(arrow_json.len(), 3); + assert!(!arrow_json.is_null(0)); + assert!(arrow_json.is_null(1)); + assert!(!arrow_json.is_null(2)); + + let string_array = arrow_json.as_any().downcast_ref::().unwrap(); + assert!(string_array.value(0).contains("Alice")); + assert!(string_array.value(2).contains("Bob")); + } + + #[test] + fn test_json_array_trait_methods() { + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#), Some(r#"{"b": 2}"#)]).unwrap(); + + // Wrapper methods + assert_eq!(json_array.len(), 2); + assert!(!json_array.is_empty()); + assert!(!json_array.is_null(0)); + + // Underlying Arrow array + assert_eq!(json_array.inner().data_type(), &DataType::LargeBinary); + assert_eq!(json_array.inner().len(), 2); + } + + #[test] + fn test_json_array_empty() { + let json_array = JsonArray::try_from_iter(Vec::>::new()).unwrap(); + assert!(json_array.is_empty()); + assert_eq!(json_array.len(), 0); + } + + #[test] + fn test_try_from_large_string_array() { + let large_string_array = LargeStringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + None, + ]); + + // Test TryFrom<&LargeStringArray> + let json_array = JsonArray::try_from(&large_string_array).unwrap(); + assert_eq!(json_array.len(), 3); + assert!(!json_array.is_null(0)); + assert!(!json_array.is_null(1)); + assert!(json_array.is_null(2)); + + // Test TryFrom (owned) + let large_string_array2 = LargeStringArray::from(vec![Some(r#"{"x": 1}"#)]); + let json_array2 = JsonArray::try_from(large_string_array2).unwrap(); + assert_eq!(json_array2.len(), 1); + } + + #[test] + fn test_try_from_array_ref() { + // Test with Utf8 + let string_array: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a": 1}"#), + Some(r#"{"b": 2}"#), + ])); + let json_array = JsonArray::try_from(string_array).unwrap(); + assert_eq!(json_array.len(), 2); + + // Test with LargeUtf8 + let large_string_array: ArrayRef = Arc::new(LargeStringArray::from(vec![ + Some(r#"{"c": 3}"#), + Some(r#"{"d": 4}"#), + ])); + let json_array2 = JsonArray::try_from(large_string_array).unwrap(); + assert_eq!(json_array2.len(), 2); + + // Test with unsupported type + let int_array: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3])); + let result = JsonArray::try_from(int_array); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Unsupported")); + } + + #[test] + fn test_arrow_json_to_lance_json_non_json_field() { + // Test that non-JSON fields are returned unchanged + let field = ArrowField::new("text", DataType::Utf8, true); + let converted = arrow_json_to_lance_json(&field); + assert_eq!(converted.data_type(), &DataType::Utf8); + assert_eq!(converted.name(), "text"); + } + + #[test] + fn test_convert_lance_json_to_arrow() { + // Create a batch with Lance JSON column (JSONB) + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + None::<&str>, + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + + let lance_json_field = json_field("data", true); + let schema = Arc::new(Schema::new(vec![lance_json_field])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(json_array.into_inner()) as ArrayRef]) + .unwrap(); + + // Convert back to Arrow JSON + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + + // Check schema + let converted_schema = converted.schema(); + let converted_field = converted_schema.field(0); + assert_eq!(converted_field.data_type(), &DataType::Utf8); + assert_eq!( + converted_field.metadata().get(ARROW_EXT_NAME_KEY), + Some(&ARROW_JSON_EXT_NAME.to_string()) + ); + + // Check data + let string_array = converted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!string_array.is_null(0)); + assert!(string_array.is_null(1)); + assert!(!string_array.is_null(2)); + assert!(string_array.value(0).contains("Alice")); + assert!(string_array.value(2).contains("Bob")); + } + + #[test] + fn test_convert_lance_json_to_arrow_empty_batch() { + // Create an empty batch with Lance JSON column + let lance_json_field = json_field("data", true); + let schema = Arc::new(Schema::new(vec![lance_json_field])); + let empty_binary = LargeBinaryBuilder::new().finish(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(empty_binary) as ArrayRef]).unwrap(); + + // Convert back to Arrow JSON + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_rows(), 0); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_lance_json_to_arrow_no_json_columns() { + // Create a batch without JSON columns + let field = ArrowField::new("text", DataType::Utf8, true); + let schema = Arc::new(Schema::new(vec![field])); + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap(); + + // Convert - should return the same batch + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_json_columns_empty_batch() { + // Create an empty batch with Arrow JSON column + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let empty_strings = arrow_array::builder::StringBuilder::new().finish(); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(empty_strings) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_rows(), 0); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + } + + #[test] + fn test_convert_json_columns_large_string() { + // Create a batch with Arrow JSON column using LargeUtf8 + let json_strings = LargeStringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]); + + let mut field = ArrowField::new("data", DataType::LargeUtf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(json_strings) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + assert_eq!(converted.num_rows(), 2); + } + + #[test] + fn test_convert_json_columns_no_json_columns() { + // Create a batch without JSON columns + let field = ArrowField::new("text", DataType::Utf8, true); + let schema = Arc::new(Schema::new(vec![field])); + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap(); + + // Convert - should return the same batch + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_json_columns_mixed_columns() { + // Create a batch with both JSON and non-JSON columns + let json_strings = StringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]); + let text_strings = StringArray::from(vec![Some("hello"), Some("world")]); + + let mut json_field = ArrowField::new("json_data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + json_field.set_metadata(metadata); + + let text_field = ArrowField::new("text_data", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![json_field, text_field])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(json_strings) as ArrayRef, + Arc::new(text_strings) as ArrayRef, + ], + ) + .unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 2); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8); + } + + #[test] + fn test_is_arrow_json_field_large_utf8() { + // Test with LargeUtf8 storage type + let mut field = ArrowField::new("data", DataType::LargeUtf8, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + assert!(is_arrow_json_field(&field)); + } + + #[test] + fn test_encode_json_invalid() { + // Test encoding invalid JSON + let result = encode_json("not valid json {"); + assert!(result.is_err()); + } + + #[test] + fn test_json_array_from_invalid_json() { + // Test creating JsonArray from invalid JSON strings + let result = JsonArray::try_from_iter(vec![Some("invalid json {")]); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Failed to encode")); + } + + #[test] + fn test_try_from_string_array_invalid_json() { + let string_array = StringArray::from(vec![Some("invalid json {")]); + let result = JsonArray::try_from(string_array); + assert!(result.is_err()); + } + + #[test] + fn test_try_from_large_string_array_invalid_json() { + let large_string_array = LargeStringArray::from(vec![Some("invalid json {")]); + let result = JsonArray::try_from(large_string_array); + assert!(result.is_err()); + } + + #[test] + fn test_convert_lance_json_to_arrow_mixed_columns() { + // Create a batch with both JSON and non-JSON columns + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + let text_strings = StringArray::from(vec![Some("hello"), Some("world")]); + + let json_f = json_field("json_data", true); + let text_field = ArrowField::new("text_data", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![json_f, text_field])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(json_array.into_inner()) as ArrayRef, + Arc::new(text_strings) as ArrayRef, + ], + ) + .unwrap(); + + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_columns(), 2); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8); + } + + #[test] + fn test_json_path_invalid_path() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + // Invalid JSONPath syntax should return error + let result = json_array.json_path(0, "invalid path without $"); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Invalid JSONPath")); + } + + #[test] + fn test_convert_json_columns_invalid_storage_type() { + // Create a batch with Arrow JSON field but wrong storage type (Int32 instead of Utf8) + let int_array = arrow_array::Int32Array::from(vec![1, 2, 3]); + + let mut field = ArrowField::new("data", DataType::Int32, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(int_array) as ArrayRef]).unwrap(); + + // This should succeed since Int32 doesn't match is_arrow_json_field check + // (is_arrow_json_field requires Utf8 or LargeUtf8) + let result = convert_json_columns(&batch); + assert!(result.is_ok()); + } + + #[test] + fn test_is_json_field_wrong_extension() { + // LargeBinary field without the correct extension metadata + let field = ArrowField::new("data", DataType::LargeBinary, true); + assert!(!is_json_field(&field)); + + // LargeBinary field with wrong extension name + let mut field2 = ArrowField::new("data", DataType::LargeBinary, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + "other.extension".to_string(), + ); + field2.set_metadata(metadata); + assert!(!is_json_field(&field2)); + } + + #[test] + fn test_is_arrow_json_field_wrong_extension() { + // Utf8 field without extension metadata + let field = ArrowField::new("data", DataType::Utf8, true); + assert!(!is_arrow_json_field(&field)); + + // Utf8 field with wrong extension name + let mut field2 = ArrowField::new("data", DataType::Utf8, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + "other.extension".to_string(), + ); + field2.set_metadata(metadata); + assert!(!is_arrow_json_field(&field2)); + + // Wrong type entirely + let field3 = ArrowField::new("data", DataType::Int32, true); + assert!(!is_arrow_json_field(&field3)); + } + + #[test] + fn test_convert_json_columns_invalid_json_utf8() { + // Test error propagation when converting invalid JSON (Utf8) + let invalid_json = StringArray::from(vec![Some("invalid json {")]); + + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap(); + + let result = convert_json_columns(&batch); + assert!(result.is_err()); + } + + #[test] + fn test_convert_json_columns_invalid_json_large_utf8() { + // Test error propagation when converting invalid JSON (LargeUtf8) + let invalid_json = LargeStringArray::from(vec![Some("invalid json {")]); + + let mut field = ArrowField::new("data", DataType::LargeUtf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap(); + + let result = convert_json_columns(&batch); + assert!(result.is_err()); + } + + #[test] + fn test_json_path_on_corrupted_jsonb() { + // Create corrupted JSONB bytes directly + let corrupted_bytes: &[u8] = &[0xFF, 0xFE, 0x00, 0x01, 0x02]; + let corrupted_binary = LargeBinaryArray::from(vec![Some(corrupted_bytes)]); + + // Wrap in JsonArray + let corrupted_json = JsonArray { + inner: corrupted_binary, + }; + + // Try to use json_path on corrupted data - the selector might fail or return unexpected results + // This exercises the code path but may not produce an error depending on jsonb library behavior + let _result = corrupted_json.json_path(0, "$.a"); + // We don't assert on the result as the behavior depends on the jsonb library + } + + #[test] + fn test_decode_json_on_various_inputs() { + // Test decode_json with various inputs + let valid_jsonb = encode_json(r#"{"key": "value"}"#).unwrap(); + let decoded = decode_json(&valid_jsonb); + assert!(decoded.contains("key")); + + // Empty bytes - jsonb library handles this gracefully + let decoded_empty = decode_json(&[]); + // Just verify it doesn't panic + let _ = decoded_empty; + + // Random bytes - jsonb library handles this gracefully + let decoded_random = decode_json(&[0xFF, 0xFE, 0x00]); + // Just verify it doesn't panic + let _ = decoded_random; + } +} diff --git a/python/vendor/lance-arrow/src/lib.rs b/python/vendor/lance-arrow/src/lib.rs new file mode 100644 index 000000000..f26e1f3e3 --- /dev/null +++ b/python/vendor/lance-arrow/src/lib.rs @@ -0,0 +1,3102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extend Arrow Functionality +//! +//! To improve Arrow-RS ergonomic + +#![warn(clippy::undocumented_unsafe_blocks)] + +// lance-arrow reinterprets value bytes as native numeric types in +// `FloatArray::as_slice` for `bf16` (rust/lance-arrow/src/bfloat16.rs), which +// requires the host byte order to match the on-disk byte order Lance writes. +// Lance writes little-endian; building on a big-endian target would silently +// produce wrong numeric values. +#[cfg(not(target_endian = "little"))] +compile_error!("lance-arrow only supports little-endian targets"); + +use std::sync::Arc; +use std::{collections::HashMap, ptr::NonNull}; + +use arrow_array::{ + Array, ArrayRef, ArrowNumericType, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray, + LargeListArray, ListArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray, + UInt8Array, UInt32Array, cast::AsArray, +}; +use arrow_array::{ + Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, new_null_array, +}; +use arrow_buffer::MutableBuffer; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema, SortOptions}; +use arrow_select::{interleave::interleave, take::take}; +use rand::prelude::*; + +pub mod deepcopy; +pub mod schema; +pub use schema::*; +pub mod bfloat16; +pub mod floats; +use crate::list::ListArrayExt; +pub use floats::*; + +pub mod ipc; +pub mod json; +pub mod list; +pub mod memory; +pub mod scalar; +pub mod stream; +pub mod r#struct; + +/// Arrow extension metadata key for extension name +pub const ARROW_EXT_NAME_KEY: &str = "ARROW:extension:name"; + +/// Arrow extension metadata key for extension metadata +pub const ARROW_EXT_META_KEY: &str = "ARROW:extension:metadata"; + +/// Key used by lance to mark a field as a blob +/// TODO: Use Arrow extension mechanism instead? +pub const BLOB_META_KEY: &str = "lance-encoding:blob"; +/// Arrow extension type name for Lance blob v2 columns +pub const BLOB_V2_EXT_NAME: &str = "lance.blob.v2"; +/// Metadata key for overriding the dedicated blob size threshold (in bytes) +pub const BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY: &str = + "lance-encoding:blob-dedicated-size-threshold"; +/// Metadata key for overriding the inline blob size threshold (in bytes) +pub const BLOB_INLINE_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-inline-size-threshold"; +/// Metadata key for overriding the maximum size (in bytes) of a packed blob sidecar file +pub const BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY: &str = + "lance-encoding:blob-pack-file-size-threshold"; + +type Result = std::result::Result; + +pub trait DataTypeExt { + /// Returns true if the data type is binary-like, such as Utf8, Binary, or the large and/or view variants. + fn is_binary_like(&self) -> bool; + + /// Returns true if the data type is a struct. + fn is_struct(&self) -> bool; + + /// Check whether the given Arrow DataType is fixed stride. + /// + /// A fixed stride type has the same byte width for all array elements + /// This includes all PrimitiveType's Boolean, FixedSizeList, FixedSizeBinary, and Decimals + fn is_fixed_stride(&self) -> bool; + + /// Returns true if the [DataType] is a dictionary type. + fn is_dictionary(&self) -> bool; + + /// Returns the byte width of the data type + /// Panics if the data type is not fixed stride. + fn byte_width(&self) -> usize; + + /// Returns the byte width of the data type, if it is fixed stride. + /// Returns None if the data type is not fixed stride. + fn byte_width_opt(&self) -> Option; +} + +impl DataTypeExt for DataType { + fn is_binary_like(&self) -> bool { + use DataType::*; + matches!( + self, + Utf8 | Binary | LargeUtf8 | LargeBinary | Utf8View | BinaryView + ) + } + + fn is_struct(&self) -> bool { + matches!(self, Self::Struct(_)) + } + + fn is_fixed_stride(&self) -> bool { + use DataType::*; + matches!( + self, + Boolean + | UInt8 + | UInt16 + | UInt32 + | UInt64 + | Int8 + | Int16 + | Int32 + | Int64 + | Float16 + | Float32 + | Float64 + | Decimal128(_, _) + | Decimal256(_, _) + | FixedSizeList(_, _) + | FixedSizeBinary(_) + | Duration(_) + | Timestamp(_, _) + | Date32 + | Date64 + | Time32(_) + | Time64(_) + ) + } + + fn is_dictionary(&self) -> bool { + matches!(self, Self::Dictionary(_, _)) + } + + fn byte_width_opt(&self) -> Option { + match self { + Self::Int8 => Some(1), + Self::Int16 => Some(2), + Self::Int32 => Some(4), + Self::Int64 => Some(8), + Self::UInt8 => Some(1), + Self::UInt16 => Some(2), + Self::UInt32 => Some(4), + Self::UInt64 => Some(8), + Self::Float16 => Some(2), + Self::Float32 => Some(4), + Self::Float64 => Some(8), + Self::Date32 => Some(4), + Self::Date64 => Some(8), + Self::Time32(_) => Some(4), + Self::Time64(_) => Some(8), + Self::Timestamp(_, _) => Some(8), + Self::Duration(_) => Some(8), + Self::Decimal128(_, _) => Some(16), + Self::Decimal256(_, _) => Some(32), + Self::Interval(unit) => match unit { + IntervalUnit::YearMonth => Some(4), + IntervalUnit::DayTime => Some(8), + IntervalUnit::MonthDayNano => Some(16), + }, + Self::FixedSizeBinary(s) => Some(*s as usize), + Self::FixedSizeList(dt, s) => dt + .data_type() + .byte_width_opt() + .map(|width| width * *s as usize), + _ => None, + } + } + + fn byte_width(&self) -> usize { + self.byte_width_opt() + .unwrap_or_else(|| panic!("Expecting fixed stride data type, found {:?}", self)) + } +} + +/// Create an [`GenericListArray`] from values and offsets. +/// +/// ``` +/// use arrow_array::{Int32Array, Int64Array, ListArray}; +/// use arrow_array::types::Int64Type; +/// use lance_arrow::try_new_generic_list_array; +/// +/// let offsets = Int32Array::from_iter([0, 2, 7, 10]); +/// let int_values = Int64Array::from_iter(0..10); +/// let list_arr = try_new_generic_list_array(int_values, &offsets).unwrap(); +/// assert_eq!(list_arr, +/// ListArray::from_iter_primitive::(vec![ +/// Some(vec![Some(0), Some(1)]), +/// Some(vec![Some(2), Some(3), Some(4), Some(5), Some(6)]), +/// Some(vec![Some(7), Some(8), Some(9)]), +/// ])) +/// ``` +pub fn try_new_generic_list_array( + values: T, + offsets: &PrimitiveArray, +) -> Result> +where + Offset::Native: OffsetSizeTrait, +{ + let data_type = if Offset::Native::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + values.data_type().clone(), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + values.data_type().clone(), + true, + ))) + }; + let data = ArrayDataBuilder::new(data_type) + .len(offsets.len() - 1) + .add_buffer(offsets.into_data().buffers()[0].clone()) + .add_child_data(values.into_data()) + .build()?; + + Ok(GenericListArray::from(data)) +} + +pub fn fixed_size_list_type(list_width: i32, inner_type: DataType) -> DataType { + DataType::FixedSizeList(Arc::new(Field::new("item", inner_type, true)), list_width) +} + +pub trait FixedSizeListArrayExt { + /// Create an [`FixedSizeListArray`] from values and list size. + /// + /// ``` + /// use arrow_array::{Int64Array, FixedSizeListArray}; + /// use arrow_array::types::Int64Type; + /// use lance_arrow::FixedSizeListArrayExt; + /// + /// let int_values = Int64Array::from_iter(0..10); + /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 2).unwrap(); + /// assert_eq!(fixed_size_list_arr, + /// FixedSizeListArray::from_iter_primitive::(vec![ + /// Some(vec![Some(0), Some(1)]), + /// Some(vec![Some(2), Some(3)]), + /// Some(vec![Some(4), Some(5)]), + /// Some(vec![Some(6), Some(7)]), + /// Some(vec![Some(8), Some(9)]) + /// ], 2)) + /// ``` + fn try_new_from_values( + values: T, + list_size: i32, + ) -> Result; + + /// Sample `n` rows from the [FixedSizeListArray] + /// + /// ``` + /// use arrow_array::{Int64Array, FixedSizeListArray, Array}; + /// use lance_arrow::FixedSizeListArrayExt; + /// + /// let int_values = Int64Array::from_iter(0..256); + /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 16).unwrap(); + /// let sampled = fixed_size_list_arr.sample(10).unwrap(); + /// assert_eq!(sampled.len(), 10); + /// assert_eq!(sampled.value_length(), 16); + /// assert_eq!(sampled.values().len(), 160); + /// ``` + fn sample(&self, n: usize) -> Result; + + /// Ensure the [FixedSizeListArray] of Float16, Float32, Float64, + /// Int8, Int16, Int32, Int64, UInt8, UInt32 type to its closest floating point type. + fn convert_to_floating_point(&self) -> Result; +} + +impl FixedSizeListArrayExt for FixedSizeListArray { + fn try_new_from_values(values: T, list_size: i32) -> Result { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let values = Arc::new(values); + + Self::try_new(field, list_size, values, None) + } + + fn sample(&self, n: usize) -> Result { + if n >= self.len() { + return Ok(self.clone()); + } + let mut rng = SmallRng::from_os_rng(); + let chosen = (0..self.len() as u32).choose_multiple(&mut rng, n); + take(self, &UInt32Array::from(chosen), None).map(|arr| arr.as_fixed_size_list().clone()) + } + + fn convert_to_floating_point(&self) -> Result { + match self.data_type() { + DataType::FixedSizeList(field, size) => match field.data_type() { + DataType::Float16 | DataType::Float32 | DataType::Float64 => Ok(self.clone()), + DataType::Int8 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int8Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int16 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int16Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int32 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int32Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int64 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int64Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + DataType::UInt8 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to UInt8Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + DataType::UInt32 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to UInt32Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + data_type => Err(ArrowError::ParseError(format!( + "Expect either floating type or integer got {:?}", + data_type + ))), + }, + data_type => Err(ArrowError::ParseError(format!( + "Expect either FixedSizeList got {:?}", + data_type + ))), + } + } +} + +/// Force downcast of an [`Array`], such as an [`ArrayRef`], to +/// [`FixedSizeListArray`], panic'ing on failure. +pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray { + arr.as_any().downcast_ref::().unwrap() +} + +pub trait FixedSizeBinaryArrayExt { + /// Create an [`FixedSizeBinaryArray`] from values and stride. + /// + /// ``` + /// use arrow_array::{UInt8Array, FixedSizeBinaryArray}; + /// use arrow_array::types::UInt8Type; + /// use lance_arrow::FixedSizeBinaryArrayExt; + /// + /// let int_values = UInt8Array::from_iter(0..10); + /// let fixed_size_list_arr = FixedSizeBinaryArray::try_new_from_values(&int_values, 2).unwrap(); + /// assert_eq!(fixed_size_list_arr, + /// FixedSizeBinaryArray::try_from_iter( + /// [[0u8, 1], [2, 3], [4, 5], [6, 7], [8, 9]].into_iter() + /// ).unwrap()) + /// ``` + fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result; +} + +impl FixedSizeBinaryArrayExt for FixedSizeBinaryArray { + fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result { + let data_type = DataType::FixedSizeBinary(stride); + let data = ArrayDataBuilder::new(data_type) + .len(values.len() / stride as usize) + .add_buffer(values.into_data().buffers()[0].clone()) + .build()?; + Ok(Self::from(data)) + } +} + +pub fn as_fixed_size_binary_array(arr: &dyn Array) -> &FixedSizeBinaryArray { + arr.as_any().downcast_ref::().unwrap() +} + +pub fn iter_str_array(arr: &dyn Array) -> Box> + Send + '_> { + match arr.data_type() { + DataType::Utf8 => Box::new(arr.as_string::().iter()), + DataType::LargeUtf8 => Box::new(arr.as_string::().iter()), + DataType::Utf8View => Box::new(arr.as_string_view().iter()), + _ => panic!( + "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}", + arr.data_type() + ), + } +} + +pub fn iter_binary_array( + arr: &dyn Array, +) -> Result> + Send + '_>> { + match arr.data_type() { + DataType::Binary => Ok(Box::new(arr.as_binary::().iter())), + DataType::LargeBinary => Ok(Box::new(arr.as_binary::().iter())), + DataType::BinaryView => Ok(Box::new(arr.as_binary_view().iter())), + DataType::FixedSizeBinary(_) => Ok(Box::new(arr.as_fixed_size_binary().iter())), + data_type => Err(ArrowError::InvalidArgumentError(format!( + "Expecting a binary type, found {data_type}" + ))), + } +} + +/// Extends Arrow's [RecordBatch]. +pub trait RecordBatchExt { + /// Append a new column to this [`RecordBatch`] and returns a new RecordBatch. + /// + /// ``` + /// use std::sync::Arc; + /// use arrow_array::{RecordBatch, Int32Array, StringArray}; + /// use arrow_schema::{Schema, Field, DataType}; + /// use lance_arrow::*; + /// + /// let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + /// let record_batch = RecordBatch::try_new(schema, vec![int_arr.clone()]).unwrap(); + /// + /// let new_field = Field::new("s", DataType::Utf8, true); + /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + /// let new_record_batch = record_batch.try_with_column(new_field, str_arr.clone()).unwrap(); + /// + /// assert_eq!( + /// new_record_batch, + /// RecordBatch::try_new( + /// Arc::new(Schema::new( + /// vec![ + /// Field::new("a", DataType::Int32, true), + /// Field::new("s", DataType::Utf8, true) + /// ]) + /// ), + /// vec![int_arr, str_arr], + /// ).unwrap() + /// ) + /// ``` + fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result; + + /// Created a new RecordBatch with column at index. + fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result; + + /// Creates a new [`RecordBatch`] from the provided [`StructArray`]. + /// + /// The fields on the [`StructArray`] need to match this [`RecordBatch`] schema + fn try_new_from_struct_array(&self, arr: StructArray) -> Result; + + /// Merge with another [`RecordBatch`] and returns a new one. + /// + /// Fields are merged based on name. First we iterate the left columns. If a matching + /// name is found in the right then we merge the two columns. If there is no match then + /// we add the left column to the output. + /// + /// To merge two columns we consider the type. If both arrays are struct arrays we recurse. + /// Otherwise we use the left array. + /// + /// Afterwards we add all non-matching right columns to the output. + /// + /// Note: This method likely does not handle nested fields correctly and you may want to consider + /// using [`Self::merge_with_schema`] instead. + /// ``` + /// use std::sync::Arc; + /// use arrow_array::*; + /// use arrow_schema::{Schema, Field, DataType}; + /// use lance_arrow::*; + /// + /// let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + /// let left = RecordBatch::try_new(left_schema, vec![int_arr.clone()]).unwrap(); + /// + /// let right_schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + /// let right = RecordBatch::try_new(right_schema, vec![str_arr.clone()]).unwrap(); + /// + /// let new_record_batch = left.merge(&right).unwrap(); + /// + /// assert_eq!( + /// new_record_batch, + /// RecordBatch::try_new( + /// Arc::new(Schema::new( + /// vec![ + /// Field::new("a", DataType::Int32, true), + /// Field::new("s", DataType::Utf8, true) + /// ]) + /// ), + /// vec![int_arr, str_arr], + /// ).unwrap() + /// ) + /// ``` + /// + /// TODO: add merge nested fields support. + fn merge(&self, other: &RecordBatch) -> Result; + + /// Create a batch by merging columns between two batches with a given schema. + /// + /// A reference schema is used to determine the proper ordering of nested fields. + /// + /// For each field in the reference schema we look for corresponding fields in + /// the left and right batches. If a field is found in both batches we recursively merge + /// it. + /// + /// If a field is only in the left or right batch we take it as it is. + fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result; + + /// Drop one column specified with the name and return the new [`RecordBatch`]. + /// + /// If the named column does not exist, it returns a copy of this [`RecordBatch`]. + fn drop_column(&self, name: &str) -> Result; + + /// Replace a column (specified by name) and return the new [`RecordBatch`]. + fn replace_column_by_name(&self, name: &str, column: Arc) -> Result; + + /// Replace a column schema (specified by name) and return the new [`RecordBatch`]. + fn replace_column_schema_by_name( + &self, + name: &str, + new_data_type: DataType, + column: Arc, + ) -> Result; + + /// Rename a column at a given index. + fn rename_column(&self, index: usize, new_name: &str) -> Result; + + /// Get (potentially nested) column by qualified name. + fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>; + + /// Project the schema over the [RecordBatch]. + fn project_by_schema(&self, schema: &Schema) -> Result; + + /// metadata of the schema. + fn metadata(&self) -> &HashMap; + + /// Add metadata to the schema. + fn add_metadata(&self, key: String, value: String) -> Result { + let mut metadata = self.metadata().clone(); + metadata.insert(key, value); + self.with_metadata(metadata) + } + + /// Replace the schema metadata with the provided one. + fn with_metadata(&self, metadata: HashMap) -> Result; + + /// Take selected rows from the [RecordBatch]. + fn take(&self, indices: &UInt32Array) -> Result; + + /// Create a new RecordBatch with compacted memory after slicing. + fn shrink_to_fit(&self) -> Result; + + /// Helper method to sort the RecordBatch by a column + fn sort_by_column(&self, column: usize, options: Option) -> Result; +} + +impl RecordBatchExt for RecordBatch { + fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result { + let new_schema = Arc::new(self.schema().as_ref().try_with_column(field)?); + let mut new_columns = self.columns().to_vec(); + new_columns.push(arr); + Self::try_new(new_schema, new_columns) + } + + fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result { + let new_schema = Arc::new(self.schema().as_ref().try_with_column_at(index, field)?); + let mut new_columns = self.columns().to_vec(); + new_columns.insert(index, arr); + Self::try_new(new_schema, new_columns) + } + + fn try_new_from_struct_array(&self, arr: StructArray) -> Result { + let schema = Arc::new(Schema::new_with_metadata( + arr.fields().to_vec(), + self.schema().metadata.clone(), + )); + let batch = Self::from(arr); + batch.with_schema(schema) + } + + fn merge(&self, other: &Self) -> Result { + if self.num_rows() != other.num_rows() { + return Err(ArrowError::InvalidArgumentError(format!( + "Attempt to merge two RecordBatch with different sizes: {} != {}", + self.num_rows(), + other.num_rows() + ))); + } + let left_struct_array: StructArray = self.clone().into(); + let right_struct_array: StructArray = other.clone().into(); + self.try_new_from_struct_array(merge(&left_struct_array, &right_struct_array)) + } + + fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result { + if self.num_rows() != other.num_rows() { + return Err(ArrowError::InvalidArgumentError(format!( + "Attempt to merge two RecordBatch with different sizes: {} != {}", + self.num_rows(), + other.num_rows() + ))); + } + let left_struct_array: StructArray = self.clone().into(); + let right_struct_array: StructArray = other.clone().into(); + self.try_new_from_struct_array(merge_with_schema( + &left_struct_array, + &right_struct_array, + schema.fields(), + )) + } + + fn drop_column(&self, name: &str) -> Result { + let mut fields = vec![]; + let mut columns = vec![]; + for i in 0..self.schema().fields.len() { + if self.schema().field(i).name() != name { + fields.push(self.schema().field(i).clone()); + columns.push(self.column(i).clone()); + } + } + Self::try_new( + Arc::new(Schema::new_with_metadata( + fields, + self.schema().metadata().clone(), + )), + columns, + ) + } + + fn rename_column(&self, index: usize, new_name: &str) -> Result { + let mut fields = self.schema().fields().to_vec(); + if index >= fields.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Index out of bounds: {}", + index + ))); + } + fields[index] = Arc::new(Field::new( + new_name, + fields[index].data_type().clone(), + fields[index].is_nullable(), + )); + Self::try_new( + Arc::new(Schema::new_with_metadata( + fields, + self.schema().metadata().clone(), + )), + self.columns().to_vec(), + ) + } + + fn replace_column_by_name(&self, name: &str, column: Arc) -> Result { + let mut columns = self.columns().to_vec(); + let field_i = self + .schema() + .fields() + .iter() + .position(|f| f.name() == name) + .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?; + columns[field_i] = column; + Self::try_new(self.schema(), columns) + } + + fn replace_column_schema_by_name( + &self, + name: &str, + new_data_type: DataType, + column: Arc, + ) -> Result { + let fields = self + .schema() + .fields() + .iter() + .map(|x| { + if x.name() != name { + x.clone() + } else { + let new_field = Field::new(name, new_data_type.clone(), x.is_nullable()); + Arc::new(new_field) + } + }) + .collect::>(); + let schema = Schema::new_with_metadata(fields, self.schema().metadata.clone()); + let mut columns = self.columns().to_vec(); + let field_i = self + .schema() + .fields() + .iter() + .position(|f| f.name() == name) + .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?; + columns[field_i] = column; + Self::try_new(Arc::new(schema), columns) + } + + fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef> { + let split = name.split('.').collect::>(); + if split.is_empty() { + return None; + } + + self.column_by_name(split[0]) + .and_then(|arr| get_sub_array(arr, &split[1..])) + } + + fn project_by_schema(&self, schema: &Schema) -> Result { + let struct_array: StructArray = self.clone().into(); + self.try_new_from_struct_array(project(&struct_array, schema.fields())?) + } + + fn metadata(&self) -> &HashMap { + self.schema_ref().metadata() + } + + fn with_metadata(&self, metadata: HashMap) -> Result { + let mut schema = self.schema_ref().as_ref().clone(); + schema.metadata = metadata; + Self::try_new(schema.into(), self.columns().into()) + } + + fn take(&self, indices: &UInt32Array) -> Result { + let struct_array: StructArray = self.clone().into(); + let taken = take(&struct_array, indices, None)?; + self.try_new_from_struct_array(taken.as_struct().clone()) + } + + fn shrink_to_fit(&self) -> Result { + // Deep copy the sliced record batch, instead of whole batch + crate::deepcopy::deep_copy_batch_sliced(self) + } + + fn sort_by_column(&self, column: usize, options: Option) -> Result { + if column >= self.num_columns() { + return Err(ArrowError::InvalidArgumentError(format!( + "Column index out of bounds: {}", + column + ))); + } + let column = self.column(column); + let sorted = arrow_ord::sort::sort_to_indices(column, options, None)?; + self.take(&sorted) + } +} + +/// Recursively projects an array to match the target field's structure. +/// This handles reordering fields inside nested List types. +fn project_array(array: &ArrayRef, target_field: &Field) -> Result { + match target_field.data_type() { + DataType::Struct(subfields) => { + let struct_arr = array.as_struct(); + let projected = project(struct_arr, subfields)?; + Ok(Arc::new(projected)) + } + DataType::List(inner_field) => { + let list_arr: &ListArray = array.as_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(ListArray::new( + inner_field.clone(), + list_arr.offsets().clone(), + projected_values, + list_arr.nulls().cloned(), + ))) + } + DataType::LargeList(inner_field) => { + let list_arr: &LargeListArray = array.as_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(LargeListArray::new( + inner_field.clone(), + list_arr.offsets().clone(), + projected_values, + list_arr.nulls().cloned(), + ))) + } + DataType::FixedSizeList(inner_field, size) => { + let list_arr = array.as_fixed_size_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(FixedSizeListArray::new( + inner_field.clone(), + *size, + projected_values, + list_arr.nulls().cloned(), + ))) + } + // A nullable entries field fails MapArray::try_new unconditionally, + // so a (schema-invalid) map declared that way keeps the clone + // fallthrough it always had rather than gaining a new error. + DataType::Map(entries_field, sorted) if !entries_field.is_nullable() => { + let map_arr = array.as_map(); + let DataType::Struct(entry_fields) = entries_field.data_type() else { + return Err(ArrowError::SchemaError(format!( + "Map entries field must be a struct, got {}", + entries_field.data_type() + ))); + }; + let projected_entries = project(map_arr.entries(), entry_fields)?; + // try_new re-checks the entries invariants (a non-null entries + // struct, two entry columns, offset bounds); null keys are ruled + // out one level down, by the struct rebuild against the + // non-nullable key field. + Ok(Arc::new(MapArray::try_new( + entries_field.clone(), + map_arr.offsets().clone(), + projected_entries, + map_arr.nulls().cloned(), + *sorted, + )?)) + } + _ => Ok(array.clone()), + } +} + +fn project(struct_array: &StructArray, fields: &Fields) -> Result { + if fields.is_empty() { + return Ok(StructArray::new_empty_fields( + struct_array.len(), + struct_array.nulls().cloned(), + )); + } + let mut columns: Vec = vec![]; + for field in fields.iter() { + if let Some(col) = struct_array.column_by_name(field.name()) { + let projected = project_array(col, field.as_ref())?; + columns.push(projected); + } else { + return Err(ArrowError::SchemaError(format!( + "field {} does not exist in the RecordBatch", + field.name() + ))); + } + } + // Preserve the struct's validity when projecting + StructArray::try_new(fields.clone(), columns, struct_array.nulls().cloned()) +} + +/// Whether `data_type` has a struct somewhere below it, through lists. +fn holds_struct(data_type: &DataType) -> bool { + match data_type { + DataType::Struct(_) => true, + DataType::List(item) | DataType::LargeList(item) => holds_struct(item.data_type()), + _ => false, + } +} + +/// Merge the values of two lists: structs field by field, nested lists +/// recursively, anything else keeps the left side. +fn merge_list_values(left: &ArrayRef, right: &ArrayRef) -> ArrayRef { + match (left.data_type(), right.data_type()) { + (DataType::Struct(_), DataType::Struct(_)) => { + Arc::new(merge(left.as_struct(), right.as_struct())) + } + (DataType::List(_), DataType::List(_)) + | (DataType::LargeList(_), DataType::LargeList(_)) + if left.data_type() != right.data_type() + && holds_struct(left.data_type()) + && holds_struct(right.data_type()) => + { + merge_list_struct(left, right) + } + _ => left.clone(), + } +} + +fn merge_list_struct_helper( + left: &dyn Array, + right: &dyn Array, +) -> Arc { + let left_list: &GenericListArray = left.as_list(); + let right_list: &GenericListArray = right.as_list(); + let item_name = match left.data_type() { + DataType::List(item) | DataType::LargeList(item) => item.name().clone(), + _ => unreachable!(), + }; + // One side all-null (a placeholder for columns this fragment lacks) has + // no values, so stand in nulls shaped like the other side's. + let null_left = left_list.null_count() == left_list.len(); + let null_right = !null_left && right_list.null_count() == right_list.len(); + let (merged_items, list, items_nullable) = if null_left || null_right { + let not_null = if null_left { right_list } else { left_list }; + let values_len = not_null.values().len(); + let values = |list: &GenericListArray, is_null: bool| { + if is_null { + new_null_array(list.values().data_type(), values_len) + } else { + list.values().clone() + } + }; + let merged = merge_list_values( + &values(left_list, null_left), + &values(right_list, null_right), + ); + (merged, not_null, true) + } else { + if left_list.offsets().inner() != right_list.offsets().inner() { + panic!("Attempt to merge list struct arrays which do not have same offsets"); + } + let merged = merge_list_values(left_list.values(), right_list.values()); + let nullable = match left.data_type() { + DataType::List(item) | DataType::LargeList(item) => item.is_nullable(), + _ => unreachable!(), + }; + (merged, left_list, nullable) + }; + let items_field = Arc::new(Field::new( + item_name, + merged_items.data_type().clone(), + items_nullable, + )); + Arc::new(GenericListArray::::new( + items_field, + list.offsets().clone(), + merged_items, + list.nulls().cloned(), + )) +} + +/// Merge two lists whose items hold structs. Only succeeds if both lists have +/// the same offsets or one of them is all-null; panics otherwise. +fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc { + match (left.data_type(), right.data_type()) { + (DataType::List(_), DataType::List(_)) => merge_list_struct_helper::(left, right), + (DataType::LargeList(_), DataType::LargeList(_)) => { + merge_list_struct_helper::(left, right) + } + _ => unreachable!(), + } +} + +/// Helper function to merge validity buffers from two struct arrays. +/// +/// A row is valid if it is valid in either input. +/// An absent validity buffer means all rows are valid, an all-null buffer acts as the identity for this merge. +fn merge_struct_validity( + left_validity: Option<&arrow_buffer::NullBuffer>, + right_validity: Option<&arrow_buffer::NullBuffer>, +) -> Option { + match (left_validity, right_validity) { + // Fast paths: no computation needed + (None, _) | (_, None) => None, + (Some(left), Some(right)) => { + if left.null_count() == 0 || right.null_count() == 0 { + return None; + } + if left.null_count() == left.len() { + return Some(right.clone()); + } + if right.null_count() == right.len() { + return Some(left.clone()); + } + + let left_buffer = left.inner(); + let right_buffer = right.inner(); + + // Perform bitwise OR directly on BooleanBuffers + // This preserves the correct semantics: 1 = valid, 0 = null + let merged_buffer = left_buffer | right_buffer; + + Some(arrow_buffer::NullBuffer::from(merged_buffer)) + } + } +} + +fn merge_list_child_values( + child_field: &Field, + left_values: ArrayRef, + right_values: ArrayRef, +) -> ArrayRef { + match child_field.data_type() { + DataType::Struct(child_fields) => Arc::new(merge_with_schema( + left_values.as_struct(), + right_values.as_struct(), + child_fields, + )) as ArrayRef, + DataType::List(grandchild) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be ListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be ListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(ListArray::new( + grandchild.clone(), + left_list.offsets().clone(), + merged_values, + merged_validity, + )) as ArrayRef + } + DataType::LargeList(grandchild) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be LargeListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be LargeListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(LargeListArray::new( + grandchild.clone(), + left_list.offsets().clone(), + merged_values, + merged_validity, + )) as ArrayRef + } + DataType::FixedSizeList(grandchild, list_size) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be FixedSizeListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be FixedSizeListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(FixedSizeListArray::new( + grandchild.clone(), + *list_size, + merged_values, + merged_validity, + )) as ArrayRef + } + _ => left_values.clone(), + } +} + +// Helper function to adjust child array validity based on parent struct validity +// When parent struct is null, propagates null to child array +// Optimized with fast paths and SIMD operations +fn adjust_child_validity( + child: &ArrayRef, + parent_validity: Option<&arrow_buffer::NullBuffer>, +) -> ArrayRef { + // Fast path: no parent validity means no adjustment needed + let parent_validity = match parent_validity { + None => return child.clone(), + Some(p) if p.null_count() == 0 => return child.clone(), // No nulls to propagate + Some(p) => p, + }; + + // Fast path: DataType::Null arrays are always entirely null by definition and cannot + // carry an explicit null bitmap (Arrow rejects it). No adjustment is needed. + if child.data_type() == &DataType::Null { + return child.clone(); + } + + let child_validity = child.nulls(); + + // Compute the new validity: child_validity AND parent_validity + let new_validity = match child_validity { + None => { + // Fast path: child has no existing validity, just use parent's + parent_validity.clone() + } + Some(child_nulls) => { + let child_buffer = child_nulls.inner(); + let parent_buffer = parent_validity.inner(); + + // Perform bitwise AND directly on BooleanBuffers + // This preserves the correct semantics: 1 = valid, 0 = null + let merged_buffer = child_buffer & parent_buffer; + + arrow_buffer::NullBuffer::from(merged_buffer) + } + }; + + // ArrayData keeps the child's offset for its value buffers. Its null bitmap + // must cover that offset too, even though the merged validity describes only + // the visible rows. Copy logical bits into an offset-aligned bitmap when needed. + let null_bit_buffer = if child.offset() == 0 && new_validity.inner().offset() == 0 { + new_validity.into_inner().into_inner() + } else { + let mut builder = + arrow_buffer::builder::BooleanBufferBuilder::new(child.offset() + child.len()); + builder.append_n(child.offset(), true); + builder.append_buffer(new_validity.inner()); + builder.build().into_inner() + }; + + // Create new array with adjusted validity + arrow_array::make_array( + arrow_data::ArrayData::try_new( + child.data_type().clone(), + child.len(), + Some(null_bit_buffer), + child.offset(), + child.to_data().buffers().to_vec(), + child.to_data().child_data().to_vec(), + ) + .unwrap(), + ) +} + +fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray { + let mut fields: Vec = vec![]; + let mut columns: Vec = vec![]; + let right_fields = right_struct_array.fields(); + let right_columns = right_struct_array.columns(); + + // Get the validity buffers from both structs + let left_validity = left_struct_array.nulls(); + let right_validity = right_struct_array.nulls(); + + // Compute merged validity + let merged_validity = merge_struct_validity(left_validity, right_validity); + + // iterate through the fields on the left hand side + for (left_field, left_column) in left_struct_array + .fields() + .iter() + .zip(left_struct_array.columns().iter()) + { + match right_fields + .iter() + .position(|f| f.name() == left_field.name()) + { + // if the field exists on the right hand side, merge them recursively if appropriate + Some(right_index) => { + let right_field = right_fields.get(right_index).unwrap(); + let right_column = right_columns.get(right_index).unwrap(); + // if both fields are struct, merge them recursively + match (left_field.data_type(), right_field.data_type()) { + (DataType::Struct(_), DataType::Struct(_)) => { + let left_sub_array = left_column.as_struct(); + let right_sub_array = right_column.as_struct(); + let merged_sub_array = merge(left_sub_array, right_sub_array); + fields.push(Field::new( + left_field.name(), + merged_sub_array.data_type().clone(), + left_field.is_nullable(), + )); + columns.push(Arc::new(merged_sub_array) as ArrayRef); + } + (DataType::List(left_list), DataType::List(right_list)) + | (DataType::LargeList(left_list), DataType::LargeList(right_list)) + if holds_struct(left_list.data_type()) + && holds_struct(right_list.data_type()) => + { + // Identical inner types: nothing to merge, keep the left column. + if left_list.data_type() == right_list.data_type() { + fields.push(left_field.as_ref().clone()); + columns.push(left_column.clone()); + } else { + // The struct fields differ, so merge them structurally. merge_list_struct + // only succeeds when both lists share offsets or one side is all-null; + // it panics otherwise. + let merged_sub_array = merge_list_struct(&left_column, &right_column); + + fields.push(Field::new( + left_field.name(), + merged_sub_array.data_type().clone(), + left_field.is_nullable(), + )); + columns.push(merged_sub_array); + } + } + // otherwise, just use the field on the left hand side + _ => { + // TODO handle list-of-struct and other types + fields.push(left_field.as_ref().clone()); + // Adjust the column validity: if left struct was null, propagate to child + let adjusted_column = adjust_child_validity(left_column, left_validity); + columns.push(adjusted_column); + } + } + } + None => { + fields.push(left_field.as_ref().clone()); + // Adjust the column validity: if left struct was null, propagate to child + let adjusted_column = adjust_child_validity(left_column, left_validity); + columns.push(adjusted_column); + } + } + } + + // now iterate through the fields on the right hand side + right_fields + .iter() + .zip(right_columns.iter()) + .for_each(|(field, column)| { + // add new columns on the right + if !left_struct_array + .fields() + .iter() + .any(|f| f.name() == field.name()) + { + fields.push(field.as_ref().clone()); + // This field doesn't exist on the left + // We use the right's column but need to adjust for struct validity + let adjusted_column = adjust_child_validity(column, right_validity); + columns.push(adjusted_column); + } + }); + + StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap() +} + +fn merge_with_schema( + left_struct_array: &StructArray, + right_struct_array: &StructArray, + fields: &Fields, +) -> StructArray { + // Helper function that returns true if both types are struct or both are non-struct + fn same_type_kind(left: &DataType, right: &DataType) -> bool { + match (left, right) { + (DataType::Struct(_), DataType::Struct(_)) => true, + (DataType::Struct(_), _) => false, + (_, DataType::Struct(_)) => false, + _ => true, + } + } + + let mut output_fields: Vec = Vec::with_capacity(fields.len()); + let mut columns: Vec = Vec::with_capacity(fields.len()); + + let left_fields = left_struct_array.fields(); + let left_columns = left_struct_array.columns(); + let right_fields = right_struct_array.fields(); + let right_columns = right_struct_array.columns(); + + // Get the validity buffers from both structs + let left_validity = left_struct_array.nulls(); + let right_validity = right_struct_array.nulls(); + + // Compute merged validity + let merged_validity = merge_struct_validity(left_validity, right_validity); + + for field in fields { + let left_match_idx = left_fields.iter().position(|f| { + f.name() == field.name() && same_type_kind(f.data_type(), field.data_type()) + }); + let right_match_idx = right_fields.iter().position(|f| { + f.name() == field.name() && same_type_kind(f.data_type(), field.data_type()) + }); + + match (left_match_idx, right_match_idx) { + (None, Some(right_idx)) => { + output_fields.push(right_fields[right_idx].as_ref().clone()); + // Adjust validity if the right struct was null + let adjusted_column = + adjust_child_validity(&right_columns[right_idx], right_validity); + columns.push(adjusted_column); + } + (Some(left_idx), None) => { + output_fields.push(left_fields[left_idx].as_ref().clone()); + // Adjust validity if the left struct was null + let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity); + columns.push(adjusted_column); + } + (Some(left_idx), Some(right_idx)) => { + match field.data_type() { + DataType::Struct(child_fields) => { + let left_sub_array = left_columns[left_idx].as_struct(); + let right_sub_array = right_columns[right_idx].as_struct(); + let merged_sub_array = + merge_with_schema(left_sub_array, right_sub_array, child_fields); + output_fields.push(Field::new( + field.name(), + merged_sub_array.data_type().clone(), + field.is_nullable(), + )); + columns.push(Arc::new(merged_sub_array) as ArrayRef); + } + DataType::List(child_field) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.trimmed_values(), + right_list.trimmed_values(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + // `trimmed_values` starts at the first used value, so offsets + // must be shifted to match or `ListArray::new` panics when the + // input list was sliced (e.g. from a filtered batch). + let merged_list = ListArray::new( + child_field.clone(), + left_list.trimmed_offsets(), + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + DataType::LargeList(child_field) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.trimmed_values(), + right_list.trimmed_values(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + let merged_list = LargeListArray::new( + child_field.clone(), + left_list.trimmed_offsets(), + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + DataType::FixedSizeList(child_field, list_size) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + let merged_list = FixedSizeListArray::new( + child_field.clone(), + *list_size, + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + _ => { + output_fields.push(left_fields[left_idx].as_ref().clone()); + // For fields that exist in both, use left but adjust validity + let adjusted_column = + adjust_child_validity(&left_columns[left_idx], left_validity); + columns.push(adjusted_column); + } + } + } + (None, None) => { + // The field will not be included in the output + } + } + } + + StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap() +} + +fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> { + if components.is_empty() { + return Some(array); + } + if !matches!(array.data_type(), DataType::Struct(_)) { + return None; + } + let struct_arr = array.as_struct(); + struct_arr + .column_by_name(components[0]) + .and_then(|arr| get_sub_array(arr, &components[1..])) +} + +/// Interleave multiple RecordBatches into a single RecordBatch. +/// +/// Behaves like [`arrow_select::interleave::interleave`], but for RecordBatches. +pub fn interleave_batches( + batches: &[RecordBatch], + indices: &[(usize, usize)], +) -> Result { + let first_batch = batches.first().ok_or_else(|| { + ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string()) + })?; + let schema = first_batch.schema(); + let num_columns = first_batch.num_columns(); + let mut columns = Vec::with_capacity(num_columns); + let mut chunks = Vec::with_capacity(batches.len()); + + for i in 0..num_columns { + for batch in batches { + chunks.push(batch.column(i).as_ref()); + } + let new_column = interleave(&chunks, indices)?; + columns.push(new_column); + chunks.clear(); + } + + RecordBatch::try_new(schema, columns) +} + +pub trait BufferExt { + /// Create an `arrow_buffer::Buffer`` from a `bytes::Bytes` object + /// + /// The alignment must be specified (as `bytes_per_value`) since we want to make + /// sure we can safely reinterpret the buffer. + /// + /// If the buffer is properly aligned this will be zero-copy. If not, a copy + /// will be made and an owned buffer returned. + /// + /// If `bytes_per_value` is not a power of two, then we assume the buffer is + /// never going to be reinterpreted into another type and we can safely + /// ignore the alignment. + /// + /// Yes, the method name is odd. It's because there is already a `from_bytes` + /// which converts from `arrow_buffer::bytes::Bytes` (not `bytes::Bytes`) + fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self; + + /// Allocates a new properly aligned arrow buffer and copies `bytes` into it + /// + /// `size_bytes` can be larger than `bytes` and, if so, the trailing bytes will + /// be zeroed out. + /// + /// # Panics + /// + /// Panics if `size_bytes` is less than `bytes.len()` + fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self; +} + +fn is_pwr_two(n: u64) -> bool { + n & (n - 1) == 0 +} + +impl BufferExt for arrow_buffer::Buffer { + fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self { + if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0 + { + // The original buffer is not aligned, cannot zero-copy + let size_bytes = bytes.len(); + Self::copy_bytes_bytes(bytes, size_bytes) + } else { + // The original buffer is aligned, can zero-copy + // SAFETY: the alignment is correct we can make this conversion + unsafe { + Self::from_custom_allocation( + NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"), + bytes.len(), + Arc::new(bytes), + ) + } + } + } + + fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self { + assert!(size_bytes >= bytes.len()); + let mut buf = MutableBuffer::with_capacity(size_bytes); + let to_fill = size_bytes - bytes.len(); + buf.extend(bytes); + buf.extend(std::iter::repeat_n(0_u8, to_fill)); + + // FIX for issue #4512: Shrink buffer to actual size before converting to immutable + // This reduces memory overhead from capacity over-allocation + buf.shrink_to_fit(); + + Self::from(buf) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ + BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, Float32Array, Int32Array, + LargeBinaryArray, ListArray, NullArray, StringArray, StructArray, new_empty_array, + new_null_array, + }; + use arrow_buffer::OffsetBuffer; + + #[test] + fn test_adjust_child_validity_with_sliced_child() { + let child: ArrayRef = Arc::new(BooleanArray::from_iter((0..1030).map(|index| { + if index == 1027 { + None + } else { + Some(index % 2 == 1) + } + }))); + let child = child.slice(1025, 4); + assert_eq!(child.offset(), 1025); + + let parent_validity = arrow_buffer::NullBuffer::from(vec![true, false, true, true]); + let adjusted = adjust_child_validity(&child, Some(&parent_validity)); + assert_eq!(adjusted.offset(), child.offset()); + let adjusted = adjusted.as_any().downcast_ref::().unwrap(); + assert_eq!( + adjusted.iter().collect::>(), + vec![Some(true), None, None, Some(false)] + ); + } + + #[test] + fn test_adjust_child_validity_with_sliced_parent_bitmap() { + let child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let parent_validity: arrow_buffer::NullBuffer = + (0..1030).map(|index| index != 1026).collect(); + let parent_validity = parent_validity.slice(1025, 4); + assert_eq!(parent_validity.inner().offset(), 1025); + + let adjusted = adjust_child_validity(&child, Some(&parent_validity)); + let adjusted = adjusted.as_any().downcast_ref::().unwrap(); + assert_eq!( + adjusted.iter().collect::>(), + vec![Some(1), None, Some(3), Some(4)] + ); + } + + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls() { + // A FixedSizeList with a null inner element must convert to a + // FixedSizeList with the null kept in place. Dropping it would + // shorten the values array and shift every later element (and, when the + // remaining count is not a multiple of the list size, panic). + let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 1.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 3.0); + assert_eq!(conv_values.value(3), 4.0); + } + + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() { + // The Float64-producing arms (Int64/UInt8/UInt32) share the same fix as the + // Float32 arms; cover one representative (UInt8 -> Float64) so both branch + // families are exercised. + let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::UInt8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 10.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 30.0); + assert_eq!(conv_values.value(3), 40.0); + } + + #[test] + fn test_merge_recursive() { + let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]); + let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]); + let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]); + let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]); + + let left_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()), + true, + ), + ]); + let left_batch = RecordBatch::try_new( + Arc::new(left_schema), + vec![ + Arc::new(a_array.clone()), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("c", DataType::Int32, true)), + Arc::new(c_array.clone()) as ArrayRef, + )])), + ], + ) + .unwrap(); + + let right_schema = Schema::new(vec![ + Field::new("e", DataType::Int32, true), + Field::new( + "b", + DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()), + true, + ), + ]); + let right_batch = RecordBatch::try_new( + Arc::new(right_schema), + vec![ + Arc::new(e_array.clone()), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("d", DataType::Utf8, true)), + Arc::new(d_array.clone()) as ArrayRef, + )])) as ArrayRef, + ], + ) + .unwrap(); + + let merged_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Struct( + vec![ + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Utf8, true), + ] + .into(), + ), + true, + ), + Field::new("e", DataType::Int32, true), + ]); + let merged_batch = RecordBatch::try_new( + Arc::new(merged_schema), + vec![ + Arc::new(a_array) as ArrayRef, + Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("c", DataType::Int32, true)), + Arc::new(c_array) as ArrayRef, + ), + ( + Arc::new(Field::new("d", DataType::Utf8, true)), + Arc::new(d_array) as ArrayRef, + ), + ])) as ArrayRef, + Arc::new(e_array) as ArrayRef, + ], + ) + .unwrap(); + + let result = left_batch.merge(&right_batch).unwrap(); + assert_eq!(result, merged_batch); + } + + #[test] + fn test_merge_with_schema() { + fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) { + let fields: Fields = names + .iter() + .zip(types) + .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false)) + .collect(); + let schema = Schema::new(vec![Field::new( + "struct", + DataType::Struct(fields.clone()), + false, + )]); + let children = types.iter().map(new_empty_array).collect::>(); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef], + ); + (schema, batch.unwrap()) + } + + let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]); + let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]); + let (output_schema, _) = test_batch( + &["b", "a", "c"], + &[DataType::Int64, DataType::Int32, DataType::Int32], + ); + + // If we use merge_with_schema the schema is respected + let merged = left_batch + .merge_with_schema(&right_batch, &output_schema) + .unwrap(); + assert_eq!(merged.schema().as_ref(), &output_schema); + + // If we use merge we get first-come first-serve based on the left batch + let (naive_schema, _) = test_batch( + &["a", "b", "c"], + &[DataType::Int32, DataType::Int64, DataType::Int32], + ); + let merged = left_batch.merge(&right_batch).unwrap(); + assert_eq!(merged.schema().as_ref(), &naive_schema); + } + + #[test] + fn test_merge_list_struct() { + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let y_field = Arc::new(Field::new("y", DataType::Int32, true)); + let x_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let y_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![y_field.clone()])), + true, + )); + let both_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])), + true, + )); + let left_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(x_struct_field.clone()), + true, + )]); + let right_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(y_struct_field.clone()), + true, + )]); + let both_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(both_struct_field.clone()), + true, + )]); + + let x = Arc::new(Int32Array::from(vec![1])); + let y = Arc::new(Int32Array::from(vec![2])); + let x_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![x.clone()], + None, + )); + let y_struct = Arc::new(StructArray::new( + Fields::from(vec![y_field.clone()]), + vec![y.clone()], + None, + )); + let both_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone(), y_field.clone()]), + vec![x.clone(), y], + None, + )); + let both_null_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field, y_field]), + vec![x, Arc::new(new_null_array(&DataType::Int32, 1))], + None, + )); + let offsets = OffsetBuffer::from_lengths([1]); + let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None); + let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None); + let both_list = ListArray::new( + both_struct_field.clone(), + offsets.clone(), + both_struct, + None, + ); + let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None); + let x_batch = + RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap(); + let y_batch = RecordBatch::try_new( + Arc::new(right_schema.clone()), + vec![Arc::new(y_s_list.clone())], + ) + .unwrap(); + let merged = x_batch.merge(&y_batch).unwrap(); + let expected = + RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap(); + assert_eq!(merged, expected); + + let y_null_list = new_null_array(y_s_list.data_type(), 1); + let y_null_batch = + RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())]) + .unwrap(); + let expected = + RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap(); + let merged = x_batch.merge(&y_null_batch).unwrap(); + assert_eq!(merged, expected); + } + + #[test] + fn test_merge_all_null_list_adds_nested_child() { + // An all-null placeholder (what a fragment reader merges in for children + // missing from its files) adds `b` to a struct below one or two lists. + let strukt = |cols: Vec<(&str, ArrayRef)>| -> ArrayRef { + let fields = cols + .iter() + .map(|(name, col)| Field::new(*name, col.data_type().clone(), true)) + .collect::>(); + Arc::new(StructArray::new( + fields.into(), + cols.into_iter().map(|(_, col)| col).collect(), + None, + )) + }; + // Two rows: `[values]` and null. + let list = |values: ArrayRef, large: bool| -> ArrayRef { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let lengths = [values.len(), 0]; + let nulls = Some(arrow_buffer::NullBuffer::from(vec![true, false])); + if large { + Arc::new(LargeListArray::new( + field, + OffsetBuffer::from_lengths(lengths), + values, + nulls, + )) + } else { + Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths(lengths), + values, + nulls, + )) + } + }; + let a = || Arc::new(Int32Array::from(vec![7])) as ArrayRef; + let b = || new_null_array(&DataType::Utf8, 1); + for large in [false, true] { + // list>>, built with or without `t`. + let under_list = |inner: ArrayRef, with_t: bool| { + let mut cols = vec![("in", inner)]; + if with_t { + cols.insert(0, ("t", Arc::new(Int32Array::from(vec![1])) as ArrayRef)); + } + list(strukt(cols), large) + }; + // list>> + let under_two_lists = |inner: ArrayRef, _: bool| { + let inner_list: ArrayRef = if large { + let field = Arc::new(Field::new("item", inner.data_type().clone(), true)); + Arc::new(LargeListArray::new( + field, + OffsetBuffer::from_lengths([1]), + inner, + None, + )) + } else { + let field = Arc::new(Field::new("item", inner.data_type().clone(), true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths([1]), + inner, + None, + )) + }; + list(inner_list, large) + }; + let shapes: [&dyn Fn(ArrayRef, bool) -> ArrayRef; 2] = [&under_list, &under_two_lists]; + for shape in shapes { + let data = shape(strukt(vec![("a", a())]), true); + let expected = shape(strukt(vec![("a", a()), ("b", b())]), true); + let placeholder = + new_null_array(shape(strukt(vec![("b", b())]), false).data_type(), 2); + let batch = |col: ArrayRef| { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "c", + col.data_type().clone(), + true, + )])), + vec![col], + ) + .unwrap() + }; + let expected = batch(expected); + // Either side may be the placeholder; field order follows the left. + for merged in [ + batch(data.clone()) + .merge(&batch(placeholder.clone())) + .unwrap(), + batch(placeholder.clone()) + .merge(&batch(data.clone())) + .unwrap(), + ] { + let merged = merged.project_by_schema(expected.schema_ref()).unwrap(); + assert_eq!(merged, expected, "large={large}"); + } + } + } + } + + #[test] + fn test_merge_list_struct_identical_schema() { + // Merging two batches whose `List` columns have identical types + // should yield a single column equal to the input (there is nothing to + // merge), not a struct with the field pushed twice. + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "list_struct", + DataType::List(item_field.clone()), + true, + )])); + + let build_list = |values: Vec| { + let len = values.len(); + let item_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![Arc::new(Int32Array::from(values))], + None, + )); + ListArray::new( + item_field.clone(), + OffsetBuffer::from_lengths([len]), + item_struct, + None, + ) + }; + + // Distinct values so the equality assertion proves the left column is kept, + // not that some column with the same values happens to be present. + let left = + RecordBatch::try_new(schema.clone(), vec![Arc::new(build_list(vec![1, 2]))]).unwrap(); + let right = RecordBatch::try_new(schema, vec![Arc::new(build_list(vec![3, 4]))]).unwrap(); + + let merged = left.merge(&right).unwrap(); + + // Exactly one column, equal to the left input: the field must not be duplicated. + assert_eq!(merged.num_columns(), 1); + assert_eq!(merged, left); + } + + #[test] + fn test_merge_nested_list_struct_identical_schema() { + // A `List` nested inside a struct reaches the identical-type + // short-circuit through the recursive `merge` call. The recursion must + // keep the field exactly once, equal to the left input. + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let companies_field = Arc::new(Field::new( + "companies", + DataType::List(item_field.clone()), + true, + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "outer", + DataType::Struct(Fields::from(vec![companies_field.clone()])), + true, + )])); + + let build_outer = |x: i32| { + // One list row: [{x}]. + let item_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![Arc::new(Int32Array::from(vec![x]))], + None, + )); + let companies = Arc::new(ListArray::new( + item_field.clone(), + OffsetBuffer::from_lengths([1]), + item_struct, + None, + )); + StructArray::new( + Fields::from(vec![companies_field.clone()]), + vec![companies], + None, + ) + }; + + // Distinct values so equality proves the left column is kept. + let left = RecordBatch::try_new(schema.clone(), vec![Arc::new(build_outer(10))]).unwrap(); + let right = RecordBatch::try_new(schema, vec![Arc::new(build_outer(20))]).unwrap(); + + let merged = left.merge(&right).unwrap(); + + // The recursive identical-type merge keeps exactly one `companies` field, + // with no double-push inside the nested struct. + assert_eq!(merged.column(0).as_struct().num_columns(), 1); + assert_eq!(merged, left); + } + + #[test] + fn test_byte_width_opt() { + assert_eq!(DataType::Int32.byte_width_opt(), Some(4)); + assert_eq!(DataType::Int64.byte_width_opt(), Some(8)); + assert_eq!(DataType::Float32.byte_width_opt(), Some(4)); + assert_eq!(DataType::Float64.byte_width_opt(), Some(8)); + assert_eq!(DataType::Utf8.byte_width_opt(), None); + assert_eq!(DataType::Binary.byte_width_opt(), None); + assert_eq!( + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(), + None + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3) + .byte_width_opt(), + Some(12) + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4) + .byte_width_opt(), + Some(16) + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5) + .byte_width_opt(), + None + ); + } + + #[test] + fn test_take_record_batch() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + ], + ) + .unwrap(); + let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap(); + assert_eq!( + taken, + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 5, 10])), + Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])), + ], + ) + .unwrap() + ) + } + + #[test] + fn test_schema_project_by_schema() { + let metadata = [("key".to_string(), "value".to_string())]; + let schema = Arc::new( + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]) + .with_metadata(metadata.clone().into()), + ); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + ], + ) + .unwrap(); + + // Empty schema + let empty_schema = Schema::empty(); + let empty_projected = batch.project_by_schema(&empty_schema).unwrap(); + let expected_schema = empty_schema.with_metadata(metadata.clone().into()); + assert_eq!( + empty_projected, + RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None)) + .with_schema(Arc::new(expected_schema)) + .unwrap() + ); + + // Re-ordered schema + let reordered_schema = Schema::new(vec![ + Field::new("b", DataType::Utf8, true), + Field::new("a", DataType::Int32, true), + ]); + let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap(); + let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into())); + assert_eq!( + reordered_projected, + RecordBatch::try_new( + expected_schema, + vec![ + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + Arc::new(Int32Array::from_iter_values(0..20)), + ], + ) + .unwrap() + ); + + // Sub schema + let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let sub_projected = batch.project_by_schema(&sub_schema).unwrap(); + let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into())); + assert_eq!( + sub_projected, + RecordBatch::try_new( + expected_schema, + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap() + ); + } + + #[test] + fn test_project_rebuilds_sliced_map() { + // A sliced MapArray keeps its full entries array behind sliced + // offsets and validity; the Map projection arm must rebuild it + // without renormalizing either. + let entry_fields = Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(entry_fields.clone()), + false, + )); + let entries = StructArray::new( + entry_fields, + vec![ + Arc::new(StringArray::from(vec!["k0", "k1", "k2"])) as ArrayRef, + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + ], + None, + ); + let map = MapArray::new( + entries_field.clone(), + OffsetBuffer::new(vec![0, 1, 1, 3].into()), + entries, + Some(arrow_buffer::NullBuffer::from(vec![true, false, true])), + false, + ); + let schema = Arc::new(Schema::new(vec![Field::new( + "m", + DataType::Map(entries_field, false), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(map) as ArrayRef]).unwrap(); + + // Rows 1..3: a null slot and a two-entry slot, offsets not zero-based. + let projected = batch + .slice(1, 2) + .project_by_schema(schema.as_ref()) + .unwrap(); + let map = projected.column(0).as_map(); + assert!(map.is_null(0)); + assert!(map.is_valid(1)); + assert_eq!(map.value_length(1), 2); + assert_eq!( + map.value(1) + .column(1) + .as_primitive::() + .values(), + &[2, 3] + ); + } + + #[test] + fn test_project_preserves_struct_validity() { + // Test that projecting a struct array preserves its validity (fix for issue #4385) + let fields = Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Float32, true), + ]); + + // Create a struct array with validity + let id_array = Int32Array::from(vec![1, 2, 3]); + let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]); + let struct_array = StructArray::new( + fields.clone(), + vec![ + Arc::new(id_array) as ArrayRef, + Arc::new(value_array) as ArrayRef, + ], + Some(vec![true, false, true].into()), // Second struct is null + ); + + // Project the struct array + let projected = project(&struct_array, &fields).unwrap(); + + // Verify the validity is preserved + assert_eq!(projected.null_count(), 1); + assert!(!projected.is_null(0)); + assert!(projected.is_null(1)); + assert!(!projected.is_null(2)); + } + + #[test] + fn test_merge_struct_with_different_validity() { + // Test case from Weston's review comment + // File 1 has height field with some nulls + let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]); + let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]); + let left_struct = StructArray::new( + left_fields, + vec![Arc::new(height_array) as ArrayRef], + Some(vec![true, false, true, false].into()), // Rows 2 and 4 are null structs + ); + + // File 2 has width field with some nulls + let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]); + let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]); + let right_struct = StructArray::new( + right_fields, + vec![Arc::new(width_array) as ArrayRef], + Some(vec![true, true, false, false].into()), // Rows 3 and 4 are null structs + ); + + // Merge the two structs + let merged = merge(&left_struct, &right_struct); + + // Expected: + // Row 1: both non-null -> {width: 300, height: 500} + // Row 2: left null, right non-null -> {width: 200, height: null} + // Row 3: left non-null, right null -> {width: null, height: 600} + // Row 4: both null -> null struct + + assert_eq!(merged.null_count(), 1); // Only row 4 is null + assert!(!merged.is_null(0)); + assert!(!merged.is_null(1)); + assert!(!merged.is_null(2)); + assert!(merged.is_null(3)); + + // Check field values + let height_col = merged.column_by_name("height").unwrap(); + let height_values = height_col.as_any().downcast_ref::().unwrap(); + assert_eq!(height_values.value(0), 500); + assert!(height_values.is_null(1)); // height is null when left struct was null + assert_eq!(height_values.value(2), 600); + + let width_col = merged.column_by_name("width").unwrap(); + let width_values = width_col.as_any().downcast_ref::().unwrap(); + assert_eq!(width_values.value(0), 300); + assert_eq!(width_values.value(1), 200); + assert!(width_values.is_null(2)); // width is null when right struct was null + + // An all-null validity buffer is data, not a placeholder meaning "this side has no + // validity": merging two of them keeps the rows null. + let all_null_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef], + Some(vec![false, false].into()), + ); + let all_null_right = StructArray::new( + Fields::from(vec![Field::new("width", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef], + Some(vec![false, false].into()), + ); + + let merged = merge(&all_null_left, &all_null_right); + assert_eq!(merged.null_count(), 2); + + // An all-null side is the identity of the merge, so the other side decides each row. + let partial_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef], + Some(vec![true, false].into()), + ); + let merged = merge(&partial_left, &all_null_right); + assert!(!merged.is_null(0)); + assert!(merged.is_null(1)); + + // A missing validity buffer means all rows are valid, which absorbs an all-null side. + let all_valid_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef], + None, + ); + let merged = merge(&all_valid_left, &all_null_right); + assert_eq!(merged.null_count(), 0); + + // An explicit all-valid buffer has the same semantics as a missing buffer. + let all_valid: arrow_buffer::NullBuffer = vec![true, true].into(); + let partial: arrow_buffer::NullBuffer = vec![true, false].into(); + assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none()); + assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none()); + } + + #[test] + fn test_merge_null_typed_column_with_parent_validity() { + // Reproduces ENT-990: panic in adjust_child_validity when a Null-typed column + // exists on one side and the parent struct has null rows. + // Arrow's Null type has no null bitmap, so passing one to ArrayData::try_new panics. + let left_struct = StructArray::new( + Fields::from(vec![Field::new("a", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef], + Some(vec![true, false].into()), + ); + let right_struct = StructArray::new( + Fields::from(vec![Field::new("b", DataType::Null, true)]), + vec![Arc::new(NullArray::new(2)) as ArrayRef], + Some(vec![true, false].into()), + ); + + // Previously panicked: "Arrays of type Null cannot contain a null bitmask" + let merged = merge(&left_struct, &right_struct); + assert_eq!(merged.len(), 2); + let b_col = merged.column_by_name("b").unwrap(); + // DataType::Null implies all-null by definition; no null bitmap is stored. + assert_eq!(b_col.data_type(), &DataType::Null); + assert_eq!(b_col.len(), 2); + } + + #[test] + fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() { + // left_list setup + let left_company_id = Arc::new(Int32Array::from(vec![None, None])); + let left_count = Arc::new(Int32Array::from(vec![None, None])); + let left_struct = Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("count", DataType::Int32, true), + ]), + vec![left_company_id, left_count], + None, + )); + let left_list = Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([2]), + left_struct, + None, + )); + + // Right List Setup + let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"])); + let right_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]), + vec![right_company_name], + None, + )); + let right_list = Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([2]), + right_struct, + None, + )); + + let target_fields = Fields::from(vec![Field::new( + "companies", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))), + true, + )]); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "companies", + left_list.data_type().clone(), + true, + )])), + vec![left_list as ArrayRef], + ) + .unwrap(); + + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "companies", + right_list.data_type().clone(), + true, + )])), + vec![right_list as ArrayRef], + ) + .unwrap(); + + let merged = left_batch + .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec())) + .unwrap(); + + // Verify the merged structure + let merged_list = merged + .column_by_name("companies") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let merged_struct = merged_list.values().as_struct(); + + // Should have all 3 fields + assert_eq!(merged_struct.num_columns(), 3); + assert!(merged_struct.column_by_name("company_id").is_some()); + assert!(merged_struct.column_by_name("company_name").is_some()); + assert!(merged_struct.column_by_name("count").is_some()); + + // Verify values + let company_id = merged_struct + .column_by_name("company_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(company_id.is_null(0)); + assert!(company_id.is_null(1)); + + let company_name = merged_struct + .column_by_name("company_name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(company_name.value(0), "Google"); + assert_eq!(company_name.value(1), "Microsoft"); + + let count = merged_struct + .column_by_name("count") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(count.is_null(0)); + assert!(count.is_null(1)); + } + + #[test] + fn test_merge_struct_lists() { + test_merge_struct_lists_generic::(); + } + + #[test] + fn test_merge_struct_large_lists() { + test_merge_struct_lists_generic::(); + } + + fn test_merge_struct_lists_generic() { + // left_list setup + let left_company_id = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + Some(11), + Some(12), + Some(13), + Some(14), + Some(15), + Some(16), + Some(17), + Some(18), + Some(19), + Some(20), + ])); + let left_count = Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + Some(30), + Some(40), + Some(50), + Some(60), + Some(70), + Some(80), + Some(90), + Some(100), + Some(110), + Some(120), + Some(130), + Some(140), + Some(150), + Some(160), + Some(170), + Some(180), + Some(190), + Some(200), + ])); + let left_struct = Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("count", DataType::Int32, true), + ]), + vec![left_company_id, left_count], + None, + )); + + let left_list = Arc::new(GenericListArray::::new( + Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([3, 1]), + left_struct.clone(), + None, + )); + + let left_list_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + ))) + }, + true, + )]), + vec![left_list as ArrayRef], + None, + )); + + // right_list setup + let right_company_name = Arc::new(StringArray::from(vec![ + "Google", + "Microsoft", + "Apple", + "Facebook", + ])); + let right_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]), + vec![right_company_name], + None, + )); + let right_list = Arc::new(GenericListArray::::new( + Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([3, 1]), + right_struct.clone(), + None, + )); + + let right_list_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + ))) + }, + true, + )]), + vec![right_list as ArrayRef], + None, + )); + + // prepare schema + let target_fields = Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))) + }, + true, + )]); + + // merge left_list and right_list + let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields); + assert_eq!(merged_array.len(), 2); + } + + #[test] + fn test_merge_with_schema_sliced_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + #[test] + fn test_merge_with_schema_sliced_large_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + // Regression for #6580: merge_with_schema panicked when the left list was a + // sliced view whose offsets did not start at zero (common after a filtered + // scan). Cloning those offsets alongside `trimmed_values` produced offsets + // larger than the trimmed child, panicking in `(Large)ListArray::new`. + fn test_merge_with_schema_sliced_list_struct_generic() { + let make_list_dtype = |item_field: Arc| { + if O::IS_LARGE { + DataType::LargeList(item_field) + } else { + DataType::List(item_field) + } + }; + + // Build a List with two rows of 5 items each, then slice away + // the first row so the remaining list's offsets start at 5, not 0. + let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]); + let left_values = Arc::new(StructArray::new( + struct_fields_a.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef], + None, + )); + let full_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)), + OffsetBuffer::::from_lengths([5, 5]), + left_values, + None, + ); + let sliced_left = full_list.slice(1, 1); + assert_eq!(sliced_left.offsets()[0].as_usize(), 5); + assert_eq!(sliced_left.offsets()[1].as_usize(), 10); + + let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]); + let right_values = Arc::new(StructArray::new( + struct_fields_b.clone(), + vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef], + None, + )); + let right_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)), + OffsetBuffer::::from_lengths([5]), + right_values, + None, + ); + + let target_item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])), + true, + )); + let target_fields = Fields::from(vec![Field::new( + "items", + make_list_dtype(target_item_field), + true, + )]); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + sliced_left.data_type().clone(), + true, + )])), + vec![Arc::new(sliced_left) as ArrayRef], + ) + .unwrap(); + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + right_list.data_type().clone(), + true, + )])), + vec![Arc::new(right_list) as ArrayRef], + ) + .unwrap(); + + let merged = left_batch + .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec())) + .unwrap(); + + let merged_list = merged + .column_by_name("items") + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(merged_list.len(), 1); + assert_eq!(merged_list.value_length(0).as_usize(), 5); + let merged_struct = merged_list.values().as_struct(); + assert_eq!(merged_struct.num_columns(), 2); + let a = merged_struct + .column_by_name("a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + // After shifting offsets to zero, values 5..10 should be first. + let a_vals: Vec = a.iter().map(|v| v.unwrap()).collect(); + assert_eq!(a_vals, vec![5, 6, 7, 8, 9]); + } + + #[test] + fn test_project_by_schema_list_struct_reorder() { + // Test that project_by_schema correctly reorders fields inside List + // This is a regression test for issue #5702 + + // Source schema with inner struct fields in order: c, b, a + let source_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("c", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + Field::new("a", DataType::Utf8, true), + ])); + let source_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "data", + DataType::List(Arc::new(Field::new( + "item", + source_inner_struct.clone(), + true, + ))), + true, + ), + ])); + + // Create source data with c, b, a order + let c_array = StringArray::from(vec!["c1", "c2"]); + let b_array = StringArray::from(vec!["b1", "b2"]); + let a_array = StringArray::from(vec!["a1", "a2"]); + let inner_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("c", DataType::Utf8, true)), + Arc::new(c_array) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + Arc::new(b_array) as ArrayRef, + ), + ( + Arc::new(Field::new("a", DataType::Utf8, true)), + Arc::new(a_array) as ArrayRef, + ), + ]); + + let list_array = ListArray::new( + Arc::new(Field::new("item", source_inner_struct, true)), + OffsetBuffer::from_lengths([1, 1]), + Arc::new(inner_struct), + None, + ); + + let batch = RecordBatch::try_new( + source_schema, + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)], + ) + .unwrap(); + + // Target schema with inner struct fields in order: a, b, c + let target_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Utf8, true), + ])); + let target_schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "data", + DataType::List(Arc::new(Field::new("item", target_inner_struct, true))), + true, + ), + ]); + + // Project should reorder the inner struct fields + let projected = batch.project_by_schema(&target_schema).unwrap(); + + // Verify the schema is correct + assert_eq!(projected.schema().as_ref(), &target_schema); + + // Verify the data is correct by checking inner struct field order + let projected_list = projected.column(1).as_list::(); + let projected_struct = projected_list.values().as_struct(); + + // Fields should now be in order: a, b, c + assert_eq!( + projected_struct.column_by_name("a").unwrap().as_ref(), + &StringArray::from(vec!["a1", "a2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column_by_name("b").unwrap().as_ref(), + &StringArray::from(vec!["b1", "b2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column_by_name("c").unwrap().as_ref(), + &StringArray::from(vec!["c1", "c2"]) as &dyn Array + ); + + // Also verify positional access matches expected order (a=0, b=1, c=2) + assert_eq!( + projected_struct.column(0).as_ref(), + &StringArray::from(vec!["a1", "a2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column(1).as_ref(), + &StringArray::from(vec!["b1", "b2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column(2).as_ref(), + &StringArray::from(vec!["c1", "c2"]) as &dyn Array + ); + } + + #[test] + fn test_project_by_schema_nested_list_struct() { + // Test deeply nested List>> projection + let inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("y", DataType::Int32, true), + Field::new("x", DataType::Int32, true), + ])); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "outer", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("b", DataType::Utf8, true), + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))), + true, + ), + Field::new("a", DataType::Utf8, true), + ])), + true, + ))), + true, + )])); + + // Create deeply nested data + let y_array = Int32Array::from(vec![1, 2]); + let x_array = Int32Array::from(vec![3, 4]); + let innermost_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("y", DataType::Int32, true)), + Arc::new(y_array) as ArrayRef, + ), + ( + Arc::new(Field::new("x", DataType::Int32, true)), + Arc::new(x_array) as ArrayRef, + ), + ]); + let inner_list = ListArray::new( + Arc::new(Field::new("item", inner_struct.clone(), true)), + OffsetBuffer::from_lengths([2]), + Arc::new(innermost_struct), + None, + ); + + let b_array = StringArray::from(vec!["b1"]); + let a_array = StringArray::from(vec!["a1"]); + let middle_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + Arc::new(b_array) as ArrayRef, + ), + ( + Arc::new(Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", inner_struct, true))), + true, + )), + Arc::new(inner_list) as ArrayRef, + ), + ( + Arc::new(Field::new("a", DataType::Utf8, true)), + Arc::new(a_array) as ArrayRef, + ), + ]); + + let outer_list = ListArray::new( + Arc::new(Field::new("item", middle_struct.data_type().clone(), true)), + OffsetBuffer::from_lengths([1]), + Arc::new(middle_struct), + None, + ); + + let batch = + RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap(); + + // Target schema with reordered fields at all levels + let target_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int32, true), // x before y now + Field::new("y", DataType::Int32, true), + ])); + let target_schema = Schema::new(vec![Field::new( + "outer", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Utf8, true), // a before b now + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", target_inner_struct, true))), + true, + ), + Field::new("b", DataType::Utf8, true), + ])), + true, + ))), + true, + )]); + + let projected = batch.project_by_schema(&target_schema).unwrap(); + + // Verify schema + assert_eq!(projected.schema().as_ref(), &target_schema); + + // Verify deeply nested data is reordered correctly + let outer_list = projected.column(0).as_list::(); + let middle_struct = outer_list.values().as_struct(); + + // Middle struct should have a first, then inner_list, then b + assert_eq!( + middle_struct.column(0).as_ref(), + &StringArray::from(vec!["a1"]) as &dyn Array + ); + assert_eq!( + middle_struct.column(2).as_ref(), + &StringArray::from(vec!["b1"]) as &dyn Array + ); + + // Inner list's struct should have x first, then y + let inner_list = middle_struct.column(1).as_list::(); + let innermost_struct = inner_list.values().as_struct(); + assert_eq!( + innermost_struct.column(0).as_ref(), + &Int32Array::from(vec![3, 4]) as &dyn Array + ); + assert_eq!( + innermost_struct.column(1).as_ref(), + &Int32Array::from(vec![1, 2]) as &dyn Array + ); + } + + #[test] + fn test_iter_binary_array_accepts_binary_variants() { + let binary = BinaryArray::from(vec![b"a".as_slice(), b"bc"]); + assert_eq!( + iter_binary_array(&binary).unwrap().collect::>(), + vec![Some(b"a".as_slice()), Some(b"bc".as_slice())] + ); + + let large_binary = LargeBinaryArray::from(vec![b"x".as_slice(), b"yz"]); + assert_eq!( + iter_binary_array(&large_binary) + .unwrap() + .collect::>(), + vec![Some(b"x".as_slice()), Some(b"yz".as_slice())] + ); + + let binary_view = BinaryViewArray::from(vec![b"1".as_slice(), b"23"]); + assert_eq!( + iter_binary_array(&binary_view).unwrap().collect::>(), + vec![Some(b"1".as_slice()), Some(b"23".as_slice())] + ); + + let fixed_size = + FixedSizeBinaryArray::try_from_iter([b"abcd", b"efgh"].into_iter()).unwrap(); + assert_eq!( + iter_binary_array(&fixed_size).unwrap().collect::>(), + vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())] + ); + } + + #[test] + fn test_iter_binary_array_rejects_non_binary() { + let int_array = Int32Array::from(vec![1, 2, 3]); + let Err(error) = iter_binary_array(&int_array) else { + panic!("expected an error for non-binary array"); + }; + assert!(error.to_string().contains("Expecting a binary type")); + } +} diff --git a/python/vendor/lance-arrow/src/list.rs b/python/vendor/lance-arrow/src/list.rs new file mode 100644 index 000000000..750b54c93 --- /dev/null +++ b/python/vendor/lance-arrow/src/list.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_array::{Array, BooleanArray, GenericListArray, OffsetSizeTrait}; +use arrow_buffer::{BooleanBufferBuilder, OffsetBuffer, ScalarBuffer}; +use arrow_schema::Field; + +pub trait ListArrayExt { + /// Filters out masked null items from the list array + /// + /// It is legal for a list array to have a null entry with a non-zero length. The + /// values inside the entry are "garbage" and should be ignored. This function + /// filters the values array to remove the garbage values. + /// + /// The output list will always have zero-length nulls. + fn filter_garbage_nulls(&self) -> Self; + /// Returns a copy of the list's values array that has been sliced to size + /// + /// It is legal for a list array's offsets to not start with zero. It's also legal + /// for a list array's offsets to not extend to the entire values array. This function + /// behaves similarly to `values()` except it slices the array so that it starts at + /// the first list offset and ends at the last list offset. + fn trimmed_values(&self) -> Arc; + /// The offset type of the underlying list array. + type Offset: OffsetSizeTrait; + /// Returns offsets shifted so the first offset is zero, matching + /// [`Self::trimmed_values`]. + /// + /// Sliced list arrays (e.g. a filtered batch) keep offsets that reference the + /// original values buffer, so combining them with trimmed values produces + /// offsets that exceed the values length. Use this together with + /// `trimmed_values` when constructing a new list array. + fn trimmed_offsets(&self) -> OffsetBuffer; +} + +impl ListArrayExt for GenericListArray { + fn filter_garbage_nulls(&self) -> Self { + if self.is_empty() { + return self.clone(); + } + let Some(validity) = self.nulls().cloned() else { + return self.clone(); + }; + + let mut should_keep = BooleanBufferBuilder::new(self.values().len()); + + // Handle case where offsets do not start at 0 + let preamble_len = self.offsets().first().unwrap().to_usize().unwrap(); + should_keep.append_n(preamble_len, false); + + let mut new_offsets: Vec = Vec::with_capacity(self.len() + 1); + new_offsets.push(OffsetSize::zero()); + let mut cur_len = OffsetSize::zero(); + for (offset, is_valid) in self.offsets().windows(2).zip(validity.iter()) { + let len = offset[1] - offset[0]; + if is_valid { + cur_len += len; + should_keep.append_n(len.to_usize().unwrap(), true); + new_offsets.push(cur_len); + } else { + should_keep.append_n(len.to_usize().unwrap(), false); + new_offsets.push(cur_len); + } + } + + // Offsets may not reference entire values buffer + let trailer = self.values().len() - should_keep.len(); + should_keep.append_n(trailer, false); + + let should_keep = should_keep.finish(); + let should_keep = BooleanArray::new(should_keep, None); + let new_values = arrow_select::filter::filter(self.values(), &should_keep).unwrap(); + let new_offsets = ScalarBuffer::from(new_offsets); + let new_offsets = OffsetBuffer::new(new_offsets); + + Self::new( + Arc::new(Field::new( + "item", + self.value_type(), + self.values().is_nullable(), + )), + new_offsets, + new_values, + Some(validity), + ) + } + + fn trimmed_values(&self) -> Arc { + let first_value = self + .offsets() + .first() + .map(|v| v.to_usize().unwrap()) + .unwrap_or(0); + let last_value = self + .offsets() + .last() + .map(|v| v.to_usize().unwrap()) + .unwrap_or(0); + self.values().slice(first_value, last_value - first_value) + } + + type Offset = OffsetSize; + + fn trimmed_offsets(&self) -> OffsetBuffer { + let offsets = self.offsets(); + let Some(&first) = offsets.first() else { + return offsets.clone(); + }; + if first == OffsetSize::zero() { + return offsets.clone(); + } + let shifted: Vec = offsets.iter().map(|&o| o - first).collect(); + OffsetBuffer::new(ScalarBuffer::from(shifted)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ListArray, UInt64Array}; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field}; + + use super::ListArrayExt; + + #[test] + fn test_filter_garbage_nulls() { + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let offsets = ScalarBuffer::::from(vec![2, 5, 8, 9]); + let offsets = OffsetBuffer::new(offsets); + let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, true)), + offsets, + Arc::new(items), + Some(list_validity.clone()), + ); + + let filtered = list_arr.filter_garbage_nulls(); + + let expected_items = UInt64Array::from(vec![2, 3, 4, 8]); + let offsets = ScalarBuffer::::from(vec![0, 3, 3, 4]); + let expected = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, false)), + OffsetBuffer::new(offsets), + Arc::new(expected_items), + Some(list_validity), + ); + + assert_eq!(filtered, expected); + } + + #[test] + fn test_trim_values() { + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let offsets = ScalarBuffer::::from(vec![2, 5, 6, 8, 9]); + let offsets = OffsetBuffer::new(offsets); + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, true)), + offsets, + Arc::new(items), + None, + ); + let list_arr = list_arr.slice(1, 2); + + let trimmed = list_arr.trimmed_values(); + + let expected_items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let expected_items = expected_items.slice(5, 3); + + assert_eq!(trimmed.as_ref(), &expected_items); + } +} diff --git a/python/vendor/lance-arrow/src/memory.rs b/python/vendor/lance-arrow/src/memory.rs new file mode 100644 index 000000000..1038a3f91 --- /dev/null +++ b/python/vendor/lance-arrow/src/memory.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashSet; + +use arrow_array::{Array, RecordBatch}; +use arrow_data::ArrayData; + +/// Counts memory used by buffers of Arrow arrays and RecordBatches. +/// +/// This is meant to capture how much memory is being used by the Arrow data +/// structures as they are. It does not represent the memory used if the data +/// were to be serialized and then deserialized. In particular: +/// +/// * This does not double count memory used by buffers shared by multiple +/// arrays or batches. Round-tripped data may use more memory because of this. +/// * This counts the **total** size of the buffers, even if the array is a slice. +/// Round-tripped data may use less memory because of this. +#[derive(Default)] +pub struct MemoryAccumulator { + seen: HashSet, + total: usize, +} + +impl MemoryAccumulator { + pub fn record_array(&mut self, array: &dyn Array) { + let data = array.to_data(); + self.record_array_data(&data); + } + + fn record_array_data(&mut self, data: &ArrayData) { + for buffer in data.buffers() { + let ptr = buffer.as_ptr(); + if self.seen.insert(ptr as usize) { + self.total += buffer.capacity(); + } + } + + if let Some(nulls) = data.nulls() { + let null_buf = nulls.inner().inner(); + let ptr = null_buf.as_ptr(); + if self.seen.insert(ptr as usize) { + self.total += null_buf.capacity(); + } + } + + for child in data.child_data() { + self.record_array_data(child); + } + } + + pub fn record_batch(&mut self, batch: &RecordBatch) { + for array in batch.columns() { + self.record_array(array); + } + } + + pub fn total(&self) -> usize { + self.total + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::Int32Array; + use arrow_schema::{DataType, Field, Schema}; + + use super::*; + + #[test] + fn test_memory_accumulator() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let slice = batch.slice(1, 2); + + let mut acc = MemoryAccumulator::default(); + + // Should record whole buffer, not just slice + acc.record_batch(&slice); + assert_eq!(acc.total(), 3 * std::mem::size_of::()); + + // Should not double count + acc.record_batch(&slice); + assert_eq!(acc.total(), 3 * std::mem::size_of::()); + } +} diff --git a/python/vendor/lance-arrow/src/scalar.rs b/python/vendor/lance-arrow/src/scalar.rs new file mode 100644 index 000000000..412253d12 --- /dev/null +++ b/python/vendor/lance-arrow/src/scalar.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_array::{ArrayRef, UInt64Array, make_array}; +use arrow_buffer::Buffer; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{ArrowError, DataType}; +use arrow_select::take::take; + +use crate::DataTypeExt; + +type Result = std::result::Result; + +pub const INLINE_VALUE_MAX_BYTES: usize = 32; + +pub fn extract_scalar_value(array: &ArrayRef, idx: usize) -> Result { + if idx >= array.len() { + return Err(ArrowError::InvalidArgumentError( + "Scalar index out of bounds".to_string(), + )); + } + + take(array.as_ref(), &UInt64Array::from(vec![idx as u64]), None) +} + +fn read_u32(buf: &[u8], offset: &mut usize) -> Result { + if *offset + 4 > buf.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: unexpected EOF".to_string(), + )); + } + let bytes = [ + buf[*offset], + buf[*offset + 1], + buf[*offset + 2], + buf[*offset + 3], + ]; + *offset += 4; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_bytes<'a>(buf: &'a [u8], offset: &mut usize, len: usize) -> Result<&'a [u8]> { + if *offset + len > buf.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: unexpected EOF".to_string(), + )); + } + let slice = &buf[*offset..*offset + len]; + *offset += len; + Ok(slice) +} + +fn write_u32(out: &mut Vec, v: u32) { + out.extend_from_slice(&v.to_le_bytes()); +} + +fn write_bytes(out: &mut Vec, bytes: &[u8]) { + out.extend_from_slice(bytes); +} + +pub fn encode_scalar_value_buffer(scalar: &ArrayRef) -> Result> { + if scalar.len() != 1 || scalar.null_count() != 0 { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer must be a single non-null value".to_string(), + )); + } + let data = scalar.to_data(); + if data.offset() != 0 { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer must have offset=0".to_string(), + )); + } + if !data.child_data().is_empty() { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer does not support nested types".to_string(), + )); + } + + // Minimal format (RFC): store the Arrow value buffers for a length-1 array. + // Null bitmap and child data are intentionally not supported here. + // + // | u32 num_buffers | + // | u32 buffer_0_len | ... | u32 buffer_{n-1}_len | + // | buffer_0 bytes | ... | buffer_{n-1} bytes | + let mut out = Vec::with_capacity(128); + let buffers = data.buffers(); + write_u32(&mut out, buffers.len() as u32); + for b in buffers { + write_u32(&mut out, b.len() as u32); + } + for b in buffers { + write_bytes(&mut out, b.as_slice()); + } + Ok(out) +} + +pub fn decode_scalar_from_value_buffer( + data_type: &DataType, + value_buffer: &[u8], +) -> Result { + if matches!( + data_type, + DataType::Struct(_) | DataType::FixedSizeList(_, _) + ) { + return Err(ArrowError::InvalidArgumentError(format!( + "Scalar value buffer does not support nested data type {:?}", + data_type + ))); + } + + let mut offset = 0; + let num_buffers = read_u32(value_buffer, &mut offset)? as usize; + let buffer_lens = (0..num_buffers) + .map(|_| read_u32(value_buffer, &mut offset).map(|l| l as usize)) + .collect::>>()?; + + let mut buffers = Vec::with_capacity(num_buffers); + for len in buffer_lens { + let bytes = read_bytes(value_buffer, &mut offset, len)?; + buffers.push(Buffer::from_vec(bytes.to_vec())); + } + + if offset != value_buffer.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: trailing bytes".to_string(), + )); + } + + let mut builder = ArrayDataBuilder::new(data_type.clone()) + .len(1) + .null_count(0); + for b in buffers { + builder = builder.add_buffer(b); + } + Ok(make_array(builder.build()?)) +} + +pub fn decode_scalar_from_inline_value( + data_type: &DataType, + inline_value: &[u8], +) -> Result { + // I expect our input to be safe here, but I added some debug_assert_eq statements just in case. + // If they are triggered, we may need to change them to return actual errors. + // + // Boolean values are bit-packed in Arrow and therefore are not "fixed-stride" in bytes. + // As a result, `byte_width_opt()` returns `None` for `DataType::Boolean`, even though a + // length-1 scalar can be represented inline using a single byte (matching `try_inline_value`). + if matches!(data_type, DataType::Boolean) { + debug_assert_eq!( + inline_value.len(), + 1, + "Invalid boolean inline scalar length (expected 1 byte, got {})", + inline_value.len() + ); + } else if let Some(byte_width) = data_type.byte_width_opt() { + debug_assert_eq!( + inline_value.len(), + byte_width, + "Inline constant length mismatch for {:?}: expected {} bytes but got {}", + data_type, + byte_width, + inline_value.len() + ); + } + + let data = ArrayDataBuilder::new(data_type.clone()) + .len(1) + .null_count(0) + .add_buffer(Buffer::from_vec(inline_value.to_vec())) + .build()?; + Ok(make_array(data)) +} + +pub fn try_inline_value(scalar: &ArrayRef) -> Option> { + if scalar.null_count() != 0 || scalar.len() != 1 { + return None; + } + let data = scalar.to_data(); + if !data.child_data().is_empty() { + return None; + } + if data.buffers().len() != 1 { + return None; + } + let bytes = data.buffers()[0].as_slice(); + if bytes.len() > INLINE_VALUE_MAX_BYTES { + return None; + } + Some(bytes.to_vec()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ + BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int8Array, Int32Array, StringArray, + cast::AsArray, types::Int8Type, + }; + + use super::*; + + #[test] + fn test_extract_scalar_value() { + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + let scalar = extract_scalar_value(&array, 2).unwrap(); + assert_eq!(scalar.len(), 1); + assert_eq!( + scalar + .as_primitive::() + .value(0), + 3 + ); + } + + #[test] + fn test_extract_scalar_value_from_full_dictionary() { + let values = Arc::new(StringArray::from( + (0..=i8::MAX) + .map(|value| format!("value-{value}")) + .collect::>(), + )); + let keys = Int8Array::from((0..=i8::MAX).collect::>()); + let array: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + + let scalar = extract_scalar_value(&array, i8::MAX as usize).unwrap(); + + let scalar = scalar.as_dictionary::(); + assert_eq!(scalar.len(), 1); + assert_eq!(scalar.key(0), Some(i8::MAX as usize)); + assert_eq!(scalar.values().len(), i8::MAX as usize + 1); + } + + #[test] + fn test_scalar_value_buffer_utf8_round_trip() { + let scalar: ArrayRef = Arc::new(StringArray::from(vec!["hello"])); + let buf = encode_scalar_value_buffer(&scalar).unwrap(); + let decoded = decode_scalar_from_value_buffer(&DataType::Utf8, &buf).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.null_count(), 0); + assert_eq!(decoded.as_string::().value(0), "hello"); + } + + #[test] + fn test_scalar_value_buffer_fixed_size_binary_round_trip() { + let val = vec![0xABu8; 33]; + let scalar: ArrayRef = Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + std::iter::once(Some(val.as_slice())), + 33, + ) + .unwrap(), + ); + let buf = encode_scalar_value_buffer(&scalar).unwrap(); + let decoded = + decode_scalar_from_value_buffer(&DataType::FixedSizeBinary(33), &buf).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.as_fixed_size_binary().value(0), val.as_slice()); + } + + #[test] + fn test_inline_value_boolean_round_trip() { + let scalar: ArrayRef = Arc::new(BooleanArray::from_iter([Some(true)])); + let inline = try_inline_value(&scalar).unwrap(); + let decoded = decode_scalar_from_inline_value(&DataType::Boolean, &inline).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.null_count(), 0); + assert!(decoded.as_boolean().value(0)); + } + + #[test] + fn test_scalar_value_buffer_rejects_nested_type() { + let field = Arc::new(arrow_schema::Field::new("item", DataType::Int32, false)); + let list: ArrayRef = Arc::new(arrow_array::FixedSizeListArray::new( + field, + 2, + Arc::new(Int32Array::from(vec![1, 2])), + None, + )); + let scalar = list.slice(0, 1); + assert!(encode_scalar_value_buffer(&scalar).is_err()); + } + + #[test] + fn test_decode_scalar_from_value_buffer_rejects_nested_type() { + let buf = Vec::::new(); + let res = + decode_scalar_from_value_buffer(&DataType::Struct(arrow_schema::Fields::empty()), &buf); + assert!(res.is_err()); + } + + #[test] + fn test_decode_scalar_from_value_buffer_trailing_bytes() { + // num_buffers = 0, plus an extra byte + let mut bytes = Vec::new(); + bytes.extend_from_slice(&0u32.to_le_bytes()); + bytes.push(1); + let res = decode_scalar_from_value_buffer(&DataType::Int32, &bytes); + assert!(res.is_err()); + } +} diff --git a/python/vendor/lance-arrow/src/schema.rs b/python/vendor/lance-arrow/src/schema.rs new file mode 100644 index 000000000..0489dd584 --- /dev/null +++ b/python/vendor/lance-arrow/src/schema.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extension to arrow schema + +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema}; + +use crate::{ARROW_EXT_NAME_KEY, BLOB_META_KEY, BLOB_V2_EXT_NAME}; + +pub enum Indentation { + OneLine, + MultiLine(u8), +} + +impl Indentation { + fn value(&self) -> String { + match self { + Self::OneLine => "".to_string(), + Self::MultiLine(spaces) => " ".repeat(*spaces as usize), + } + } + + fn deepen(&self) -> Self { + match self { + Self::OneLine => Self::OneLine, + Self::MultiLine(spaces) => Self::MultiLine(spaces + 2), + } + } +} + +/// Extends the functionality of [arrow_schema::Field]. +pub trait FieldExt { + /// Create a compact string representation of the field + /// + /// This is intended for display purposes and not for serialization + fn to_compact_string(&self, indent: Indentation) -> String; + + /// Check if the field is marked as a packed struct + fn is_packed_struct(&self) -> bool; + + /// Check if the field is marked as a blob + fn is_blob(&self) -> bool; + + /// Check if the field is marked as a blob + fn is_blob_v2(&self) -> bool; +} + +impl FieldExt for Field { + fn to_compact_string(&self, indent: Indentation) -> String { + let mut result = format!("{}: ", self.name().clone()); + match self.data_type() { + DataType::Struct(fields) => { + result += "{"; + result += &indent.value(); + for (field_idx, field) in fields.iter().enumerate() { + result += field.to_compact_string(indent.deepen()).as_str(); + if field_idx < fields.len() - 1 { + result += ","; + } + result += indent.value().as_str(); + } + result += "}"; + } + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) => { + result += "["; + result += field.to_compact_string(indent.deepen()).as_str(); + result += "]"; + } + DataType::FixedSizeList(child, dimension) => { + result += &format!( + "[{}; {}]", + child.to_compact_string(indent.deepen()), + dimension + ); + } + DataType::Dictionary(key_type, value_type) => { + result += &value_type.to_string(); + result += "@"; + result += &key_type.to_string(); + } + _ => { + result += &self.data_type().to_string(); + } + } + if self.is_nullable() { + result += "?"; + } + result + } + + // Check if field has metadata `packed` set to true, this check is case insensitive. + fn is_packed_struct(&self) -> bool { + let field_metadata = self.metadata(); + const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"]; + PACKED_KEYS.iter().any(|key| { + field_metadata + .get(*key) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(false) + }) + } + + fn is_blob(&self) -> bool { + let field_metadata = self.metadata(); + field_metadata.get(BLOB_META_KEY).is_some() + || field_metadata + .get(ARROW_EXT_NAME_KEY) + .map(|value| value == BLOB_V2_EXT_NAME) + .unwrap_or(false) + } + + fn is_blob_v2(&self) -> bool { + let field_metadata = self.metadata(); + field_metadata + .get(ARROW_EXT_NAME_KEY) + .map(|value| value == BLOB_V2_EXT_NAME) + .unwrap_or(false) + } +} + +/// Extends the functionality of [arrow_schema::Schema]. +pub trait SchemaExt { + /// Create a new [`Schema`] with one extra field. + fn try_with_column(&self, field: Field) -> std::result::Result; + + fn try_with_column_at( + &self, + index: usize, + field: Field, + ) -> std::result::Result; + + fn field_names(&self) -> Vec<&String>; + + fn without_column(&self, column_name: &str) -> Schema; + + /// Create a compact string representation of the schema + /// + /// This is intended for display purposes and not for serialization + fn to_compact_string(&self, indent: Indentation) -> String; +} + +impl SchemaExt for Schema { + fn try_with_column(&self, field: Field) -> std::result::Result { + if self.column_with_name(field.name()).is_some() { + return Err(ArrowError::SchemaError(format!( + "Can not append column {} on schema: {:?}", + field.name(), + self + ))); + }; + let mut fields: Vec = self.fields().iter().cloned().collect(); + fields.push(FieldRef::new(field)); + Ok(Self::new_with_metadata(fields, self.metadata.clone())) + } + + fn try_with_column_at( + &self, + index: usize, + field: Field, + ) -> std::result::Result { + if self.column_with_name(field.name()).is_some() { + return Err(ArrowError::SchemaError(format!( + "Failed to modify schema: Inserting column {} would create a duplicate column in schema: {:?}", + field.name(), + self + ))); + }; + let mut fields: Vec = self.fields().iter().cloned().collect(); + fields.insert(index, FieldRef::new(field)); + Ok(Self::new_with_metadata(fields, self.metadata.clone())) + } + + /// Project the schema to remove the given column. + /// + /// This only works on top-level fields right now. If a field does not exist, + /// the schema will be returned as is. + fn without_column(&self, column_name: &str) -> Schema { + let fields: Vec = self + .fields() + .iter() + .filter(|f| f.name() != column_name) + .cloned() + .collect(); + Self::new_with_metadata(fields, self.metadata.clone()) + } + + fn field_names(&self) -> Vec<&String> { + self.fields().iter().map(|f| f.name()).collect() + } + + fn to_compact_string(&self, indent: Indentation) -> String { + let mut result = "{".to_string(); + result += &indent.value(); + for (field_idx, field) in self.fields.iter().enumerate() { + result += field.to_compact_string(indent.deepen()).as_str(); + if field_idx < self.fields.len() - 1 { + result += ","; + } + result += indent.value().as_str(); + } + result += "}"; + result + } +} diff --git a/python/vendor/lance-arrow/src/stream.rs b/python/vendor/lance-arrow/src/stream.rs new file mode 100644 index 000000000..c5e55f093 --- /dev/null +++ b/python/vendor/lance-arrow/src/stream.rs @@ -0,0 +1,554 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities for working with streams of [`RecordBatch`]. + +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, SchemaRef}; +use futures::stream::{self, Stream, StreamExt}; +use std::pin::Pin; + +use crate::deepcopy::deep_copy_batch_sliced; + +/// Rechunks a stream of [`RecordBatch`] so that each output batch has +/// approximately `target_bytes` of array data. +/// +/// Small input batches are accumulated (by concatenation) until at least +/// `min_bytes` of data has been collected. If the resulting batch exceeds +/// `max_bytes`, it is sliced into roughly equal pieces of ~`max_bytes` +/// (assuming uniform row sizes). +pub fn rechunk_stream_by_size( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, false) +} + +/// Like [`rechunk_stream_by_size`] but deep-copies slices so that +/// `get_array_memory_size` reflects the true size of each output batch. +/// +/// After a normal `RecordBatch::slice`, the backing buffers are shared with +/// the original batch, so `get_array_memory_size` still reports the full +/// parent size. This variant deep-copies every slice produced during the +/// splitting phase, which allows the stream to detect and re-split slices +/// that still exceed `max_bytes` (e.g. because a single row is much larger +/// than average). +/// +/// The deep copy is a last resort and potentially expensive for large +/// batches. However, it is only performed when a batch actually needs to be +/// sliced — batches that are already within the target range pass through at +/// zero cost. Use this only when the hard cap on `max_bytes` is a +/// correctness requirement, not merely a performance hint. +pub fn rechunk_stream_by_size_deep_copy( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, true) +} + +fn rechunk_stream_by_size_inner( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, + deep_copy: bool, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + stream::try_unfold( + RechunkState { + input: Box::pin(input), + accumulated: Vec::new(), + acc_bytes: 0, + done: false, + input_schema, + min_bytes, + max_bytes, + deep_copy, + }, + |mut state| async move { + if state.done && state.accumulated.is_empty() { + return Ok(None); + } + + // Pull batches until we reach the byte target or exhaust input. + // Always pull at least one batch so that min_bytes=0 works. + while !state.done && (state.accumulated.is_empty() || state.acc_bytes < state.min_bytes) + { + match state.input.next().await { + Some(Ok(batch)) => { + state.acc_bytes += batch.get_array_memory_size(); + state.accumulated.push(batch); + } + Some(Err(e)) => return Err(e), + None => { + state.done = true; + } + } + } + + if state.accumulated.is_empty() { + return Ok(None); + } + + // Fast path: if the first accumulated batch already meets the + // byte threshold, deliver it directly instead of concatenating + // everything together (which would just get sliced back apart). + if state.accumulated.len() > 1 + && state.accumulated[0].get_array_memory_size() >= state.min_bytes + { + let b = state.accumulated.remove(0); + state.acc_bytes -= b.get_array_memory_size(); + return Ok(Some((b, state))); + } + + let batch = if state.accumulated.len() == 1 { + state.accumulated.pop().unwrap() + } else { + let b = + arrow_select::concat::concat_batches(&state.input_schema, &state.accumulated) + .map_err(E::from)?; + state.accumulated.clear(); + b + }; + state.acc_bytes = 0; + + // Slice the batch into ~max_bytes pieces assuming uniform row sizes. + let mut slices = + slice_batch(batch, state.max_bytes, state.deep_copy).map_err(E::from)?; + + if slices.len() == 1 { + Ok(Some((slices.pop().unwrap(), state))) + } else { + let first = slices.remove(0); + + // Stash leftover slices for subsequent iterations. + for a in &slices { + state.acc_bytes += a.get_array_memory_size(); + } + state.accumulated = slices; + + Ok(Some((first, state))) + } + }, + ) +} + +/// Slice a batch into pieces of at most `max_bytes`. +/// +/// When `deep_copy` is false, slices share buffers with the original batch +/// and `get_array_memory_size` will still report the parent buffer size. +/// This is fine when the caller only needs approximate sizing. +/// +/// When `deep_copy` is true, each slice is deep-copied so that +/// `get_array_memory_size` reflects the true size. If a deep-copied slice +/// still exceeds `max_bytes` (due to non-uniform row sizes), it is +/// recursively split until every piece is within budget or contains only a +/// single row. +fn slice_batch( + batch: RecordBatch, + max_bytes: usize, + deep_copy: bool, +) -> Result, ArrowError> { + let batch_bytes = batch.get_array_memory_size(); + let num_rows = batch.num_rows(); + + if batch_bytes <= max_bytes { + return Ok(vec![batch]); + } + + if num_rows <= 1 { + // A single row cannot be split further, but the size just measured may + // belong to its source buffer rather than to the row: + // `get_array_memory_size` reports whole buffers, and a slice shares + // them. Returning it uncopied hands the caller that inflated figure -- + // a 100 KB row inside a 200 MB buffer measures as 200 MB, and a caller + // budgeting on it concludes the row cannot fit anywhere. + if deep_copy { + return Ok(vec![deep_copy_batch_sliced(&batch)?]); + } + return Ok(vec![batch]); + } + + let rows_per_chunk = (max_bytes as u64 * num_rows as u64 / batch_bytes as u64).max(1) as usize; + + let mut result = Vec::new(); + let mut offset = 0; + while offset < num_rows { + let len = rows_per_chunk.min(num_rows - offset); + let slice = batch.slice(offset, len); + if deep_copy { + let copied = deep_copy_batch_sliced(&slice)?; + // Recurse: the deep-copied slice has accurate sizes, so if it + // still exceeds max_bytes we can split further. + result.extend(slice_batch(copied, max_bytes, true)?); + } else { + result.push(slice); + } + offset += len; + } + + Ok(result) +} + +/// Internal state for [`rechunk_stream`]. +/// +/// Kept as a named struct so the `try_unfold` closure stays readable. +struct RechunkState { + input: Pin>, + accumulated: Vec, + acc_bytes: usize, + done: bool, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, + deep_copy: bool, +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow_array::Int32Array; + use arrow_schema::{DataType, Field, Schema}; + use futures::executor::block_on; + + /// A single row that shares a large source buffer must report its own + /// size, not the buffer's. Callers budget on this figure, and an inflated + /// one makes a small row look impossible to place. + #[test] + fn a_single_row_slice_reports_its_own_size() { + use arrow_array::LargeStringArray; + + let row = "x".repeat(100 * 1024); + let values: Vec<&str> = (0..2000).map(|_| row.as_str()).collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "text", + DataType::LargeUtf8, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(LargeStringArray::from(values))]).unwrap(); + let whole = batch.get_array_memory_size(); + + let one_row = batch.slice(7, 1); + assert_eq!( + one_row.get_array_memory_size(), + whole, + "a slice shares its source buffer, which is the reason this test exists" + ); + + // Budget far below the row so the split recurses to the single-row floor. + let chunks = slice_batch(one_row, 1024, true).unwrap(); + assert_eq!(chunks.len(), 1); + let measured = chunks[0].get_array_memory_size(); + assert_eq!(chunks[0].num_rows(), 1); + assert!( + measured < whole / 100, + "single row still measured as {measured} bytes against a {whole} byte source" + ); + } + + fn make_batch(num_rows: usize) -> RecordBatch { + let schema = test_schema(); + let values: Vec = (0..num_rows as i32).collect(); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn collect_rechunked( + batches: Vec, + min_bytes: usize, + max_bytes: usize, + ) -> Vec { + let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>)); + let rechunked = rechunk_stream_by_size(input, test_schema(), min_bytes, max_bytes); + block_on(rechunked.collect::>()) + .into_iter() + .map(|r| r.unwrap()) + .collect() + } + + fn total_rows(batches: &[RecordBatch]) -> usize { + batches.iter().map(|b| b.num_rows()).sum() + } + + #[test] + fn test_empty_stream() { + let result = collect_rechunked(vec![], 100, 200); + assert!(result.is_empty()); + } + + #[test] + fn test_single_batch_passthrough() { + let batch = make_batch(100); + let bytes = batch.get_array_memory_size(); + // Batch is between min and max — should pass through as-is. + let result = collect_rechunked(vec![batch], bytes / 2, bytes * 2); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 100); + } + + #[test] + fn test_small_batches_concatenated() { + let one_batch_bytes = make_batch(10).get_array_memory_size(); + let batches: Vec<_> = (0..8).map(|_| make_batch(10)).collect(); + // min = 5 batches worth, max = 10 batches worth. + let result = collect_rechunked(batches, one_batch_bytes * 5, one_batch_bytes * 10); + assert_eq!(total_rows(&result), 80); + // Should have been concatenated into fewer batches than the 8 inputs. + assert!( + result.len() < 8, + "expected fewer output batches, got {}", + result.len() + ); + } + + #[test] + fn test_large_batch_sliced() { + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4); + assert_eq!(total_rows(&result), 1000); + assert!( + result.len() >= 4, + "expected at least 4 slices, got {}", + result.len() + ); + } + + #[test] + fn test_sliced_leftovers_are_not_recombined() { + // Key test for the fast-path optimisation. When a large batch is + // sliced, leftover slices should be delivered one-at-a-time without + // being concatenated back together. We verify this by checking that + // every output buffer pointer falls inside the original batch's + // allocation (i.e. they are all zero-copy slices, not fresh copies). + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let orig_data = batch.column(0).to_data(); + let orig_buf = &orig_data.buffers()[0]; + let orig_start = orig_buf.as_ptr() as usize; + let orig_end = orig_start + orig_buf.len(); + + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4); + + assert_eq!(total_rows(&result), 1000); + assert!(result.len() >= 4); + + for (i, b) in result.iter().enumerate() { + let ptr = b.column(0).to_data().buffers()[0].as_ptr() as usize; + assert!( + ptr >= orig_start && ptr < orig_end, + "slice {i} buffer at {ptr:#x} is outside the original allocation \ + [{orig_start:#x}, {orig_end:#x}) — it was re-concatenated" + ); + } + } + + #[test] + fn test_flush_remainder_on_stream_end() { + // Data below min_bytes should still be flushed when the stream ends. + let batch = make_batch(10); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes * 100, bytes * 200); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 10); + } + + #[test] + fn test_large_then_small_batches() { + // After a large batch is fully drained, subsequent small batches + // should be accumulated normally. + let large = make_batch(1000); + let small_bytes = make_batch(10).get_array_memory_size(); + let batches = vec![ + large, + make_batch(10), + make_batch(10), + make_batch(10), + make_batch(10), + make_batch(10), + ]; + let result = collect_rechunked(batches, small_bytes * 3, small_bytes * 100); + assert_eq!(total_rows(&result), 1050); + // The large batch should appear (possibly sliced) followed by + // concatenated small batches, so we should have fewer output batches + // than the 6 inputs. + assert!(result.len() < 6); + } + + #[test] + fn test_row_preservation_across_slicing() { + // Verify that every input row appears exactly once in the output + // and in the correct order after slicing. + let batch = make_batch(237); // odd count to exercise remainder slice + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 5); + + assert_eq!(total_rows(&result), 237); + + let values: Vec = result + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + let expected: Vec = (0..237).collect(); + assert_eq!(values, expected); + } + + #[test] + fn test_min_bytes_zero_still_yields_all_rows() { + // When min_bytes=0, the stream should still yield every batch. + // This is the "chop only, don't coalesce" use case. + let batches: Vec<_> = (0..5).map(|_| make_batch(100)).collect(); + let batch_bytes = batches[0].get_array_memory_size(); + let result = collect_rechunked(batches, 0, batch_bytes * 2); + assert_eq!(total_rows(&result), 500); + } + + #[test] + fn test_min_bytes_zero_slices_oversized() { + // min_bytes=0 with a small max_bytes should still slice large batches. + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], 0, bytes / 4); + assert_eq!(total_rows(&result), 1000); + assert!( + result.len() >= 4, + "expected at least 4 slices, got {}", + result.len() + ); + } + + /// Build a batch with one variable-length string column. + /// Every row is `small_size` bytes except the row at index `big_row_idx` + /// which is `big_size` bytes. + fn make_variable_batch( + num_rows: usize, + small_size: usize, + big_row_idx: usize, + big_size: usize, + ) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let values: Vec = (0..num_rows) + .map(|i| { + if i == big_row_idx { + "X".repeat(big_size) + } else { + "x".repeat(small_size) + } + }) + .collect(); + let array = arrow_array::StringArray::from(values); + RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap() + } + + fn variable_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])) + } + + fn collect_rechunked_variable( + batches: Vec, + min_bytes: usize, + max_bytes: usize, + ) -> Vec { + let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>)); + let rechunked = + rechunk_stream_by_size_deep_copy(input, variable_schema(), min_bytes, max_bytes); + block_on(rechunked.collect::>()) + .into_iter() + .map(|r| r.unwrap()) + .collect() + } + + #[test] + fn test_oversized_row_at_end() { + // 100 rows: 99 small (64 bytes each) + 1 large (100KiB) at the end. + let batch = make_variable_batch(100, 64, 99, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_oversized_row_at_start() { + // 100 rows: 1 large (100KiB) at the start + 99 small (64 bytes each). + let batch = make_variable_batch(100, 64, 0, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_oversized_row_in_middle() { + // 100 rows: 1 large (100KiB) in the middle + 99 small (64 bytes each). + let batch = make_variable_batch(100, 64, 50, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_error_propagation() { + let input = stream::iter(vec![ + Ok(make_batch(10)), + Err(ArrowError::ComputeError("boom".into())), + Ok(make_batch(10)), + ]); + let rechunked = rechunk_stream_by_size(input, test_schema(), 1, usize::MAX); + let results: Vec> = block_on(rechunked.collect()); + assert!(results.iter().any(|r| r.is_err())); + } +} diff --git a/python/vendor/lance-arrow/src/struct.rs b/python/vendor/lance-arrow/src/struct.rs new file mode 100644 index 000000000..ed4b6288e --- /dev/null +++ b/python/vendor/lance-arrow/src/struct.rs @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extension to arrow struct arrays + +use arrow_array::{Array, StructArray, cast::AsArray, make_array}; +use arrow_buffer::NullBuffer; +use arrow_data::{ArrayData, ArrayDataBuilder}; +use arrow_schema::{ArrowError, DataType}; + +pub trait StructArrayExt { + /// Removes the offset / length of the struct array by pushing it into the children + /// + /// In arrow-rs when slice is called it recursively slices the children. + /// In arrow-cpp when slice is called it just sets the offset/length of + /// the struct array and leaves the children as-is + /// + /// Both are legal approaches (â•Ĩīšâ•Ĩ) + /// + /// This method helps reduce complexity by folding into the arrow-rs approach + fn normalize_slicing(&self) -> Result + where + Self: Sized; + + /// Structs are allowed to mask valid items. For example, a struct array might be: + /// + /// [ {"items": [1, 2, 3]}, NULL, {"items": NULL}] + /// + /// However, the underlying items array might be: [[1, 2, 3], [4, 5], NULL] + /// + /// The [4, 5] list is masked out because the struct array is null. + /// + /// The struct validity would be [true, false, true] and the list validity would be [true, true, false] + /// + /// This method pushes nulls down into all children. In the above example the list validity would become + /// [true, false, false]. + /// + /// This method is not recursive. If a child is a struct array it will not push that child's nulls down. + /// + /// This method does not remove garbage lists. It only updates the validity so a future call to + /// [crate::list::ListArrayExt::filter_garbage_nulls] will remove the garbage lists (without + /// this pushdown it would not) + fn pushdown_nulls(&self) -> Result + where + Self: Sized; +} + +fn normalized_struct_array_data(data: ArrayData) -> Result { + let parent_offset = data.offset(); + let parent_len = data.len(); + let modified_children = data + .child_data() + .iter() + .map(|d| { + let d = normalized_struct_array_data(d.clone())?; + let offset = d.offset(); + let len = d.len(); + if len < parent_len + parent_offset { + return Err(ArrowError::InvalidArgumentError(format!( + "Child array {} has length {} which is less than the parent length {} plus the parent offset {}", + d.data_type(), + len, + parent_len, + parent_offset + ))); + } + let new_offset = offset + parent_offset; + d.into_builder().offset(new_offset) + .len(parent_len) + .build() + }) + .collect::, _>>()?; + ArrayDataBuilder::new(data.data_type().clone()) + .len(parent_len) + .offset(0) + .buffers(data.buffers().to_vec()) + .child_data(modified_children) + .build() +} + +impl StructArrayExt for StructArray { + fn normalize_slicing(&self) -> Result + where + Self: Sized, + { + if self.offset() == 0 && self.columns().iter().all(|c| c.len() == self.len()) { + return Ok(self.clone()); + } + + let data = normalized_struct_array_data(self.to_data())?; + Ok(Self::from(data)) + } + + fn pushdown_nulls(&self) -> Result + where + Self: Sized, + { + let Some(validity) = self.nulls() else { + return Ok(self.clone()); + }; + let data = self.to_data(); + let children = data + .child_data() + .iter() + .map(|c| { + if c.data_type() == &DataType::Null { + // A Null array is already null in every slot, and arrow-rs rejects + // a validity buffer on it. + Ok(c.clone()) + } else if let Some(child_validity) = c.nulls() { + let new_validity = child_validity.inner() & validity.inner(); + c.clone() + .into_builder() + .nulls(Some(NullBuffer::from(new_validity))) + .build() + } else { + Ok(c.clone() + .into_builder() + .nulls(Some(validity.clone())) + .build()?) + } + }) + .collect::, _>>()?; + let arr = make_array(data.into_builder().child_data(children).build()?); + Ok(arr.as_struct().clone()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::{ + Array, Int32Array, NullArray, StructArray, cast::AsArray, make_array, types::Int32Type, + }; + use arrow_buffer::NullBuffer; + use arrow_schema::{DataType, Field, Fields}; + use std::sync::Arc; + + use crate::r#struct::StructArrayExt; + + #[test] + fn test_normalize_slicing_no_offset() { + let x = Int32Array::from(vec![1, 2, 3]); + let y = Int32Array::from(vec![4, 5, 6]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + let normalized = struct_array.normalize_slicing().unwrap(); + assert_eq!(normalized, struct_array); + } + + #[test] + fn test_arrow_rs_slicing() { + let x = Int32Array::from(vec![1, 2, 3, 4]); + let y = Int32Array::from(vec![5, 6, 7, 8]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + // Slicing with arrow-rs propagates the slicing to the children so there should + // be no change needed to the struct array + let sliced = struct_array.slice(1, 2); + let normalized = sliced.normalize_slicing().unwrap(); + + assert_eq!(normalized, sliced); + } + + #[test] + fn test_arrow_cpp_slicing() { + let x = Int32Array::from(vec![1, 2, 3, 4]); + let y = Int32Array::from(vec![5, 6, 7, 8]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + let data = struct_array.to_data(); + let sliced = data.into_builder().offset(1).len(2).build().unwrap(); + let sliced = make_array(sliced); + let normalized = sliced.as_struct().clone().normalize_slicing().unwrap(); + + assert_eq!(normalized, struct_array.slice(1, 2)); + } + + #[test] + fn test_pushdown_nulls_with_null_typed_child() { + // `pa.Table.from_pylist([{"s": {"a": 1, "b": None}}, {"s": None}])` infers a + // struct with a Null-typed child. Arrow forbids a validity buffer on a Null + // array, so the struct's nulls must not be pushed into it. + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Null, true), + ]), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(NullArray::new(3)), + ], + Some(NullBuffer::from(vec![true, false, true])), + ); + + let pushed = struct_array.pushdown_nulls().unwrap(); + + assert_eq!(pushed.nulls(), struct_array.nulls()); + assert_eq!( + pushed.column(0).as_primitive::(), + &Int32Array::from(vec![Some(1), None, Some(3)]) + ); + assert_eq!(pushed.column(1).data_type(), &DataType::Null); + assert_eq!(pushed.column(1).logical_null_count(), 3); + } +} diff --git a/rust/lancedb/tests/nullable_struct_merge.rs b/rust/lancedb/tests/nullable_struct_merge.rs new file mode 100644 index 000000000..c245b1c2c --- /dev/null +++ b/rust/lancedb/tests/nullable_struct_merge.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, RecordBatch, StructArray}; +use arrow_buffer::NullBuffer; +use arrow_schema::{DataType, Field, Schema}; +use lance_arrow::RecordBatchExt; + +#[test] +fn merge_nullable_struct_with_sliced_boolean_child() { + let left_struct = StructArray::new( + vec![Field::new("id", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef], + None, + ); + + // BooleanArray retains its offset when sliced. The merged child bitmap must + // cover that offset as well as the four visible rows. + let boolean_child: ArrayRef = Arc::new(BooleanArray::from_iter((0..1030).map(|index| { + if index == 1027 { + None + } else { + Some(index % 2 == 1) + } + }))); + let boolean_child = boolean_child.slice(1025, 4); + assert_eq!(boolean_child.offset(), 1025); + let right_struct = StructArray::new( + vec![Field::new("flag", DataType::Boolean, true)].into(), + vec![boolean_child], + Some(NullBuffer::from(vec![true, false, true, true])), + ); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "nested", + left_struct.data_type().clone(), + true, + )])), + vec![Arc::new(left_struct)], + ) + .unwrap(); + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "nested", + right_struct.data_type().clone(), + true, + )])), + vec![Arc::new(right_struct)], + ) + .unwrap(); + + let merged = left_batch.merge(&right_batch).unwrap(); + let nested = merged.column_by_name("nested").unwrap(); + let nested = nested.as_any().downcast_ref::().unwrap(); + let flag = nested.column_by_name("flag").unwrap(); + let flag = flag.as_any().downcast_ref::().unwrap(); + assert_eq!( + flag.iter().collect::>(), + vec![Some(true), None, None, Some(false)] + ); +} From 55897d7c3b03c6b94b461161abc1b28d34903fde Mon Sep 17 00:00:00 2001 From: Yingzi Zhang <164733623+LindseyZ1205@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:26:43 -0400 Subject: [PATCH 32/33] fix(python): preserve Arrow schema when merging reranker scores (#4332) Fixes #4331 With `return_score="all"`, model rerankers reconstruct merged rows without their Arrow schema. This changes fixed-size float32 vectors into variable-length float64 lists, loses metadata, and can overflow valid uint64 row IDs. Unify the original input schemas before inserting placeholder scores, then use the unified schema when reconstructing the merged table. Row merging and reranking behavior stay the same. Capturing the schema first also preserves float64 scores without conflicting with the float32 null placeholders. Validation: - New regression cases: 5 failures and 5 passing controls before the fix; all 10 pass afterward. - Reranker, hybrid-query and TypeSafe reranker tests: 150 passed, 17 skipped. Optional live model/API tests were not run. - The fake-client tests that use random metadata passed again with seeds 7 and 41. - Repository-wide Ruff lint/format checks, configured Pyright check and applicable pre-commit hooks pass. - Tests used the locally built Rust/Python extension and the repository's locked uv environment. No Rust code changed. AI disclosure: I used Codex for investigation, implementation, tests, local checks, and drafting this description. --- python/python/lancedb/rerankers/base.py | 5 +- python/python/tests/test_rerankers.py | 108 ++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/rerankers/base.py b/python/python/lancedb/rerankers/base.py index 4af938a1e..e1bc7f0dd 100644 --- a/python/python/lancedb/rerankers/base.py +++ b/python/python/lancedb/rerankers/base.py @@ -241,6 +241,7 @@ class Reranker(ABC): This op is slower than just keeping relevance score but can be useful for debugging. """ + combined_schema = pa.unify_schemas([vector_results.schema, fts_results.schema]) # add nulls to fts results for _distance if "_distance" not in fts_results.column_names: fts_results = fts_results.append_column( @@ -265,7 +266,9 @@ class Reranker(ABC): else: vector_results_dict[key] = value - combined = pa.Table.from_pylist(list(vector_results_dict.values())) + combined = pa.Table.from_pylist( + list(vector_results_dict.values()), schema=combined_schema + ) return combined def _keep_relevance_score(self, combined_results: pa.Table): diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 850b477a0..d5ee6b6a4 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,114 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@pytest.mark.parametrize("return_score", ["relevance", "all"]) +@pytest.mark.parametrize("first_row_id", [1, 2**63]) +@pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()]) +def test_model_reranker_preserves_arrow_schema( + mocker, return_score, first_row_id, score_type +): + common_schema = pa.schema( + [ + pa.field("_rowid", pa.uint64()), + pa.field("text", pa.string()), + pa.field("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id, + "text": "vector", + "vector": [1, 0], + "_distance": 0.1, + }, + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_distance": 0.2, + }, + ], + schema=common_schema.append(pa.field("_distance", score_type)), + ) + fts_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_score": 0.9, + }, + { + "_rowid": first_row_id + 2, + "text": "fts", + "vector": [1, 1], + "_score": 0.8, + }, + ], + schema=common_schema.append(pa.field("_score", score_type)), + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + reranker._client.rerank.return_value.results = [ + mocker.Mock(index=index, relevance_score=score) + for index, score in [(2, 0.9), (1, 0.8), (0, 0.7)] + ] + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result["_rowid"].to_pylist() == [ + first_row_id + 2, + first_row_id + 1, + first_row_id, + ] + assert result["text"].to_pylist() == ["fts", "both", "vector"] + if return_score == "all": + assert ( + result["_distance"] + .combine_chunks() + .equals(pa.array([None, 0.2, 0.1], type=score_type)) + ) + assert ( + result["_score"] + .combine_chunks() + .equals(pa.array([0.8, 0.9, None], type=score_type)) + ) + + +@pytest.mark.parametrize("return_score", ["relevance", "all"]) +def test_model_reranker_empty_results_preserve_schema(mocker, return_score): + common_schema = pa.schema( + [ + ("_rowid", pa.uint64()), + ("text", pa.string()), + ("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_distance", pa.float32())) + ) + fts_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_score", pa.float32())) + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert len(result) == 0 + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result.schema.field("_relevance_score").type == pa.float32() + reranker._client.rerank.assert_not_called() + + def test_cross_encoder_reranker_return_all(tmp_path): pytest.importorskip("sentence_transformers") reranker = CrossEncoderReranker(return_score="all") From 7f932b8552c22355902829be1c5615a867d1bd45 Mon Sep 17 00:00:00 2001 From: Prince Jain Date: Thu, 22 Oct 2026 23:15:26 +0530 Subject: [PATCH 33/33] fix(nodejs): reject listNamespaces limit above i32::MAX instead of wrapping (#4584) ## Summary Fixes #4583. `Connection::list_namespaces` narrowed `Option` to `i32` with an unchecked `as` cast, so a JavaScript `limit` between 2^31 and 2^32 - 1 (e.g. `3000000000`) wrapped negative before reaching `ListNamespacesRequest`. The server received a negative page size instead of a clear client-side error. Replaced the cast with a `checked_limit` helper using `i32::try_from`, returning `napi::Error::from_reason` on overflow, matching the existing `try_from` validation used for `limit` elsewhere in the same file. ## Tests Unit tests cover the wrap case (`u32` values above `i32::MAX` reject), the boundary (`i32::MAX` passes), and `None`/zero passthrough. `cargo test -p lancedb-nodejs` passes 8/8 and `pnpm build` compiles the binding and TypeScript definitions cleanly. --- nodejs/src/connection.rs | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 58ec9859a..67c4706e4 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -673,7 +673,7 @@ impl Connection { let req = ListNamespacesRequest { id: namespace_path, page_token, - limit: limit.map(|l| l as i32), + limit: checked_limit(limit)?, ..Default::default() }; let resp = self @@ -787,3 +787,33 @@ impl Connection { .default_error() } } + +fn checked_limit(limit: Option) -> napi::Result> { + limit + .map(i32::try_from) + .transpose() + .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) +} + +#[cfg(test)] +mod tests { + use super::checked_limit; + + #[test] + fn checked_limit_rejects_values_above_i32_max() { + // u32 values above i32::MAX wrapped negative through `as i32`, + // sending a negative page size to the server. + assert!(checked_limit(Some(3_000_000_000)).is_err()); + assert!(checked_limit(Some(u32::MAX)).is_err()); + } + + #[test] + fn checked_limit_passes_boundary_and_missing_values() { + assert_eq!( + checked_limit(Some(i32::MAX as u32)).unwrap(), + Some(i32::MAX) + ); + assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); + assert_eq!(checked_limit(None).unwrap(), None); + } +}