diff --git a/rust/lancedb/CHANGELOG.md b/rust/lancedb/CHANGELOG.md new file mode 100644 index 000000000..028f67128 --- /dev/null +++ b/rust/lancedb/CHANGELOG.md @@ -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. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 061ae8bf2..214e06d4e 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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::>(); + 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::(), 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::>() + .await + .unwrap(); + let row_ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("_rowid") + .unwrap() + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + 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))] diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index 973ab0123..8f187d465 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -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, + // Only the read schema includes opted-in system/FTS columns. schema: Arc, + write_schema: Arc, + blob_schema: BlobSchema, fts_query: Option, use_rte: Option, } @@ -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 { + /// 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 { - 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, ) -> DataFusionResult> { - // 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 = 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, insert_op: InsertOp, ) -> DataFusionResult> { + 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::>() + .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::>(), + ["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}; diff --git a/rust/lancedb/src/table/datafusion/blob_nulls.rs b/rust/lancedb/src/table/datafusion/blob_nulls.rs index 7c0cbda42..5befee5c3 100644 --- a/rust/lancedb/src/table/datafusion/blob_nulls.rs +++ b/rust/lancedb/src/table/datafusion/blob_nulls.rs @@ -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)) diff --git a/rust/lancedb/tests/datafusion_blob_schema.rs b/rust/lancedb/tests/datafusion_blob_schema.rs new file mode 100644 index 000000000..d51c182f6 --- /dev/null +++ b/rust/lancedb/tests/datafusion_blob_schema.rs @@ -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) -> (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 { + 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::() + .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::(), + &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::() + .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::() + .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>) -> 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::() + .iter() + .flatten() + .collect::>() + .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::(), + &StringArray::from(vec![Some("a"), Some("b"), None]) + ); + assert_eq!( + sql(&ctx, "SELECT count(extra.note) FROM docs") + .await + .column(0) + .as_primitive::() + .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::(), + &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::() + .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::() + .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::() + .value(0), + 2 + ); + assert_eq!( + sql(&ctx, "SELECT count(*) FROM docs") + .await + .column(0) + .as_primitive::() + .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]>]) + ); +}