fix!: restore blob SQL inserts with explicit adapter schemas (#4687)

This commit is contained in:
lancedb-gatefixer[bot]
2026-10-27 09:31:23 -07:00
committed by GitHub
parent 4b373b2070
commit fe079eb421
5 changed files with 1083 additions and 8 deletions
+15
View File
@@ -0,0 +1,15 @@
# Changelog
## Unreleased
- `BaseTableAdapter` now defaults to `BlobSchema::Read`, exposing the Blob v2
descriptors returned by scans. SQL INSERT into Blob v2 tables requires opting
into `BlobSchema::Write` with `with_blob_schema`. Write-mode queries support
scalar projections that exclude blobs, scalar filters, and `COUNT(*)`.
Scans must exactly match the projected provider schema, including metadata
and nullability. Blob v2 descriptor scans fail during planning; legacy v1
blobs expose `LargeBinary` write inputs and reject descriptor scans with a
schema mismatch. A scalar sibling inside a blob-containing struct also needs
a separate read adapter because DataFusion projects the whole parent struct.
Tables without blob columns behave the same in both modes. Generated `_rowid`,
`_rowaddr`, and `_score` columns are exposed only in read mode.
+221
View File
@@ -5712,6 +5712,227 @@ mod tests {
assert!(fields.find("data").is_none());
}
#[tokio::test]
async fn test_adapter_blob_read_schema_matches_remote_scan() {
use datafusion::prelude::SessionContext;
use datafusion_catalog::TableProvider;
use lance_core::datatypes::Projection;
let write_schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
crate::blob("image", true),
]);
let read_schema = Arc::new(
Projection::full(Arc::new(LanceSchema::try_from(&write_schema).unwrap()))
.to_arrow_schema(),
);
let descriptor = arrow_array::new_null_array(read_schema.field(1).data_type(), 2);
let response = RecordBatch::try_new(
read_schema,
vec![Arc::new(Int64Array::from(vec![1, 2])), descriptor],
)
.unwrap();
let table =
Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => {
schema_describe_response(&write_schema).map(String::into_bytes)
}
"/v1/table/my_table/query/" => {
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
.unwrap();
let columns = body["columns"]
.as_array()
.unwrap()
.iter()
.map(|column| {
response
.schema()
.index_of(column.as_str().unwrap())
.unwrap()
})
.collect::<Vec<_>>();
http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&response.project(&columns).unwrap()))
.unwrap()
}
path => panic!("unexpected request: {path}"),
});
let adapter =
crate::table::datafusion::BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
let ctx = SessionContext::new();
let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap();
assert_eq!(plan.schema(), adapter.schema());
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
assert!(
batches
.iter()
.all(|batch| batch.schema() == adapter.schema())
);
let adapter = adapter.with_blob_schema(crate::table::datafusion::BlobSchema::Write);
let plan = adapter
.scan(&ctx.state(), Some(&vec![0]), &[], None)
.await
.unwrap();
let expected = Arc::new(adapter.schema().project(&[0]).unwrap());
assert_eq!(plan.schema(), expected);
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap();
assert!(batches.iter().all(|batch| batch.schema() == expected));
}
#[tokio::test]
async fn test_adapter_write_scan_checks_metadata_and_nullability() {
use crate::table::datafusion::{BaseTableAdapter, BlobSchema};
use datafusion::prelude::SessionContext;
use datafusion_catalog::TableProvider;
let declared = Field::new("id", DataType::Int64, false)
.with_metadata(HashMap::from([("meaning".into(), "identifier".into())]));
for returned in [
declared.clone(),
declared.clone().with_nullable(true),
declared.clone().with_metadata(HashMap::new()),
] {
let matches = returned == declared;
let write_schema = Schema::new(vec![declared.clone()]);
let response = RecordBatch::try_new(
Arc::new(Schema::new(vec![returned])),
vec![Arc::new(Int64Array::from(vec![1]))],
)
.unwrap();
let table =
Table::new_with_handler("my_table", move |request| match request.url().path() {
"/v1/table/my_table/describe/" => {
schema_describe_response(&write_schema).map(String::into_bytes)
}
"/v1/table/my_table/query/" => http::Response::builder()
.status(200)
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
.body(write_ipc_file(&response))
.unwrap(),
path => panic!("unexpected request: {path}"),
});
let adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_blob_schema(BlobSchema::Write);
let ctx = SessionContext::new();
let result = adapter.scan(&ctx.state(), None, &[], None).await;
if matches {
let plan = result.unwrap();
assert_eq!(plan.schema(), adapter.schema());
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap();
assert_eq!(batches[0].schema(), adapter.schema());
} else {
let error = result.unwrap_err();
assert!(
error.to_string().contains("Scan schema mismatch"),
"{error}"
);
}
}
}
#[tokio::test]
async fn test_adapter_blob_write_sql_remote_payload_round_trip() {
use crate::table::datafusion::{BaseTableAdapter, BlobSchema};
use datafusion::prelude::SessionContext;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
crate::blob("image", true),
]));
let described_schema = schema.clone();
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/" => schema_describe_response(&described_schema),
"/v1/table/my_table/insert/" => {
let body = request.body_mut().take().unwrap();
sender.send(body).unwrap();
http::Response::builder()
.status(200)
.body(r#"{"version": 2}"#.to_string())
.unwrap()
}
path => panic!("unexpected request: {path}"),
});
let adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_blob_schema(BlobSchema::Write);
let ctx = SessionContext::new();
ctx.register_table("docs", Arc::new(adapter)).unwrap();
// Persist the exact Arrow stream sent to the mocked remote endpoint to
// a real native sink, then check storage bytes independently of schemas.
let tmp = tempfile::tempdir().unwrap();
let db = crate::connect(tmp.path().to_str().unwrap())
.execute()
.await
.unwrap();
let sink = db
.create_empty_table("sink", schema)
.execute()
.await
.unwrap();
for query in [
"INSERT INTO docs SELECT 1, named_struct('data', X'6869', 'uri', NULL)",
"INSERT INTO docs VALUES (2, NULL)",
"INSERT INTO docs (id) VALUES (3)",
"INSERT INTO docs SELECT 4, named_struct('data', X'', 'uri', NULL)",
] {
ctx.sql(query).await.unwrap().collect().await.unwrap();
let body = collect_body(receiver.recv().unwrap()).await;
let reader = StreamReader::try_new(Cursor::new(body), None).unwrap();
for batch in reader {
sink.add(batch.unwrap()).execute().await.unwrap();
}
}
let batches = sink
.query()
.with_row_id()
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let row_ids = batches
.iter()
.flat_map(|batch| {
batch
.column_by_name("_rowid")
.unwrap()
.as_primitive::<arrow_array::types::UInt64Type>()
.values()
.iter()
.copied()
})
.collect::<Vec<_>>();
let actual = sink.fetch_blobs("image", &row_ids).await.unwrap();
assert_eq!(
actual,
LargeBinaryArray::from(vec![
Some(b"hi".as_slice()),
None,
None,
Some(b"".as_slice())
])
);
}
#[rstest]
#[case(semver::Version::new(0, 1, 0))]
#[case(semver::Version::new(0, 5, 0))]
+210 -7
View File
@@ -166,10 +166,40 @@ impl ExecutionPlan for MetadataEraserExec {
}
}
/// The blob schema exposed by a [`BaseTableAdapter`].
///
/// DataFusion uses one provider schema for both scans and INSERT coercion.
/// Blob descriptors and blob write inputs have different types, so applications
/// that need both can register separate read and write adapters for one table.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum BlobSchema {
/// Expose Blob v2 descriptors (`kind`, `position`, `size`, `blob_id`, `blob_uri`)
/// and legacy v1 descriptors (`position`, `size`). Scans support all blob fields;
/// SQL INSERT into Blob v2 tables is unsupported.
#[default]
Read,
/// Expose Blob v2 write inputs (`data`, `uri`, `position`, `size`) and legacy
/// v1 `LargeBinary` inputs for SQL INSERT. Scans must match the projected
/// provider schema exactly; use [`Self::Read`] to scan blob descriptors or
/// scalar siblings whose projected parent contains a blob.
Write,
}
/// A DataFusion provider for a LanceDB table.
///
/// Defaults to [`BlobSchema::Read`], matching the descriptor schema produced by
/// scans without fetching blob bytes. Use [`Self::with_blob_schema`] to enable
/// SQL INSERT into blob columns. Legacy v1 blobs expose descriptors in read mode
/// and `LargeBinary` inputs in write mode; write-mode descriptor scans fail with
/// a schema mismatch. Tables without blob columns behave the same in both modes.
#[derive(Debug)]
pub struct BaseTableAdapter {
table: Arc<dyn BaseTable>,
// Only the read schema includes opted-in system/FTS columns.
schema: Arc<ArrowSchema>,
write_schema: Arc<ArrowSchema>,
blob_schema: BlobSchema,
fts_query: Option<FullTextSearchQuery>,
use_rte: Option<bool>,
}
@@ -183,15 +213,59 @@ impl BaseTableAdapter {
let schema = Projection::full(Arc::new(LanceSchema::try_from(write_schema.as_ref())?))
.to_arrow_schema();
let schema = blob_nulls::read_schema(&schema).with_metadata(HashMap::default());
let write_schema = Arc::new(
write_schema
.as_ref()
.clone()
.with_metadata(HashMap::default()),
);
Ok(Self {
table,
schema: Arc::new(schema),
write_schema,
blob_schema: BlobSchema::Read,
fts_query: None,
use_rte: None,
})
}
/// Choose the schema used by DataFusion for blob columns.
///
/// [`BlobSchema::Write`] enables INSERT with blob write structs, NULLs, and
/// omitted nullable blobs. Scans requiring blobs fail with
/// "Blob v2 scans require BlobSchema::Read" before execution. Legacy v1 blobs
/// expose `LargeBinary` write inputs, so their descriptor scans fail with a
/// schema mismatch. Write-mode scans are accepted only when the actual schema
/// exactly matches the projected provider schema, including metadata and
/// nullability. No blob payloads are fetched to emulate reads.
///
/// Scalar projections that exclude blob fields, `COUNT(*)`, and scalar filters
/// remain supported. DataFusion projects whole top-level fields, so selecting
/// `extra.note` also requests its parent `extra`. If that struct contains a blob,
/// even this scalar projection requires [`BlobSchema::Read`], as do computed
/// scalar and window queries that project the same parent.
///
/// Register a separate read adapter for blob descriptor queries and scalar
/// queries rejected in write mode. A self-insert whose SELECT accesses blobs
/// also fails before writing any rows.
/// Other adapter builders preserve the selected mode. Generated `_rowid`,
/// `_rowaddr`, and `_score` columns are exposed only in read mode.
///
/// ```
/// use lancedb::table::datafusion::{BaseTableAdapter, BlobSchema};
///
/// # async fn adapter(table: &lancedb::Table) -> lancedb::Result<BaseTableAdapter> {
/// Ok(BaseTableAdapter::try_new(table.base_table().clone())
/// .await?
/// .with_blob_schema(BlobSchema::Write))
/// # }
/// ```
pub fn with_blob_schema(mut self, blob_schema: BlobSchema) -> Self {
self.blob_schema = blob_schema;
self
}
/// Expose the requested Lance system columns to DataFusion.
///
/// Supports [`ROW_ID`] (`_rowid`) and [`ROW_ADDR`] (`_rowaddr`), both nullable
@@ -199,7 +273,8 @@ impl BaseTableAdapter {
/// already exposed by this adapter. Other names return [`Error::InvalidInput`].
///
/// System columns are hidden by default so existing `SELECT *` queries keep
/// their data-only schema. After opting in, `SELECT *` includes these columns.
/// their data-only schema. After opting in, `SELECT *` includes these columns
/// in [`BlobSchema::Read`] mode. The write schema remains data-only for INSERT.
///
/// ```
/// use lance_core::{ROW_ADDR, ROW_ID};
@@ -239,6 +314,8 @@ impl BaseTableAdapter {
}
/// Create a new adapter with an FTS query applied.
///
/// The generated `_score` column is exposed only in [`BlobSchema::Read`] mode.
pub fn with_fts_query(&self, fts_query: FullTextSearchQuery) -> Self {
// Add _score column to the schema
let score_field = Field::new("_score", DataType::Float32, true);
@@ -249,16 +326,28 @@ impl BaseTableAdapter {
Self {
table: self.table.clone(),
schema,
write_schema: self.write_schema.clone(),
blob_schema: self.blob_schema,
fts_query: Some(fts_query),
use_rte: self.use_rte,
}
}
}
fn blob_scan_error() -> DataFusionError {
DataFusionError::Plan(
"Blob v2 scans require BlobSchema::Read; use a separate read adapter for blob fields"
.into(),
)
}
#[async_trait]
impl TableProvider for BaseTableAdapter {
fn schema(&self) -> Arc<ArrowSchema> {
self.schema.clone()
match self.blob_schema {
BlobSchema::Read => self.schema.clone(),
BlobSchema::Write => self.write_schema.clone(),
}
}
fn table_type(&self) -> TableType {
@@ -272,8 +361,18 @@ impl TableProvider for BaseTableAdapter {
filters: &[Expr],
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
// For FTS queries, disable auto-projection of _score to match DataFusion expectations
let disable_scoring = self.fts_query.is_some() && projection.is_some();
// Check directly supplied filters too, not just optimizer pushdown calls.
if self.blob_schema == BlobSchema::Write {
for filter in filters {
if blob_nulls::references_descriptor(filter, &self.schema)? {
return Err(blob_scan_error());
}
}
}
// The writer never advertises _score. For readers, an explicit projection
// controls whether it is included, rather than FTS auto-projection.
let disable_scoring = self.fts_query.is_some()
&& (projection.is_some() || self.blob_schema == BlobSchema::Write);
let mut query = QueryRequest {
full_text_search: self.fts_query.clone(),
@@ -282,16 +381,17 @@ impl TableProvider for BaseTableAdapter {
..Default::default()
};
let provider_schema = self.schema();
if let Some(projection) = projection {
let field_names = projection
.iter()
.map(|i| self.schema.field(*i).name().clone())
.map(|i| provider_schema.field(*i).name().clone())
.collect();
query.select = Select::Columns(field_names);
} else {
// Select::All excludes system columns, which may be part of our schema.
query.select = Select::Columns(
self.schema
provider_schema
.fields()
.iter()
.map(|field| field.name().clone())
@@ -323,7 +423,32 @@ impl TableProvider for BaseTableAdapter {
.map_err(|err| DataFusionError::External(err.into()))
.await?;
let plan = blob_nulls::mask_plan(plan)?;
Ok(Arc::new(MetadataEraserExec::new(plan)))
let plan: Arc<dyn ExecutionPlan> = Arc::new(MetadataEraserExec::new(plan));
if self.blob_schema == BlobSchema::Write {
let schema = match projection {
Some(projection) => Arc::new(provider_schema.project(projection)?),
None => provider_schema,
};
// Compare the full schemas, including field metadata and nullability.
// A scan must satisfy its advertised contract before it is returned
// to DataFusion; later projections cannot repair a mismatched scan.
if plan.schema() != schema {
if plan
.schema()
.fields()
.iter()
.any(|field| blob_nulls::has_descriptor(field))
{
return Err(blob_scan_error());
}
return Err(arrow_schema::ArrowError::SchemaError(format!(
"Scan schema mismatch: expected {schema:?}, got {:?}",
plan.schema()
))
.into());
}
}
Ok(plan)
}
fn supports_filters_pushdown(
@@ -338,6 +463,9 @@ impl TableProvider for BaseTableAdapter {
.map(|filter| {
Ok(
if blob_nulls::references_descriptor(filter, &self.schema)? {
if self.blob_schema == BlobSchema::Write {
return Err(blob_scan_error());
}
TableProviderFilterPushDown::Unsupported
} else {
TableProviderFilterPushDown::Exact
@@ -358,6 +486,17 @@ impl TableProvider for BaseTableAdapter {
input: Arc<dyn ExecutionPlan>,
insert_op: InsertOp,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
if self.blob_schema == BlobSchema::Read
&& self
.schema
.fields()
.iter()
.any(|field| blob_nulls::has_descriptor(field))
{
return Err(DataFusionError::NotImplemented(
"SQL INSERT into Blob v2 tables requires BlobSchema::Write".into(),
));
}
let mode = match insert_op {
InsertOp::Append => WriteMode::Append,
InsertOp::Overwrite => WriteMode::Overwrite,
@@ -603,6 +742,70 @@ pub mod tests {
))
}
#[tokio::test]
async fn test_blob_schema_builders_preserve_modes_and_system_columns() {
use super::BlobSchema;
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap())
.execute()
.await
.unwrap();
let table = db
.create_empty_table(
"docs",
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
crate::blob("payload", true),
])),
)
.execute()
.await
.unwrap();
let adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_blob_schema(BlobSchema::Write)
.with_system_columns(&[ROW_ID])
.unwrap()
.with_use_rte(Some(false))
.with_fts_query(lance_index::scalar::FullTextSearchQuery::new(
"hello".into(),
))
.with_system_columns(&[ROW_ADDR])
.unwrap();
assert_eq!(adapter.blob_schema, BlobSchema::Write);
assert_eq!(adapter.use_rte, Some(false));
assert!(adapter.fts_query.is_some());
let names = adapter
.schema()
.fields()
.iter()
.map(|field| field.name().as_str())
.collect::<Vec<_>>()
.join(",");
assert_eq!(names, "id,payload");
let DataType::Struct(fields) = adapter.schema().field(1).data_type().clone() else {
unreachable!()
};
assert!(fields.find("data").is_some());
let adapter = adapter.with_blob_schema(BlobSchema::Read);
assert_eq!(
adapter
.schema()
.fields()
.iter()
.map(|field| field.name().as_str())
.collect::<Vec<_>>(),
["id", "payload", "_rowid", "_score", "_rowaddr"]
);
let DataType::Struct(fields) = adapter.schema().field(1).data_type().clone() else {
unreachable!()
};
assert!(fields.find("kind").is_some());
assert_eq!(adapter.schema().fields().len(), 5);
}
#[tokio::test]
async fn test_sql_blob_descriptor_nulls() {
use arrow_array::{UInt8Array, UInt64Array};
@@ -24,7 +24,7 @@ use datafusion_physical_plan::{ExecutionPlan, projection::ProjectionExec};
use lance_arrow::FieldExt;
use lance_core::datatypes::BlobV2Layout;
fn is_descriptor(field: &Field) -> bool {
pub(super) fn is_descriptor(field: &Field) -> bool {
field.is_blob()
&& matches!(field.data_type(), DataType::Struct(fields)
if BlobV2Layout::classify(fields) == Some(BlobV2Layout::Descriptor))
@@ -0,0 +1,636 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use arrow_array::{
Array, ArrayRef, Int64Array, LargeBinaryArray, ListArray, RecordBatch, StringArray,
StructArray, UInt64Array, cast::AsArray,
};
use arrow_buffer::{NullBuffer, OffsetBuffer};
use arrow_schema::{DataType, Field, Schema};
use arrow_select::concat::concat_batches;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_catalog::TableProvider;
use lance_core::{ROW_ADDR, ROW_ID};
use lance_index::scalar::FullTextSearchQuery;
use lancedb::{
Table, blob, connect,
database::listing::OPT_NEW_TABLE_STORAGE_VERSION,
table::datafusion::{BaseTableAdapter, BlobSchema},
};
use tempfile::{TempDir, tempdir};
async fn empty_table(fields: Vec<Field>) -> (TempDir, Table) {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap())
.execute()
.await
.unwrap();
let table = db
.create_empty_table("docs", Arc::new(Schema::new(fields)))
.execute()
.await
.unwrap();
(tmp, table)
}
async fn register(ctx: &SessionContext, table: &Table, mode: BlobSchema) {
let adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_blob_schema(mode);
ctx.deregister_table("docs").unwrap();
ctx.register_table("docs", Arc::new(adapter)).unwrap();
}
async fn sql(ctx: &SessionContext, query: &str) -> RecordBatch {
let batches = ctx.sql(query).await.unwrap().collect().await.unwrap();
concat_batches(&batches[0].schema(), &batches).unwrap()
}
async fn planning_error(ctx: &SessionContext, query: &str) -> String {
let error = match ctx.sql(query).await {
Ok(df) => df.create_physical_plan().await.unwrap_err(),
Err(error) => error,
};
error.to_string()
}
async fn row_ids(ctx: &SessionContext, table: &Table) -> Vec<u64> {
let reader = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_system_columns(&[ROW_ID])
.unwrap();
ctx.register_table("read_docs", Arc::new(reader)).unwrap();
sql(ctx, "SELECT _rowid FROM read_docs ORDER BY id")
.await
.column(0)
.as_primitive::<arrow_array::types::UInt64Type>()
.values()
.to_vec()
}
#[tokio::test]
async fn write_mode_inserts_blobs_nulls_omitted_and_reordered_columns() {
let (_tmp, table) = empty_table(vec![
Field::new("id", DataType::Int64, false),
blob("payload", true),
])
.await;
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Write).await;
for query in [
"INSERT INTO docs SELECT 1, named_struct('data', X'6869', 'uri', NULL)",
"INSERT INTO docs VALUES (2, NULL)",
"INSERT INTO docs (id) VALUES (3)",
"INSERT INTO docs SELECT 4, named_struct('data', X'', 'uri', NULL)",
"INSERT INTO docs (payload, id) SELECT named_struct('data', X'627965', 'uri', NULL), 5",
] {
sql(&ctx, query).await;
}
let ids = row_ids(&ctx, &table).await;
assert_eq!(
table.fetch_blobs("payload", &ids).await.unwrap(),
LargeBinaryArray::from(vec![
Some(b"hi".as_slice()),
None,
None,
Some(b"".as_slice()),
Some(b"bye".as_slice())
])
);
let batch = sql(&ctx, "SELECT id, payload.size FROM read_docs ORDER BY id").await;
assert_eq!(
batch
.column(1)
.as_primitive::<arrow_array::types::UInt64Type>(),
&UInt64Array::from(vec![Some(2), None, None, Some(0), Some(3)])
);
assert_eq!(
sql(&ctx, "SELECT count(payload) FROM read_docs")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
3
);
}
#[tokio::test]
async fn write_mode_inserts_two_blobs_in_explicit_column_order() {
let (_tmp, table) = empty_table(vec![
Field::new("id", DataType::Int64, false),
blob("payload", true),
blob("thumbnail", true),
])
.await;
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Write).await;
for query in [
"INSERT INTO docs (thumbnail, id, payload) SELECT named_struct('data', X'', 'uri', NULL), 1, named_struct('data', X'6869', 'uri', NULL)",
"INSERT INTO docs (id) VALUES (2)",
"INSERT INTO docs VALUES (3, NULL, NULL)",
] {
sql(&ctx, query).await;
}
let ids = row_ids(&ctx, &table).await;
assert_eq!(
table.fetch_blobs("payload", &ids).await.unwrap(),
LargeBinaryArray::from(vec![Some(b"hi".as_slice()), None, None])
);
assert_eq!(
table.fetch_blobs("thumbnail", &ids).await.unwrap(),
LargeBinaryArray::from(vec![Some(b"".as_slice()), None, None])
);
}
#[tokio::test]
async fn write_mode_rejects_blob_reads_before_execution_and_supports_scalar_self_insert() {
let (_tmp, table) = empty_table(vec![
Field::new("id", DataType::Int64, false),
blob("payload", true),
])
.await;
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Write).await;
sql(
&ctx,
"INSERT INTO docs SELECT 1, named_struct('data', X'6869', 'uri', NULL)",
)
.await;
sql(&ctx, "INSERT INTO docs VALUES (2, NULL)").await;
let system_adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap()
.with_system_columns(&[ROW_ID, lance_core::ROW_ADDR])
.unwrap()
.with_blob_schema(BlobSchema::Read);
ctx.register_table("system_docs", Arc::new(system_adapter))
.unwrap();
assert_eq!(
sql(
&ctx,
"SELECT id, _rowid, _rowaddr FROM system_docs WHERE id > 0"
)
.await
.num_rows(),
2
);
for query in [
"SELECT * FROM docs",
"SELECT payload FROM docs",
"SELECT payload.size FROM docs",
"SELECT payload.data FROM docs",
"SELECT count(payload) FROM docs",
"SELECT id FROM docs WHERE payload IS NULL",
"SELECT id FROM docs WHERE payload.size = 0",
"INSERT INTO docs SELECT id, payload FROM docs",
] {
let error = planning_error(&ctx, query).await;
assert!(
error.contains("Blob v2 scans require BlobSchema::Read"),
"{query}: {error}"
);
assert_eq!(table.count_rows(None).await.unwrap(), 2);
}
assert_eq!(
sql(&ctx, "SELECT count(*) FROM docs WHERE id > 0")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
2
);
assert_eq!(
sql(&ctx, "SELECT id FROM docs WHERE id = 1")
.await
.num_rows(),
1
);
sql(&ctx, "INSERT INTO docs (id) SELECT id + 100 FROM docs").await;
let ids = row_ids(&ctx, &table).await;
assert_eq!(ids.len(), 4);
assert_eq!(
table.fetch_blobs("payload", &ids).await.unwrap(),
LargeBinaryArray::from(vec![Some(b"hi".as_slice()), None, None, None])
);
}
fn blob_inputs(values: Vec<Option<&[u8]>>) -> ArrayRef {
let data = LargeBinaryArray::from(values);
let DataType::Struct(fields) = blob("clip", true).data_type().clone() else {
unreachable!()
};
let len = data.len();
let nulls = data.nulls().cloned();
Arc::new(StructArray::new(
fields,
vec![
Arc::new(data),
Arc::new(StringArray::new_null(len)),
Arc::new(UInt64Array::new_null(len)),
Arc::new(UInt64Array::new_null(len)),
],
nulls,
))
}
#[tokio::test]
async fn nested_blob_projections_require_a_read_adapter() {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap())
.execute()
.await
.unwrap();
let fields = vec![
Arc::new(Field::new("note", DataType::Utf8, true)),
Arc::new(blob("clip", true)),
];
let extra = StructArray::new(
fields.clone().into(),
vec![
Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])),
blob_inputs(vec![Some(b"hi"), None, Some(b"")]),
],
Some(NullBuffer::from(vec![true, true, false])),
);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("extra", DataType::Struct(fields.into()), true),
])),
vec![Arc::new(Int64Array::from(vec![1, 2, 3])), Arc::new(extra)],
)
.unwrap();
let table = db.create_table("docs", batch).execute().await.unwrap();
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Write).await;
// DataFusion requests the same whole parent struct for extra.note and
// extra.clip. Neither can satisfy the writer's advertised scan schema.
for query in [
"SELECT extra.note FROM docs ORDER BY id",
"SELECT count(extra.note) FROM docs",
"SELECT extra.note || 'x' FROM docs",
"SELECT row_number() OVER (ORDER BY extra.note) FROM docs",
"SELECT extra FROM docs",
"SELECT extra.clip FROM docs",
"SELECT extra.clip.size FROM docs",
"SELECT count(extra.clip) FROM docs",
] {
let error = planning_error(&ctx, query).await;
assert!(
error.contains("Blob v2 scans require BlobSchema::Read"),
"{query}: {error}"
);
}
let explain = sql(&ctx, "EXPLAIN SELECT id FROM docs WHERE extra.note = 'a'").await;
let explain = explain
.column(1)
.as_string::<i32>()
.iter()
.flatten()
.collect::<Vec<_>>()
.join("\n");
assert!(explain.contains("full_filters=["), "{explain}");
assert!(!explain.contains("FilterExec"), "{explain}");
assert_eq!(
sql(&ctx, "SELECT id FROM docs WHERE extra.note = 'a'")
.await
.num_rows(),
1
);
register(&ctx, &table, BlobSchema::Read).await;
let batch = sql(&ctx, "SELECT extra.note FROM docs ORDER BY id").await;
assert_eq!(
batch.column(0).as_string::<i32>(),
&StringArray::from(vec![Some("a"), Some("b"), None])
);
assert_eq!(
sql(&ctx, "SELECT count(extra.note) FROM docs")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
2
);
assert_eq!(
sql(&ctx, "SELECT extra.note FROM docs WHERE extra.note = 'a'")
.await
.num_rows(),
1
);
let reader = Arc::new(
BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap(),
);
let plan = reader.scan(&ctx.state(), None, &[], None).await.unwrap();
assert_eq!(plan.schema(), reader.schema());
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap();
assert!(
batches
.iter()
.all(|batch| batch.schema() == reader.schema())
);
let schema = reader.schema();
ctx.register_table("read_docs", reader).unwrap();
assert_eq!(sql(&ctx, "SELECT * FROM read_docs").await.schema(), schema);
assert_eq!(
sql(&ctx, "SELECT extra.clip.size FROM read_docs ORDER BY id")
.await
.column(0)
.as_primitive::<arrow_array::types::UInt64Type>(),
&UInt64Array::from(vec![Some(2), None, None])
);
}
#[tokio::test]
async fn list_descriptors_match_declared_schema_and_cannot_be_scanned_in_write_mode() {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap())
.execute()
.await
.unwrap();
let list = ListArray::new(
Arc::new(blob("item", true)),
OffsetBuffer::from_lengths([1, 2]),
blob_inputs(vec![Some(b"hi"), None, Some(b"")]),
None,
);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"payloads",
list.data_type().clone(),
true,
)])),
vec![Arc::new(list)],
)
.unwrap();
let table = db.create_table("docs", batch).execute().await.unwrap();
let ctx = SessionContext::new();
let reader = Arc::new(
BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap(),
);
let plan = reader.scan(&ctx.state(), None, &[], None).await.unwrap();
assert_eq!(plan.schema(), reader.schema());
let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap();
assert!(
batches
.iter()
.all(|batch| batch.schema() == reader.schema())
);
ctx.register_table("docs", reader).unwrap();
assert_eq!(
sql(&ctx, "SELECT count(payloads) FROM docs")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
2
);
register(&ctx, &table, BlobSchema::Write).await;
assert!(
planning_error(&ctx, "SELECT payloads FROM docs")
.await
.contains("Blob v2 scans require BlobSchema::Read")
);
assert_eq!(
sql(&ctx, "SELECT count(*) FROM docs")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
2
);
}
#[tokio::test]
async fn tables_without_blobs_support_inserts_and_full_scans_in_both_modes() {
let (_tmp, table) = empty_table(vec![Field::new("id", DataType::Int64, false)]).await;
let ctx = SessionContext::new();
for (mode, id) in [(BlobSchema::Read, 1), (BlobSchema::Write, 2)] {
register(&ctx, &table, mode).await;
sql(&ctx, &format!("INSERT INTO docs VALUES ({id})")).await;
let provider = ctx.table_provider("docs").await.unwrap();
let plan = provider.scan(&ctx.state(), None, &[], None).await.unwrap();
assert_eq!(plan.schema(), provider.schema());
assert_eq!(
sql(&ctx, "SELECT * FROM docs").await.num_rows(),
id as usize
);
}
}
#[tokio::test]
async fn read_mode_rejects_null_and_omitted_blob_inserts_without_writing() {
let (_tmp, table) = empty_table(vec![
Field::new("id", DataType::Int64, false),
blob("payload", true),
])
.await;
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Read).await;
for query in [
"INSERT INTO docs VALUES (2, NULL)",
"INSERT INTO docs (id) VALUES (3)",
] {
let error = planning_error(&ctx, query).await;
assert!(
error.contains("SQL INSERT into Blob v2 tables requires BlobSchema::Write"),
"{query}: {error}"
);
assert_eq!(table.count_rows(None).await.unwrap(), 0);
}
}
#[tokio::test]
async fn write_mode_inserts_exclude_read_only_builder_columns() {
for mode_first in [false, true] {
for system_first in [false, true] {
let (_tmp, table) = empty_table(vec![
Field::new("id", DataType::Int64, false),
Field::new("text", DataType::Utf8, true),
blob("payload", true),
])
.await;
let mut adapter = BaseTableAdapter::try_new(table.base_table().clone())
.await
.unwrap();
if mode_first {
adapter = adapter.with_blob_schema(BlobSchema::Write);
}
let fts = FullTextSearchQuery::new("hello".into());
adapter = if system_first {
adapter
.with_system_columns(&[ROW_ID, ROW_ADDR])
.unwrap()
.with_fts_query(fts)
} else {
adapter
.with_fts_query(fts)
.with_system_columns(&[ROW_ID, ROW_ADDR])
.unwrap()
};
let writer = adapter.with_blob_schema(BlobSchema::Write);
let ctx = SessionContext::new();
ctx.register_table("docs", Arc::new(writer)).unwrap();
sql(
&ctx,
"INSERT INTO docs SELECT 1, 'hello', named_struct('data', X'6869', 'uri', NULL)",
)
.await;
sql(&ctx, "INSERT INTO docs (text, id) VALUES ('hello', 2)").await;
let ids = row_ids(&ctx, &table).await;
assert_eq!(
table.fetch_blobs("payload", &ids).await.unwrap(),
LargeBinaryArray::from(vec![Some(b"hi".as_slice()), None])
);
}
}
}
#[tokio::test]
async fn legacy_blob_write_mode_rejects_descriptor_scans_and_keeps_scalar_scans() {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().to_str().unwrap())
.storage_options([(OPT_NEW_TABLE_STORAGE_VERSION, "2.1")])
.execute()
.await
.unwrap();
let legacy = Field::new("payload", DataType::LargeBinary, true).with_metadata(
std::collections::HashMap::from([("lance-encoding:blob".into(), "true".into())]),
);
let table = db
.create_empty_table(
"docs",
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
legacy,
])),
)
.execute()
.await
.unwrap();
let ctx = SessionContext::new();
register(&ctx, &table, BlobSchema::Write).await;
sql(&ctx, "INSERT INTO docs VALUES (1, X'6869')").await;
let writer = ctx.table_provider("docs").await.unwrap();
assert_eq!(writer.schema().field(1).data_type(), &DataType::LargeBinary);
let error = writer
.scan(&ctx.state(), None, &[], None)
.await
.unwrap_err();
assert!(
error.to_string().contains("Scan schema mismatch"),
"{error}"
);
assert!(!error.to_string().contains("Blob v2"), "{error}");
let plan = writer
.scan(&ctx.state(), Some(&vec![0]), &[], None)
.await
.unwrap();
assert_eq!(
plan.schema(),
Arc::new(writer.schema().project(&[0]).unwrap())
);
assert_eq!(
sql(&ctx, "SELECT id + 1 FROM docs WHERE id = 1")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
2
);
assert_eq!(
sql(&ctx, "SELECT count(*) FROM docs")
.await
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
1
);
register(&ctx, &table, BlobSchema::Read).await;
let reader = ctx.table_provider("docs").await.unwrap();
let plan = reader.scan(&ctx.state(), None, &[], None).await.unwrap();
assert_eq!(plan.schema(), reader.schema());
let DataType::Struct(fields) = reader.schema().field(1).data_type().clone() else {
panic!("expected a legacy descriptor");
};
assert_eq!(fields.len(), 2);
assert!(fields.find("position").is_some());
assert!(fields.find("size").is_some());
assert_eq!(sql(&ctx, "SELECT * FROM docs").await.num_rows(), 1);
}
#[tokio::test]
async fn copying_non_null_descriptors_between_tables_commits_no_rows() {
let (_tmp, source) = empty_table(vec![
Field::new("id", DataType::Int64, false),
blob("payload", true),
])
.await;
let ctx = SessionContext::new_with_config(
SessionConfig::new()
.with_batch_size(1)
.with_target_partitions(1),
);
register(&ctx, &source, BlobSchema::Write).await;
sql(
&ctx,
"INSERT INTO docs SELECT 1, named_struct('data', X'6869', 'uri', NULL)",
)
.await;
sql(&ctx, "INSERT INTO docs VALUES (2, NULL)").await;
register(&ctx, &source, BlobSchema::Read).await;
let (_dest_tmp, destination) = empty_table(
source
.schema()
.await
.unwrap()
.fields()
.iter()
.map(|f| f.as_ref().clone())
.collect(),
)
.await;
let writer = BaseTableAdapter::try_new(destination.base_table().clone())
.await
.unwrap()
.with_blob_schema(BlobSchema::Write);
ctx.register_table("destination", Arc::new(writer)).unwrap();
let version = destination.version().await.unwrap();
// The null row is valid and reaches the writer before the non-null
// descriptor. Even if it is written to a fragment, it must not be committed.
let plan = ctx
.sql("INSERT INTO destination SELECT id, payload FROM docs ORDER BY id DESC")
.await
.unwrap()
.create_physical_plan()
.await
.unwrap();
let error = datafusion_physical_plan::collect(plan, ctx.task_ctx())
.await
.unwrap_err();
assert!(error.to_string().contains("blob"), "{error}");
assert_eq!(destination.count_rows(None).await.unwrap(), 0);
assert_eq!(destination.version().await.unwrap(), version);
sql(
&ctx,
"INSERT INTO destination SELECT id, payload FROM docs WHERE payload IS NULL",
)
.await;
let ids = row_ids(&ctx, &destination).await;
assert_eq!(ids.len(), 1);
assert_eq!(
destination.fetch_blobs("payload", &ids).await.unwrap(),
LargeBinaryArray::from(vec![None::<&[u8]>])
);
}