rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-26 13:36:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix!: normalize blob layouts and support external byte ranges (#4589)

Browse Source
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-19 16:44:00 -07:00
committed by GitHub
parent 6507615a6f
commit 4a03859b2a
20 changed files with 2837 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/blob.rs
+58 -5
View File
@@ -244,8 +244,34 @@ impl BlobRangeRequest {
/// Creates an Arrow field for a Lance blob v2 column.
///
/// `Struct<data, uri>` with the `lance.blob.v2` marker. Same layout Lance
/// expects on write.
/// `Struct<data, uri, position, size>` with nullable children and the
/// `lance.blob.v2` marker. Set both range fields to read a byte range from a
/// URI, or leave both null for a whole file. Inline bytes are also supported.
///
/// # Migration
///
/// This helper previously returned only `data` and `uri`. Rust callers building
/// a `StructArray` from this field must now supply four children: append typed
/// null UInt64 arrays for `position` and `size` when no range is needed:
///
/// ```
/// use std::sync::Arc;
/// use arrow_array::{ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt64Array};
/// use arrow_schema::DataType;
/// let field = lancedb::blob("image", true);
/// let DataType::Struct(children) = field.data_type() else { unreachable!() };
/// let image = StructArray::new(children.clone(), vec![
/// Arc::new(LargeBinaryArray::from(vec![Some(b"image".as_slice())])) as ArrayRef,
/// Arc::new(StringArray::from(vec![None::<&str>])),
/// Arc::new(UInt64Array::new_null(1)),
/// Arc::new(UInt64Array::new_null(1)),
/// ], None);
/// ```
///
/// Existing two-field tables remain writable: append and merge-insert fill
/// missing range fields or drop all-null range fields to match the table.
/// Non-null ranges cannot be written to a two-field table. Overwrite replaces
/// the table schema with the input schema rather than normalizing its layout.
///
/// A blob column may be top-level or nested inside a struct or list. Blobs
/// nested inside structs use dotted paths (e.g. `info.blob`) in the read APIs.
@@ -261,9 +287,21 @@ impl BlobRangeRequest {
/// ]);
/// ```
pub fn blob(name: impl AsRef<str>, nullable: bool) -> Field {
lance::blob::blob_field(name.as_ref(), nullable)
blob_with_options(name, nullable, BlobOptions::default())
}
static BLOB_STORAGE_TYPE: std::sync::LazyLock<DataType> = std::sync::LazyLock::new(|| {
DataType::Struct(
vec![
Field::new("data", DataType::LargeBinary, true),
Field::new("uri", DataType::Utf8, true),
Field::new("position", DataType::UInt64, true),
Field::new("size", DataType::UInt64, true),
]
.into(),
)
});
/// Options for creating a blob v2 column field with tier thresholds.
///
/// Re-exported from lance so the thresholds are built exactly the way
@@ -293,6 +331,7 @@ pub use lance::blob::BlobFieldOptions as BlobOptions;
/// ```
pub fn blob_with_options(name: impl AsRef<str>, nullable: bool, options: BlobOptions) -> Field {
lance::blob::blob_field_with_options(name.as_ref(), nullable, options)
.with_data_type(BLOB_STORAGE_TYPE.clone())
}
/// Returns true if `field` is a blob v2 column.
@@ -306,7 +345,7 @@ pub fn is_blob(field: &Field) -> bool {
}
/// Returns true if `field`, or any field nested under it, is a blob v2 column.
fn field_tree_has_blob_v2(field: &Field) -> bool {
pub(crate) fn field_tree_has_blob_v2(field: &Field) -> bool {
if field.is_blob_v2() {
return true;
}
@@ -673,7 +712,21 @@ mod tests {
field.metadata().get(ARROW_EXT_NAME_KEY).map(String::as_str),
Some("lance.blob.v2")
);
assert!(matches!(field.data_type(), DataType::Struct(_)));
let DataType::Struct(children) = field.data_type() else {
panic!("expected struct")
};
assert_eq!(
children
.iter()
.map(|f| (f.name().as_str(), f.data_type(), f.is_nullable()))
.collect::<Vec<_>>(),
vec![
("data", &DataType::LargeBinary, true),
("uri", &DataType::Utf8, true),
("position", &DataType::UInt64, true),
("size", &DataType::UInt64, true),
]
);
}
#[test]
rust/lancedb/src/remote/table.rs
+302
View File
@@ -3154,6 +3154,17 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
let rescannable = source.rescannable();
let input: Arc<dyn ExecutionPlan> =
Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None));
let table_schema = self.schema().await?;
let input = if crate::blob::has_blob_columns(&table_schema) {
crate::table::merge::validate_merge_source_fields(
input.schema().fields(),
table_schema.fields(),
"",
)?;
crate::table::datafusion::cast::cast_to_table_schema(input, &table_schema)?
} else {
input
};
let freshness_request = self.snapshot_freshness_headers();
let mut merge: Arc<dyn ExecutionPlan> = Arc::new(
@@ -4499,6 +4510,253 @@ mod tests {
assert!(image.column_by_name("data").unwrap().is_null(0));
}
#[tokio::test]
async fn add_external_blob_ranges_preserves_remote_wire_values() {
let field = crate::blob("image", true);
let DataType::Struct(children) = field.data_type().clone() else {
unreachable!();
};
let schema = Arc::new(Schema::new(vec![field]));
// Registered external bases are enforced by the server, so remote
// clients send the URI and range without the local-only opt-in.
let uri = "s3://registered/payload.bin";
let image = StructArray::new(
children,
vec![
arrow_array::new_null_array(&DataType::LargeBinary, 2),
Arc::new(StringArray::from(vec![uri, uri])),
Arc::new(arrow_array::UInt64Array::from(vec![0, 4])),
Arc::new(arrow_array::UInt64Array::from(vec![6, 6])),
],
None,
);
let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(image)]).unwrap();
let describe_body = describe_response(&schema).into_bytes();
let payload = b"0123456789abcdef";
let fetched = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"image",
DataType::LargeBinary,
true,
)])),
vec![Arc::new(LargeBinaryArray::from(vec![
&payload[..6],
&payload[4..10],
]))],
)
.unwrap();
let fetched_body = write_ipc_stream_uncompressed(&fetched);
let (sender, receiver) = std::sync::mpsc::channel();
let table = Table::new_with_handler_version(
"my_table",
semver::Version::new(0, 5, 0),
move |mut request| match request.url().path() {
"/v1/table/my_table/describe/" => http::Response::builder()
.status(200)
.body(describe_body.clone())
.unwrap(),
"/v1/table/my_table/insert/" => {
sender.send(request.body_mut().take().unwrap()).unwrap();
http::Response::builder()
.status(200)
.body(br#"{"version": 2}"#.to_vec())
.unwrap()
}
"/v1/table/my_table/fetch_blobs/" => {
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
.unwrap();
assert_eq!(body["column"], "image");
assert_eq!(body["row_ids"], json!([0, 1]));
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
.body(fetched_body.clone())
.unwrap()
}
path => panic!("unexpected path: {path}"),
},
);
table.add(input).execute().await.unwrap();
let body = collect_body(receiver.recv().unwrap()).await;
let batch = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None)
.unwrap()
.next()
.unwrap()
.unwrap();
assert_eq!(batch.schema(), schema);
let image = batch.column(0).as_struct();
let positions = image
.column(2)
.as_primitive::<arrow::datatypes::UInt64Type>();
let sizes = image
.column(3)
.as_primitive::<arrow::datatypes::UInt64Type>();
assert_eq!(positions.values().as_ref(), &[0, 4]);
assert_eq!(sizes.values().as_ref(), &[6, 6]);
assert_eq!(image.column(1).as_string::<i32>().value(0), uri);
assert_eq!(image.column(1).as_string::<i32>().value(1), uri);
let bytes = table.fetch_blobs("image", &[0, 1]).await.unwrap();
assert_eq!(bytes.value(0), &payload[..6]);
assert_eq!(bytes.value(1), &payload[4..10]);
}
#[tokio::test]
async fn remote_blob_layout_normalization_and_overwrite() {
use arrow_array::{ArrayRef, UInt64Array};
for wide_target in [false, true] {
let field = |wide| {
if wide {
crate::blob("image", true)
} else {
lance::blob::blob_field("image", true)
}
};
let target = Schema::new(vec![
Field::new("id", DataType::Int64, false),
field(wide_target),
]);
for operation in ["add", "merge", "overwrite"] {
for range in [
None,
Some((Some(4), None)),
Some((None, Some(6))),
Some((Some(4), Some(6))),
] {
if range.is_some() && (wide_target || operation == "overwrite") {
continue;
}
// Range input always uses the full layout. Otherwise use the
// opposite layout to exercise both normalization directions.
let wide_input = range.is_some() || !wide_target;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
field(wide_input),
]));
let DataType::Struct(children) = schema.field(1).data_type().clone() else {
unreachable!()
};
let mut arrays: Vec<ArrayRef> = vec![
Arc::new(LargeBinaryArray::from(vec![Some(b"inline".as_slice())])),
Arc::new(StringArray::from(vec![None::<&str>])),
];
if wide_input {
arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.0)])));
arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.1)])));
}
let image = Arc::new(StructArray::new(children, arrays, None));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int64Array::from(vec![1])), image],
)
.unwrap();
let describe = describe_response(&target);
let (sender, receiver) = std::sync::mpsc::channel();
let table = Table::new_with_handler("my_table", move |mut request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(describe.clone())
.unwrap();
}
assert!(
request.url().path().ends_with("/insert/")
|| request.url().path().ends_with("/merge_insert/")
);
if operation == "overwrite" {
assert!(
request
.url()
.query_pairs()
.any(|(k, v)| k == "mode" && v == "overwrite")
);
}
sender.send(request.body_mut().take().unwrap()).unwrap();
http::Response::builder()
.status(200)
.body(r#"{"version": 2, "num_inserted_rows": 1}"#.to_string())
.unwrap()
});
// A bad later batch must abort the request, including after
// an earlier valid batch was encoded. No write handler runs.
let mut batches = Vec::new();
if range.is_some() {
let empty_range_arrays: Vec<ArrayRef> = vec![
Arc::new(LargeBinaryArray::from(vec![Some(b"first".as_slice())])),
Arc::new(StringArray::from(vec![None::<&str>])),
Arc::new(UInt64Array::new_null(1)),
Arc::new(UInt64Array::new_null(1)),
];
let DataType::Struct(children) = schema.field(1).data_type().clone() else {
unreachable!()
};
batches.push(Ok(RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(vec![0])),
Arc::new(StructArray::new(children, empty_range_arrays, None)),
],
)
.unwrap()));
}
batches.push(Ok(batch));
let reader: Box<dyn RecordBatchReader + Send> =
Box::new(RecordBatchIterator::new(batches, schema.clone()));
let result = match operation {
"merge" => {
let mut merge = table.merge_insert(&["id"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
merge.execute(reader).await.map(|_| ())
}
"overwrite" => table
.add(reader)
.mode(AddDataMode::Overwrite)
.execute()
.await
.map(|_| ()),
_ => table.add(reader).execute().await.map(|_| ()),
};
if !wide_target && range.is_some() && operation != "overwrite" {
assert!(
result
.unwrap_err()
.to_string()
.contains("cannot store byte ranges")
);
assert!(
receiver.try_recv().is_err(),
"invalid stream must not reach a committed write"
);
continue;
}
result.unwrap();
let body = collect_body(receiver.recv().unwrap()).await;
let reader = StreamReader::try_new(Cursor::new(body), None).unwrap();
let expected = if operation == "overwrite" {
schema.as_ref()
} else {
&target
};
assert_eq!(reader.schema().as_ref(), expected);
for batch in reader {
let batch = batch.unwrap();
let image = batch.column(1).as_struct();
assert!(
image.column(0).as_binary::<i64>().value(0) == b"inline"
|| image.column(0).as_binary::<i64>().value(0) == b"first"
);
if image.num_columns() == 4 && range.is_none() {
assert!(image.column(2).is_null(0));
assert!(image.column(3).is_null(0));
}
}
}
}
}
}
#[rstest]
#[case(true)]
#[case(false)]
@@ -4852,6 +5110,12 @@ mod tests {
));
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
.unwrap();
}
if request.url().path() == "/v1/table/my_table/merge_insert/" {
assert_eq!(request.method(), "POST");
@@ -4908,6 +5172,12 @@ mod tests {
));
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
.unwrap();
}
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
@@ -4993,6 +5263,12 @@ mod tests {
let num_deleted_rows = if delete_unmatched { 3 } else { 0 };
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(describe_response(&Schema::empty()))
.unwrap();
}
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE);
@@ -5063,6 +5339,12 @@ mod tests {
));
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
.unwrap();
}
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
// One repeated `on` per column, in the order the caller gave them.
@@ -5134,6 +5416,12 @@ mod tests {
// Default parameters
let table = Table::new_with_handler("my_table", |request| {
if request.url().path().ends_with("/describe/") {
return http::Response::builder()
.status(200)
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
.unwrap();
}
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
@@ -14405,6 +14693,20 @@ mod tests {
));
let table = Table::new_with_handler("my_table", move |request| {
match request.url().path() {
"/v1/table/my_table/describe/" => {
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
.unwrap();
assert_eq!(body["branch"], "exp");
http::Response::builder()
.status(200)
.body(describe_response(&Schema::new(vec![Field::new(
"id",
DataType::Int32,
false,
)])))
.unwrap()
}
"/v1/table/my_table/branches/create/" => http::Response::builder()
.status(200)
.body("{}".to_string())
rust/lancedb/src/table/add_data.rs
+106 -6
View File
@@ -1,10 +1,12 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashSet;
use std::sync::Arc;
use arrow_schema::{DataType, Fields, Schema};
use arrow_schema::{DataType, Field, Fields, Schema};
use lance::dataset::WriteMode;
use lance_arrow::FieldExt;
use serde::{Deserialize, Serialize};
use crate::data::scannable::Scannable;
@@ -233,6 +235,7 @@ pub struct PreprocessingOutput {
/// Fields can be in different orders, so match by name.
///
/// If a column exists in input but not in table, error (no extra columns allowed).
/// Input field names must be unique within each struct.
///
/// If a column exists in table but not in input, that is okay - it may be filled with nulls.
///
@@ -244,7 +247,13 @@ fn validate_schema(input: &Schema, table: &Schema) -> Result<()> {
}
fn validate_fields(input: &Fields, table: &Fields) -> Result<()> {
let mut names = HashSet::with_capacity(input.len());
for field in input {
if !names.insert(field.name()) {
return Err(Error::InvalidInput {
message: format!("field '{}' is specified more than once", field.name()),
});
}
match table.iter().find(|f| f.name() == field.name()) {
None => {
return Err(Error::InvalidInput {
@@ -252,17 +261,33 @@ fn validate_fields(input: &Fields, table: &Fields) -> Result<()> {
});
}
Some(table_field) => {
if let (DataType::Struct(in_children), DataType::Struct(tbl_children)) =
(field.data_type(), table_field.data_type())
{
validate_fields(in_children, tbl_children)?;
}
validate_field(field, table_field)?;
}
}
}
Ok(())
}
fn validate_field(input: &Field, table: &Field) -> Result<()> {
// Blob coercion validates its children and permits dropping only
// all-null range fields when targeting a legacy two-field layout.
if table.is_blob_v2() {
return Ok(());
}
match (input.data_type(), table.data_type()) {
(DataType::Struct(input), DataType::Struct(table)) => validate_fields(input, table),
(
DataType::List(input) | DataType::LargeList(input) | DataType::FixedSizeList(input, _),
DataType::List(table) | DataType::LargeList(table) | DataType::FixedSizeList(table, _),
) => {
// Preserve the item's extension metadata, and match nested struct
// children by name rather than comparing collection item names.
validate_field(input, table)
}
_ => Ok(()),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
@@ -299,6 +324,81 @@ mod tests {
conn.create_table("test", batch).execute().await.unwrap()
}
#[rstest::rstest]
fn test_validate_append_collection_fields(
#[values("list", "large_list", "fixed_size_list")] source_kind: &str,
#[values("list", "large_list", "fixed_size_list")] target_kind: &str,
) {
fn schema(kind: &str, item: Field) -> Schema {
let item = Arc::new(item);
let data_type = match kind {
"list" => DataType::List(item),
"large_list" => DataType::LargeList(item),
"fixed_size_list" => DataType::FixedSizeList(item, 1),
_ => unreachable!(),
};
Schema::new(vec![Field::new("items", data_type, true)])
}
let image = crate::blob("image", true);
let legacy_image = lance::blob::blob_field("image", true);
let caption = Field::new("caption", DataType::Utf8, true);
let input = schema(
source_kind,
Field::new(
"source_item",
DataType::Struct(vec![image.clone(), caption.clone()].into()),
true,
),
);
let target = schema(
target_kind,
Field::new(
"target_item",
DataType::Struct(vec![legacy_image.clone()].into()),
true,
),
);
let error = super::validate_schema(&input, &target)
.unwrap_err()
.to_string();
assert!(
error.contains("caption") && error.contains("does not exist"),
"{error}"
);
let target = schema(
target_kind,
Field::new(
"target_item",
DataType::Struct(vec![caption.clone(), legacy_image.clone()].into()),
true,
),
);
// Declared siblings can be reordered, but must not occur more than once.
super::validate_schema(&input, &target).unwrap();
let duplicate = schema(
source_kind,
Field::new(
"source_item",
DataType::Struct(vec![image.clone(), caption.clone(), caption].into()),
true,
),
);
let error = super::validate_schema(&duplicate, &target)
.unwrap_err()
.to_string();
assert!(
error.contains("caption") && error.contains("specified more than once"),
"{error}"
);
// Item names need not match, and blob children still reach range coercion.
super::validate_schema(
&schema(source_kind, image.with_name("source_item")),
&schema(target_kind, legacy_image.with_name("target_item")),
)
.unwrap();
}
async fn test_add_with_data<T>(data: T)
where
T: Scannable + 'static,
rust/lancedb/src/table/datafusion/blob_coerce.rs
+1119 -32
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/table/datafusion/cast.rs
+38 -23
View File
@@ -33,9 +33,12 @@ pub fn cast_to_table_schema(
return Ok(input);
}
let exprs = build_field_exprs(input_schema.fields(), table_schema.fields(), &|idx| {
Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc<dyn PhysicalExpr>
})?;
let exprs = build_field_exprs(
input_schema.fields(),
table_schema.fields(),
&|idx| Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc<dyn PhysicalExpr>,
&[],
)?;
let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = exprs
.into_iter()
@@ -52,10 +55,13 @@ pub fn cast_to_table_schema(
/// For each table field that exists in the input, produce an expression that
/// reads from the input and casts if needed. Fields in the table but not in the
/// input are omitted (the storage layer handles missing columns).
fn build_field_exprs(
/// `ancestors` carries validity for blob range validation without nullifying
/// non-nullable children hidden beneath a null parent.
pub(super) fn build_field_exprs(
input_fields: &Fields,
table_fields: &Fields,
get_input_expr: &dyn Fn(usize) -> Arc<dyn PhysicalExpr>,
ancestors: &[Arc<dyn PhysicalExpr>],
) -> Result<Vec<(Arc<dyn PhysicalExpr>, FieldRef)>> {
let config = Arc::new(ConfigOptions::default());
let mut result = Vec::new();
@@ -109,13 +115,14 @@ fn build_field_exprs(
input_field,
table_field,
&config,
ancestors,
)?);
continue;
}
// Blob values in lists need the same coercion as top-level blob columns.
if let Some(coerced) =
coerce_blob_list_expr(input_expr.clone(), input_field, table_field, &config)?
coerce_blob_list_expr(input_expr.clone(), input_field, table_field, ancestors)?
{
result.push(coerced);
continue;
@@ -138,30 +145,38 @@ fn build_field_exprs(
(DataType::Struct(in_children), DataType::Struct(tbl_children))
if in_children != tbl_children =>
{
let sub_exprs = build_field_exprs(in_children, tbl_children, &|child_idx| {
let child_name = in_children[child_idx].name();
Arc::new(ScalarFunctionExpr::new(
&format!("get_field({child_name})"),
get_field(),
vec![
input_expr.clone(),
Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))),
],
Arc::new(in_children[child_idx].as_ref().clone()),
config.clone(),
)) as Arc<dyn PhysicalExpr>
})?;
let mut child_ancestors = ancestors.to_vec();
child_ancestors.push(input_expr.clone());
let sub_exprs = build_field_exprs(
in_children,
tbl_children,
&|child_idx| {
let child_name = in_children[child_idx].name();
Arc::new(ScalarFunctionExpr::new(
&format!("get_field({child_name})"),
get_field(),
vec![
input_expr.clone(),
Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))),
],
Arc::new(in_children[child_idx].as_ref().clone()),
config.clone(),
)) as Arc<dyn PhysicalExpr>
},
&child_ancestors,
)?;
let output_struct_fields: Fields = sub_exprs
.iter()
.map(|(_, f)| f.clone())
.collect::<Vec<_>>()
.into();
let output_field: FieldRef = Arc::new(Field::new(
table_field.name(),
DataType::Struct(output_struct_fields),
table_field.is_nullable(),
));
let output_field: FieldRef = Arc::new(
table_field
.as_ref()
.clone()
.with_data_type(DataType::Struct(output_struct_fields)),
);
// Build named_struct(lit("a"), expr_a, lit("b"), expr_b, ...)
let mut ns_args: Vec<Arc<dyn PhysicalExpr>> = Vec::new();
rust/lancedb/src/table/merge.rs
+42 -29
View File
@@ -254,7 +254,11 @@ fn canonicalize_merge_filter(filter: Option<MergeFilter>) -> Result<Option<Merge
// The JSON projection iterates target fields, so reject unknown source fields
// before it can erase them. Missing and reordered fields remain valid merge
// inputs; type compatibility is still checked by the cast and Lance writer.
fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) -> Result<()> {
pub(crate) fn validate_merge_source_fields(
input: &Fields,
target: &Fields,
parent: &str,
) -> Result<()> {
let mut names = HashSet::with_capacity(input.len());
for field in input {
let path = if parent.is_empty() {
@@ -273,7 +277,9 @@ fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) -
.ok_or_else(|| Error::InvalidInput {
message: format!("merge source field '{path}' is not present in the table schema"),
})?;
validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?;
if !crate::blob::is_blob(target_field) {
validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?;
}
}
Ok(())
}
@@ -296,7 +302,13 @@ fn validate_merge_source_type(input: &DataType, target: &DataType, path: &str) -
| DataType::LargeListView(target),
)
| (DataType::Map(input, _), DataType::Map(target, _)) => {
validate_merge_source_type(input.data_type(), target.data_type(), path)
// Collection items carry extension metadata just like named fields.
// Blob coercion, rather than struct validation, checks range narrowing.
if crate::blob::is_blob(target) {
Ok(())
} else {
validate_merge_source_type(input.data_type(), target.data_type(), path)
}
}
(DataType::Dictionary(_, input), _) => validate_merge_source_type(input, target, path),
(_, DataType::Dictionary(_, target)) => validate_merge_source_type(input, target, path),
@@ -313,35 +325,18 @@ pub(crate) async fn execute_merge_insert(
new_data: Box<dyn RecordBatchReader + Send>,
) -> Result<MergeResult> {
params.canonicalize_filters()?;
let schema = table.schema().await?;
super::computed_columns::ensure_no_function_bindings_for_mutation(
table.schema().await?.as_ref(),
schema.as_ref(),
"merge_insert",
)?;
match lsm::lsm_dispatch_decision(table, &params).await? {
lsm::LsmDispatch::Lsm(plan) => {
let future =
lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, new_data);
return match params.timeout {
Some(timeout) => match tokio::time::timeout(timeout, future).await {
Ok(result) => result,
Err(_) => Err(Error::Runtime {
message: "merge insert timed out".to_string(),
}),
},
None => future.await,
};
}
lsm::LsmDispatch::Standard => {}
}
let dataset = table.dataset.get().await?;
let schema = arrow_schema::Schema::from(dataset.schema());
// JSON source fields must carry the stored extension metadata just as on
// append. Keep arrow.json text labelled until Lance encodes it as JSONB.
let source = if schema
.fields()
.iter()
.any(|field| lance_arrow::json::has_json_fields(field))
// Normalize blob layouts and JSON extension metadata just as on append.
// Keep arrow.json text labelled until Lance encodes it as JSONB.
let source = if crate::blob::has_blob_columns(&schema)
|| schema
.fields()
.iter()
.any(|field| lance_arrow::json::has_json_fields(field))
{
validate_merge_source_fields(new_data.schema().fields(), schema.fields(), "")?;
let plan = Arc::new(super::datafusion::scannable_exec::ScannableExec::new(
@@ -356,6 +351,24 @@ pub(crate) async fn execute_merge_insert(
} else {
new_data.into_stream()
};
match lsm::lsm_dispatch_decision(table, &params).await? {
lsm::LsmDispatch::Lsm(plan) => {
let future =
lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, source);
return match params.timeout {
Some(timeout) => match tokio::time::timeout(timeout, future).await {
Ok(result) => result,
Err(_) => Err(Error::Runtime {
message: "merge insert timed out".to_string(),
}),
},
None => future.await,
};
}
lsm::LsmDispatch::Standard => {}
}
let dataset = table.dataset.get().await?;
let mut builder = LanceMergeInsertBuilder::try_new(dataset.clone(), params.on)?;
match (
params.when_matched_update_all,
rust/lancedb/src/table/merge/lsm.rs
+5 -4
View File
@@ -24,8 +24,10 @@ use arrow_array::cast::AsArray;
use arrow_array::types::{
Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
};
use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, RecordBatchReader};
use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch};
use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef};
use datafusion_physical_plan::SendableRecordBatchStream;
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::mem_wal::{
DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec,
@@ -707,7 +709,7 @@ pub(crate) async fn execute_lsm_merge_insert(
table: &NativeTable,
plan: LsmPlan,
validate_single_shard: bool,
new_data: Box<dyn RecordBatchReader + Send>,
mut new_data: SendableRecordBatchStream,
) -> Result<MergeResult> {
let dataset = table.dataset.get().await?;
let target_schema: SchemaRef = Arc::new(ArrowSchema::from(dataset.schema()));
@@ -718,8 +720,7 @@ pub(crate) async fn execute_lsm_merge_insert(
let mut batches: Vec<RecordBatch> = Vec::new();
let mut total_rows: u64 = 0;
for batch in new_data {
let batch = batch.map_err(|e| Error::Arrow { source: e })?;
while let Some(batch) = new_data.try_next().await? {
if batch.num_rows() == 0 {
continue;
}
rust/lancedb/src/table/refresh.rs
+3 -3
View File
@@ -1747,7 +1747,7 @@ mod tests {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", arrow_schema::DataType::Int32, false),
crate::blob("image", true),
lance::blob::blob_field("image", true),
])),
vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()],
)
@@ -1901,7 +1901,7 @@ mod tests {
let mut blob_builder = lance::blob::BlobArrayBuilder::new(2);
blob_builder.push_bytes(b"nested").unwrap();
blob_builder.push_null().unwrap();
let blob_field = crate::blob("image", true);
let blob_field = lance::blob::blob_field("image", true);
let metadata_fields = Fields::from(vec![blob_field.clone()]);
let metadata = StructArray::new(
metadata_fields.clone(),
@@ -1968,7 +1968,7 @@ mod tests {
blob_builder.push_bytes(b"a").unwrap();
blob_builder.push_bytes(b"bb").unwrap();
blob_builder.push_null().unwrap();
let item = Arc::new(crate::blob("item", true));
let item = Arc::new(lance::blob::blob_field("item", true));
let images = ListArray::new(
item.clone(),
OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
rust/lancedb/tests/blob_integration.rs
+269 -1
View File
@@ -10,12 +10,13 @@ use arrow_array::{
use arrow_schema::{DataType, Field, Fields, Schema};
use futures::TryStreamExt;
use lance::Dataset;
use lance::blob::blob_field as blob;
use lance::dataset::WriteParams;
use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
use lance_table::format::BasePath;
use lancedb::{
Connection, Error, Result, Table,
blob::{BlobRangeRequest, blob},
blob::BlobRangeRequest,
connect, connect_namespace,
database::listing::{
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
@@ -1610,6 +1611,88 @@ async fn external_uri_struct_round_trips_with_flag() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn external_ranges_round_trip_with_registered_base_and_opt_in() -> Result<()> {
let tmp = tempdir().unwrap();
let payload = b"0123456789abcdef";
let uri = write_payload_file_uri(tmp.path(), "payload.bin", payload);
let base_uri = url::Url::from_directory_path(tmp.path())
.unwrap()
.to_string();
let db = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await?;
for registered_base in [false, true] {
let field = lancedb::blob("image", true);
let DataType::Struct(children) = field.data_type().clone() else {
unreachable!("blob field is a struct");
};
let schema = Arc::new(Schema::new(vec![field]));
let name = if registered_base {
"registered"
} else {
"opt_in"
};
let table = db
.create_empty_table(name, schema.clone())
.write_options(WriteOptions {
lance_write_params: Some(WriteParams {
initial_bases: registered_base.then(|| {
vec![BasePath {
id: 1,
name: Some("external".to_string()),
path: base_uri.clone(),
is_dataset_root: false,
}]
}),
..Default::default()
}),
})
.execute()
.await?;
let values = StructArray::new(
children,
vec![
Arc::new(LargeBinaryArray::from(vec![
None,
None,
None,
Some(b"inline".as_slice()),
])),
Arc::new(StringArray::from(vec![
Some(uri.as_str()),
Some(uri.as_str()),
Some(uri.as_str()),
None,
])),
Arc::new(UInt64Array::from(vec![Some(0), Some(4), None, None])),
Arc::new(UInt64Array::from(vec![Some(6), Some(6), None, None])),
],
None,
);
table
.add(RecordBatch::try_new(schema, vec![Arc::new(values)])?)
.allow_external_blob_outside_bases(!registered_base)
.execute()
.await?;
let reopened = db.open_table(name).execute().await?;
let ids = collect_row_ids(&reopened).await?;
let bytes = reopened.fetch_blobs("image", &ids).await?;
assert_eq!(bytes.value(0), &payload[..6]);
assert_eq!(bytes.value(1), &payload[4..10]);
assert_eq!(bytes.value(2), payload);
assert_eq!(bytes.value(3), b"inline");
let files = reopened.fetch_blob_files("image", &ids).await?;
assert_eq!(files[1].as_ref().unwrap().size(), 6);
assert_eq!(
files[1].as_ref().unwrap().read_range(1..3).await?,
&payload[5..7]
);
}
Ok(())
}
#[tokio::test]
async fn external_uri_add_requires_opt_in() -> Result<()> {
let tmp = tempdir().unwrap();
@@ -1761,3 +1844,188 @@ async fn malformed_string_uri_is_rejected_at_write() -> Result<()> {
assert_eq!(table.count_rows(None).await?, 0);
Ok(())
}
fn layout_batch(wide: bool, id: i64, position: Option<u64>, size: Option<u64>) -> RecordBatch {
let field = if wide {
lancedb::blob("image", true)
} else {
blob("image", true)
};
let DataType::Struct(children) = field.data_type().clone() else {
unreachable!()
};
let mut columns: Vec<arrow_array::ArrayRef> = vec![
Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_slice())])),
new_null_array(&DataType::Utf8, 1),
];
if wide {
columns.push(Arc::new(UInt64Array::from(vec![position])));
columns.push(Arc::new(UInt64Array::from(vec![size])));
}
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
field,
])),
vec![
Arc::new(Int64Array::from(vec![id])),
Arc::new(StructArray::new(children, columns, None)),
],
)
.unwrap()
}
fn layout_reader(batches: Vec<RecordBatch>) -> Box<dyn arrow_array::RecordBatchReader + Send> {
let schema = batches[0].schema();
Box::new(arrow_array::RecordBatchIterator::new(
batches.into_iter().map(Ok),
schema,
))
}
#[tokio::test]
async fn blob_layout_append_and_merge_normalize_both_directions() -> Result<()> {
for wide in [false, true] {
for merge in [false, true] {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let initial = layout_batch(wide, 1, None, None);
let schema = initial.schema();
let table = db.create_table("layouts", initial).execute().await?;
let input = layout_batch(!wide, if merge { 1 } else { 2 }, None, None);
if merge {
let mut builder = table.merge_insert(&["id"]);
builder
.when_matched_update_all(None)
.when_not_matched_insert_all();
let result = builder
.execute(layout_reader(vec![
input,
layout_batch(!wide, 2, None, None),
]))
.await?;
assert_eq!(result.num_updated_rows, 1);
assert_eq!(result.num_inserted_rows, 1);
} else {
table.add(input).execute().await?;
}
assert_eq!(table.schema().await?, schema);
assert_eq!(table.count_rows(None).await?, 2);
let ids = collect_row_ids(&table).await?;
let values = table.fetch_blobs("image", &ids).await?;
assert_eq!(values.value(0), b"payload");
assert_eq!(values.value(1), b"payload");
}
}
Ok(())
}
#[tokio::test]
async fn blob_layout_rejects_non_null_ranges_atomically_in_later_batches() -> Result<()> {
for merge in [false, true] {
for (position, size) in [(Some(0), None), (None, Some(6)), (Some(0), Some(6))] {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let table = db
.create_table("layouts", layout_batch(false, 1, None, None))
.execute()
.await?;
let version = table.version().await?;
// A successful first batch must not commit when a later batch has a range.
let input = layout_reader(vec![
layout_batch(true, 2, None, None),
layout_batch(true, 3, position, size),
]);
let result = if merge {
let mut builder = table.merge_insert(&["id"]);
builder
.when_matched_update_all(None)
.when_not_matched_insert_all();
builder.execute(input).await.map(|_| ())
} else {
table.add(input).execute().await.map(|_| ())
};
let error = result.unwrap_err().to_string();
assert!(error.contains("cannot store byte ranges"), "{error}");
assert_eq!(table.version().await?, version);
assert_eq!(table.count_rows(None).await?, 1);
let ids = collect_row_ids(&table).await?;
assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload");
}
}
Ok(())
}
#[tokio::test]
async fn blob_layout_overwrite_replaces_schema_in_both_directions() -> Result<()> {
for wide in [false, true] {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let table = db
.create_table("layouts", layout_batch(wide, 1, None, None))
.execute()
.await?;
let input = layout_batch(!wide, 2, None, None);
let expected_schema = input.schema();
table
.add(input)
.mode(AddDataMode::Overwrite)
.execute()
.await?;
assert_eq!(table.schema().await?, expected_schema);
assert_eq!(table.count_rows(None).await?, 1);
let ids = collect_row_ids(&table).await?;
assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload");
}
Ok(())
}
#[tokio::test]
async fn blob_layout_memwal_merge_normalizes_before_writing() -> Result<()> {
use lancedb::{query::Select, table::MemWalWriteSpec};
for wide in [false, true] {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
let table = db
.create_table("layouts", layout_batch(wide, 1, None, None))
.execute()
.await?;
table.set_unenforced_primary_key(["id"]).await?;
table.enable_memwal(MemWalWriteSpec::unsharded()).await?;
let mut merge = table.merge_insert(&["id"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
merge
.clone()
.execute(layout_reader(vec![layout_batch(!wide, 2, None, None)]))
.await?;
let batches = table
.query()
.select(Select::columns(&["id"]))
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
if !wide {
let error = merge
.execute(layout_reader(vec![
layout_batch(true, 3, None, None),
layout_batch(true, 4, Some(0), Some(6)),
]))
.await
.unwrap_err();
assert!(error.to_string().contains("cannot store byte ranges"));
let batches = table
.query()
.select(Select::columns(&["id"]))
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
}
}
Ok(())
}
rust/lancedb/tests/duplicate_pairs.rs
+2
View File
@@ -410,6 +410,8 @@ async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Re
(0..6).map(|id| vec![id; 4096]),
)),
Arc::new(StringArray::from(vec![None::<&str>; 6])),
Arc::new(UInt64Array::new_null(6)),
Arc::new(UInt64Array::new_null(6)),
],
None,
);