mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-28 22:46:14 +00:00
fix: preserve null blob descriptor fields in SQL (#4571)
SQL descriptor access returns NULL for every field of a NULL Blob v2 value, so `payload.size = 0` matches only empty blobs and `payload.size IS NULL` matches NULL blobs. DataFusion 54.1.0 extracts struct children without combining their validity with the parent struct. Lance descriptors keep nullness in the parent and can store zero-filled children beneath NULL blobs. The adapter derives the descriptor schema from Lance's default projection, marks descriptor children nullable, and masks them before SQL evaluation. Schema derivation requires no query request, and the mask handles struct and list nesting while preserving existing child nulls and value buffers without fetching blob bytes. Filter classification follows the referenced struct field path. Blob predicates stay above the mask; scalar siblings such as `extra.note` retain exact filter pushdown even when `extra.clip` is a blob. **Merge hold:** as requested in the [review discussion](https://github.com/lancedb/lancedb/pull/4571#issuecomment-6253443077), blob SQL INSERT compatibility remains pending a separate write-mode design. DataFusion uses the provider's single schema for reads and INSERT coercion, so the descriptor schema still rejects blob write structs. This follow-up does not restore those INSERTs; the merge should wait until the write contract is settled. Validation: - Reproduced the original NULL blob returning `kind = 0` before enabling the mask. - Reproduced both the extra remote query and lost nested sibling pushdown before the follow-up fixes. - All 441 DataFusion adapter tests passed after merging main at `e6789a41`, including SQL projections, filters, mixed predicates, LIMIT, COUNT, schema agreement, sliced arrays, nested structs, list variants, and EXPLAIN assertions for scalar sibling pushdown. The blob fixtures use the current four-field write format, and NULL-blob filtering is also covered with row IDs and row addresses enabled. - Remote constructor regression passed: one describe request and no query request, with nullable descriptor fields preserved. - `cargo fmt --all` - `cargo clippy --locked --profile ci -p lancedb --features remote --tests --examples -- -D warnings` Fixes #4569 <!-- lance-gatekeeper-fix:v1 agent=2ca7389efd48d11a8debe9111abeb00e generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
@@ -5687,6 +5687,31 @@ mod tests {
|
||||
schema_describe_response(&schema)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_adapter_construction_only_describes_blob_schema() {
|
||||
use datafusion_catalog::TableProvider;
|
||||
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let observed = requests.clone();
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
observed.fetch_add(1, Ordering::SeqCst);
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/describe/");
|
||||
blob_describe_response()
|
||||
});
|
||||
|
||||
let adapter =
|
||||
crate::table::datafusion::BaseTableAdapter::try_new(table.base_table().clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(requests.load(Ordering::SeqCst), 1);
|
||||
let schema = adapter.schema();
|
||||
let DataType::Struct(fields) = schema.field_with_name("image").unwrap().data_type() else {
|
||||
panic!("expected a blob descriptor");
|
||||
};
|
||||
assert!(fields.find("size").unwrap().1.is_nullable());
|
||||
assert!(fields.find("data").is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(semver::Version::new(0, 1, 0))]
|
||||
#[case(semver::Version::new(0, 5, 0))]
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
//! This module contains adapters to allow LanceDB tables to be used as DataFusion table providers.
|
||||
|
||||
mod blob_coerce;
|
||||
mod blob_nulls;
|
||||
pub mod cast;
|
||||
pub mod insert;
|
||||
pub mod reject_nan;
|
||||
@@ -25,7 +26,10 @@ use datafusion_physical_plan::{
|
||||
};
|
||||
use futures::{TryFutureExt, TryStreamExt};
|
||||
use lance::dataset::{WriteMode, WriteParams};
|
||||
use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD};
|
||||
use lance_core::{
|
||||
ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD,
|
||||
datatypes::{Projection, Schema as LanceSchema},
|
||||
};
|
||||
|
||||
use super::{AnyQuery, BaseTable};
|
||||
use crate::{
|
||||
@@ -171,13 +175,14 @@ pub struct BaseTableAdapter {
|
||||
}
|
||||
|
||||
impl BaseTableAdapter {
|
||||
/// Create an adapter whose schema matches default scans, including blob descriptors.
|
||||
///
|
||||
/// Deriving the descriptor schema does not execute a query or read blob payloads.
|
||||
pub async fn try_new(table: Arc<dyn BaseTable>) -> Result<Self> {
|
||||
let schema = table
|
||||
.schema()
|
||||
.await?
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_metadata(HashMap::default());
|
||||
let write_schema = table.schema().await?;
|
||||
let schema = Projection::full(Arc::new(LanceSchema::try_from(write_schema.as_ref())?))
|
||||
.to_arrow_schema();
|
||||
let schema = blob_nulls::read_schema(&schema).with_metadata(HashMap::default());
|
||||
|
||||
Ok(Self {
|
||||
table,
|
||||
@@ -317,6 +322,7 @@ impl TableProvider for BaseTableAdapter {
|
||||
.create_plan(&AnyQuery::Query(query), options)
|
||||
.map_err(|err| DataFusionError::External(err.into()))
|
||||
.await?;
|
||||
let plan = blob_nulls::mask_plan(plan)?;
|
||||
Ok(Arc::new(MetadataEraserExec::new(plan)))
|
||||
}
|
||||
|
||||
@@ -324,7 +330,21 @@ impl TableProvider for BaseTableAdapter {
|
||||
&self,
|
||||
filters: &[&Expr],
|
||||
) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
|
||||
Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
|
||||
// Lance evaluates pushed filters before the descriptor validity mask.
|
||||
// Keep predicates touching blobs above that mask, including IS NULL:
|
||||
// the stored descriptor's children are not independently nullable.
|
||||
filters
|
||||
.iter()
|
||||
.map(|filter| {
|
||||
Ok(
|
||||
if blob_nulls::references_descriptor(filter, &self.schema)? {
|
||||
TableProviderFilterPushDown::Unsupported
|
||||
} else {
|
||||
TableProviderFilterPushDown::Exact
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn statistics(&self) -> Option<Statistics> {
|
||||
@@ -562,6 +582,287 @@ pub mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn inline_blob_input(values: Vec<Option<&[u8]>>) -> arrow_array::ArrayRef {
|
||||
use arrow_array::{Array, LargeBinaryArray, StructArray, UInt64Array};
|
||||
|
||||
let data = LargeBinaryArray::from(values);
|
||||
let len = data.len();
|
||||
let nulls = data.nulls().cloned();
|
||||
let DataType::Struct(fields) = crate::blob("payload", true).data_type().clone() else {
|
||||
unreachable!()
|
||||
};
|
||||
Arc::new(StructArray::new(
|
||||
fields,
|
||||
vec![
|
||||
Arc::new(data),
|
||||
Arc::new(StringArray::new_null(len)),
|
||||
Arc::new(UInt64Array::new_null(len)),
|
||||
Arc::new(UInt64Array::new_null(len)),
|
||||
],
|
||||
nulls,
|
||||
))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sql_blob_descriptor_nulls() {
|
||||
use arrow_array::{UInt8Array, UInt64Array};
|
||||
use arrow_select::concat::concat_batches;
|
||||
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let blobs = inline_blob_input(vec![Some(b"hello"), Some(b""), None]);
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
crate::blob("payload", true),
|
||||
])),
|
||||
vec![Arc::new(Int32Array::from(vec![1, 2, 3])), blobs],
|
||||
)
|
||||
.unwrap();
|
||||
let table = db.create_table("docs", batch).execute().await.unwrap();
|
||||
let adapter = Arc::new(
|
||||
BaseTableAdapter::try_new(table.base_table().clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(2));
|
||||
ctx.register_table("docs", adapter.clone()).unwrap();
|
||||
|
||||
let batches = ctx
|
||||
.sql("SELECT id, payload.kind, payload.position, payload.size, payload.blob_id, payload.blob_uri FROM docs ORDER BY id")
|
||||
.await.unwrap().collect().await.unwrap();
|
||||
let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
|
||||
assert_eq!(batch.num_rows(), 3);
|
||||
assert_eq!(
|
||||
batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<UInt8Array>()
|
||||
.unwrap(),
|
||||
&UInt8Array::from(vec![Some(0), Some(0), None])
|
||||
);
|
||||
assert_eq!(
|
||||
batch
|
||||
.column(3)
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap(),
|
||||
&UInt64Array::from(vec![Some(5), Some(0), None])
|
||||
);
|
||||
for column in &batch.columns()[1..] {
|
||||
assert!(!column.is_null(0));
|
||||
assert!(!column.is_null(1));
|
||||
assert!(column.is_null(2));
|
||||
}
|
||||
|
||||
for (predicate, expected) in [
|
||||
("payload.size = 0", vec![2]),
|
||||
("payload.kind = 0", vec![1, 2]),
|
||||
("payload.blob_uri = ''", vec![1, 2]),
|
||||
("payload.size IS NULL", vec![3]),
|
||||
("payload.size IS NOT NULL", vec![1, 2]),
|
||||
("payload.kind IS NULL", vec![3]),
|
||||
("payload.position IS NULL", vec![3]),
|
||||
("payload.blob_id IS NULL", vec![3]),
|
||||
("payload.blob_uri IS NULL", vec![3]),
|
||||
("payload IS NULL", vec![3]),
|
||||
("payload.size = 0 OR payload.size IS NULL", vec![2, 3]),
|
||||
("id > 1 AND payload.size = 0", vec![2]),
|
||||
] {
|
||||
let batches = ctx
|
||||
.sql(&format!(
|
||||
"SELECT id FROM docs WHERE {predicate} ORDER BY id"
|
||||
))
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
let ids: Vec<i32> = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.values()
|
||||
.to_vec()
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(ids, expected, "{predicate}");
|
||||
}
|
||||
|
||||
let batches = ctx
|
||||
.sql("SELECT id FROM docs WHERE payload.size IS NULL LIMIT 1")
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
batches[0]
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.value(0),
|
||||
3
|
||||
);
|
||||
let batches = ctx
|
||||
.sql("SELECT COUNT(*), COUNT(payload.size) FROM docs")
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
batches[0]
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap()
|
||||
.value(0),
|
||||
3
|
||||
);
|
||||
assert_eq!(
|
||||
batches[0]
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap()
|
||||
.value(0),
|
||||
2
|
||||
);
|
||||
assert!(ctx.sql("SELECT payload.data FROM docs").await.is_err());
|
||||
|
||||
// Whole-struct scans must agree with the declared schema too.
|
||||
let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap();
|
||||
assert_eq!(plan.schema(), adapter.schema());
|
||||
|
||||
// Opted-in system columns must compose with the descriptor schema and mask.
|
||||
let adapter = BaseTableAdapter::try_new(table.base_table().clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.with_system_columns(&[ROW_ID, ROW_ADDR])
|
||||
.unwrap();
|
||||
let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap();
|
||||
assert_eq!(plan.schema(), adapter.schema());
|
||||
let expected = datafusion_physical_plan::collect(plan, ctx.task_ctx())
|
||||
.await
|
||||
.unwrap();
|
||||
let expected = concat_batches(&adapter.schema(), &expected).unwrap();
|
||||
ctx.register_table("docs_with_system_columns", Arc::new(adapter))
|
||||
.unwrap();
|
||||
let batches = ctx
|
||||
.sql("SELECT _rowaddr, _rowid, id FROM docs_with_system_columns WHERE payload.size IS NULL")
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
let actual = concat_batches(&batches[0].schema(), &batches).unwrap();
|
||||
let columns =
|
||||
[ROW_ADDR, ROW_ID, "id"].map(|name| expected.schema().index_of(name).unwrap());
|
||||
assert_eq!(actual, expected.project(&columns).unwrap().slice(2, 1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_nested_blob_sibling_filter_pushdown() {
|
||||
use arrow_array::{Array, StructArray};
|
||||
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let blobs = inline_blob_input(vec![Some(b"hello"), None, Some(b"")]);
|
||||
let fields = vec![
|
||||
Arc::new(Field::new("note", DataType::Utf8, false)),
|
||||
Arc::new(crate::blob("clip", true)),
|
||||
Arc::new(Field::new("note.with.dot", DataType::Utf8, false)),
|
||||
];
|
||||
let extra = StructArray::new(
|
||||
fields.into(),
|
||||
vec![
|
||||
Arc::new(StringArray::from(vec!["note 7", "note 8", "note 9"])),
|
||||
blobs,
|
||||
Arc::new(StringArray::from(vec!["a", "b", "c"])),
|
||||
],
|
||||
None,
|
||||
);
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
Field::new("extra", extra.data_type().clone(), false),
|
||||
])),
|
||||
vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(extra)],
|
||||
)
|
||||
.unwrap();
|
||||
let table = db.create_table("docs", batch).execute().await.unwrap();
|
||||
let adapter = Arc::new(
|
||||
BaseTableAdapter::try_new(table.base_table().clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let ctx = SessionContext::new();
|
||||
ctx.register_table("docs", adapter).unwrap();
|
||||
|
||||
for (predicate, pushed, expected) in [
|
||||
("extra.note = 'note 7'", true, vec![1]),
|
||||
("extra.\"note.with.dot\" = 'b'", true, vec![2]),
|
||||
("extra.clip.size IS NULL", false, vec![2]),
|
||||
("get_field(extra, 'clip', 'size') = 0", false, vec![3]),
|
||||
("extra.clip IS NULL", false, vec![2]),
|
||||
(
|
||||
"extra.note = 'note 7' OR extra.clip.size IS NULL",
|
||||
false,
|
||||
vec![1, 2],
|
||||
),
|
||||
] {
|
||||
let sql = format!("SELECT id FROM docs WHERE {predicate}");
|
||||
let explain = ctx
|
||||
.sql(&format!("EXPLAIN {sql}"))
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
let explain = explain
|
||||
.iter()
|
||||
.flat_map(|batch| batch.column(1).as_string::<i32>().iter())
|
||||
.flatten()
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
assert_eq!(explain.contains("FilterExec"), !pushed, "{explain}");
|
||||
if pushed {
|
||||
assert!(explain.contains("full_filters=["), "{explain}");
|
||||
}
|
||||
let batches = ctx
|
||||
.sql(&format!("{sql} ORDER BY id"))
|
||||
.await
|
||||
.unwrap()
|
||||
.collect()
|
||||
.await
|
||||
.unwrap();
|
||||
let ids = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(0)
|
||||
.as_primitive::<arrow_array::types::Int32Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.copied()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(ids, expected, "{predicate}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_system_columns_opt_in() {
|
||||
let fixture = TestFixture::new().await;
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Apply Blob v2 struct validity before DataFusion extracts descriptor fields.
|
||||
|
||||
use std::fmt;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, RecordBatch, StructArray,
|
||||
cast::AsArray, make_array,
|
||||
};
|
||||
use arrow_buffer::NullBuffer;
|
||||
use arrow_schema::{DataType, Field, FieldRef, Schema};
|
||||
use datafusion_common::{
|
||||
Result,
|
||||
tree_node::{TreeNode, TreeNodeRecursion},
|
||||
};
|
||||
use datafusion_expr::{ColumnarValue, Expr};
|
||||
use datafusion_functions::core::getfield::GetFieldFunc;
|
||||
use datafusion_physical_expr::{PhysicalExpr, expressions::Column};
|
||||
use datafusion_physical_plan::{ExecutionPlan, projection::ProjectionExec};
|
||||
use lance_arrow::FieldExt;
|
||||
use lance_core::datatypes::BlobV2Layout;
|
||||
|
||||
fn is_descriptor(field: &Field) -> bool {
|
||||
field.is_blob()
|
||||
&& matches!(field.data_type(), DataType::Struct(fields)
|
||||
if BlobV2Layout::classify(fields) == Some(BlobV2Layout::Descriptor))
|
||||
}
|
||||
|
||||
pub(super) fn has_descriptor(field: &Field) -> bool {
|
||||
if is_descriptor(field) {
|
||||
return true;
|
||||
}
|
||||
match field.data_type() {
|
||||
DataType::Struct(fields) => fields.iter().any(|field| has_descriptor(field)),
|
||||
DataType::List(item) | DataType::LargeList(item) | DataType::FixedSizeList(item, _) => {
|
||||
has_descriptor(item)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve static struct accesses without confusing a scalar sibling with the
|
||||
/// entire root column. Track descriptor ancestors too: `blob.size` is scalar,
|
||||
/// but still requires the descriptor validity mask before evaluation.
|
||||
fn referenced_field<'a>(expr: &Expr, schema: &'a Schema) -> Option<(&'a Field, bool)> {
|
||||
match expr {
|
||||
Expr::Column(column) => Some((schema.field_with_name(&column.name).ok()?, false)),
|
||||
Expr::ScalarFunction(call)
|
||||
if call.func.inner().downcast_ref::<GetFieldFunc>().is_some() =>
|
||||
{
|
||||
let (base, path) = call.args.split_first()?;
|
||||
let (mut field, mut within_descriptor) = referenced_field(base, schema)?;
|
||||
for name in path {
|
||||
let Expr::Literal(name, _) = name else {
|
||||
return None;
|
||||
};
|
||||
within_descriptor |= is_descriptor(field);
|
||||
let DataType::Struct(fields) = field.data_type() else {
|
||||
return None;
|
||||
};
|
||||
field = fields.find(name.try_as_str().flatten()?)?.1;
|
||||
}
|
||||
Some((field, within_descriptor))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn references_descriptor(expr: &Expr, schema: &Schema) -> Result<bool> {
|
||||
let mut found = false;
|
||||
expr.apply(|expr| {
|
||||
if let Some((field, within_descriptor)) = referenced_field(expr, schema) {
|
||||
if within_descriptor || has_descriptor(field) {
|
||||
found = true;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
}
|
||||
// The complete path is safe. Do not revisit its root as a whole
|
||||
// struct, which could contain an unrelated blob sibling.
|
||||
return Ok(TreeNodeRecursion::Jump);
|
||||
}
|
||||
// Unknown/computed paths conservatively visit their arguments.
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
})?;
|
||||
Ok(found)
|
||||
}
|
||||
|
||||
fn nullable_descriptor_fields(field: &FieldRef) -> FieldRef {
|
||||
let data_type = match field.data_type() {
|
||||
DataType::Struct(fields) if is_descriptor(field) => DataType::Struct(
|
||||
fields
|
||||
.iter()
|
||||
.map(|field| Arc::new(field.as_ref().clone().with_nullable(true)))
|
||||
.collect(),
|
||||
),
|
||||
DataType::Struct(fields) => {
|
||||
DataType::Struct(fields.iter().map(nullable_descriptor_fields).collect())
|
||||
}
|
||||
DataType::List(item) => DataType::List(nullable_descriptor_fields(item)),
|
||||
DataType::LargeList(item) => DataType::LargeList(nullable_descriptor_fields(item)),
|
||||
DataType::FixedSizeList(item, size) => {
|
||||
DataType::FixedSizeList(nullable_descriptor_fields(item), *size)
|
||||
}
|
||||
_ => return field.clone(),
|
||||
};
|
||||
Arc::new(field.as_ref().clone().with_data_type(data_type))
|
||||
}
|
||||
|
||||
pub(super) fn read_schema(schema: &Schema) -> Schema {
|
||||
Schema::new_with_metadata(
|
||||
schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(nullable_descriptor_fields)
|
||||
.collect::<Vec<_>>(),
|
||||
schema.metadata().clone(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Preserve value buffers, offsets, metadata, and existing child nulls. Only
|
||||
/// validity buffers are combined; this never materializes blob payloads.
|
||||
fn mask_array(array: &ArrayRef, field: &Field, ancestors: Option<&NullBuffer>) -> Result<ArrayRef> {
|
||||
if !has_descriptor(field) {
|
||||
return Ok(array.clone());
|
||||
}
|
||||
match field.data_type() {
|
||||
DataType::Struct(fields) => {
|
||||
let values = array.as_struct();
|
||||
let nulls = NullBuffer::union(values.nulls(), ancestors);
|
||||
let columns = if is_descriptor(field) {
|
||||
values
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|column| {
|
||||
let nulls = NullBuffer::union(column.nulls(), nulls.as_ref());
|
||||
Ok(make_array(
|
||||
column.to_data().into_builder().nulls(nulls).build()?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
} else {
|
||||
values
|
||||
.columns()
|
||||
.iter()
|
||||
.zip(fields)
|
||||
.map(|(column, field)| mask_array(column, field, nulls.as_ref()))
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
};
|
||||
Ok(Arc::new(StructArray::try_new_with_length(
|
||||
fields.clone(),
|
||||
columns,
|
||||
nulls,
|
||||
values.len(),
|
||||
)?))
|
||||
}
|
||||
DataType::List(item) => {
|
||||
let list = array.as_list::<i32>();
|
||||
Ok(Arc::new(ListArray::try_new(
|
||||
item.clone(),
|
||||
list.offsets().clone(),
|
||||
mask_array(list.values(), item, None)?,
|
||||
list.nulls().cloned(),
|
||||
)?))
|
||||
}
|
||||
DataType::LargeList(item) => {
|
||||
let list = array.as_list::<i64>();
|
||||
Ok(Arc::new(LargeListArray::try_new(
|
||||
item.clone(),
|
||||
list.offsets().clone(),
|
||||
mask_array(list.values(), item, None)?,
|
||||
list.nulls().cloned(),
|
||||
)?))
|
||||
}
|
||||
DataType::FixedSizeList(item, size) => {
|
||||
let list = array.as_fixed_size_list();
|
||||
Ok(Arc::new(FixedSizeListArray::try_new(
|
||||
item.clone(),
|
||||
*size,
|
||||
mask_array(list.values(), item, None)?,
|
||||
list.nulls().cloned(),
|
||||
)?))
|
||||
}
|
||||
_ => Ok(array.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MaskBlobNulls {
|
||||
source: Arc<dyn PhysicalExpr>,
|
||||
field: FieldRef,
|
||||
}
|
||||
|
||||
impl fmt::Display for MaskBlobNulls {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "mask_blob_nulls({})", self.source)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for MaskBlobNulls {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.source.eq(&other.source) && self.field == other.field
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for MaskBlobNulls {}
|
||||
|
||||
impl Hash for MaskBlobNulls {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.source.hash(state);
|
||||
self.field.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
impl PhysicalExpr for MaskBlobNulls {
|
||||
fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
|
||||
Ok(self.field.clone())
|
||||
}
|
||||
|
||||
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
|
||||
let source = self.source.evaluate(batch)?.into_array(batch.num_rows())?;
|
||||
Ok(ColumnarValue::Array(mask_array(
|
||||
&source,
|
||||
&self.field,
|
||||
None,
|
||||
)?))
|
||||
}
|
||||
|
||||
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
|
||||
vec![&self.source]
|
||||
}
|
||||
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn PhysicalExpr>>,
|
||||
) -> Result<Arc<dyn PhysicalExpr>> {
|
||||
assert_eq!(children.len(), 1);
|
||||
Ok(Arc::new(Self {
|
||||
source: children[0].clone(),
|
||||
field: self.field.clone(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{self}")
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn mask_plan(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let schema = input.schema();
|
||||
if !schema.fields().iter().any(|field| has_descriptor(field)) {
|
||||
return Ok(input);
|
||||
}
|
||||
let expressions = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, field)| {
|
||||
let source: Arc<dyn PhysicalExpr> = Arc::new(Column::new(field.name(), index));
|
||||
let expression: Arc<dyn PhysicalExpr> = if has_descriptor(field) {
|
||||
Arc::new(MaskBlobNulls {
|
||||
source,
|
||||
field: nullable_descriptor_fields(field),
|
||||
})
|
||||
} else {
|
||||
source
|
||||
};
|
||||
(expression, field.name().clone())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(Arc::new(ProjectionExec::try_new(expressions, input)?))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use arrow_array::{StringArray, UInt8Array, UInt32Array, UInt64Array};
|
||||
use arrow_buffer::OffsetBuffer;
|
||||
use lance_core::datatypes::BLOB_V2_DESC_FIELD;
|
||||
|
||||
fn descriptor() -> (FieldRef, ArrayRef) {
|
||||
let field = Arc::new(BLOB_V2_DESC_FIELD.clone().with_nullable(true));
|
||||
let field = nullable_descriptor_fields(&field);
|
||||
let DataType::Struct(fields) = field.data_type() else {
|
||||
unreachable!()
|
||||
};
|
||||
let array = StructArray::new(
|
||||
fields.clone(),
|
||||
vec![
|
||||
Arc::new(UInt8Array::from(vec![0; 5])),
|
||||
Arc::new(UInt64Array::from(vec![10; 5])),
|
||||
Arc::new(UInt64Array::from(vec![
|
||||
Some(9),
|
||||
Some(5),
|
||||
Some(0),
|
||||
Some(0),
|
||||
None,
|
||||
])),
|
||||
Arc::new(UInt32Array::from(vec![0; 5])),
|
||||
Arc::new(StringArray::from(vec![""; 5])),
|
||||
],
|
||||
Some(NullBuffer::from(vec![true, true, true, false, true])),
|
||||
);
|
||||
(field, Arc::new(array))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sliced_descriptor_preserves_child_nulls_and_buffers() {
|
||||
let (field, array) = descriptor();
|
||||
let array = array.slice(1, 4);
|
||||
let masked = mask_array(&array, &field, None).unwrap();
|
||||
let masked = masked.as_struct();
|
||||
let size = masked
|
||||
.column_by_name("size")
|
||||
.unwrap()
|
||||
.as_primitive::<arrow_array::types::UInt64Type>();
|
||||
assert_eq!(size, &UInt64Array::from(vec![Some(5), Some(0), None, None]));
|
||||
assert_eq!(masked.nulls(), array.nulls());
|
||||
for (input, output) in array.as_struct().columns().iter().zip(masked.columns()) {
|
||||
assert_eq!(input.to_data().buffers(), output.to_data().buffers());
|
||||
assert!(output.is_null(2));
|
||||
}
|
||||
|
||||
let empty = mask_array(&array.slice(0, 0), &field, None).unwrap();
|
||||
assert_eq!(empty.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nested_descriptor_inherits_struct_nulls() {
|
||||
let (blob_field, array) = descriptor();
|
||||
let parent = Arc::new(Field::new(
|
||||
"parent",
|
||||
DataType::Struct(vec![blob_field.clone()].into()),
|
||||
true,
|
||||
));
|
||||
let array: ArrayRef = Arc::new(StructArray::new(
|
||||
vec![blob_field].into(),
|
||||
vec![array],
|
||||
Some(NullBuffer::from(vec![true, false, true, true, true])),
|
||||
));
|
||||
let array = array.slice(1, 3);
|
||||
let masked = mask_array(&array, &parent, None).unwrap();
|
||||
let blob = masked.as_struct().column(0).as_struct();
|
||||
assert_eq!(
|
||||
blob.nulls().unwrap(),
|
||||
&NullBuffer::from(vec![false, true, false])
|
||||
);
|
||||
for column in blob.columns() {
|
||||
assert!(column.is_null(0));
|
||||
assert!(!column.is_null(1));
|
||||
assert!(column.is_null(2));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_descriptor_masks_preserve_offsets_and_validity() {
|
||||
let (item, values) = descriptor();
|
||||
let nulls = Some(NullBuffer::from(vec![true, false, true]));
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(ListArray::new(
|
||||
item.clone(),
|
||||
OffsetBuffer::new(vec![0, 1, 3, 5].into()),
|
||||
values.clone(),
|
||||
nulls.clone(),
|
||||
)),
|
||||
Arc::new(LargeListArray::new(
|
||||
item.clone(),
|
||||
OffsetBuffer::new(vec![0_i64, 1, 3, 5].into()),
|
||||
values.clone(),
|
||||
nulls,
|
||||
)),
|
||||
Arc::new(FixedSizeListArray::new(item.clone(), 1, values, None)),
|
||||
];
|
||||
for array in arrays {
|
||||
let field = Field::new("blobs", array.data_type().clone(), true);
|
||||
let masked = mask_array(&array, &field, None).unwrap();
|
||||
assert_eq!(masked.nulls(), array.nulls());
|
||||
assert_eq!(masked.to_data().buffers(), array.to_data().buffers());
|
||||
let values = match masked.data_type() {
|
||||
DataType::List(_) => masked.as_list::<i32>().values(),
|
||||
DataType::LargeList(_) => masked.as_list::<i64>().values(),
|
||||
DataType::FixedSizeList(_, _) => masked.as_fixed_size_list().values(),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
for column in values.as_struct().columns() {
|
||||
assert!(column.is_null(3));
|
||||
}
|
||||
assert_eq!(
|
||||
mask_array(&array.slice(1, 1), &field, None).unwrap().len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user