diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index 81df7d300..22ec5caeb 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -4,6 +4,7 @@ //! This module contains adapters to allow LanceDB tables to be used as DataFusion table providers. mod blob_coerce; +mod blob_materialize; mod blob_nulls; pub mod cast; pub mod insert; @@ -31,7 +32,7 @@ use lance_core::{ datatypes::{Projection, Schema as LanceSchema}, }; -use super::{AnyQuery, BaseTable}; +use super::{AnyQuery, BaseTable, NativeTableExt}; use crate::{ Error, Result, query::{QueryExecutionOptions, QueryFilter, QueryRequest, Select}, @@ -199,7 +200,8 @@ impl ExecutionPlan for MetadataEraserExec { 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. + /// SQL INSERT SELECT materializes descriptors from a single native-table + /// snapshot; remote and Realtime/MemWAL sources are not supported. #[default] Read, /// Expose Blob v2 write inputs (`data`, `uri`, `position`, `size`) and legacy @@ -213,9 +215,33 @@ pub enum BlobSchema { /// /// 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. +/// SQL INSERT with logical blob inputs. In read mode, INSERT SELECT can copy +/// Blob v2 descriptors from a single native-table snapshot into target-owned +/// payloads, including projections, filters, sorting, and limits. JOIN, UNION, +/// aggregation, and computed descriptors fail during planning because their +/// source row addresses cannot be established safely. 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. +/// Descriptor copying requires a native source table; remote and Realtime/MemWAL +/// sources are not supported by this materialization path. +/// +/// With read adapters registered as `source` and `target`, copying blobs needs +/// no payload UDF or explicit system columns: +/// +/// ```sql +/// INSERT INTO target SELECT * FROM source WHERE id > 10; +/// ``` +/// +/// The target owns the copied payloads independently of subsequent source +/// compaction or deletion. External URI references retain their external +/// storage semantics and remain subject to the target's registered external +/// bases; copying a descriptor does not authorize a new external location. +/// +/// Materialization allocates payload bytes for each input batch outside the +/// DataFusion memory pool. For large blobs, configure a conservative +/// `SessionConfig::with_batch_size`; the pool limit alone does not bound these +/// allocations, and a single blob can still exceed that limit. #[derive(Debug)] pub struct BaseTableAdapter { table: Arc, @@ -440,13 +466,55 @@ impl TableProvider for BaseTableAdapter { ..Default::default() }; - let plan = self - .table - .create_plan(&AnyQuery::Query(query), options) + // The scan and any later INSERT rewrite must use the same snapshot, + // even if the live table advances between planning and execution. + let projects_descriptors = match projection { + Some(indices) => indices + .iter() + .any(|index| blob_nulls::has_descriptor(provider_schema.field(*index))), + None => provider_schema + .fields() + .iter() + .any(|field| blob_nulls::has_descriptor(field)), + }; + let snapshot = if self.blob_schema == BlobSchema::Read + && self.table.as_native().is_some() + && projects_descriptors + { + Some( + self.table + .query_snapshot() + .await + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ) + } else { + None + }; + let source = snapshot.as_ref().unwrap_or(&self.table); + let plan = source + .create_plan(&AnyQuery::Query(query.clone()), options.clone()) .map_err(|err| DataFusionError::External(err.into())) .await?; let plan = blob_nulls::mask_plan(plan)?; - let plan: Arc = Arc::new(MetadataEraserExec::new(plan)); + let mut plan: Arc = Arc::new(MetadataEraserExec::new(plan)); + if let Some(snapshot) = snapshot + && plan + .schema() + .fields() + .iter() + .any(|field| blob_nulls::has_descriptor(field)) + { + let dataset = snapshot + .as_native() + .unwrap() + .dataset + .get() + .await + .map_err(|error| DataFusionError::External(Box::new(error)))?; + plan = Arc::new(blob_materialize::BlobScanExec::new( + plan, snapshot, dataset, query, options, + )); + } if self.blob_schema == BlobSchema::Write { let schema = match projection { Some(projection) => Arc::new(provider_schema.project(projection)?), @@ -509,17 +577,30 @@ impl TableProvider for BaseTableAdapter { input: Arc, insert_op: InsertOp, ) -> DataFusionResult> { + let has_descriptors = input + .schema() + .fields() + .iter() + .any(|field| blob_nulls::has_descriptor(field)) + || (blob_materialize::has_source_scan(&input) + && input + .schema() + .fields() + .iter() + .any(|field| blob_materialize::has_descriptor(field))); if self.blob_schema == BlobSchema::Read && self .schema .fields() .iter() .any(|field| blob_nulls::has_descriptor(field)) + && !has_descriptors { return Err(DataFusionError::NotImplemented( "SQL INSERT into Blob v2 tables requires BlobSchema::Write".into(), )); } + let input = blob_materialize::materialize(input, self.write_schema.as_ref()).await?; let mode = match insert_op { InsertOp::Append => WriteMode::Append, InsertOp::Overwrite => WriteMode::Overwrite, diff --git a/rust/lancedb/src/table/datafusion/blob_materialize.rs b/rust/lancedb/src/table/datafusion/blob_materialize.rs new file mode 100644 index 000000000..22d0a023d --- /dev/null +++ b/rust/lancedb/src/table/datafusion/blob_materialize.rs @@ -0,0 +1,615 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Preserve scan provenance until INSERT can materialize source-owned blobs. +//! +//! External descriptors remain references and must satisfy the destination's +//! existing external-base policy; source authorization grants no target trust. + +use std::fmt; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::{DataFusionError, Result, tree_node::TreeNodeRecursion}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_functions::core::getfield::GetFieldFunc; +use datafusion_physical_expr::{ + PhysicalExpr, ScalarFunctionExpr, + expressions::{CastExpr, Column, Literal}, +}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + coalesce_partitions::CoalescePartitionsExec, + execution_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}, + filter::FilterExec, + limit::{GlobalLimitExec, LocalLimitExec}, + projection::ProjectionExec, + repartition::RepartitionExec, + sorts::{sort::SortExec, sort_preserving_merge::SortPreservingMergeExec}, + stream::RecordBatchStreamAdapter, +}; +use futures::{FutureExt, TryStreamExt, future::BoxFuture}; +use lance::{Dataset, dataset::BlobV2BatchRewritePlan}; +use lance_core::{ + ROW_ADDR, + datatypes::{BlobV2Layout, Schema as LanceSchema}, +}; + +use super::{BaseTable, MetadataEraserExec, blob_nulls}; +use crate::{ + query::{QueryExecutionOptions, QueryRequest, Select}, + table::AnyQuery, +}; + +const HIDDEN_ROW_ADDR: &str = "__lancedb_blob_rowaddr"; + +// Expressions such as CASE can retain the descriptor layout while dropping +// extension metadata. They still require proven provenance before writing. +pub(super) fn has_descriptor(field: &Field) -> bool { + contains_layout(field.data_type(), BlobV2Layout::Descriptor) +} + +fn descriptor_for_target(input: &Field, target: &Field) -> bool { + if crate::blob::is_blob(target) { + return matches!(input.data_type(), DataType::Struct(fields) + if BlobV2Layout::classify(fields) == Some(BlobV2Layout::Descriptor)); + } + match (input.data_type(), target.data_type()) { + (DataType::Struct(input), DataType::Struct(target)) => target.iter().any(|target| { + input + .find(target.name()) + .is_some_and(|(_, input)| descriptor_for_target(input, target)) + }), + ( + DataType::List(input) | DataType::LargeList(input) | DataType::FixedSizeList(input, _), + DataType::List(target) + | DataType::LargeList(target) + | DataType::FixedSizeList(target, _), + ) => descriptor_for_target(input, target), + _ => false, + } +} + +fn contains_layout(data_type: &DataType, layout: BlobV2Layout) -> bool { + match data_type { + DataType::Struct(fields) => { + BlobV2Layout::classify(fields) == Some(layout) + || fields + .iter() + .any(|field| contains_layout(field.data_type(), layout)) + } + DataType::List(item) | DataType::LargeList(item) | DataType::FixedSizeList(item, _) => { + contains_layout(item.data_type(), layout) + } + _ => false, + } +} + +pub(super) fn has_source_scan(input: &Arc) -> bool { + input.downcast_ref::().is_some() + || input.children().iter().any(|child| has_source_scan(child)) +} + +#[expect( + deprecated, + reason = "DataFusion 55 still accepts CoalesceBatchesExec plans" +)] +fn preserves_row_identity(input: &Arc) -> bool { + input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input.downcast_ref::().is_some() + || input + .downcast_ref::() + .is_some() + || input.downcast_ref::().is_some() +} + +fn provenance_error(detail: impl fmt::Display) -> DataFusionError { + DataFusionError::Plan(format!( + "Blob INSERT SELECT requires descriptors from one source table snapshot with an \ + unambiguous row address; {detail}. JOIN, UNION, aggregation, and computed blob \ + descriptors are not supported" + )) +} + +/// A transparent scan boundary. It retains the snapshot and scan request, rather +/// than exposing a system column to SELECT * or fetching any payloads on reads. +#[derive(Debug)] +pub(super) struct BlobScanExec { + input: Arc, + source: Arc, + dataset: Arc, + query: QueryRequest, + options: QueryExecutionOptions, +} + +impl BlobScanExec { + pub(super) fn new( + input: Arc, + source: Arc, + dataset: Arc, + query: QueryRequest, + options: QueryExecutionOptions, + ) -> Self { + Self { + input, + source, + dataset, + query, + options, + } + } + + async fn with_row_address(&self) -> Result { + let mut query = self.query.clone(); + let Select::Columns(columns) = &mut query.select else { + return Err(provenance_error( + "source scan has no explicit column projection", + )); + }; + if !columns.iter().any(|name| name == ROW_ADDR) { + columns.push(ROW_ADDR.to_owned()); + } + let plan = self + .source + .create_plan(&AnyQuery::Query(query), self.options.clone()) + .await + .map_err(|error| DataFusionError::External(Box::new(error)))?; + let plan = blob_nulls::mask_plan(plan)?; + let plan: Arc = Arc::new(MetadataEraserExec::new(plan)); + let schema = plan.schema(); + let mut projection: Vec<(Arc, String)> = self + .input + .schema() + .fields() + .iter() + .map(|field| { + Ok(( + Arc::new(Column::new(field.name(), schema.index_of(field.name())?)) as _, + field.name().clone(), + )) + }) + .collect::>()?; + projection.push(( + Arc::new(Column::new(ROW_ADDR, schema.index_of(ROW_ADDR)?)), + HIDDEN_ROW_ADDR.into(), + )); + Ok(Provenance { + plan: Arc::new(ProjectionExec::try_new(projection, plan)?), + dataset: self.dataset.clone(), + fields: self + .input + .schema() + .fields() + .iter() + .map(|field| Some(field.name().clone())) + .collect(), + }) + } +} + +impl DisplayAs for BlobScanExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "BlobScanExec: version={}", + self.dataset.version().version + ) + } +} + +impl ExecutionPlan for BlobScanExec { + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn name(&self) -> &str { + "BlobScanExec" + } + fn properties(&self) -> &Arc { + self.input.properties() + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn maintains_input_order(&self) -> Vec { + vec![true] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input]: [_; 1] = children + .try_into() + .map_err(|_| provenance_error("invalid scan children"))?; + Ok(Arc::new(Self::new( + input, + self.source.clone(), + self.dataset.clone(), + self.query.clone(), + self.options.clone(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } +} + +struct Provenance { + plan: Arc, + dataset: Arc, + // Source field path for each output column, before the hidden address. + fields: Vec>, +} + +fn source_field(expr: &Arc, fields: &[Option]) -> Option { + if let Some(column) = expr.downcast_ref::() { + return fields.get(column.index()).cloned().flatten(); + } + let call = expr.downcast_ref::()?; + call.fun().inner().downcast_ref::()?; + let (root, path) = call.args().split_first()?; + let mut source = source_field(root, fields)?; + for component in path { + let name = component + .downcast_ref::()? + .value() + .try_as_str() + .flatten()?; + source.push('.'); + source.push_str(name); + } + Some(source) +} + +fn with_provenance(input: Arc) -> BoxFuture<'static, Result> { + async move { + if let Some(scan) = input.downcast_ref::() { + return scan.with_row_address().await; + } + let children = input.children(); + // An explicit allowlist keeps row identity from silently propagating + // through a future operator that changes cardinality or column meaning. + if children.len() != 1 || !preserves_row_identity(&input) { + return Err(provenance_error(format!( + "cannot establish provenance through {}", + input.name() + ))); + } + let mut provenance = with_provenance(children[0].clone()).await?; + if let Some(projection) = input.downcast_ref::() { + let fields = projection + .expr() + .iter() + .map(|expr| source_field(&expr.expr, &provenance.fields)) + .collect(); + let mut expressions: Vec<(Arc, String)> = projection + .expr() + .iter() + .map(|expr| (expr.expr.clone(), expr.alias.clone())) + .collect(); + expressions.push(( + Arc::new(Column::new( + HIDDEN_ROW_ADDR, + provenance.plan.schema().fields().len() - 1, + )), + HIDDEN_ROW_ADDR.into(), + )); + provenance.plan = Arc::new(ProjectionExec::try_new(expressions, provenance.plan)?); + provenance.fields = fields; + } else { + if input.schema() != children[0].schema() { + return Err(provenance_error(format!( + "{} changes the source schema", + input.name() + ))); + } + provenance.plan = input.replace_children( + vec![provenance.plan], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + } + Ok(provenance) + } + .boxed() +} + +/// Defer only a container coercion whose Blob leaves retain their named paths. +/// Descriptor leaves are atomic: casting their members is never provenance. +fn can_defer_blob_cast(source: &Field, cast_type: &DataType, target: &Field) -> bool { + if crate::blob::is_blob(target) { + let (DataType::Struct(source), DataType::Struct(cast)) = (source.data_type(), cast_type) + else { + return false; + }; + return BlobV2Layout::classify(source) == Some(BlobV2Layout::Descriptor) + && (BlobV2Layout::classify(cast) == Some(BlobV2Layout::Logical) + || (BlobV2Layout::classify(cast) == Some(BlobV2Layout::Descriptor) + && source.len() == cast.len() + && source.iter().zip(cast).all(|(source, cast)| { + source.name() == cast.name() && source.data_type() == cast.data_type() + }))); + } + match (source.data_type(), cast_type, target.data_type()) { + (DataType::Struct(source), DataType::Struct(cast), DataType::Struct(target)) => { + source.len() == cast.len() + && source.len() == target.len() + && target.iter().all(|target| { + source + .find(target.name()) + .zip(cast.find(target.name())) + .is_some_and(|((_, source), (_, cast))| { + can_defer_blob_cast(source, cast.data_type(), target) + }) + }) + } + (DataType::List(source), DataType::List(cast), DataType::List(target)) + | (DataType::LargeList(source), DataType::LargeList(cast), DataType::LargeList(target)) => { + can_defer_blob_cast(source, cast.data_type(), target) + } + ( + DataType::FixedSizeList(source, source_len), + DataType::FixedSizeList(cast, cast_len), + DataType::FixedSizeList(target, target_len), + ) => { + source_len == cast_len + && source_len == target_len + && can_defer_blob_cast(source, cast.data_type(), target) + } + _ => !has_descriptor(source) && cast_type == target.data_type(), + } +} + +pub(super) async fn materialize( + input: Arc, + target_schema: &Schema, +) -> Result> { + if !target_schema + .fields() + .iter() + .any(|field| crate::blob::field_tree_has_blob_v2(field)) + { + return Ok(input); + } + // INSERT may cast containers for metadata, field order, scalar siblings, + // or a Write adapter's logical layout. Bind the original Blob leaves and + // materialize before applying the target coercion. Other expressions still + // have to pass the Column/GetField provenance checks below. + let input = if let Some(projection) = input.downcast_ref::() { + let mut changed = false; + let expressions = projection + .expr() + .iter() + .map(|projected| { + let mut expr = projected.expr.clone(); + if let Some(cast) = expr.downcast_ref::() + && let Ok(target) = target_schema.field_with_name(&projected.alias) + && crate::blob::field_tree_has_blob_v2(target) + && let source = cast + .expr() + .return_field(projection.input().schema().as_ref())? + && has_descriptor(&source) + { + if !can_defer_blob_cast( + &source, + &cast.data_type(projection.input().schema().as_ref())?, + target, + ) { + return Err(provenance_error(format!( + "cast of column '{}' does not preserve Blob leaf identity", + projected.alias + ))); + } + expr = cast.expr().clone(); + changed = true; + } + Ok((expr, projected.alias.clone())) + }) + .collect::>>()?; + if changed { + Arc::new(ProjectionExec::try_new( + expressions, + projection.input().clone(), + )?) as Arc + } else { + input + } + } else { + input + }; + let rewrite_columns: Vec = input + .schema() + .fields() + .iter() + .map(|field| { + target_schema + .field_with_name(field.name()) + .is_ok_and(|target| descriptor_for_target(field, target)) + }) + .collect(); + if !rewrite_columns.iter().any(|rewrite| *rewrite) { + return Ok(input); + } + let provenance = with_provenance(input.clone()).await?; + let mut source_schema = LanceSchema::default(); + let mut expressions: Vec<(Arc, String)> = Vec::new(); + for (index, (field, source)) in input + .schema() + .fields() + .iter() + .zip(&provenance.fields) + .enumerate() + { + let name = format!("__lancedb_blob_column_{index}"); + if rewrite_columns[index] { + let source = source.as_ref().ok_or_else(|| { + provenance_error(format!("column '{}' has no source", field.name())) + })?; + let mut source_field = provenance + .dataset + .schema() + .field(source) + .ok_or_else(|| provenance_error(format!("source field '{source}' is missing")))? + .clone(); + if !crate::blob::field_tree_has_blob_v2(&Field::from(&source_field)) { + return Err(provenance_error(format!( + "source field '{source}' is not a Blob v2 column" + ))); + } + source_field.name = name.clone(); + source_schema.fields.push(source_field); + } + expressions.push((Arc::new(Column::new(field.name(), index)), name)); + } + expressions.push(( + Arc::new(Column::new(HIDDEN_ROW_ADDR, provenance.fields.len())), + ROW_ADDR.into(), + )); + let plan: Arc = + Arc::new(ProjectionExec::try_new(expressions, provenance.plan)?); + // preserve_managed=false forces even managed source payloads to become + // target-owned bytes instead of retaining source storage references. + let rewrite = Arc::new(BlobV2BatchRewritePlan::try_new( + &source_schema, + plan.schema().as_ref(), + false, + false, + )?); + let schema = Arc::new(Schema::new( + rewrite + .output_schema() + .fields() + .iter() + .zip(input.schema().fields()) + .map(|(rewritten, original)| rewritten.as_ref().clone().with_name(original.name())) + .collect::>(), + )); + let plan = Arc::new(MaterializeBlobsExec::new( + plan, + provenance.dataset, + rewrite, + schema, + )); + super::cast::cast_to_table_schema(plan, target_schema) + .map_err(|error| DataFusionError::External(Box::new(error))) +} + +struct MaterializeBlobsExec { + input: Arc, + dataset: Arc, + rewrite: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl fmt::Debug for MaterializeBlobsExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MaterializeBlobsExec") + .field("input", &self.input) + .finish_non_exhaustive() + } +} + +impl MaterializeBlobsExec { + fn new( + input: Arc, + dataset: Arc, + rewrite: Arc, + schema: SchemaRef, + ) -> Self { + // Rewriting changes column types; descriptor equivalences/orderings do + // not describe the resulting payload structs. + let properties = Arc::new(PlanProperties::new( + datafusion_physical_expr::EquivalenceProperties::new(schema.clone()), + datafusion_physical_expr::Partitioning::UnknownPartitioning( + input.output_partitioning().partition_count(), + ), + input.properties().emission_type, + input.properties().boundedness, + )); + Self { + input, + dataset, + rewrite, + schema, + properties, + } + } +} + +impl DisplayAs for MaterializeBlobsExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "MaterializeBlobsExec") + } +} + +impl ExecutionPlan for MaterializeBlobsExec { + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn name(&self) -> &str { + "MaterializeBlobsExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn maintains_input_order(&self) -> Vec { + vec![true] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input]: [_; 1] = children + .try_into() + .map_err(|_| provenance_error("invalid rewrite children"))?; + Ok(Arc::new(Self::new( + input, + self.dataset.clone(), + self.rewrite.clone(), + self.schema.clone(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let dataset = self.dataset.clone(); + let rewrite = self.rewrite.clone(); + let schema = self.schema.clone(); + let stream = self + .input + .execute(partition, context)? + .and_then(move |batch| { + let dataset = dataset.clone(); + let rewrite = rewrite.clone(); + let schema = schema.clone(); + async move { + let batch = rewrite.transform_batch(&dataset, batch).await?; + Ok(RecordBatch::try_new(schema, batch.columns().to_vec())?) + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } +} diff --git a/rust/lancedb/src/table/datafusion/insert.rs b/rust/lancedb/src/table/datafusion/insert.rs index ec1e603d2..a2d5f34c9 100644 --- a/rust/lancedb/src/table/datafusion/insert.rs +++ b/rust/lancedb/src/table/datafusion/insert.rs @@ -195,6 +195,7 @@ impl ExecutionPlan for InsertExec { context: Arc, ) -> DataFusionResult { let input_stream = self.input.execute(partition, context)?; + let input_plan = self.input.clone(); let dataset = self.dataset.clone(); let mut write_params = self.write_params.clone(); let partial_transactions = self.partial_transactions.clone(); @@ -227,6 +228,9 @@ impl ExecutionPlan for InsertExec { )); let stream = futures::stream::once(async move { + // Source blob snapshots must outlive the final target commit, not + // only the consumption of the input stream. + let _input_plan = input_plan; if let Some(tracker) = tracker && write_params.write_progress.is_none() { diff --git a/rust/lancedb/tests/datafusion_blob_schema.rs b/rust/lancedb/tests/datafusion_blob_schema.rs index d51c182f6..5af3122ab 100644 --- a/rust/lancedb/tests/datafusion_blob_schema.rs +++ b/rust/lancedb/tests/datafusion_blob_schema.rs @@ -571,66 +571,701 @@ async fn legacy_blob_write_mode_rejects_descriptor_scans_and_keeps_scalar_scans( } #[tokio::test] -async fn copying_non_null_descriptors_between_tables_commits_no_rows() { +async fn copying_descriptors_materializes_payloads_and_preserves_snapshot() { + use std::collections::HashMap; + use std::num::NonZeroUsize; + + use lancedb::blob::{BlobOptions, blob_with_options}; + use lancedb::table::OptimizeAction; + + for managed in [false, true] { + for mode in [BlobSchema::Read, BlobSchema::Write] { + let source_tmp = tempdir().unwrap(); + let db = connect(source_tmp.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let payload = blob_with_options( + "payload", + true, + BlobOptions { + inline_size_threshold: Some(8), + dedicated_size_threshold: NonZeroUsize::new(32), + }, + ); + let values: Vec> = vec![ + Some(b"hi"), + Some(&[2; 16]), + Some(&[3; 64]), + None, + Some(b""), + Some(b"skip"), + ]; + let schema = Arc::new(Schema::new_with_metadata( + vec![Field::new("id", DataType::Int64, false), payload], + HashMap::from([("lance.blob.managed".into(), managed.to_string())]), + )); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5, 6])), + blob_inputs(values.clone()), + ], + ) + .unwrap(); + let source = db.create_table("docs", batch).execute().await.unwrap(); + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .with_batch_size(2) + .with_target_partitions(2), + ); + register(&ctx, &source, BlobSchema::Read).await; + // Descriptor queries and SELECT * retain their original read schema. + let sizes = sql(&ctx, "SELECT payload.size FROM docs ORDER BY id").await; + assert_eq!( + sizes + .column(0) + .as_primitive::() + .value(2), + 64 + ); + assert_eq!(sql(&ctx, "SELECT * FROM docs").await.num_columns(), 2); + let kinds = sql(&ctx, "SELECT payload.kind FROM docs ORDER BY id").await; + let kinds = kinds + .column(0) + .as_primitive::(); + if !managed { + assert_eq!( + kinds.value(0), + lance_core::datatypes::BlobKind::Inline as u8 + ); + assert_eq!( + kinds.value(1), + lance_core::datatypes::BlobKind::Packed as u8 + ); + assert_eq!( + kinds.value(2), + lance_core::datatypes::BlobKind::Dedicated as u8 + ); + } else { + assert_eq!( + kinds.value(2), + lance_core::datatypes::BlobKind::Managed as u8 + ); + } + let (_destination_tmp, destination) = empty_table(vec![ + blob("copied", true), + Field::new("id", DataType::Int64, false), + ]) + .await; + let writer = BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap() + .with_blob_schema(mode); + ctx.register_table("destination", Arc::new(writer)).unwrap(); + let plan = ctx.sql("INSERT INTO destination (id, copied) SELECT id + 10, payload AS copied FROM docs WHERE id > 0 AND (payload.size IS NULL OR payload.size < 100) ORDER BY id DESC LIMIT 5 OFFSET 1") + .await.unwrap().create_physical_plan().await.unwrap(); + // A live-table change after planning must not change the source + // version used to decode the descriptors in this insert. + source.delete("id = 1").await.unwrap(); + datafusion_physical_plan::collect(plan, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!(destination.count_rows(None).await.unwrap(), 5); + source.optimize(OptimizeAction::All).await.unwrap(); + std::fs::remove_dir_all(source_tmp.path()).unwrap(); + let reader = BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ID]) + .unwrap(); + ctx.register_table("copied_rows", Arc::new(reader)).unwrap(); + let ids = sql(&ctx, "SELECT _rowid FROM copied_rows ORDER BY id").await; + let ids = ids + .column(0) + .as_primitive::() + .values(); + assert_eq!( + destination.fetch_blobs("copied", ids).await.unwrap(), + LargeBinaryArray::from(values[..5].to_vec()) + ); + } + } +} + +#[tokio::test] +async fn copying_nested_descriptors_preserves_parent_nulls_and_field_identity() { + use arrow_array::Int32Array; + use lancedb::blob::{BlobOptions, blob_with_options}; + + for difference in ["metadata", "order", "scalar"] { + for mode in [BlobSchema::Read, BlobSchema::Write] { + let source_tmp = tempdir().unwrap(); + let db = connect(source_tmp.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let children = vec![ + Arc::new(Field::new("note", DataType::Utf8, true)), + Arc::new(Field::new("rank", DataType::Int32, true)), + Arc::new(blob_with_options( + "clip", + true, + BlobOptions { + inline_size_threshold: Some(8), + ..Default::default() + }, + )), + ]; + let extra_field = Field::new("extra", DataType::Struct(children.clone().into()), true); + let extra = StructArray::new( + children.into(), + vec![ + Arc::new(StringArray::from(vec![ + Some("a"), + Some("hidden"), + Some("c"), + ])), + Arc::new(Int32Array::from(vec![1, 2, 3])), + blob_inputs(vec![Some(b"hi"), Some(b"hidden"), Some(b"")]), + ], + Some(NullBuffer::from(vec![true, false, true])), + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + extra_field, + ])), + vec![Arc::new(Int64Array::from(vec![1, 2, 3])), Arc::new(extra)], + ) + .unwrap(); + let source = db.create_table("docs", batch).execute().await.unwrap(); + let ctx = SessionContext::new(); + register(&ctx, &source, BlobSchema::Read).await; + let mut target_children = vec![ + Field::new("note", DataType::Utf8, true), + Field::new( + "rank", + if difference == "scalar" { + DataType::Int64 + } else { + DataType::Int32 + }, + true, + ), + blob("clip", true), + ]; + if difference == "order" { + target_children.swap(0, 2); + } + let (_destination_tmp, destination) = empty_table(vec![ + Field::new("id", DataType::Int64, false), + Field::new("extra", DataType::Struct(target_children.into()), true), + blob("flattened", true), + ]) + .await; + ctx.register_table( + "destination", + Arc::new( + BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap() + .with_blob_schema(mode), + ), + ) + .unwrap(); + let version = destination.version().await.unwrap(); + for expression in [ + "CASE WHEN id = 1 THEN extra ELSE NULL END", + "named_struct('note', extra.note, 'rank', extra.rank, 'clip', extra.clip)", + ] { + let query = format!( + "INSERT INTO destination SELECT id, {expression} AS extra, extra.clip AS flattened FROM docs" + ); + let error = planning_error(&ctx, &query).await; + assert!( + error.contains("one source table snapshot"), + "{difference}, {mode:?}: {error}" + ); + assert_eq!(destination.version().await.unwrap(), version); + assert_eq!(destination.count_rows(None).await.unwrap(), 0); + } + sql(&ctx, "INSERT INTO destination SELECT id, extra, extra.clip AS flattened FROM docs ORDER BY id DESC").await; + std::fs::remove_dir_all(source_tmp.path()).unwrap(); + let ids = row_ids(&ctx, &destination).await; + let expected = + LargeBinaryArray::from(vec![Some(b"hi".as_slice()), None, Some(b"".as_slice())]); + assert_eq!( + destination.fetch_blobs("extra.clip", &ids).await.unwrap(), + expected, + "{difference}, {mode:?}" + ); + assert_eq!( + destination.fetch_blobs("flattened", &ids).await.unwrap(), + expected + ); + let nested = sql(&ctx, "SELECT extra FROM read_docs ORDER BY id").await; + let extra = nested.column(0).as_struct(); + assert!(extra.is_null(1)); + assert_eq!( + extra + .column_by_name("note") + .unwrap() + .as_string::() + .value(2), + "c" + ); + let rank = extra.column_by_name("rank").unwrap(); + if difference == "scalar" { + assert_eq!( + rank.as_primitive::() + .value(2), + 3 + ); + } else { + assert_eq!( + rank.as_primitive::() + .value(2), + 3 + ); + } + } + } +} + +#[tokio::test] +async fn ordinary_descriptor_shaped_structs_remain_plain_data() { + let fields = vec![ + Field::new("id", DataType::Int64, false), + Field::new( + "raw", + DataType::Struct( + vec![ + Field::new("kind", DataType::UInt8, true), + Field::new("position", DataType::UInt64, true), + Field::new("size", DataType::UInt64, true), + Field::new("blob_id", DataType::UInt32, true), + Field::new("blob_uri", DataType::Utf8, true), + ] + .into(), + ), + true, + ), + ]; + let (source_tmp, source) = empty_table(fields.clone()).await; + let (_destination_tmp, destination) = empty_table(fields).await; + let ctx = SessionContext::new(); + register(&ctx, &source, BlobSchema::Read).await; + sql(&ctx, "INSERT INTO docs VALUES (1, named_struct('kind', 1, 'position', 2, 'size', 3, 'blob_id', 4, 'blob_uri', 'plain data'))").await; + let expected = sql(&ctx, "SELECT * FROM docs").await; + ctx.register_table( + "destination", + Arc::new( + BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap(), + ), + ) + .unwrap(); + sql(&ctx, "INSERT INTO destination SELECT * FROM docs").await; + std::fs::remove_dir_all(source_tmp.path()).unwrap(); + let actual = sql(&ctx, "SELECT * FROM destination").await; + assert_eq!(actual.columns(), expected.columns()); +} + +#[tokio::test] +async fn ambiguous_descriptor_sources_fail_during_planning_without_writing() { 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), - ); + let ctx = SessionContext::new(); 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(), - ) + let (_destination_tmp, destination) = empty_table(vec![ + Field::new("id", DataType::Int64, false), + blob("payload", true), + ]) .await; let writer = BaseTableAdapter::try_new(destination.base_table().clone()) .await - .unwrap() - .with_blob_schema(BlobSchema::Write); + .unwrap(); 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") + for select in [ + "SELECT a.id, a.payload FROM docs a JOIN docs b ON a.id = b.id", + "SELECT id, payload FROM docs UNION ALL SELECT id, payload FROM docs", + "SELECT id, payload FROM docs GROUP BY id, payload", + "SELECT id, CASE WHEN id = 1 THEN payload ELSE NULL END AS payload FROM docs", + ] { + let query = format!("INSERT INTO destination {select}"); + let error = planning_error(&ctx, &query).await; + assert!( + error.contains("one source table snapshot"), + "{query}: {error}" + ); + } + assert_eq!(destination.count_rows(None).await.unwrap(), 0); + assert_eq!(destination.version().await.unwrap(), version); +} + +#[tokio::test] +async fn external_descriptor_copy_obeys_destination_bases_and_preserves_ranges() { + use lance::dataset::WriteParams; + use lance_core::datatypes::BlobKind; + use lance_table::format::BasePath; + use lancedb::table::WriteOptions; + + let external_tmp = tempdir().unwrap(); + let payload = b"0123456789abcdef"; + let path = external_tmp.path().join("payload.bin"); + std::fs::write(&path, payload).unwrap(); + let uri = url::Url::from_file_path(&path).unwrap().to_string(); + let base_uri = url::Url::from_directory_path(external_tmp.path()) + .unwrap() + .to_string(); + let base = |id| BasePath { + id, + name: Some("external".into()), + path: base_uri.clone(), + is_dataset_root: false, + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + blob("payload", true), + ])); + let DataType::Struct(children) = schema.field(1).data_type().clone() else { + unreachable!() + }; + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StructArray::new( + children, + vec![ + Arc::new(LargeBinaryArray::from(vec![ + None, + Some(b"inline".as_slice()), + ])), + Arc::new(StringArray::from(vec![Some(uri.as_str()), None])), + Arc::new(UInt64Array::from(vec![Some(4), None])), + Arc::new(UInt64Array::from(vec![Some(6), None])), + ], + None, + )), + ], + ) + .unwrap(); + for source_registered in [false, true] { + let source_tmp = tempdir().unwrap(); + let db = connect(source_tmp.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let source = db + .create_empty_table("source", schema.clone()) + .write_options(WriteOptions { + lance_write_params: Some(WriteParams { + initial_bases: source_registered.then(|| vec![base(7)]), + ..Default::default() + }), + }) + .execute() + .await + .unwrap(); + source + .add(input.clone()) + .allow_external_blob_outside_bases(!source_registered) + .execute() + .await + .unwrap(); + for mode in [BlobSchema::Read, BlobSchema::Write] { + let ctx = SessionContext::new(); + register(&ctx, &source, BlobSchema::Read).await; + let descriptor = sql(&ctx, "SELECT payload.blob_id FROM docs WHERE id = 1").await; + assert_eq!( + descriptor + .column(0) + .as_primitive::() + .value(0), + if source_registered { 7 } else { 0 } + ); + for authorized in [false, true] { + let destination_tmp = tempdir().unwrap(); + let destination_db = connect(destination_tmp.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let destination = destination_db + .create_empty_table("destination", schema.clone()) + .write_options(WriteOptions { + lance_write_params: Some(WriteParams { + initial_bases: authorized.then(|| vec![base(11)]), + ..Default::default() + }), + }) + .execute() + .await + .unwrap(); + ctx.deregister_table("destination").unwrap(); + ctx.register_table( + "destination", + Arc::new( + BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap() + .with_blob_schema(mode), + ), + ) + .unwrap(); + let version = destination.version().await.unwrap(); + let result = ctx + .sql("INSERT INTO destination SELECT * FROM docs ORDER BY id") + .await + .unwrap() + .collect() + .await; + if !authorized { + let error = result.unwrap_err().to_string(); + assert!( + error.contains("outside registered external bases"), + "{error}" + ); + let reopened = destination_db + .open_table("destination") + .execute() + .await + .unwrap(); + assert_eq!(reopened.version().await.unwrap(), version); + assert_eq!(reopened.count_rows(None).await.unwrap(), 0); + // A logical URI input has the same destination authorization boundary. + ctx.register_table( + "logical_target", + Arc::new( + BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap() + .with_blob_schema(BlobSchema::Write), + ), + ) + .unwrap(); + let error = ctx.sql(&format!("INSERT INTO logical_target SELECT 3, named_struct('data', NULL, 'uri', '{uri}')")).await.unwrap().collect().await.unwrap_err().to_string(); + assert!( + error.contains("allow_external_blob_outside_bases"), + "{error}" + ); + destination.checkout_latest().await.unwrap(); + assert_eq!(destination.version().await.unwrap(), version); + assert_eq!(destination.count_rows(None).await.unwrap(), 0); + } else { + result.unwrap(); + let ids = row_ids(&ctx, &destination).await; + let descriptor = sql(&ctx, "SELECT payload.kind, payload.blob_id, payload.blob_uri, payload.position, payload.size FROM read_docs WHERE id = 1").await; + assert_eq!( + descriptor + .column(0) + .as_primitive::() + .value(0), + BlobKind::External as u8 + ); + assert_eq!( + descriptor + .column(1) + .as_primitive::() + .value(0), + 11 + ); + assert_eq!( + descriptor.column(2).as_string::().value(0), + "payload.bin" + ); + assert_eq!( + descriptor + .column(3) + .as_primitive::() + .value(0), + 4 + ); + assert_eq!( + descriptor + .column(4) + .as_primitive::() + .value(0), + 6 + ); + assert_eq!( + destination.fetch_blobs("payload", &ids).await.unwrap(), + LargeBinaryArray::from(vec![ + Some(&payload[4..10]), + Some(b"inline".as_slice()) + ]) + ); + let files = destination.fetch_blob_files("payload", &ids).await.unwrap(); + let file = files[0].as_ref().unwrap(); + assert_eq!(file.uri(), Some("payload.bin")); + let dataset = destination.dataset().unwrap().get().await.unwrap(); + let registered = &dataset.manifest().base_paths[&11]; + let resolved = + url::Url::parse(&format!("{}/", registered.path.trim_end_matches('/'))) + .unwrap() + .join(file.uri().unwrap()) + .unwrap(); + assert_eq!(resolved.as_str(), uri); + assert_eq!(file.position(), Some(4)); + assert_eq!(file.size(), 6); + } + } + } + } +} + +#[tokio::test] +async fn source_payload_failure_after_written_partition_does_not_commit() { + use datafusion_physical_plan::{ + ExecutionPlanProperties, + execution_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}, + union::UnionExec, + }; + use futures::TryStreamExt; + use lancedb::blob::{BlobOptions, blob_with_options}; + use std::num::NonZeroUsize; + + let source_tmp = tempdir().unwrap(); + let db = connect(source_tmp.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + blob_with_options( + "payload", + true, + BlobOptions { + inline_size_threshold: Some(0), + dedicated_size_threshold: NonZeroUsize::new(1), + }, + ), + ])); + let source = db + .create_table( + "docs", + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + blob_inputs(vec![Some(&[1; 64]), Some(&[2; 64])]), + ], + ) + .unwrap(), + ) + .execute() + .await + .unwrap(); + let (_destination_tmp, destination) = empty_table(vec![ + Field::new("id", DataType::Int64, false), + blob("payload", true), + ]) + .await; + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(1), + ); + register(&ctx, &source, BlobSchema::Read).await; + ctx.register_table( + "destination", + Arc::new( + BaseTableAdapter::try_new(destination.base_table().clone()) + .await + .unwrap(), + ), + ) + .unwrap(); + let first = ctx + .sql("INSERT INTO destination SELECT * FROM docs WHERE id = 1") .await .unwrap() .create_physical_plan() .await .unwrap(); - let error = datafusion_physical_plan::collect(plan, ctx.task_ctx()) + let second = ctx + .sql("INSERT INTO destination SELECT * FROM docs WHERE id = 2") .await - .unwrap_err(); - assert!(error.to_string().contains("blob"), "{error}"); + .unwrap() + .create_physical_plan() + .await + .unwrap(); + assert_eq!(first.name(), "InsertExec"); + assert_eq!(second.name(), "InsertExec"); + // Combine already validated SQL inputs only to control partition scheduling: + // finish one real uncommitted write before faulting the next payload read. + let input = UnionExec::try_new(vec![ + first.children()[0].clone(), + second.children()[0].clone(), + ]) + .unwrap(); + assert_eq!(input.output_partitioning().partition_count(), 2); + let insert = first + .replace_children( + vec![input], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + let version = destination.version().await.unwrap(); + let batches: Vec = insert + .execute(0, ctx.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches[0] + .column(0) + .as_primitive::() + .value(0), + 1 + ); 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]>]) + let dataset = destination.dataset().unwrap().get().await.unwrap(); + let data_dir = _destination_tmp.path().join("docs.lance").join("data"); + assert!( + std::fs::read_dir(data_dir).unwrap().next().is_some(), + "first partition must have written a data file" ); + let ids = row_ids(&ctx, &source).await; + let files = source.fetch_blob_files("payload", &ids).await.unwrap(); + let file = files[1].as_ref().unwrap(); + assert_eq!( + file.kind(), + Some(lance_core::datatypes::BlobKind::Dedicated) + ); + let missing_path = std::path::Path::new("/").join(file.data_path().unwrap().as_ref()); + drop(files); + std::fs::remove_file(&missing_path).unwrap(); + let error = insert + .execute(1, ctx.task_ctx()) + .unwrap() + .try_collect::>() + .await + .unwrap_err() + .to_string(); + assert!( + error.contains(missing_path.file_name().unwrap().to_str().unwrap()), + "{error}" + ); + // Check both the live handle and an independently reopened manifest. + destination.checkout_latest().await.unwrap(); + assert_eq!(destination.version().await.unwrap(), version); + assert_eq!(destination.count_rows(None).await.unwrap(), 0); + let reopened = lance::Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!(reopened.version().version, version); + assert_eq!(reopened.count_rows(None).await.unwrap(), 0); }