mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(rust): support cross-table blob copies in DataFusion (#4976)
`INSERT INTO target SELECT ... FROM source` currently cannot copy Blob v2 descriptors into an independent table. This enables copies from a single native source snapshot through projections, filters, sorting, and limits, while preserving descriptor queries such as `SELECT payload.size`. Source-owned payloads become target-owned data, so subsequent source mutation, compaction, or deletion does not invalidate the copy; nested fields retain their named identity across metadata differences, child reordering, scalar casts, and nullable parents. External blobs remain URI/range references under the destination's existing authorization policy: a destination that independently registers the matching external base accepts the copy, while an unregistered destination rejects it without committing. Source base registrations are not transferred, and `allow_external_blob_outside_bases` is not relaxed. JOIN, UNION, aggregation, and computed descriptors are rejected during planning. This addresses the DataFusion cross-table copy portion of OSS-696; remote/cloud execution and Realtime/MemWAL remain outside this support boundary. Local acceptance covers 15 passing DataFusion Blob integration cases, including byte reads after physical source-directory deletion, authorized and unauthorized External ranges, and a source-read failure after an earlier partition has written an uncommitted data file, with both the live destination and an independently reopened manifest unchanged. Existing regression evidence also includes 195 DataFusion unit tests and 59 Blob integration tests. Rust formatting and whitespace checks passed. Remote-feature builds, cloud/S3 execution, clippy, and the full workspace suite were not run.
This commit is contained in:
Xuanwo
@@ -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<dyn BaseTable>,
|
||||
@@ -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<dyn ExecutionPlan> = Arc::new(MetadataEraserExec::new(plan));
|
||||
let mut plan: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan>,
|
||||
insert_op: InsertOp,
|
||||
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
||||
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,
|
||||
|
||||
@@ -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<dyn ExecutionPlan>) -> bool {
|
||||
input.downcast_ref::<BlobScanExec>().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<dyn ExecutionPlan>) -> bool {
|
||||
input.downcast_ref::<ProjectionExec>().is_some()
|
||||
|| input.downcast_ref::<FilterExec>().is_some()
|
||||
|| input.downcast_ref::<SortExec>().is_some()
|
||||
|| input.downcast_ref::<SortPreservingMergeExec>().is_some()
|
||||
|| input.downcast_ref::<GlobalLimitExec>().is_some()
|
||||
|| input.downcast_ref::<LocalLimitExec>().is_some()
|
||||
|| input.downcast_ref::<CoalescePartitionsExec>().is_some()
|
||||
|| input
|
||||
.downcast_ref::<datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec>()
|
||||
.is_some()
|
||||
|| input.downcast_ref::<RepartitionExec>().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<dyn ExecutionPlan>,
|
||||
source: Arc<dyn BaseTable>,
|
||||
dataset: Arc<Dataset>,
|
||||
query: QueryRequest,
|
||||
options: QueryExecutionOptions,
|
||||
}
|
||||
|
||||
impl BlobScanExec {
|
||||
pub(super) fn new(
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
source: Arc<dyn BaseTable>,
|
||||
dataset: Arc<Dataset>,
|
||||
query: QueryRequest,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Self {
|
||||
Self {
|
||||
input,
|
||||
source,
|
||||
dataset,
|
||||
query,
|
||||
options,
|
||||
}
|
||||
}
|
||||
|
||||
async fn with_row_address(&self) -> Result<Provenance> {
|
||||
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<dyn ExecutionPlan> = Arc::new(MetadataEraserExec::new(plan));
|
||||
let schema = plan.schema();
|
||||
let mut projection: Vec<(Arc<dyn PhysicalExpr>, 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::<Result<_>>()?;
|
||||
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<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
|
||||
) -> Result<TreeNodeRecursion> {
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
}
|
||||
fn name(&self) -> &str {
|
||||
"BlobScanExec"
|
||||
}
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
self.input.properties()
|
||||
}
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![&self.input]
|
||||
}
|
||||
fn maintains_input_order(&self) -> Vec<bool> {
|
||||
vec![true]
|
||||
}
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
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<TaskContext>,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
self.input.execute(partition, context)
|
||||
}
|
||||
}
|
||||
|
||||
struct Provenance {
|
||||
plan: Arc<dyn ExecutionPlan>,
|
||||
dataset: Arc<Dataset>,
|
||||
// Source field path for each output column, before the hidden address.
|
||||
fields: Vec<Option<String>>,
|
||||
}
|
||||
|
||||
fn source_field(expr: &Arc<dyn PhysicalExpr>, fields: &[Option<String>]) -> Option<String> {
|
||||
if let Some(column) = expr.downcast_ref::<Column>() {
|
||||
return fields.get(column.index()).cloned().flatten();
|
||||
}
|
||||
let call = expr.downcast_ref::<ScalarFunctionExpr>()?;
|
||||
call.fun().inner().downcast_ref::<GetFieldFunc>()?;
|
||||
let (root, path) = call.args().split_first()?;
|
||||
let mut source = source_field(root, fields)?;
|
||||
for component in path {
|
||||
let name = component
|
||||
.downcast_ref::<Literal>()?
|
||||
.value()
|
||||
.try_as_str()
|
||||
.flatten()?;
|
||||
source.push('.');
|
||||
source.push_str(name);
|
||||
}
|
||||
Some(source)
|
||||
}
|
||||
|
||||
fn with_provenance(input: Arc<dyn ExecutionPlan>) -> BoxFuture<'static, Result<Provenance>> {
|
||||
async move {
|
||||
if let Some(scan) = input.downcast_ref::<BlobScanExec>() {
|
||||
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::<ProjectionExec>() {
|
||||
let fields = projection
|
||||
.expr()
|
||||
.iter()
|
||||
.map(|expr| source_field(&expr.expr, &provenance.fields))
|
||||
.collect();
|
||||
let mut expressions: Vec<(Arc<dyn PhysicalExpr>, 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<dyn ExecutionPlan>,
|
||||
target_schema: &Schema,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
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::<ProjectionExec>() {
|
||||
let mut changed = false;
|
||||
let expressions = projection
|
||||
.expr()
|
||||
.iter()
|
||||
.map(|projected| {
|
||||
let mut expr = projected.expr.clone();
|
||||
if let Some(cast) = expr.downcast_ref::<CastExpr>()
|
||||
&& 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::<Result<Vec<_>>>()?;
|
||||
if changed {
|
||||
Arc::new(ProjectionExec::try_new(
|
||||
expressions,
|
||||
projection.input().clone(),
|
||||
)?) as Arc<dyn ExecutionPlan>
|
||||
} else {
|
||||
input
|
||||
}
|
||||
} else {
|
||||
input
|
||||
};
|
||||
let rewrite_columns: Vec<bool> = 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<dyn PhysicalExpr>, 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<dyn ExecutionPlan> =
|
||||
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::<Vec<_>>(),
|
||||
));
|
||||
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<dyn ExecutionPlan>,
|
||||
dataset: Arc<Dataset>,
|
||||
rewrite: Arc<BlobV2BatchRewritePlan>,
|
||||
schema: SchemaRef,
|
||||
properties: Arc<PlanProperties>,
|
||||
}
|
||||
|
||||
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<dyn ExecutionPlan>,
|
||||
dataset: Arc<Dataset>,
|
||||
rewrite: Arc<BlobV2BatchRewritePlan>,
|
||||
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<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
|
||||
) -> Result<TreeNodeRecursion> {
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
}
|
||||
fn name(&self) -> &str {
|
||||
"MaterializeBlobsExec"
|
||||
}
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![&self.input]
|
||||
}
|
||||
fn maintains_input_order(&self) -> Vec<bool> {
|
||||
vec![true]
|
||||
}
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
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<TaskContext>,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
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,
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -195,6 +195,7 @@ impl ExecutionPlan for InsertExec {
|
||||
context: Arc<TaskContext>,
|
||||
) -> DataFusionResult<SendableRecordBatchStream> {
|
||||
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()
|
||||
{
|
||||
|
||||
@@ -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<Option<&[u8]>> = 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::<arrow_array::types::UInt64Type>()
|
||||
.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::<arrow_array::types::UInt8Type>();
|
||||
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::<arrow_array::types::UInt64Type>()
|
||||
.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::<i32>()
|
||||
.value(2),
|
||||
"c"
|
||||
);
|
||||
let rank = extra.column_by_name("rank").unwrap();
|
||||
if difference == "scalar" {
|
||||
assert_eq!(
|
||||
rank.as_primitive::<arrow_array::types::Int64Type>()
|
||||
.value(2),
|
||||
3
|
||||
);
|
||||
} else {
|
||||
assert_eq!(
|
||||
rank.as_primitive::<arrow_array::types::Int32Type>()
|
||||
.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::<arrow_array::types::UInt32Type>()
|
||||
.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::<arrow_array::types::UInt8Type>()
|
||||
.value(0),
|
||||
BlobKind::External as u8
|
||||
);
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.column(1)
|
||||
.as_primitive::<arrow_array::types::UInt32Type>()
|
||||
.value(0),
|
||||
11
|
||||
);
|
||||
assert_eq!(
|
||||
descriptor.column(2).as_string::<i32>().value(0),
|
||||
"payload.bin"
|
||||
);
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.column(3)
|
||||
.as_primitive::<arrow_array::types::UInt64Type>()
|
||||
.value(0),
|
||||
4
|
||||
);
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.column(4)
|
||||
.as_primitive::<arrow_array::types::UInt64Type>()
|
||||
.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<RecordBatch> = insert
|
||||
.execute(0, ctx.task_ctx())
|
||||
.unwrap()
|
||||
.try_collect()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
batches[0]
|
||||
.column(0)
|
||||
.as_primitive::<arrow_array::types::UInt64Type>()
|
||||
.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::<Vec<_>>()
|
||||
.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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user