mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-30 15:26:52 +00:00
feat: accept blob URI writes (#3954)
#3528 added blob declarations and binary coercion. String values were still rejected. They now coerce to the blob `uri` child. ```python table.add([{"id": 1, "image": "s3://bucket/media/cat.jpg"}]) payload = table.fetch_blobs("image", table.search().to_arrow()) ``` A URI under a registered base writes with no extra options. An unregistered URI fails. `allow_external_blob_outside_bases` is a local escape hatch that stores an absolute URI. It does not register a base. Remote `add` rejects that flag before making a request. String input still coerces and is sent as a `uri` struct. `add_bases` is a follow-up. `merge_insert` does not coerce string blob input. ### Testing - `cargo test -p lancedb --test blob_integration` - `cargo test -p lancedb blob_coerce` - `cargo test -p lancedb --features remote --lib add_rejects_external_blob_flag add_string_blob_becomes_uri_struct` - `cd python && uv run --extra tests pytest python/tests/test_blob.py -k uri -q` Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
@@ -283,6 +283,7 @@ class Table:
|
||||
mode: Literal["append", "overwrite"],
|
||||
progress: Optional[Any] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult: ...
|
||||
async def update(
|
||||
self, updates: Dict[str, str], where: Optional[str]
|
||||
|
||||
@@ -610,6 +610,7 @@ class RemoteTable(Table):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -642,6 +643,8 @@ class RemoteTable(Table):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Not supported on LanceDB Cloud. Setting this raises.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -658,6 +661,7 @@ class RemoteTable(Table):
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
|
||||
@@ -1269,6 +1269,7 @@ class Table(ABC):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -1320,6 +1321,10 @@ class Table(ABC):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Store blob URIs that sit outside registered blob bases. The row
|
||||
keeps a reference, so the object has to stay readable. Local
|
||||
tables only.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -3409,6 +3414,7 @@ class LanceTable(Table):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add data to the table.
|
||||
If vector columns are missing and the table
|
||||
@@ -3436,6 +3442,9 @@ class LanceTable(Table):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
Local tables only.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -3452,6 +3461,7 @@ class LanceTable(Table):
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -5365,6 +5375,7 @@ class AsyncTable:
|
||||
fill_value: Optional[float] = None,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
|
||||
|
||||
@@ -5395,6 +5406,9 @@ class AsyncTable:
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
Local tables only.
|
||||
|
||||
"""
|
||||
schema = await self.schema()
|
||||
@@ -5431,6 +5445,7 @@ class AsyncTable:
|
||||
mode or "append",
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Cast error" in str(e):
|
||||
|
||||
@@ -617,3 +617,71 @@ def test_fetch_blobs_nested_path_survives_sort_after_query():
|
||||
def _identifiable_payload(size: int) -> bytes:
|
||||
block = 256
|
||||
return b"".join(bytes([i % 256]) * block for i in range(size // block))
|
||||
|
||||
|
||||
def _external_uri_blob_array(uris):
|
||||
blob_type = lancedb.blob("image").type
|
||||
storage_type = blob_type.storage_type
|
||||
child_names = [field.name for field in storage_type]
|
||||
assert "uri" in child_names, "blob layout no longer has a uri child"
|
||||
children = [
|
||||
pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
|
||||
for field in storage_type
|
||||
]
|
||||
storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
|
||||
return pa.ExtensionArray.from_storage(blob_type, storage)
|
||||
|
||||
|
||||
def _external_uri_table_and_rows(name, uris):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
rows = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array(range(len(uris)), type=pa.int64()),
|
||||
_external_uri_blob_array(uris),
|
||||
],
|
||||
schema=schema,
|
||||
)
|
||||
return table, rows
|
||||
|
||||
|
||||
def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
|
||||
payload = b"external-uri-bytes"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
|
||||
table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
|
||||
table.add(rows, allow_external_blob_outside_bases=True)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert blobs[0].as_py() == payload
|
||||
|
||||
|
||||
def test_add_external_uri_without_flag_raises(tmp_path):
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(b"unreachable")
|
||||
|
||||
table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
|
||||
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
|
||||
table.add(rows)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_add_external_uri_string_round_trips_with_flag(tmp_path):
|
||||
payload = b"external-uri-bytes"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("external_string", schema=schema)
|
||||
table.add(
|
||||
[{"id": 1, "image": blob_path.as_uri()}],
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert blobs[0].as_py() == payload
|
||||
|
||||
+6
-2
@@ -780,15 +780,19 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
|
||||
pub fn add<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
data: PyScannable,
|
||||
mode: String,
|
||||
progress: Option<Py<PyAny>>,
|
||||
write_parallelism: Option<usize>,
|
||||
allow_external_blob_outside_bases: bool,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let mut op = self_.inner_ref()?.add(data);
|
||||
let mut op = self_
|
||||
.inner_ref()?
|
||||
.add(data)
|
||||
.allow_external_blob_outside_bases(allow_external_blob_outside_bases);
|
||||
if mode == "append" {
|
||||
op = op.mode(AddDataMode::Append);
|
||||
} else if mode == "overwrite" {
|
||||
|
||||
@@ -2161,6 +2161,15 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
async fn add(&self, mut add: AddDataBuilder) -> Result<AddResult> {
|
||||
self.check_mutable().await?;
|
||||
|
||||
if add.allow_external_blob_outside_bases {
|
||||
return Err(Error::NotSupported {
|
||||
message: "allow_external_blob_outside_bases is only supported on local tables"
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
// String blob values still coerce to the uri child in into_plan.
|
||||
// Remote and local share that input shape.
|
||||
|
||||
let table_schema = self.schema().await?;
|
||||
crate::table::computed_columns::ensure_supported_function_metadata(table_schema.as_ref())?;
|
||||
let table_def = TableDefinition::try_from_rich_schema(table_schema.clone())?;
|
||||
@@ -3269,7 +3278,10 @@ mod tests {
|
||||
use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type};
|
||||
use arrow_array::Array;
|
||||
use arrow_array::builder::LargeBinaryBuilder;
|
||||
use arrow_array::{BinaryArray, Int32Array, RecordBatch, RecordBatchIterator, record_batch};
|
||||
use arrow_array::{
|
||||
BinaryArray, Int32Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray,
|
||||
StructArray, record_batch,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt, TryFutureExt, future::BoxFuture};
|
||||
@@ -3621,6 +3633,88 @@ mod tests {
|
||||
assert_eq!(&body, &expected_body);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn add_rejects_external_blob_flag_before_any_request() {
|
||||
let table = Table::new_with_handler::<String>("my_table", |request| {
|
||||
panic!("Unexpected request: {}", request.url().path())
|
||||
});
|
||||
let data = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(vec![1]))],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let err = table
|
||||
.add(data)
|
||||
.allow_external_blob_outside_bases(true)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(err, Error::NotSupported { .. }), "got {err:?}");
|
||||
assert!(err.to_string().contains("local tables"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn add_string_blob_becomes_uri_struct_without_the_local_flag() {
|
||||
let table_schema = Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
crate::blob("image", true),
|
||||
]);
|
||||
let describe_body = describe_response(&table_schema);
|
||||
let input = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("image", DataType::Utf8, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![1])),
|
||||
Arc::new(StringArray::from(vec![Some("s3://bucket/key")])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let (sender, receiver) = std::sync::mpsc::channel();
|
||||
let table =
|
||||
Table::new_with_handler("my_table", 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/" => {
|
||||
let mut body_out = reqwest::Body::from(Vec::new());
|
||||
std::mem::swap(request.body_mut().as_mut().unwrap(), &mut body_out);
|
||||
sender.send(body_out).unwrap();
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 2}"#.to_string())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {path}"),
|
||||
});
|
||||
|
||||
table.add(input).execute().await.unwrap();
|
||||
|
||||
let body = collect_body(receiver.recv().unwrap()).await;
|
||||
let mut reader =
|
||||
arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None).unwrap();
|
||||
let batch = reader.next().unwrap().unwrap();
|
||||
let image = batch
|
||||
.column_by_name("image")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.expect("remote add should send the coerced blob struct");
|
||||
let uri: &StringArray = image
|
||||
.column_by_name("uri")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref()
|
||||
.unwrap();
|
||||
assert_eq!(uri.value(0), "s3://bucket/key");
|
||||
assert!(image.column_by_name("data").unwrap().is_null(0));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(true)]
|
||||
#[case(false)]
|
||||
|
||||
@@ -3208,7 +3208,7 @@ impl BaseTable for NativeTable {
|
||||
|
||||
let output = add.into_plan(&table_schema, &table_def)?;
|
||||
|
||||
let lance_params = output
|
||||
let mut lance_params = output
|
||||
.write_options
|
||||
.lance_write_params
|
||||
.unwrap_or(WriteParams {
|
||||
@@ -3218,6 +3218,9 @@ impl BaseTable for NativeTable {
|
||||
},
|
||||
..Default::default()
|
||||
});
|
||||
if output.allow_external_blob_outside_bases {
|
||||
lance_params.allow_external_blob_outside_bases = true;
|
||||
}
|
||||
|
||||
// Repartition for write parallelism if beneficial.
|
||||
let plan = if num_partitions > 1 {
|
||||
|
||||
@@ -60,6 +60,7 @@ pub struct AddDataBuilder {
|
||||
pub(crate) embedding_registry: Option<Arc<dyn EmbeddingRegistry>>,
|
||||
pub(crate) progress_callback: Option<ProgressCallback>,
|
||||
pub(crate) write_parallelism: Option<usize>,
|
||||
pub(crate) allow_external_blob_outside_bases: bool,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AddDataBuilder {
|
||||
@@ -87,6 +88,7 @@ impl AddDataBuilder {
|
||||
embedding_registry,
|
||||
progress_callback: None,
|
||||
write_parallelism: None,
|
||||
allow_external_blob_outside_bases: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,6 +143,16 @@ impl AddDataBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Store blob URIs that sit outside registered blob bases.
|
||||
///
|
||||
/// The row keeps a reference, so the object has to stay readable.
|
||||
/// [`crate::table::Table::fetch_blobs`] reads from that location.
|
||||
/// Defaults to `false`. Local tables only.
|
||||
pub fn allow_external_blob_outside_bases(mut self, allow: bool) -> Self {
|
||||
self.allow_external_blob_outside_bases = allow;
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn execute(self) -> Result<AddResult> {
|
||||
if self.write_parallelism.map(|p| p == 0).unwrap_or(false) {
|
||||
return Err(Error::InvalidInput {
|
||||
@@ -199,6 +211,7 @@ impl AddDataBuilder {
|
||||
write_options: self.write_options,
|
||||
mode: self.mode,
|
||||
tracker,
|
||||
allow_external_blob_outside_bases: self.allow_external_blob_outside_bases,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -212,6 +225,7 @@ pub struct PreprocessingOutput {
|
||||
pub write_options: WriteOptions,
|
||||
pub mode: AddDataMode,
|
||||
pub tracker: Option<Arc<WriteProgressTracker>>,
|
||||
pub allow_external_blob_outside_bases: bool,
|
||||
}
|
||||
|
||||
/// Check that the input schema is valid for insert.
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::{DataType, Field, FieldRef};
|
||||
use arrow_schema::{DataType, Field, FieldRef, Fields};
|
||||
use datafusion::functions::core::{get_field, named_struct};
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::config::ConfigOptions;
|
||||
@@ -35,8 +35,9 @@ pub(super) fn coerce_blob_expr(
|
||||
});
|
||||
};
|
||||
|
||||
let input_struct_children = match input_field.data_type() {
|
||||
DataType::Binary | DataType::LargeBinary | DataType::BinaryView => None,
|
||||
let input_shape = match input_field.data_type() {
|
||||
DataType::Binary | DataType::LargeBinary | DataType::BinaryView => BlobInputShape::Bytes,
|
||||
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => BlobInputShape::String,
|
||||
DataType::Struct(children) => {
|
||||
if !children
|
||||
.iter()
|
||||
@@ -49,13 +50,15 @@ pub(super) fn coerce_blob_expr(
|
||||
),
|
||||
});
|
||||
}
|
||||
Some(children)
|
||||
BlobInputShape::Struct(children)
|
||||
}
|
||||
other => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"cannot coerce column '{}' with type {} into a blob v2 struct. \
|
||||
expected Binary, LargeBinary, BinaryView, or a Struct with a 'data' or 'uri' child",
|
||||
expected binary bytes (Binary, LargeBinary, BinaryView), \
|
||||
strings (Utf8, LargeUtf8, Utf8View), \
|
||||
or a Struct with a 'data' or 'uri' child",
|
||||
table_field.name(),
|
||||
other,
|
||||
),
|
||||
@@ -69,9 +72,8 @@ pub(super) fn coerce_blob_expr(
|
||||
declared.name().as_str(),
|
||||
))));
|
||||
|
||||
let value: Arc<dyn PhysicalExpr> = match input_struct_children {
|
||||
// Raw binary lands in `data` and everything else is a typed null.
|
||||
None => {
|
||||
let value: Arc<dyn PhysicalExpr> = match &input_shape {
|
||||
BlobInputShape::Bytes => {
|
||||
if declared.name() == "data" {
|
||||
Arc::new(CastExpr::new(
|
||||
input_expr.clone(),
|
||||
@@ -82,30 +84,43 @@ pub(super) fn coerce_blob_expr(
|
||||
typed_null(declared.data_type())?
|
||||
}
|
||||
}
|
||||
Some(children) => match children.iter().find(|c| c.name() == declared.name()) {
|
||||
Some(child) => {
|
||||
let field_expr: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
|
||||
&format!("get_field({})", declared.name()),
|
||||
get_field(),
|
||||
vec![
|
||||
input_expr.clone(),
|
||||
Arc::new(Literal::new(ScalarValue::from(declared.name().as_str()))),
|
||||
],
|
||||
Arc::new(child.as_ref().clone()),
|
||||
config.clone(),
|
||||
));
|
||||
if child.data_type() == declared.data_type() {
|
||||
field_expr
|
||||
} else {
|
||||
Arc::new(CastExpr::new(
|
||||
field_expr,
|
||||
declared.data_type().clone(),
|
||||
None,
|
||||
))
|
||||
}
|
||||
BlobInputShape::String => {
|
||||
if declared.name() == "uri" {
|
||||
Arc::new(CastExpr::new(
|
||||
input_expr.clone(),
|
||||
declared.data_type().clone(),
|
||||
None,
|
||||
))
|
||||
} else {
|
||||
typed_null(declared.data_type())?
|
||||
}
|
||||
None => typed_null(declared.data_type())?,
|
||||
},
|
||||
}
|
||||
BlobInputShape::Struct(children) => {
|
||||
match children.iter().find(|c| c.name() == declared.name()) {
|
||||
Some(child) => {
|
||||
let field_expr: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
|
||||
&format!("get_field({})", declared.name()),
|
||||
get_field(),
|
||||
vec![
|
||||
input_expr.clone(),
|
||||
Arc::new(Literal::new(ScalarValue::from(declared.name().as_str()))),
|
||||
],
|
||||
Arc::new(child.as_ref().clone()),
|
||||
config.clone(),
|
||||
));
|
||||
if child.data_type() == declared.data_type() {
|
||||
field_expr
|
||||
} else {
|
||||
Arc::new(CastExpr::new(
|
||||
field_expr,
|
||||
declared.data_type().clone(),
|
||||
None,
|
||||
))
|
||||
}
|
||||
}
|
||||
None => typed_null(declared.data_type())?,
|
||||
}
|
||||
}
|
||||
};
|
||||
ns_args.push(value);
|
||||
}
|
||||
@@ -120,6 +135,12 @@ pub(super) fn coerce_blob_expr(
|
||||
Ok((expr, table_field.clone()))
|
||||
}
|
||||
|
||||
enum BlobInputShape<'a> {
|
||||
Bytes,
|
||||
String,
|
||||
Struct(&'a Fields),
|
||||
}
|
||||
|
||||
fn typed_null(data_type: &DataType) -> Result<Arc<dyn PhysicalExpr>> {
|
||||
let scalar = ScalarValue::try_from(data_type).map_err(|e| Error::InvalidInput {
|
||||
message: format!("cannot build null literal for blob child type {data_type}: {e}"),
|
||||
@@ -134,7 +155,7 @@ mod tests {
|
||||
use crate::blob::blob;
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, LargeBinaryArray,
|
||||
RecordBatch, StringArray, StructArray, UInt8Array, UInt64Array,
|
||||
RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array,
|
||||
};
|
||||
use arrow_schema::Schema;
|
||||
use datafusion::prelude::SessionContext;
|
||||
@@ -436,14 +457,78 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn unsupported_input_type_is_rejected_with_column_name() {
|
||||
let batch = batch_with_image(
|
||||
Field::new("image", DataType::Utf8, true),
|
||||
Arc::new(StringArray::from(vec!["not bytes"])),
|
||||
Field::new("image", DataType::Int64, true),
|
||||
Arc::new(Int64Array::from(vec![42])),
|
||||
);
|
||||
let err = coerce_err(batch, &blob_table_schema()).await;
|
||||
assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}");
|
||||
assert!(err.to_string().contains("image"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn utf8_string_coerces_to_uri_child() {
|
||||
let batch = batch_with_image(
|
||||
Field::new("image", DataType::Utf8, true),
|
||||
Arc::new(StringArray::from(vec![Some("s3://bucket/key"), None])),
|
||||
);
|
||||
let coerced = coerce(batch, &blob_table_schema()).await;
|
||||
let image = image_struct(&coerced);
|
||||
let uri: &StringArray = image
|
||||
.column_by_name("uri")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref()
|
||||
.unwrap();
|
||||
assert_eq!(uri.value(0), "s3://bucket/key");
|
||||
assert!(image.column_by_name("data").unwrap().is_null(0));
|
||||
assert!(uri.is_null(1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn large_utf8_string_coerces_into_four_child_blob_layout() {
|
||||
use arrow_array::LargeStringArray;
|
||||
|
||||
let table_schema = Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
wide_blob_field("image"),
|
||||
]);
|
||||
let batch = batch_with_image(
|
||||
Field::new("image", DataType::LargeUtf8, true),
|
||||
Arc::new(LargeStringArray::from(vec!["file:///tmp/blob.bin"])),
|
||||
);
|
||||
let coerced = coerce(batch, &table_schema).await;
|
||||
let image = image_struct(&coerced);
|
||||
assert_eq!(image.num_columns(), 4);
|
||||
let uri: &StringArray = image
|
||||
.column_by_name("uri")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref()
|
||||
.unwrap();
|
||||
assert_eq!(uri.value(0), "file:///tmp/blob.bin");
|
||||
assert!(image.column_by_name("data").unwrap().is_null(0));
|
||||
assert!(image.column_by_name("position").unwrap().is_null(0));
|
||||
assert!(image.column_by_name("size").unwrap().is_null(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn utf8_view_string_coerces_to_uri_child() {
|
||||
let batch = batch_with_image(
|
||||
Field::new("image", DataType::Utf8View, true),
|
||||
Arc::new(StringViewArray::from(vec![Some("s3://bucket/view-key")])),
|
||||
);
|
||||
let coerced = coerce(batch, &blob_table_schema()).await;
|
||||
let image = image_struct(&coerced);
|
||||
let uri: &StringArray = image
|
||||
.column_by_name("uri")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref()
|
||||
.unwrap();
|
||||
assert_eq!(uri.value(0), "s3://bucket/view-key");
|
||||
assert!(image.column_by_name("data").unwrap().is_null(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn blob_metadata_survives_cast_of_sibling_column() {
|
||||
let batch = RecordBatch::try_new(
|
||||
|
||||
@@ -5,12 +5,14 @@ use std::sync::Arc;
|
||||
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, BinaryArray, Int64Array, LargeBinaryArray, RecordBatch, StringArray,
|
||||
StructArray, UInt64Array,
|
||||
StructArray, UInt64Array, new_null_array,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance::dataset::WriteParams;
|
||||
use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
|
||||
use lance_table::format::BasePath;
|
||||
use lancedb::{
|
||||
Connection, Error, Result, Table,
|
||||
blob::{BlobRangeRequest, blob},
|
||||
@@ -19,7 +21,7 @@ use lancedb::{
|
||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
},
|
||||
query::{ExecutableQuery, QueryBase},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats, WriteOptions},
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
@@ -261,11 +263,11 @@ async fn add_rejects_uncoercible_blob_input() -> Result<()> {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("image", DataType::Utf8, true),
|
||||
Field::new("image", DataType::Int64, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![1])),
|
||||
Arc::new(StringArray::from(vec!["not bytes"])),
|
||||
Arc::new(Int64Array::from(vec![42])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
@@ -1332,3 +1334,223 @@ async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn uri_struct_batch(id: i64, uri: &str) -> RecordBatch {
|
||||
let image_field = blob("image", true);
|
||||
let DataType::Struct(child_fields) = image_field.data_type().clone() else {
|
||||
unreachable!("blob field is a struct");
|
||||
};
|
||||
let children: Vec<ArrayRef> = child_fields
|
||||
.iter()
|
||||
.map(|field| match field.name().as_str() {
|
||||
"uri" => Arc::new(StringArray::from(vec![Some(uri)])) as ArrayRef,
|
||||
_ => new_null_array(field.data_type(), 1),
|
||||
})
|
||||
.collect();
|
||||
let image = StructArray::new(child_fields, children, None);
|
||||
RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
image_field,
|
||||
])),
|
||||
vec![Arc::new(Int64Array::from(vec![id])), Arc::new(image)],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn uri_string_batch(id: i64, uri: &str) -> RecordBatch {
|
||||
RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("image", DataType::Utf8, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![id])),
|
||||
Arc::new(StringArray::from(vec![Some(uri)])),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn write_payload_file_uri(dir: &std::path::Path, name: &str, payload: &[u8]) -> String {
|
||||
let path = dir.join(name);
|
||||
std::fs::write(&path, payload).unwrap();
|
||||
url::Url::from_file_path(&path).unwrap().to_string()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn external_uri_struct_round_trips_with_flag() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let payload: &[u8] = b"external-struct-payload";
|
||||
let uri = write_payload_file_uri(tmp.path(), "payload.bin", payload);
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
table
|
||||
.add(uri_struct_batch(1, &uri))
|
||||
.allow_external_blob_outside_bases(true)
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
let bytes = table.fetch_blobs("image", &ids).await?;
|
||||
assert_eq!(bytes.value(0), payload);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn external_uri_add_requires_opt_in() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let uri = write_payload_file_uri(tmp.path(), "payload.bin", b"unreachable");
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let err = table
|
||||
.add(uri_struct_batch(1, &uri))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("allow_external_blob_outside_bases"),
|
||||
"got: {err}"
|
||||
);
|
||||
assert_eq!(table.count_rows(None).await?, 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn string_uri_input_round_trips_as_external_reference() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let payload: &[u8] = b"external-string-payload";
|
||||
let uri = write_payload_file_uri(tmp.path(), "payload.bin", payload);
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
table
|
||||
.add(uri_string_batch(1, &uri))
|
||||
.allow_external_blob_outside_bases(true)
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
let bytes = table.fetch_blobs("image", &ids).await?;
|
||||
assert_eq!(bytes.value(0), payload);
|
||||
|
||||
let files = table.fetch_blob_files("image", &ids).await?;
|
||||
let file = files[0].as_ref().expect("missing blob file");
|
||||
assert_eq!(file.uri(), Some(uri.as_str()));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn string_uri_inside_registered_base_does_not_need_the_flag() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db_path = tmp.path().join("db");
|
||||
let external_base = tmp.path().join("external_base");
|
||||
let object_dir = external_base.join("objects");
|
||||
std::fs::create_dir_all(&object_dir).unwrap();
|
||||
let payload: &[u8] = b"mapped-in-base";
|
||||
let object_path = object_dir.join("mapped.bin");
|
||||
std::fs::write(&object_path, payload).unwrap();
|
||||
let object_uri = url::Url::from_file_path(&object_path).unwrap().to_string();
|
||||
let base_uri = url::Url::from_file_path(&external_base)
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
let db = connect(db_path.to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.write_options(WriteOptions {
|
||||
lance_write_params: Some(WriteParams {
|
||||
initial_bases: Some(vec![BasePath {
|
||||
id: 1,
|
||||
name: Some("external".to_string()),
|
||||
path: base_uri,
|
||||
is_dataset_root: false,
|
||||
}]),
|
||||
..Default::default()
|
||||
}),
|
||||
})
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
table
|
||||
.add(uri_string_batch(1, &object_uri))
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
let bytes = table.fetch_blobs("image", &ids).await?;
|
||||
assert_eq!(bytes.value(0), payload);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn external_uri_rows_mix_with_inline_rows() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let external_payload: &[u8] = b"external-bytes";
|
||||
let uri = write_payload_file_uri(tmp.path(), "payload.bin", external_payload);
|
||||
let table =
|
||||
create_inline_blob_table(&db, "t", &[1], &[Some(b"inline-bytes".as_slice())]).await?;
|
||||
|
||||
table
|
||||
.add(uri_string_batch(2, &uri))
|
||||
.allow_external_blob_outside_bases(true)
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let pairs = collect_id_rowid(&table).await?;
|
||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, r)| *r).collect();
|
||||
let bytes = table.fetch_blobs("image", &row_ids).await?;
|
||||
for (i, (id, _)) in pairs.iter().enumerate() {
|
||||
match id {
|
||||
1 => assert_eq!(bytes.value(i), b"inline-bytes"),
|
||||
2 => assert_eq!(bytes.value(i), external_payload),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_string_uri_is_rejected_at_write() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
let err = table
|
||||
.add(uri_string_batch(1, "not a uri"))
|
||||
.allow_external_blob_outside_bases(true)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("not a uri"), "got: {err}");
|
||||
assert_eq!(table.count_rows(None).await?, 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user