mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-28 22:46:14 +00:00
fix!: restore blob SQL inserts with explicit adapter schemas (#4687)
This commit is contained in:
@@ -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.
|
||||
@@ -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))]
|
||||
|
||||
@@ -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]>])
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user