Compare commits

...

3 Commits

Author SHA1 Message Date
Gatefixer b61e14dd92 fix(rust): infer vector dimensions from list data 2026-08-05 20:52:47 +00:00
Justin Miller c7ea91f3ea test: cover blob null/empty preservation across Table::optimize (#3774)
## Description

`Table::optimize()` compacts through
`lance::dataset::optimize::compact_files`
(`rust/lancedb/src/table/optimize.rs:155`). Until
lance-format/lance#7965 that rewrite corrupted blob columns holding null
or empty values, which is what #3744 reports:

- **storage 2.0** (legacy v1 `lance-encoding:blob` descriptors): every
payload following a null or empty row in the same fragment was rewritten
as `{position: 0, size: 0}`, so it read back as `b""` and the new
fragment no longer referenced the bytes — silent payload loss,
unrecoverable once the pre-optimize versions are pruned.
- **storage 2.2** (blob v2): a valid empty value was rewritten as null,
destroying the null-vs-empty distinction.

Both manifestations share one root cause: `is_inline_null_blob`
classified any inline blob with `position == 0 && size == 0` as null,
which is also exactly what a *valid empty value* looks like. Such rows
were dropped from `blob_read_addrs`, misaligning every payload that
followed.

The behaviour is already correct on `main`: the vendored lance crate
first carried the fix at `v10.0.0-beta.3` (#3710) and is now
`v10.1.0-beta.1` (#3757). What was missing is coverage — nothing in this
repo exercised a blob column containing a null or empty value through
`optimize()`, which is why this shipped unnoticed. This PR adds that
guard.

## Tests

Two tests in `rust/lancedb/tests/blob_integration.rs`, reusing the
file's existing 64 KiB dedicated-blob helpers and a delete-triggered
fragment rewrite. After `id IN (1, 4)` is deleted the surviving rows are
`2` (null), `3` (valid empty), `5` and `6` (payloads) — payloads sit
immediately after the null/empty, which is where the misalignment
landed.

- `optimize_preserves_v1_blob_payloads_with_null_and_empty` — storage
2.0; asserts the **payload bytes** are unchanged across
`OptimizeAction::All` (what the Python/Node `optimize()` bindings
invoke). Payloads are read through `lance::Dataset::take_blobs`, since
`Table::fetch_blobs` rejects legacy v1 columns. The before/after
descriptors are reported on failure but deliberately *not* asserted:
compaction repacks the blob file, so they shift legitimately (id 5
`(131072, 65536)` → `(0, 65536)`, id 6 `(196608, 65536)` → `(65536,
65536)`). Note that a post-compaction `position: 0` is both the
legitimate first-payload offset and the bug's signature, so asserting
descriptors would be actively misleading.
- `optimize_preserves_blob_v2_null_and_empty_distinction` — storage >=
2.2; asserts a null stays null and a valid empty value stays non-null
empty.

Both assert the pre-optimize state first, so a setup change that stops
producing the null/empty/payload mix fails loudly instead of passing
vacuously.

Both also assert the returned `CompactionMetrics` show a fragment was
actually rewritten. These tests depend on `delete("id IN (1, 4)")`
pushing the fragment past lance's `materialize_deletions_threshold` (0.1
by default; 2 of 6 rows here). That coupling is invisible and unasserted
otherwise: against a forced no-op (`materialize_deletions_threshold:
1.5`) the metrics come back all zeroes and *every payload assertion
still passes*. Since the whole point of these tests is to survive
dependency changes, they check that the rewrite happened rather than
trusting the planner to keep selecting the fragment.

Guard verified against a pre-fix lance: with the published
`lancedb==0.36.0` wheel (vendors lance 9.0.0), `Table.optimize()` on the
same data rewrites the descriptors of the two rows following the
null/empty from `(131072, 65536)` and `(196608, 65536)` to `(0, 0)`, and
the payloads read back empty. Against the pinned `v10.1.0-beta.1`, all
39 tests in the file pass, adding roughly 10–20 ms to the file's
runtime.

## Not addressed here

- **No released artifact has the fix yet.** PyPI `lancedb` 0.36.0
(2026-07-29) vendors lance 9.0.0; npm `@lancedb/lancedb` 0.37.1-beta.0
predates the bump. No 9.x lance tag carries the fix: `v10.0.0-beta.3` is
the first tag containing it, every `v9.1.0-beta.1`…`beta.8` is behind
it, and `v9.0.0` / `v9.0.1-rc.1` sit on a diverged branch without it. A
stable lancedb release needs a stable lance >= 10.
- **The version skew #3744 flagged is still live.**
`python/pyproject.toml` pins `pylance==9.0.0rc1` for the `tests` extra
against a vendored `10.1.0-beta.1`, so Python CI still cannot observe
this class of divergence.
- **Only the single-fragment rewrite shape is covered.** Both tests
rewrite one fragment by materializing deletions. lance's own
`test_compact_blob_v1/v2_preserves_null_empty_and_payload_order` cover
the multi-fragment merge shape (3 fragments → 1) at unit level, so this
PR is complementary rather than redundant — it covers the binding-level
path through `Table::optimize` — but it would not catch a regression
that only appears when *merging* fragments.
`multi_fragment_dedicated_blob_table` in the same file makes that a
cheap follow-up.

Closes #3744

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-04 12:23:03 -07:00
Wyatt Alt 8e24dd3828 feat(rust)!: make add_columns a builder (#3778)
Table::add_columns now takes no arguments and returns AddColumnsBuilder,
so calls become .add_columns().transform(t).execute().

read_columns was the second positional argument but reaches only one of
the five transform variants. In lance's add_columns_to_fragments only
BatchUDF receives the caller's value: SqlExpressions replaces it with
the columns its expressions reference, Stream and Reader pass None, and
AllNulls reads nothing. So it was mandatory on every call -- all
eighteen call sites here passed None -- and silently discarded four
times out of five. As a builder method it is optional, and setting it
where lance would discard it is now an error, which does reject a call
that previously succeeded while ignoring the argument.

Matches the builders add, update, and merge_insert already use.
2026-08-04 11:18:22 -07:00
14 changed files with 747 additions and 69 deletions
Generated
+8 -8
View File
@@ -7583,7 +7583,7 @@ version = "0.14.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
dependencies = [
"heck 0.5.0",
"heck 0.4.1",
"itertools 0.14.0",
"log",
"multimap",
@@ -8498,9 +8498,9 @@ dependencies = [
[[package]]
name = "rkyv"
version = "0.8.16"
version = "0.8.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73389e0c99e664f919275ab5b5b0471391fe9a8de61e1dff9b1eaf56a90f16e3"
checksum = "815cc8a37159a463064825246cadb07961e25cd9885908606f6d08a98d8f8874"
dependencies = [
"bytecheck",
"bytes",
@@ -8517,9 +8517,9 @@ dependencies = [
[[package]]
name = "rkyv_derive"
version = "0.8.16"
version = "0.8.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d2ed0b54125315fb36bd021e82d314d1c126548f871634b483f46b31d13cac6"
checksum = "c0ed1a78a1b19d184b0daa629dd9a024573173ec7d485b287cb369fb3607cc1c"
dependencies = [
"proc-macro2",
"quote",
@@ -9277,7 +9277,7 @@ version = "0.8.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451"
dependencies = [
"heck 0.5.0",
"heck 0.4.1",
"proc-macro2",
"quote",
"syn 2.0.117",
@@ -9289,7 +9289,7 @@ version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40"
dependencies = [
"heck 0.5.0",
"heck 0.4.1",
"proc-macro2",
"quote",
"syn 2.0.117",
@@ -9730,7 +9730,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.2",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
+6 -2
View File
@@ -339,7 +339,9 @@ impl Table {
let transforms = NewColumnTransform::SqlExpressions(transforms);
let res = self
.inner_ref()?
.add_columns(transforms, None)
.add_columns()
.transform(transforms)
.execute()
.await
.default_error()?;
Ok(res.into())
@@ -356,7 +358,9 @@ impl Table {
let transforms = NewColumnTransform::AllNulls(schema);
let res = self
.inner_ref()?
.add_columns(transforms, None)
.add_columns()
.transform(transforms)
.execute()
.await
.default_error()?;
Ok(res.into())
+12 -2
View File
@@ -1375,7 +1375,12 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner.add_columns(definitions, None).await.infer_error()?;
let result = inner
.add_columns()
.transform(definitions)
.execute()
.await
.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
@@ -1389,7 +1394,12 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner.add_columns(transform, None).await.infer_error()?;
let result = inner
.add_columns()
.transform(transform)
.execute()
.await
.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
+5
View File
@@ -420,6 +420,11 @@ impl Connection {
///
/// * `name` - The name of the table
/// * `initial_data` - The initial data to write to the table
///
/// Floating-point `List` columns named `vec`, or with `vector` or `embedding`
/// in their name, are inferred as vector columns when the first batch has a
/// uniform, non-zero list length. The inferred dimension is validated for all
/// subsequent batches and stored as a `FixedSizeList`.
pub fn create_table<T: Scannable + 'static>(
&self,
name: impl Into<String>,
+3 -1
View File
@@ -8,7 +8,7 @@ use lance_io::object_store::StorageOptionsProvider;
use crate::{
Error, Result, Table,
connection::{merge_storage_options, set_storage_options_provider},
data::scannable::{Scannable, WithEmbeddingsScannable},
data::scannable::{Scannable, WithEmbeddingsScannable, maybe_infer_vector_schema},
database::{CreateTableMode, CreateTableRequest, Database},
embeddings::{EmbeddingDefinition, EmbeddingFunction, EmbeddingRegistry},
table::WriteOptions,
@@ -147,6 +147,8 @@ impl CreateTableBuilder {
let embedding_registry = self.embedding_registry.clone();
let parent = self.parent.clone();
self.request.data = maybe_infer_vector_schema(self.request.data).await?;
// If embeddings were configured via add_embedding(), wrap the data
if !self.embeddings.is_empty() {
let wrapped_data: Box<dyn Scannable> = Box::new(WithEmbeddingsScannable::try_new(
+145 -2
View File
@@ -18,13 +18,16 @@ use crate::embeddings::{
};
use crate::table::{ColumnDefinition, ColumnKind, TableDefinition};
use crate::{Error, Result};
use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader};
use arrow_schema::{ArrowError, SchemaRef};
use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, cast::AsArray};
use arrow_cast::{CastOptions, cast_with_options};
use arrow_schema::{ArrowError, DataType, Schema, SchemaRef};
use async_trait::async_trait;
use futures::StreamExt;
use futures::stream::once;
use lance_datafusion::utils::StreamingWriteSource;
use super::inspect::infer_dimension;
pub trait Scannable: Send {
/// Returns the schema of the data.
fn schema(&self) -> SchemaRef;
@@ -497,6 +500,146 @@ impl Scannable for PeekedScannable {
}
}
fn name_suggests_vector_column(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "vec" || name.contains("vector") || name.contains("embedding")
}
/// Infer fixed dimensions for vector-like floating-point list columns.
///
/// Lance vector search requires `FixedSizeList` columns, but Arrow data assembled
/// from runtime embedding models is often represented as `List`. For vector-like
/// column names, inspect the first batch and convert uniform, non-empty lists to a
/// fixed-size schema. Every subsequent batch is cast with strict length checking.
pub(crate) async fn maybe_infer_vector_schema(
data: Box<dyn Scannable>,
) -> Result<Box<dyn Scannable>> {
let input_schema = data.schema();
let candidates = input_schema
.fields()
.iter()
.enumerate()
.filter(|(_, field)| {
name_suggests_vector_column(field.name())
&& matches!(
field.data_type(),
DataType::List(item) | DataType::LargeList(item)
if item.data_type().is_floating()
)
})
.map(|(index, _)| index)
.collect::<Vec<_>>();
if candidates.is_empty() {
return Ok(data);
}
let mut peeked = PeekedScannable::new(data);
let Some(first_batch) = peeked.peek().await else {
return Ok(Box::new(peeked));
};
let mut fields = input_schema.fields().iter().cloned().collect::<Vec<_>>();
let mut changed = false;
for index in candidates {
let array = first_batch.column(index);
let dimension = match array.data_type() {
DataType::List(_) => {
infer_dimension::<arrow_array::types::Int32Type>(array.as_list::<i32>())?
.map(i64::from)
}
DataType::LargeList(_) => {
infer_dimension::<arrow_array::types::Int64Type>(array.as_list::<i64>())?
}
_ => unreachable!(),
};
let Some(dimension) = dimension.filter(|dimension| *dimension > 0) else {
continue;
};
let dimension = i32::try_from(dimension).map_err(|_| Error::InvalidInput {
message: format!(
"Vector column '{}' has a dimension larger than i32::MAX",
fields[index].name()
),
})?;
let item = match fields[index].data_type() {
DataType::List(item) | DataType::LargeList(item) => item.clone(),
_ => unreachable!(),
};
fields[index] = Arc::new(
fields[index]
.as_ref()
.clone()
.with_data_type(DataType::FixedSizeList(item, dimension)),
);
changed = true;
}
if !changed {
return Ok(Box::new(peeked));
}
let output_schema = Arc::new(Schema::new_with_metadata(
fields,
input_schema.metadata().clone(),
));
Ok(Box::new(InferredVectorScannable {
inner: peeked,
output_schema,
}))
}
struct InferredVectorScannable {
inner: PeekedScannable,
output_schema: SchemaRef,
}
impl Scannable for InferredVectorScannable {
fn schema(&self) -> SchemaRef {
self.output_schema.clone()
}
fn scan_as_stream(&mut self) -> SendableRecordBatchStream {
let output_schema = self.output_schema.clone();
let stream_schema = output_schema.clone();
let stream = self.inner.scan_as_stream().map(move |batch| {
let batch = batch?;
let columns = batch
.columns()
.iter()
.zip(output_schema.fields())
.map(|(array, field)| {
if array.data_type() == field.data_type() {
Ok(array.clone())
} else {
cast_with_options(
array,
field.data_type(),
&CastOptions {
safe: false,
..Default::default()
},
)
.map_err(Error::from)
}
})
.collect::<Result<Vec<_>>>()?;
Ok(RecordBatch::try_new(output_schema.clone(), columns)?)
});
Box::pin(SimpleRecordBatchStream {
schema: stream_schema,
stream,
})
}
fn num_rows(&self) -> Option<usize> {
self.inner.num_rows()
}
fn rescannable(&self) -> bool {
self.inner.rescannable()
}
}
/// Compute the number of write partitions based on data size estimates.
///
/// `sample_bytes` and `sample_rows` come from a representative batch and are
+5 -1
View File
@@ -72,7 +72,9 @@
//!
//! LanceDB uses [arrow-rs](https://github.com/apache/arrow-rs) to define schema, data types and array itself.
//! It treats [`FixedSizeList<Float16/Float32>`](https://docs.rs/arrow/latest/arrow/array/struct.FixedSizeListArray.html)
//! columns as vector columns.
//! columns as vector columns. When creating a table with a floating-point `List`
//! column named `vec`, or with `vector` or `embedding` in its name, LanceDB infers
//! a uniform dimension from the first batch and stores it as a `FixedSizeList`.
//!
//! For more details, please refer to the [LanceDB documentation](https://docs.lancedb.com).
//!
@@ -82,6 +84,8 @@
//! schema of the `RecordBatch` determines the schema of the table.
//!
//! Vector columns should be represented as `FixedSizeList<Float16/Float32>` data type.
//! A vector-like `List<Float16/Float32>` input is also accepted when every vector
//! has the same runtime dimension.
//!
//! ```rust
//! # use std::sync::Arc;
+87 -2
View File
@@ -1648,8 +1648,8 @@ mod tests {
use super::*;
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
use arrow_array::{
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
types::Float32Type,
FixedSizeListArray, Float32Array, Int32Array, ListArray, RecordBatch, StringArray,
cast::AsArray, types::Float32Type,
};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use futures::{StreamExt, TryStreamExt};
@@ -2282,6 +2282,91 @@ mod tests {
);
}
#[tokio::test]
async fn vector_search_infers_dimension_from_list_array() {
let tmp_dir = tempdir().unwrap();
let schema = Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, false),
ArrowField::new(
"vec",
DataType::List(Arc::new(ArrowField::new("item", DataType::Float32, true))),
true,
),
]));
let vectors = ListArray::from_iter_primitive::<Float32Type, _, _>([
Some([Some(0.0), Some(0.0)]),
Some([Some(1.0), Some(1.0)]),
]);
let batch = RecordBatch::try_new(
schema,
vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)],
)
.unwrap();
let table = connect(tmp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap()
.create_table("vectors", batch)
.execute()
.await
.unwrap();
assert!(matches!(
table.schema().await.unwrap().field(1).data_type(),
DataType::FixedSizeList(_, 2)
));
let results = table
.vector_search(&[0.0, 0.0])
.unwrap()
.limit(1)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(results[0]["id"].as_primitive::<Int32Type>().value(0), 0);
}
#[tokio::test]
async fn inferred_vector_dimension_is_validated_across_batches() {
let tmp_dir = tempdir().unwrap();
let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
"vec",
DataType::List(Arc::new(ArrowField::new("item", DataType::Float32, true))),
true,
)]));
let first = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(
ListArray::from_iter_primitive::<Float32Type, _, _>([Some([Some(0.0), Some(0.0)])]),
)],
)
.unwrap();
let wrong_dimension = RecordBatch::try_new(
schema,
vec![Arc::new(
ListArray::from_iter_primitive::<Float32Type, _, _>([Some([
Some(1.0),
Some(1.0),
Some(1.0),
])]),
)],
)
.unwrap();
let result = connect(tmp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap()
.create_table("vectors", vec![first, wrong_dimension])
.execute()
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_fast_search_plan() {
let tmp_dir = tempdir().unwrap();
+24 -19
View File
@@ -3089,10 +3089,12 @@ mod tests {
Box::pin(table.delete("false").map_ok(|_| ())),
Box::pin(
table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("x".into(), "y".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"x".into(),
"y".into(),
)]))
.execute()
.map_ok(|_| ()),
),
Box::pin(async {
@@ -6388,13 +6390,12 @@ mod tests {
});
let result = table
.add_columns(
NewColumnTransform::SqlExpressions(vec![
("b".into(), "a + 1".into()),
("x".into(), "cast(NULL as int32)".into()),
]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![
("b".into(), "a + 1".into()),
("x".into(), "cast(NULL as int32)".into()),
]))
.execute()
.await
.unwrap();
@@ -7119,10 +7120,12 @@ mod tests {
}
"add_columns" => {
let _ = table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + 1".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"c".into(),
"a + 1".into(),
)]))
.execute()
.await;
}
"drop_columns" => {
@@ -9880,10 +9883,12 @@ mod tests {
.await
.unwrap();
branch
.add_columns(
NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"b".into(),
"a + 1".into(),
)]))
.execute()
.await
.unwrap();
branch
+4 -6
View File
@@ -65,6 +65,7 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
use self::dataset::DatasetConsistencyWrapper;
use self::merge::MergeInsertBuilder;
pub mod add_columns;
mod add_data;
pub mod branch_merge;
mod create_index;
@@ -79,6 +80,7 @@ pub mod schema_evolution;
pub mod update;
pub mod write_progress;
use crate::index::waiter::wait_for_index;
pub use add_columns::AddColumnsBuilder;
#[cfg(feature = "remote")]
pub(crate) use add_data::PreprocessingOutput;
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
@@ -1620,12 +1622,8 @@ impl Table {
}
/// Add new columns to the table, providing values to fill in.
pub async fn add_columns(
&self,
transforms: NewColumnTransform,
read_columns: Option<Vec<String>>,
) -> Result<AddColumnsResult> {
self.inner.add_columns(transforms, read_columns).await
pub fn add_columns(&self) -> AddColumnsBuilder {
AddColumnsBuilder::new(self.inner.clone())
}
/// Change a column's name or nullability.
+161
View File
@@ -0,0 +1,161 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Builder for adding columns to a table.
use std::sync::Arc;
use lance::dataset::NewColumnTransform;
use super::BaseTable;
use super::schema_evolution::AddColumnsResult;
use crate::{Error, Result};
/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns).
pub struct AddColumnsBuilder {
parent: Arc<dyn BaseTable>,
transform: Option<NewColumnTransform>,
read_columns: Option<Vec<String>>,
}
impl std::fmt::Debug for AddColumnsBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AddColumnsBuilder")
.field("parent", &self.parent)
.field("has_transform", &self.transform.is_some())
.field("read_columns", &self.read_columns)
.finish()
}
}
impl AddColumnsBuilder {
pub(crate) fn new(parent: Arc<dyn BaseTable>) -> Self {
Self {
parent,
transform: None,
read_columns: None,
}
}
/// Set how the new columns' values are produced. Required.
pub fn transform(mut self, transform: NewColumnTransform) -> Self {
self.transform = Some(transform);
self
}
/// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper
/// receives. Every other transform determines what it reads, so setting
/// this alongside one is an error rather than a silent no-op.
pub fn read_columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.read_columns = Some(columns.into_iter().map(Into::into).collect());
self
}
/// Add the columns.
pub async fn execute(self) -> Result<AddColumnsResult> {
let Self {
parent,
transform,
read_columns,
} = self;
let Some(transform) = transform else {
return Err(Error::InvalidInput {
message: "add_columns requires a transform".into(),
});
};
if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) {
return Err(Error::InvalidInput {
message: "read_columns applies only to a BatchUDF transform; \
every other transform determines what it reads"
.into(),
});
}
parent.add_columns(transform, read_columns).await
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{Int32Array, RecordBatch, record_batch};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{BatchUDF, NewColumnTransform};
use crate::Table;
use crate::connect;
async fn table_with_two_columns(name: &str) -> Table {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap();
conn.create_table(name, batch).execute().await.unwrap()
}
#[tokio::test]
async fn test_requires_a_transform() {
let table = table_with_two_columns("no_transform").await;
let err = table.add_columns().execute().await.unwrap_err();
assert!(
err.to_string().contains("requires a transform"),
"got: {err}"
);
}
#[tokio::test]
async fn test_read_columns_with_sql_expressions_is_rejected() {
let table = table_with_two_columns("read_cols_sql").await;
let err = table
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"doubled".into(),
"x * 2".into(),
)]))
.read_columns(["x"])
.execute()
.await
.unwrap_err();
assert!(err.to_string().contains("BatchUDF"), "got: {err}");
let schema = table.schema().await.unwrap();
assert!(
schema.field_with_name("doubled").is_err(),
"a rejected call must not commit"
);
}
#[tokio::test]
async fn test_read_columns_limits_what_a_batch_udf_sees() {
let table = table_with_two_columns("read_cols_udf").await;
let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)]));
let mapper_schema = output_schema.clone();
let udf = BatchUDF {
mapper: Box::new(move |batch: &RecordBatch| {
assert!(batch.column_by_name("x").is_some());
assert!(batch.column_by_name("y").is_none(), "y was not requested");
let x = batch["x"].as_any().downcast_ref::<Int32Array>().unwrap();
let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect();
Ok(RecordBatch::try_new(
mapper_schema.clone(),
vec![Arc::new(doubled)],
)?)
}),
output_schema,
result_checkpoint: None,
};
table
.add_columns()
.transform(NewColumnTransform::BatchUDF(udf))
.read_columns(["x"])
.execute()
.await
.unwrap();
let schema = table.schema().await.unwrap();
assert!(schema.field_with_name("sum").is_ok());
}
}
+9 -5
View File
@@ -576,10 +576,12 @@ mod tests {
// Add a new physical column AFTER the embedding column.
table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"score".into(),
"42.0".into(),
)]))
.execute()
.await
.unwrap();
@@ -683,7 +685,9 @@ mod tests {
true,
)]));
table
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
.add_columns()
.transform(NewColumnTransform::AllNulls(nested_schema))
.execute()
.await
.unwrap();
+24 -19
View File
@@ -193,10 +193,12 @@ mod tests {
// Add a computed column
let result = table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"doubled".into(),
"id * 2".into(),
)]))
.execute()
.await
.unwrap();
@@ -251,13 +253,12 @@ mod tests {
// Add multiple columns at once
table
.add_columns(
NewColumnTransform::SqlExpressions(vec![
("y".into(), "x + 1".into()),
("z".into(), "x * x".into()),
]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![
("y".into(), "x + 1".into()),
("z".into(), "x * x".into()),
]))
.execute()
.await
.unwrap();
@@ -283,10 +284,12 @@ mod tests {
// Add a column with a constant value
table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"constant".into(),
"42".into(),
)]))
.execute()
.await
.unwrap();
@@ -659,10 +662,12 @@ mod tests {
// Add column increments version
let add_result = table
.add_columns(
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
None,
)
.add_columns()
.transform(NewColumnTransform::SqlExpressions(vec![(
"c".into(),
"a + b".into(),
)]))
.execute()
.await
.unwrap();
assert!(add_result.version > v1);
+254 -2
View File
@@ -9,14 +9,17 @@ use arrow_array::{
};
use arrow_schema::{DataType, Field, Fields, Schema};
use futures::TryStreamExt;
use lance::Dataset;
use lance_encoding::version::LanceFileVersion;
use lancedb::{
Connection, Error, Result, Table,
blob::{BlobRangeRequest, blob},
connect, connect_namespace,
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
database::listing::{
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
},
query::{ExecutableQuery, QueryBase},
table::{AddDataMode, CompactionOptions, OptimizeAction},
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
};
use tempfile::tempdir;
@@ -1075,3 +1078,252 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
}
Ok(())
}
/// Rows exercising the null/empty interleavings from
/// <https://github.com/lancedb/lancedb/issues/3744>: a payload, a null, a valid
/// empty value, then payloads whose descriptors a fragment rewrite used to zero.
fn null_empty_input_batch() -> RecordBatch {
let owned = [
Some(dedicated_blob_bytes(1)),
None,
Some(Vec::new()),
Some(dedicated_blob_bytes(4)),
Some(dedicated_blob_bytes(5)),
Some(dedicated_blob_bytes(6)),
];
let payloads: Vec<Option<&[u8]>> = owned.iter().map(|payload| payload.as_deref()).collect();
binary_input_batch(&[1, 2, 3, 4, 5, 6], &payloads)
}
/// One `(id, Some((payload length, first byte)))` per live row, or `(id, None)`
/// for a null blob. Comparing lengths and first bytes keeps failure output
/// readable where comparing whole payloads would not.
type BlobSummary = Vec<(i64, Option<(usize, Option<u8>)>)>;
/// The rows [`null_empty_input_batch`] leaves behind after `id IN (1, 4)` is
/// deleted: a null, a valid empty value, and the two payloads that follow them.
fn expected_null_empty_survivors() -> BlobSummary {
vec![
(2, None),
(3, Some((0, None))),
(5, Some((DEDICATED_BLOB_LEN, Some(5)))),
(6, Some((DEDICATED_BLOB_LEN, Some(6)))),
]
}
/// `optimize()` only rewrites a fragment when lance's compaction planner selects
/// it — here because the delete pushes the fragment past
/// `materialize_deletions_threshold` (0.1 by default; these tests delete 2 of 6
/// rows). Without this check, a planner or threshold change upstream would leave
/// both regression tests green while no rewrite happened at all.
fn assert_compacted(stats: &OptimizeStats) {
let metrics = stats
.compaction
.as_ref()
.expect("OptimizeAction::All runs compaction");
assert!(
metrics.fragments_removed >= 1,
"optimize() rewrote no fragment, so this test proves nothing: {metrics:?}"
);
}
fn summarize(rows: &[(i64, Option<Vec<u8>>)]) -> BlobSummary {
rows.iter()
.map(|(id, payload)| {
(
*id,
payload
.as_ref()
.map(|bytes| (bytes.len(), bytes.first().copied())),
)
})
.collect()
}
async fn sorted_id_rowid(table: &Table) -> Result<Vec<(i64, u64)>> {
let mut pairs = collect_id_rowid(table).await?;
pairs.sort_by_key(|(id, _)| *id);
Ok(pairs)
}
/// `{position, size}` descriptors of a legacy v1 blob column, keyed by `id`.
async fn v1_blob_descriptors(table: &Table) -> Result<Vec<(i64, Option<(u64, u64)>)>> {
let batches = table
.query()
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
let ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let descriptors = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<StructArray>()
.expect("v1 blob column reads back as a descriptor struct");
let position = descriptors
.column_by_name("position")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let size = descriptors
.column_by_name("size")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let mut rows: Vec<(i64, Option<(u64, u64)>)> = (0..batch.num_rows())
.map(|row| {
let descriptor =
(!descriptors.is_null(row)).then(|| (position.value(row), size.value(row)));
(ids.value(row), descriptor)
})
.collect();
rows.sort_by_key(|(id, _)| *id);
Ok(rows)
}
/// Payload bytes of every live row of a legacy v1 blob column, keyed by `id`.
/// [`Table::fetch_blobs`] rejects v1 columns, so read them through lance.
async fn v1_blob_payloads(dataset_uri: &str, table: &Table) -> Result<Vec<(i64, Option<Vec<u8>>)>> {
let pairs = sorted_id_rowid(table).await?;
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
let dataset = Arc::new(Dataset::open(dataset_uri).await?);
let files = dataset.take_blobs(&row_ids, "image").await?;
assert_eq!(
files.len(),
pairs.len(),
"take_blobs returned {} handles for {} live rows",
files.len(),
pairs.len()
);
let mut rows = Vec::with_capacity(pairs.len());
for ((id, _), file) in pairs.iter().zip(files) {
let payload = match file {
Some(file) => Some(file.read().await?.to_vec()),
None => None,
};
rows.push((*id, payload));
}
Ok(rows)
}
/// Length and first byte of every live blob v2 value, keyed by `id`.
async fn blob_v2_values(table: &Table) -> Result<BlobSummary> {
let pairs = sorted_id_rowid(table).await?;
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
let bytes = table.fetch_blobs("image", &row_ids).await?;
Ok(pairs
.iter()
.enumerate()
.map(|(slot, (id, _))| {
let value = (!bytes.is_null(slot))
.then(|| (bytes.value(slot).len(), bytes.value(slot).first().copied()));
(*id, value)
})
.collect())
}
/// Regression test for [#3744]: on storage 2.0 (legacy v1 descriptors),
/// compaction rewrote every payload following a null or empty value in the same
/// fragment as `{position: 0, size: 0}`, so the payload bytes read back as `b""`
/// and the new fragment no longer referenced them at all.
///
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
#[tokio::test]
async fn optimize_preserves_v1_blob_payloads_with_null_and_empty() -> Result<()> {
let tmp = tempdir().unwrap();
let db_uri = tmp.path().to_str().unwrap().to_string();
let db = connect(&db_uri)
.database_options(&ListingDatabaseOptions {
new_table_config: NewTableConfig {
data_storage_version: Some(LanceFileVersion::V2_0),
..Default::default()
},
..Default::default()
})
.execute()
.await?;
let legacy = Field::new("image", DataType::LargeBinary, true).with_metadata(
std::collections::HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]),
);
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
legacy,
]));
let table = db.create_empty_table("t", schema).execute().await?;
table.add(null_empty_input_batch()).execute().await?;
assert_eq!(
storage_format_version(&table).await,
LanceFileVersion::V2_0.resolve(),
"v1 blob descriptors only exist below storage 2.2"
);
let dataset_uri = table.uri().await?;
// Any rewrite triggers it; deleting rows is the shape from the issue.
table.delete("id IN (1, 4)").await?;
let descriptors_before = v1_blob_descriptors(&table).await?;
let before = v1_blob_payloads(&dataset_uri, &table).await?;
assert_eq!(
summarize(&before),
expected_null_empty_survivors(),
"test setup no longer produces the null/empty/payload mix"
);
let stats = table.optimize(OptimizeAction::All).await?;
assert_compacted(&stats);
let descriptors_after = v1_blob_descriptors(&table).await?;
let after = v1_blob_payloads(&dataset_uri, &table).await?;
assert_eq!(
summarize(&after),
summarize(&before),
"optimize() lost blob payloads; descriptors before={descriptors_before:?} after={descriptors_after:?}"
);
assert!(after == before, "optimize() changed blob payload bytes");
Ok(())
}
/// Regression test for the blob v2 half of [#3744]: compaction rewrote a valid
/// empty value as null, destroying the null-vs-empty distinction.
///
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
#[tokio::test]
async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let table = db
.create_empty_table("t", blob_table_schema())
.execute()
.await?;
table.add(null_empty_input_batch()).execute().await?;
assert!(
storage_format_version(&table).await >= LanceFileVersion::V2_2,
"blob v2 columns require storage >= 2.2"
);
table.delete("id IN (1, 4)").await?;
let before = blob_v2_values(&table).await?;
assert_eq!(
before,
expected_null_empty_survivors(),
"test setup no longer produces the null/empty/payload mix"
);
let stats = table.optimize(OptimizeAction::All).await?;
assert_compacted(&stats);
assert_eq!(
blob_v2_values(&table).await?,
before,
"optimize() changed blob v2 values"
);
Ok(())
}