use std::collections::{HashMap, HashSet};
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use chrono::{DateTime, Utc};
use crate::Result;
use crate::column_rename::ColumnRenameExec;
use crate::delete_filter::DeleteFilterExec;
use crate::inlined_filter::translate_inlined_filters;
use crate::metadata_provider::{
DuckLakeFileColumnStatistics, DuckLakeFileData, DuckLakeFileMetadata, DuckLakeNameMapping,
DuckLakeStatistics, DuckLakeTableColumn, DuckLakeTableColumnStatistics, DuckLakeTableField,
DuckLakeTableFile, FILE_METADATA_BATCH_SIZE, INLINED_DATA_REMEDIATION, MetadataProvider,
};
use crate::nan_pruning_barrier::NanPruningBarrierExec;
use crate::partition::PartitionSpec;
use crate::path_resolver::resolve_path;
use crate::row_id::{
ROW_ID_PARQUET_FIELD_ID, ROWID_COLUMN_NAME, RowIdExec, SNAPSHOT_ID_PARQUET_FIELD_ID,
positional_table_schema, positional_table_schema_reserving, rowid_field,
};
use crate::snapshot_filter::SnapshotFilterExec;
use crate::stats_filter::{self, StatsFilter};
use crate::types::{
DuckLakeDefaultExprAdapterFactory, INITIAL_DEFAULT_METADATA_KEY,
build_arrow_schema_from_fields, build_read_schema_with_field_id_mapping,
build_read_schema_with_field_id_mapping_from_schema, build_read_schema_with_name_mapping,
ducklake_to_arrow_type, extract_parquet_field_ids, parse_ducklake_default_scalar,
parse_ducklake_scalar,
};
#[cfg(feature = "write")]
use crate::delete_exec::DuckLakeDeleteExec;
#[cfg(feature = "write")]
use crate::insert_exec::DuckLakeInsertExec;
#[cfg(feature = "write")]
use crate::metadata_writer::{MetadataWriter, WriteMode};
#[cfg(feature = "write")]
use crate::update_exec::DuckLakeUpdateExec;
use datafusion::common::DFSchema;
use datafusion::common::pruning::{PrunableStatistics, PruningStatistics};
#[cfg(feature = "write")]
use datafusion::logical_expr::Operator;
use datafusion::physical_expr::PhysicalExpr;
#[cfg(feature = "write")]
use datafusion::physical_expr::expressions::BinaryExpr;
use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningPredicateBuilder};
#[cfg(feature = "encryption")]
use crate::encryption::DuckLakeEncryptionFactory;
use arrow::array::{Array, ArrayRef, BooleanArray, Int64Array};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::stats::Precision;
use datafusion::common::{Column, ColumnStatistics, ScalarValue, Statistics};
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::physical_plan::parquet::CachedParquetFileReaderFactory;
use datafusion::datasource::physical_plan::{
FileGroup, FileScanConfigBuilder, FileSource, ParquetFileReaderFactory, ParquetSource,
};
use datafusion::datasource::source::DataSourceExec;
use datafusion::datasource::table_schema::TableSchema;
use datafusion::error::{DataFusionError, Result as DataFusionResult};
use datafusion::execution::object_store::ObjectStoreUrl;
#[cfg(feature = "write")]
use datafusion::logical_expr::dml::InsertOp;
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
use datafusion::physical_plan::ExecutionPlan;
use datafusion_datasource_parquet::metadata::DFParquetMetadata;
use futures::StreamExt;
use object_store::ObjectMeta;
use object_store::ObjectStoreExt;
use object_store::path::Path as ObjectPath;
use percent_encoding::percent_decode_str;
#[cfg(feature = "encryption")]
use datafusion::execution::parquet_encryption::EncryptionFactory;
pub const DELETE_FILE_PATH_COL: &str = "file_path";
pub const DELETE_POS_COL: &str = "pos";
pub const DELETE_FILE_PATH_FIELD_ID: i32 = 2_147_483_646;
pub const DELETE_POS_FIELD_ID: i32 = 2_147_483_645;
struct FileMetadataPages<'a> {
provider: &'a dyn MetadataProvider,
table_id: i64,
snapshot_id: i64,
after_data_file_id: Option<i64>,
filter: Option<&'a StatsFilter>,
page_name: &'static str,
finished: bool,
}
impl Iterator for FileMetadataPages<'_> {
type Item = Result<Vec<DuckLakeFileMetadata>>;
fn next(&mut self) -> Option<Self::Item> {
if self.finished {
return None;
}
let metadata = match self.provider.get_table_file_metadata_page_filtered(
self.table_id,
self.snapshot_id,
self.after_data_file_id,
FILE_METADATA_BATCH_SIZE,
self.filter,
) {
Ok(metadata) => metadata,
Err(error) => {
self.finished = true;
return Some(Err(error));
},
};
if metadata.is_empty() {
self.finished = true;
return None;
}
if metadata.len() > FILE_METADATA_BATCH_SIZE {
self.finished = true;
return Some(Err(crate::DuckLakeError::InvalidConfig(format!(
"metadata provider returned {} files for a {}-file {} page",
metadata.len(),
FILE_METADATA_BATCH_SIZE,
self.page_name,
))));
}
let next_after = metadata.last().unwrap().file.data_file_id;
if self
.after_data_file_id
.is_some_and(|after| next_after <= after)
{
self.finished = true;
return Some(Err(crate::DuckLakeError::InvalidConfig(
"metadata provider returned a non-advancing file page".to_string(),
)));
}
self.after_data_file_id = Some(next_after);
self.finished = metadata.len() < FILE_METADATA_BATCH_SIZE;
Some(Ok(metadata))
}
}
fn parquet_field_id_metadata(field_id: i32) -> HashMap<String, String> {
HashMap::from([("PARQUET:field_id".to_string(), field_id.to_string())])
}
pub(crate) fn validated_file_size(file_size_bytes: i64, file_path: &str) -> DataFusionResult<u64> {
u64::try_from(file_size_bytes).map_err(|_| {
DataFusionError::Execution(format!(
"Invalid file_size_bytes ({}) for file '{}': value must be non-negative",
file_size_bytes, file_path
))
})
}
pub(crate) fn validated_record_count(record_count: i64, file_path: &str) -> DataFusionResult<u64> {
u64::try_from(record_count).map_err(|_| {
DataFusionError::Execution(format!(
"Invalid record_count ({}) for file '{}': value must be non-negative",
record_count, file_path
))
})
}
fn statistic_usize(value: i64, statistic: &str) -> Option<usize> {
match usize::try_from(value) {
Ok(value) => Some(value),
Err(_) => {
tracing::warn!(
value,
statistic,
"Ignoring invalid negative DuckLake statistic"
);
None
},
}
}
fn parse_statistic_scalar(
value: &str,
column: &DuckLakeTableColumn,
data_type: &DataType,
) -> Option<ScalarValue> {
let ducklake_type = column.column_type.trim().to_ascii_lowercase();
if ducklake_type.starts_with("list")
|| ducklake_type.starts_with("array")
|| ducklake_type.starts_with("struct")
|| ducklake_type.starts_with("map")
|| matches!(
ducklake_type.as_str(),
"geometry"
| "point"
| "linestring"
| "polygon"
| "multipoint"
| "multilinestring"
| "multipolygon"
| "geometrycollection"
| "linestring z"
| "timetz"
| "time with time zone"
| "interval"
)
{
return None;
}
if matches!(
value.to_ascii_lowercase().as_str(),
"infinity" | "-infinity"
) {
return None;
}
let parsed = parse_ducklake_scalar(value, data_type);
if parsed.is_none() {
tracing::debug!(
column = %column.column_name,
ducklake_type = %column.column_type,
value,
"Ignoring DuckLake statistic that could not be decoded"
);
}
parsed
}
fn validate_column_defaults(columns: &[DuckLakeTableColumn]) -> Result<HashMap<String, Expr>> {
let mut defaults = HashMap::new();
for column in columns {
let data_type = ducklake_to_arrow_type(&column.column_type)?;
if let Some(value) = &column.initial_default
&& parse_ducklake_default_scalar(value, &data_type).is_none()
{
return Err(crate::DuckLakeError::InvalidConfig(format!(
"Cannot decode initial_default '{value}' for column '{}' as {}",
column.column_name, data_type
)));
}
let Some(value) = &column.default_value else {
continue;
};
match column.default_value_type.as_deref() {
None | Some("literal") => {
let scalar = parse_ducklake_default_scalar(value, &data_type).ok_or_else(|| {
crate::DuckLakeError::InvalidConfig(format!(
"Cannot decode default_value '{value}' for column '{}' as {}",
column.column_name, data_type
))
})?;
defaults.insert(column.column_name.clone(), Expr::Literal(scalar, None));
},
Some("expression") => {
tracing::debug!(
column = %column.column_name,
dialect = column.default_value_dialect.as_deref().unwrap_or("unknown"),
"Ignoring unevaluated DuckLake expression default"
);
},
Some(value_type) => {
return Err(crate::DuckLakeError::Unsupported(format!(
"Default value type '{value_type}' for column '{}' is not supported",
column.column_name
)));
},
}
}
Ok(defaults)
}
pub(crate) fn apply_initial_default_metadata(
schema: &Schema,
columns: &[DuckLakeTableColumn],
name_mapping: &HashMap<String, String>,
) -> Schema {
let fields = schema
.fields()
.iter()
.map(|field| {
let logical_name = name_mapping
.get(field.name())
.map_or(field.name().as_str(), String::as_str);
let Some(initial_default) = columns
.iter()
.find(|column| column.column_name == logical_name)
.and_then(|column| column.initial_default.as_deref())
.filter(|value| !value.eq_ignore_ascii_case("NULL"))
else {
return Arc::clone(field);
};
let mut metadata = field.metadata().clone();
metadata.insert(
INITIAL_DEFAULT_METADATA_KEY.to_string(),
initial_default.to_string(),
);
Arc::new(field.as_ref().clone().with_metadata(metadata))
})
.collect::<Vec<_>>();
Schema::new_with_metadata(fields, schema.metadata().clone())
}
fn float_bound_is_usable(data_type: &DataType, contains_nan: Option<bool>) -> bool {
!matches!(
data_type,
DataType::Float16 | DataType::Float32 | DataType::Float64
) || contains_nan == Some(false)
}
fn scalar_precision(
value: Option<&str>,
column: &DuckLakeTableColumn,
data_type: &DataType,
exact: bool,
) -> Precision<ScalarValue> {
match value.and_then(|value| parse_statistic_scalar(value, column, data_type)) {
Some(value) if exact => Precision::Exact(value),
Some(value) => Precision::Inexact(value),
None => Precision::Absent,
}
}
fn file_row_count(
file: &DuckLakeTableFile,
file_columns: Option<&HashMap<i64, DuckLakeFileColumnStatistics>>,
) -> Precision<usize> {
let gross = file.max_row_count.or_else(|| {
file_columns.and_then(|columns| columns.values().find_map(|stats| stats.value_count))
});
let Some(gross) = gross.and_then(|value| statistic_usize(value, "record_count")) else {
return Precision::Absent;
};
if file.delete_file.is_some() {
let Some(deleted) = file
.delete_count
.and_then(|value| statistic_usize(value, "delete_count"))
else {
return Precision::Absent;
};
gross
.checked_sub(deleted)
.map(Precision::Exact)
.unwrap_or(Precision::Absent)
} else {
Precision::Exact(gross)
}
}
fn file_summary_row_count(file: &DuckLakeTableFile) -> Precision<usize> {
if file.delete_file.is_some() || file.delete_count.is_some_and(|count| count > 0) {
return Precision::Absent;
}
file.max_row_count
.and_then(|value| statistic_usize(value, "record_count"))
.map(Precision::Exact)
.unwrap_or(Precision::Absent)
}
fn bound_may_be_widened(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Utf8
| DataType::LargeUtf8
| DataType::Utf8View
| DataType::Binary
| DataType::LargeBinary
| DataType::BinaryView
| DataType::FixedSizeBinary(_)
)
}
fn file_is_known_empty(file: &DuckLakeTableFile) -> bool {
file.max_row_count == Some(0)
}
fn build_datafusion_statistics(
schema: &Schema,
columns: &[DuckLakeTableColumn],
table_files: &[DuckLakeTableFile],
catalog: DuckLakeStatistics,
use_current_table_statistics: bool,
file_metadata_complete: bool,
) -> (Statistics, HashMap<i64, Arc<Statistics>>) {
let table_column_rows: HashMap<i64, DuckLakeTableColumnStatistics> = catalog
.columns
.into_iter()
.map(|stats| (stats.column_id, stats))
.collect();
let mut file_column_rows: HashMap<i64, HashMap<i64, DuckLakeFileColumnStatistics>> =
HashMap::new();
for stats in catalog.files {
file_column_rows
.entry(stats.data_file_id)
.or_default()
.insert(stats.column_id, stats);
}
let mut file_statistics = HashMap::with_capacity(table_files.len());
for file in table_files {
let raw_columns = file_column_rows.get(&file.data_file_id);
let has_deletes = file.delete_file.is_some();
let mut statistics = Statistics::new_unknown(schema);
statistics.num_rows = file_row_count(file, raw_columns);
statistics.total_byte_size =
statistic_usize(file.file.file_size_bytes, "data_file.file_size_bytes")
.map(Precision::Inexact)
.unwrap_or(Precision::Absent);
for (index, column) in columns.iter().enumerate() {
let Some(raw) = raw_columns.and_then(|stats| stats.get(&column.column_id)) else {
continue;
};
let field_type = schema.field(index).data_type();
let exact = !has_deletes;
let column_statistics = &mut statistics.column_statistics[index];
column_statistics.null_count = raw
.null_count
.and_then(|value| statistic_usize(value, "file_column_stats.null_count"))
.map(|value| {
if exact {
Precision::Exact(value)
} else {
Precision::Inexact(value)
}
})
.unwrap_or(Precision::Absent);
let bounds_usable = float_bound_is_usable(field_type, raw.contains_nan);
column_statistics.min_value = if bounds_usable {
scalar_precision(raw.min_value.as_deref(), column, field_type, exact)
} else {
Precision::Absent
};
column_statistics.max_value = if bounds_usable {
scalar_precision(raw.max_value.as_deref(), column, field_type, exact)
} else {
Precision::Absent
};
column_statistics.byte_size = raw
.column_size_bytes
.and_then(|value| statistic_usize(value, "file_column_stats.column_size_bytes"))
.map(Precision::Inexact)
.unwrap_or(Precision::Absent);
}
file_statistics.insert(file.data_file_id, Arc::new(statistics));
}
let mut table_statistics = Statistics::new_unknown(schema);
let mut row_total = Some(0usize);
for file in table_files {
let value = file_row_count(file, file_column_rows.get(&file.data_file_id));
row_total = match (row_total, value.get_value()) {
(Some(total), Some(value)) => total.checked_add(*value),
_ => None,
};
}
table_statistics.num_rows = if file_metadata_complete && let Some(rows) = row_total {
Precision::Exact(rows)
} else if use_current_table_statistics {
catalog
.table
.as_ref()
.and_then(|stats| stats.record_count)
.and_then(|value| statistic_usize(value, "table_stats.record_count"))
.map(Precision::Exact)
.unwrap_or(Precision::Absent)
} else {
Precision::Absent
};
table_statistics.total_byte_size = if use_current_table_statistics {
catalog
.table
.as_ref()
.and_then(|stats| stats.file_size_bytes)
.and_then(|value| statistic_usize(value, "table_stats.file_size_bytes"))
.map(Precision::Inexact)
.unwrap_or_else(|| fallback_table_byte_size(table_files))
} else {
fallback_table_byte_size(table_files)
};
let any_deletes = table_files.iter().any(|file| file.delete_file.is_some());
for (index, column) in columns.iter().enumerate() {
let field_type = schema.field(index).data_type();
let output = &mut table_statistics.column_statistics[index];
if use_current_table_statistics && let Some(raw) = table_column_rows.get(&column.column_id)
{
if raw.contains_null == Some(false) {
output.null_count = Precision::Exact(0);
}
let bounds_usable = float_bound_is_usable(field_type, raw.contains_nan);
output.min_value = if bounds_usable {
scalar_precision(
raw.min_value.as_deref(),
column,
field_type,
raw.bounds_are_exact,
)
} else {
Precision::Absent
};
output.max_value = if bounds_usable {
scalar_precision(
raw.max_value.as_deref(),
column,
field_type,
raw.bounds_are_exact,
)
} else {
Precision::Absent
};
output.byte_size = raw
.column_size_bytes
.and_then(|value| statistic_usize(value, "file_column_stats.column_size_bytes"))
.map(Precision::Inexact)
.unwrap_or(Precision::Absent);
}
if !file_metadata_complete {
continue;
}
if table_files.is_empty() {
output.null_count = Precision::Exact(0);
output.byte_size = Precision::Exact(0);
continue;
}
let mut null_total = Some(0usize);
let mut byte_total = Some(0usize);
let mut min_value: Option<ScalarValue> = None;
let mut max_value: Option<ScalarValue> = None;
let mut min_complete = true;
let mut max_complete = true;
for file in table_files {
let Some(raw) = file_column_rows
.get(&file.data_file_id)
.and_then(|stats| stats.get(&column.column_id))
else {
null_total = None;
byte_total = None;
min_complete = false;
max_complete = false;
continue;
};
null_total = match (
null_total,
raw.null_count
.and_then(|value| statistic_usize(value, "file_column_stats.null_count")),
) {
(Some(total), Some(value)) => total.checked_add(value),
_ => None,
};
byte_total = match (
byte_total,
raw.column_size_bytes.and_then(|value| {
statistic_usize(value, "file_column_stats.column_size_bytes")
}),
) {
(Some(total), Some(value)) => total.checked_add(value),
_ => None,
};
let all_null =
matches!((raw.value_count, raw.null_count), (Some(v), Some(n)) if v == n);
let usable_min = raw
.min_value
.as_deref()
.filter(|_| float_bound_is_usable(field_type, raw.contains_nan));
match usable_min.and_then(|value| parse_statistic_scalar(value, column, field_type)) {
Some(value) => {
min_value = match min_value {
Some(current) => current.partial_cmp(&value).map(|ordering| {
if ordering.is_le() {
current
} else {
value
}
}),
None => Some(value),
};
min_complete &= min_value.is_some();
},
None if all_null => {},
None => min_complete = false,
}
let usable_max = raw
.max_value
.as_deref()
.filter(|_| float_bound_is_usable(field_type, raw.contains_nan));
match usable_max.and_then(|value| parse_statistic_scalar(value, column, field_type)) {
Some(value) => {
max_value = match max_value {
Some(current) => current.partial_cmp(&value).map(|ordering| {
if ordering.is_ge() {
current
} else {
value
}
}),
None => Some(value),
};
max_complete &= max_value.is_some();
},
None if all_null => {},
None => max_complete = false,
}
}
if let Some(value) = null_total {
output.null_count = if any_deletes {
Precision::Inexact(value)
} else {
Precision::Exact(value)
};
}
if let Some(value) = byte_total {
output.byte_size = Precision::Inexact(value);
}
if min_complete && let Some(value) = min_value {
output.min_value = if any_deletes {
Precision::Inexact(value)
} else {
Precision::Exact(value)
};
}
if max_complete && let Some(value) = max_value {
output.max_value = if any_deletes {
Precision::Inexact(value)
} else {
Precision::Exact(value)
};
}
}
(table_statistics, file_statistics)
}
fn fallback_table_byte_size(table_files: &[DuckLakeTableFile]) -> Precision<usize> {
let data_bytes: i128 = table_files
.iter()
.map(|file| i128::from(file.file.file_size_bytes))
.sum();
let delete_bytes: i128 = table_files
.iter()
.filter_map(|file| file.delete_file.as_ref())
.map(|file| i128::from(file.file_size_bytes))
.sum();
usize::try_from((data_bytes - delete_bytes).max(0))
.map(Precision::Inexact)
.unwrap_or(Precision::Absent)
}
pub fn delete_file_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new(DELETE_FILE_PATH_COL, DataType::Utf8, false)
.with_metadata(parquet_field_id_metadata(DELETE_FILE_PATH_FIELD_ID)),
Field::new(DELETE_POS_COL, DataType::Int64, false)
.with_metadata(parquet_field_id_metadata(DELETE_POS_FIELD_ID)),
]))
}
type PagedFileStatistics = (
Vec<DuckLakeTableFile>,
HashMap<i64, Arc<Statistics>>,
HashMap<i64, Arc<Statistics>>,
);
type SchemaMapping = (
SchemaRef,
HashMap<String, String>,
HashMap<String, ScalarValue>,
);
#[derive(Debug, Clone)]
struct FileReadConfig {
read_schema: SchemaRef,
name_mapping: HashMap<String, String>,
row_group_count: usize,
renamed_column_indices: Arc<HashSet<usize>>,
constants: HashMap<String, ScalarValue>,
embedded_rowid_parquet_name: Option<String>,
embedded_snapshot_parquet_name: Option<String>,
drops_current_columns: bool,
}
pub(crate) struct ParquetFooterFacts {
pub(crate) field_ids: HashMap<i32, String>,
pub(crate) arrow_schema: SchemaRef,
pub(crate) row_group_count: usize,
}
#[derive(Debug, Clone)]
pub(crate) struct ParquetFileLayout {
pub(crate) read_schema: SchemaRef,
pub(crate) name_mapping: HashMap<String, String>,
pub(crate) constants: HashMap<String, ScalarValue>,
pub(crate) embedded_rowid_parquet_name: Option<String>,
pub(crate) embedded_snapshot_parquet_name: Option<String>,
pub(crate) drops_current_columns: bool,
pub(crate) row_group_count: usize,
}
fn epoch_object_meta(location: ObjectPath, size: u64) -> ObjectMeta {
ObjectMeta {
location,
last_modified: DateTime::<Utc>::from_timestamp_nanos(0),
size,
e_tag: None,
version: None,
}
}
pub(crate) async fn read_parquet_footer_facts(
state: &dyn Session,
object_store_url: &ObjectStoreUrl,
resolved_path: &str,
file_size_bytes: Option<i64>,
encryption_key: Option<&str>,
) -> DataFusionResult<ParquetFooterFacts> {
let object_store = state.runtime_env().object_store(object_store_url)?;
let object_path = ObjectPath::from(resolved_path);
let size = match file_size_bytes {
Some(size) => validated_file_size(size, resolved_path)?,
None => object_store.head(&object_path).await?.size,
};
let object_meta = epoch_object_meta(object_path, size);
#[cfg(feature = "encryption")]
let decryption_properties = match encryption_key {
Some(key) if !key.is_empty() => {
let key_bytes = crate::encryption::DuckLakeEncryptionFactory::decode_key(key)?;
let props = parquet::encryption::decrypt::FileDecryptionProperties::builder(key_bytes)
.build()
.map_err(|e| {
DataFusionError::Execution(format!(
"Failed to create decryption properties: {}",
e
))
})?;
Some(props)
},
_ => None,
};
#[cfg(not(feature = "encryption"))]
let decryption_properties = {
let _ = encryption_key;
None
};
let metadata = DFParquetMetadata::new(object_store.as_ref(), &object_meta)
.with_file_metadata_cache(Some(
state.runtime_env().cache_manager.get_file_metadata_cache(),
))
.with_decryption_properties(decryption_properties)
.fetch_metadata()
.await?;
let field_ids = extract_parquet_field_ids(&metadata);
let arrow_schema = Arc::new(
parquet::arrow::parquet_to_arrow_schema(
metadata.file_metadata().schema_descr(),
metadata.file_metadata().key_value_metadata(),
)
.map_err(|e| DataFusionError::External(Box::new(e)))?,
);
Ok(ParquetFooterFacts {
field_ids,
arrow_schema,
row_group_count: metadata.row_groups().len(),
})
}
pub(crate) fn cached_parquet_reader_factory(
state: &dyn Session,
object_store_url: &ObjectStoreUrl,
) -> DataFusionResult<Arc<dyn ParquetFileReaderFactory>> {
let object_store = state.runtime_env().object_store(object_store_url)?;
let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
Ok(Arc::new(CachedParquetFileReaderFactory::new(
object_store,
metadata_cache,
)))
}
pub(crate) async fn read_parquet_file_layout(
state: &dyn Session,
object_store_url: &ObjectStoreUrl,
resolved_path: &str,
file_size_bytes: Option<i64>,
encryption_key: Option<&str>,
columns: &[DuckLakeTableColumn],
fallback_schema: &SchemaRef,
) -> DataFusionResult<Arc<ParquetFileLayout>> {
let facts = read_parquet_footer_facts(
state,
object_store_url,
resolved_path,
file_size_bytes,
encryption_key,
)
.await?;
let (read_schema, name_mapping) = if facts.field_ids.is_empty() {
(fallback_schema.clone(), HashMap::new())
} else {
let (schema, mapping) = build_read_schema_with_field_id_mapping(
columns,
&facts.field_ids,
Some(facts.arrow_schema.as_ref()),
)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
(Arc::new(schema), mapping)
};
let current_column_ids: HashSet<i32> = columns
.iter()
.flat_map(|column| {
std::iter::once(column.column_id).chain(column.nested_column_ids.iter().copied())
})
.map(|column_id| column_id as i32)
.collect();
let drops_current_columns = facts.field_ids.keys().any(|fid| {
*fid != ROW_ID_PARQUET_FIELD_ID
&& *fid != SNAPSHOT_ID_PARQUET_FIELD_ID
&& !current_column_ids.contains(fid)
});
Ok(Arc::new(ParquetFileLayout {
read_schema,
name_mapping,
constants: HashMap::new(),
embedded_rowid_parquet_name: facts.field_ids.get(&ROW_ID_PARQUET_FIELD_ID).cloned(),
embedded_snapshot_parquet_name: facts.field_ids.get(&SNAPSHOT_ID_PARQUET_FIELD_ID).cloned(),
drops_current_columns,
row_group_count: facts.row_group_count,
}))
}
pub(crate) fn apply_name_mapping_to_layout(
provider: &dyn MetadataProvider,
table_id: i64,
snapshot_id: i64,
columns: &[DuckLakeTableColumn],
resolved_path: &str,
mapping_id: Option<i64>,
layout: Arc<ParquetFileLayout>,
) -> DataFusionResult<Arc<ParquetFileLayout>> {
let Some(mapping_id) = mapping_id else {
return Ok(layout);
};
let mapping = provider
.get_name_mapping(mapping_id)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
if mapping.table_id != table_id {
return Err(DataFusionError::Execution(format!(
"DuckLake name mapping {mapping_id} belongs to table {}, not table {table_id}",
mapping.table_id
)));
}
let fields = provider
.get_table_fields(table_id, snapshot_id)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let mapped = build_read_schema_with_name_mapping(&fields, &mapping)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let constants = mapped
.partitions
.into_iter()
.map(|partition| {
DuckLakeTable::hive_partition_value(
resolved_path,
&partition.source_name,
&partition.data_type,
)
.map(|value| (partition.logical_name, value))
})
.collect::<DataFusionResult<HashMap<_, _>>>()?;
let mut layout = layout.as_ref().clone();
layout.read_schema = Arc::new(apply_initial_default_metadata(
&mapped.schema,
columns,
&mapped.names,
));
layout.name_mapping = mapped.names;
layout.constants = constants;
Ok(Arc::new(layout))
}
#[derive(Clone)]
pub struct DuckLakeTable {
#[allow(dead_code)]
table_id: i64,
table_name: String,
#[allow(dead_code)]
provider: Arc<dyn MetadataProvider>,
#[cfg_attr(not(feature = "write"), allow(dead_code))]
snapshot_id: i64,
object_store_url: Arc<ObjectStoreUrl>,
table_path: String,
schema: SchemaRef,
physical_schema: SchemaRef,
row_lineage: bool,
columns: Vec<DuckLakeTableColumn>,
column_defaults: HashMap<String, Expr>,
fields: Vec<DuckLakeTableField>,
name_mapping_cache: Arc<std::sync::Mutex<HashMap<i64, Arc<DuckLakeNameMapping>>>>,
table_statistics: Statistics,
partition_spec: Option<PartitionSpec>,
file_read_config_cache: Arc<std::sync::Mutex<HashMap<String, Arc<FileReadConfig>>>>,
inlined_materialized_row_count: Arc<AtomicUsize>,
#[cfg(feature = "encryption")]
encryption_keys: Arc<InstalledEncryptionKeys>,
#[cfg(feature = "write")]
schema_name: Option<String>,
#[cfg(feature = "write")]
writer: Option<Arc<dyn MetadataWriter>>,
#[cfg(feature = "write")]
pub(crate) write_options: crate::table_writer::DuckLakeWriteOptions,
}
impl std::fmt::Debug for DuckLakeTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DuckLakeTable")
.field("table_id", &self.table_id)
.field("table_name", &self.table_name)
.field("table_path", &self.table_path)
.field("schema", &self.schema)
.field("columns", &self.columns)
.finish_non_exhaustive()
}
}
#[cfg(feature = "encryption")]
#[derive(Default)]
struct InstalledEncryptionKeys {
keys: Arc<std::sync::RwLock<HashMap<String, String>>>,
factory: std::sync::OnceLock<Arc<dyn EncryptionFactory>>,
}
#[cfg(feature = "encryption")]
fn record_encryption_key(keys: &mut HashMap<String, String>, path: &str, key: Option<&str>) {
if let Some(key) = key.filter(|key| !key.is_empty()) {
keys.insert(path.to_string(), key.to_string());
}
}
impl DuckLakeTable {
pub fn new(
table_id: i64,
table_name: impl Into<String>,
provider: Arc<dyn MetadataProvider>,
snapshot_id: i64, object_store_url: Arc<ObjectStoreUrl>,
table_path: String,
) -> Result<Self> {
let fields = provider.get_table_fields(table_id, snapshot_id)?;
let columns = provider.get_table_structure(table_id, snapshot_id)?;
let column_defaults = validate_column_defaults(&columns)?;
let partition_spec = provider.get_partition_spec(table_id, snapshot_id)?;
let physical_schema = Arc::new(apply_initial_default_metadata(
&build_arrow_schema_from_fields(&fields)?,
&columns,
&HashMap::new(),
));
let schema = physical_schema.clone();
let catalog_statistics = provider.get_table_summary_statistics(table_id, snapshot_id)?;
let use_current_table_statistics = provider.get_current_snapshot()? == snapshot_id;
let (table_statistics, _) = build_datafusion_statistics(
physical_schema.as_ref(),
&columns,
&[],
catalog_statistics,
use_current_table_statistics,
false,
);
#[cfg(feature = "encryption")]
let encryption_keys = Arc::new(InstalledEncryptionKeys::default());
Ok(Self {
table_id,
table_name: table_name.into(),
provider,
snapshot_id,
object_store_url,
table_path,
schema,
physical_schema,
row_lineage: false,
columns,
column_defaults,
fields,
name_mapping_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
table_statistics,
partition_spec,
#[cfg(feature = "encryption")]
encryption_keys,
file_read_config_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
inlined_materialized_row_count: Arc::new(AtomicUsize::new(0)),
#[cfg(feature = "write")]
schema_name: None,
#[cfg(feature = "write")]
writer: None,
#[cfg(feature = "write")]
write_options: crate::table_writer::DuckLakeWriteOptions::default(),
})
}
pub fn with_row_lineage(mut self, enabled: bool) -> Self {
self.row_lineage = enabled;
self.schema = if enabled {
let mut fields: Vec<Arc<Field>> =
self.physical_schema.fields().iter().cloned().collect();
fields.push(Arc::new(rowid_field()));
Arc::new(Schema::new(fields))
} else {
self.physical_schema.clone()
};
self
}
#[must_use]
pub fn inlined_materialized_row_count(&self) -> usize {
self.inlined_materialized_row_count.load(Ordering::Relaxed)
}
fn rowid_index(&self) -> Option<usize> {
self.row_lineage
.then(|| self.physical_schema.fields().len())
}
pub fn files(&self) -> Result<Vec<DuckLakeTableFile>> {
let files = self
.provider
.get_table_files_for_select(self.table_id, self.snapshot_id)?;
#[cfg(feature = "encryption")]
self.configure_encryption_factory(&files)?;
Ok(files)
}
pub fn files_matching(
&self,
predicate: &Arc<dyn PhysicalExpr>,
) -> Result<Vec<DuckLakeTableFile>> {
let conjuncts = datafusion::physical_expr::split_conjunction(predicate)
.into_iter()
.cloned();
let pruning = match self.pruning_predicates(conjuncts) {
Ok(pruning) => pruning,
Err(error) => {
tracing::debug!(%error, "skipping predicate-based file pruning");
Vec::new()
},
};
let stats_filter =
stats_filter::lower_predicate(predicate, &self.physical_schema, &self.columns);
#[cfg(feature = "encryption")]
let mut encryption_keys = HashMap::new();
let mut matching = Vec::new();
for metadata in self.file_metadata_pages("file matching", stats_filter.as_ref()) {
let (table_files, mut file_statistics) = self.page_files_with_statistics(metadata?);
restate_in_physical_row_space(&mut file_statistics, &table_files);
#[cfg(feature = "encryption")]
self.collect_encryption_keys(&mut encryption_keys, &table_files)?;
matching.extend(
self.prune_table_files_iteratively(&pruning, &table_files, &file_statistics)
.into_iter()
.cloned(),
);
}
#[cfg(feature = "encryption")]
self.install_encryption_keys(encryption_keys);
Ok(matching)
}
pub(crate) fn inlined_data_with_row_ids(
&self,
) -> Result<Vec<crate::metadata_provider::DuckLakeInlinedData>> {
self.provider
.get_inlined_data_with_row_ids(self.table_id, self.snapshot_id, &self.columns)
}
fn file_metadata_pages<'a>(
&'a self,
page_name: &'static str,
filter: Option<&'a StatsFilter>,
) -> FileMetadataPages<'a> {
FileMetadataPages {
provider: self.provider.as_ref(),
table_id: self.table_id,
snapshot_id: self.snapshot_id,
after_data_file_id: None,
filter,
page_name,
finished: false,
}
}
fn resolve_file_path(&self, file: &DuckLakeFileData) -> DataFusionResult<String> {
resolve_path(&self.table_path, &file.path, file.path_is_relative)
.map_err(|e| DataFusionError::External(Box::new(e)))
}
fn hive_partition_raw_value<'a>(resolved_path: &'a str, source_name: &str) -> Option<&'a str> {
let mut partition_start = 0;
let mut equality_sign = 0;
let mut candidate_partition = true;
for (index, byte) in resolved_path.bytes().enumerate() {
if matches!(byte, b'?' | b'\n') {
candidate_partition = false;
}
if matches!(byte, b'/' | b'\\') {
if candidate_partition && equality_sign > partition_start {
let key = &resolved_path[partition_start..equality_sign];
if key == source_name {
return Some(&resolved_path[equality_sign + 1..index]);
}
}
partition_start = index + 1;
candidate_partition = true;
} else if byte == b'=' {
if equality_sign > partition_start {
candidate_partition = false;
}
equality_sign = index;
}
}
None
}
fn normalize_hive_hex_integer(value: &str, data_type: &DataType) -> Option<String> {
if !matches!(
data_type,
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
) {
return None;
}
let digits = value
.strip_prefix("0x")
.or_else(|| value.strip_prefix("0X"))?;
if digits.is_empty()
|| digits
.split('_')
.any(|part| part.is_empty() || !part.bytes().all(|byte| byte.is_ascii_hexdigit()))
{
return None;
}
let digits = digits.replace('_', "");
u128::from_str_radix(&digits, 16)
.ok()
.map(|value| value.to_string())
}
fn normalize_hive_exponent_decimal(value: &str, data_type: &DataType) -> Option<String> {
let (precision, scale) = match data_type {
DataType::Decimal128(precision, scale) | DataType::Decimal256(precision, scale) => {
(*precision, *scale)
},
_ => return None,
};
let exponent_index = value.find(['e', 'E'])?;
if value[exponent_index + 1..].contains(['e', 'E']) {
return None;
}
let exponent = value[exponent_index + 1..].parse::<i32>().ok()?;
let mantissa = &value[..exponent_index];
let (negative, mantissa) = if let Some(mantissa) = mantissa.strip_prefix('-') {
(true, mantissa)
} else if let Some(mantissa) = mantissa.strip_prefix('+') {
(false, mantissa)
} else {
(false, mantissa)
};
let mut components = mantissa.split('.');
let whole = components.next()?;
let fraction = components.next().unwrap_or_default();
if components.next().is_some()
|| (whole.is_empty() && fraction.is_empty())
|| !whole
.bytes()
.chain(fraction.bytes())
.all(|byte| byte.is_ascii_digit())
{
return None;
}
let digits = format!("{whole}{fraction}");
let decimal_index = i64::try_from(whole.len()).ok()? + i64::from(exponent);
let max_expansion = usize::from(precision)
.saturating_add(usize::from(scale.unsigned_abs()))
.saturating_add(digits.len())
.saturating_add(2);
let leading_zeros = usize::try_from((-decimal_index).max(0)).ok()?;
let trailing_zeros =
usize::try_from((decimal_index - i64::try_from(digits.len()).ok()?).max(0)).ok()?;
if leading_zeros > max_expansion || trailing_zeros > max_expansion {
return None;
}
let mut normalized = String::with_capacity(digits.len() + max_expansion.min(32));
if negative {
normalized.push('-');
}
if decimal_index <= 0 {
normalized.push_str("0.");
normalized.extend(std::iter::repeat_n('0', leading_zeros));
normalized.push_str(&digits);
} else if usize::try_from(decimal_index).ok()? >= digits.len() {
normalized.push_str(&digits);
normalized.extend(std::iter::repeat_n('0', trailing_zeros));
} else {
let decimal_index = usize::try_from(decimal_index).ok()?;
normalized.push_str(&digits[..decimal_index]);
normalized.push('.');
normalized.push_str(&digits[decimal_index..]);
}
Some(normalized)
}
pub(crate) fn hive_partition_value(
resolved_path: &str,
source_name: &str,
data_type: &DataType,
) -> DataFusionResult<ScalarValue> {
let raw_value = Self::hive_partition_raw_value(resolved_path, source_name).ok_or_else(|| {
DataFusionError::Execution(format!(
"column '{source_name}' uses a DuckLake Hive name mapping but is absent from file path '{resolved_path}'"
))
})?;
let value = percent_decode_str(raw_value).decode_utf8().map_err(|e| {
DataFusionError::Execution(format!(
"invalid percent encoding in Hive partition value '{raw_value}': {e}"
))
})?;
let value = value.into_owned();
let is_varchar = matches!(
data_type,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
);
if value == "__HIVE_DEFAULT_PARTITION__"
|| value.eq_ignore_ascii_case("NULL")
|| (value.is_empty() && !is_varchar)
{
return ScalarValue::try_from(data_type);
}
let scalar_input = Self::normalize_hive_hex_integer(&value, data_type)
.or_else(|| Self::normalize_hive_exponent_decimal(&value, data_type))
.unwrap_or_else(|| value.clone());
ScalarValue::try_from_string(scalar_input, data_type).map_err(|e| {
DataFusionError::Execution(format!(
"failed to parse Hive partition '{source_name}={value}' as {data_type}: {e}"
))
})
}
fn mapped_schema(
&self,
mapping_id: i64,
resolved_path: &str,
) -> DataFusionResult<SchemaMapping> {
let mapping = {
let cached = self.name_mapping_cache.lock().unwrap();
cached.get(&mapping_id).cloned()
};
let mapping = match mapping {
Some(mapping) => mapping,
None => {
let mapping = Arc::new(
self.provider
.get_name_mapping(mapping_id)
.map_err(|e| DataFusionError::External(Box::new(e)))?,
);
self.name_mapping_cache
.lock()
.unwrap()
.insert(mapping_id, Arc::clone(&mapping));
mapping
},
};
if mapping.table_id != self.table_id {
return Err(DataFusionError::Execution(format!(
"DuckLake name mapping {mapping_id} belongs to table {}, not table {}",
mapping.table_id, self.table_id
)));
}
let mapped = build_read_schema_with_name_mapping(&self.fields, mapping.as_ref())
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let constants = mapped
.partitions
.into_iter()
.map(|partition| {
Self::hive_partition_value(
resolved_path,
&partition.source_name,
&partition.data_type,
)
.map(|value| (partition.logical_name, value))
})
.collect::<DataFusionResult<HashMap<_, _>>>()?;
let schema = apply_initial_default_metadata(&mapped.schema, &self.columns, &mapped.names);
Ok((Arc::new(schema), mapped.names, constants))
}
fn partitioned_file(&self, file: &DuckLakeFileData) -> DataFusionResult<PartitionedFile> {
let resolved_path = self.resolve_file_path(file)?;
let mut pf = PartitionedFile::new(
&resolved_path,
validated_file_size(file.file_size_bytes, &resolved_path)?,
);
if let Some(footer_size) = file.footer_size
&& footer_size > 0
&& let Ok(hint) = usize::try_from(footer_size)
{
pf = pf.with_metadata_size_hint(hint);
}
Ok(pf)
}
fn partitioned_data_file(
&self,
table_file: &DuckLakeTableFile,
include_rowid: bool,
file_statistics: &HashMap<i64, Arc<Statistics>>,
) -> DataFusionResult<PartitionedFile> {
let mut file = self.partitioned_file(&table_file.file)?;
if let Some(statistics) = file_statistics.get(&table_file.data_file_id) {
let statistics = if include_rowid {
let mut statistics = statistics.as_ref().clone();
statistics
.column_statistics
.push(ColumnStatistics::new_unknown());
Arc::new(statistics)
} else {
Arc::clone(statistics)
};
file = file.with_statistics(statistics);
}
Ok(file)
}
fn page_files_with_statistics(
&self,
metadata: Vec<DuckLakeFileMetadata>,
) -> (Vec<DuckLakeTableFile>, HashMap<i64, Arc<Statistics>>) {
let (table_files, pruning, _) = self.page_files_resolved(metadata, false);
(table_files, pruning)
}
fn page_files_with_both_statistics(
&self,
metadata: Vec<DuckLakeFileMetadata>,
) -> PagedFileStatistics {
self.page_files_resolved(metadata, true)
}
fn page_files_resolved(
&self,
metadata: Vec<DuckLakeFileMetadata>,
want_summary: bool,
) -> PagedFileStatistics {
let mut catalog_file_statistics = Vec::new();
let mut table_files = Vec::with_capacity(metadata.len());
for DuckLakeFileMetadata {
file,
column_statistics,
} in metadata
{
table_files.push(file);
catalog_file_statistics.extend(column_statistics);
}
let (_, mut file_statistics) = build_datafusion_statistics(
self.physical_schema.as_ref(),
&self.columns,
&table_files,
DuckLakeStatistics {
files: catalog_file_statistics,
..Default::default()
},
false,
true,
);
let summary_statistics = if want_summary {
self.summary_statistics_from(&table_files, &file_statistics)
} else {
HashMap::new()
};
self.apply_partition_bounds(&table_files, &mut file_statistics);
(table_files, file_statistics, summary_statistics)
}
fn summary_statistics_from(
&self,
table_files: &[DuckLakeTableFile],
file_statistics: &HashMap<i64, Arc<Statistics>>,
) -> HashMap<i64, Arc<Statistics>> {
let mut summary = HashMap::with_capacity(table_files.len());
for file in table_files {
let Some(pruning) = file_statistics.get(&file.data_file_id) else {
continue;
};
let mut statistics = pruning.as_ref().clone();
statistics.num_rows = file_summary_row_count(file);
let has_deletes = file.delete_file.is_some();
for (index, column) in statistics.column_statistics.iter_mut().enumerate() {
let widened = self
.physical_schema
.fields()
.get(index)
.is_none_or(|field| bound_may_be_widened(field.data_type()));
if has_deletes || widened {
column.min_value = Precision::Absent;
column.max_value = Precision::Absent;
}
}
summary.insert(file.data_file_id, Arc::new(statistics));
}
summary
}
fn apply_partition_bounds(
&self,
table_files: &[DuckLakeTableFile],
file_statistics: &mut HashMap<i64, Arc<Statistics>>,
) {
let Some(spec) = self.partition_spec.as_ref() else {
return;
};
if !spec.prune_safe {
return;
}
for file in table_files {
if file.partition_values.is_empty() {
continue;
}
let mut updates: Vec<(usize, ScalarValue, ScalarValue)> = Vec::new();
for (key_index, value) in &file.partition_values {
let Some(value) = value.as_deref() else {
continue; };
let Some(column) = spec
.columns
.iter()
.find(|c| c.partition_key_index == *key_index)
else {
continue;
};
let Some(index) = self
.columns
.iter()
.position(|c| c.column_id == column.column_id)
else {
continue;
};
let data_type = self.physical_schema.field(index).data_type();
if let Some((min, max)) = column.transform.source_bounds(value, data_type) {
updates.push((index, min, max));
}
}
if updates.is_empty() {
continue;
}
let Some(stats) = file_statistics.get_mut(&file.data_file_id) else {
continue;
};
let stats = Arc::make_mut(stats);
for (index, min, max) in updates {
let Some(column_statistics) = stats.column_statistics.get_mut(index) else {
continue;
};
let exact = min == max;
if matches!(column_statistics.min_value, Precision::Absent) {
column_statistics.min_value = if exact {
Precision::Exact(min)
} else {
Precision::Inexact(min)
};
}
if matches!(column_statistics.max_value, Precision::Absent) {
column_statistics.max_value = if exact {
Precision::Exact(max)
} else {
Precision::Inexact(max)
};
}
}
}
}
fn prune_table_files_iteratively<'a>(
&self,
predicates: &[PruningPredicate],
table_files: &'a [DuckLakeTableFile],
file_statistics: &HashMap<i64, Arc<Statistics>>,
) -> Vec<&'a DuckLakeTableFile> {
let candidates: Vec<&'a DuckLakeTableFile> = table_files
.iter()
.filter(|file| !file_is_known_empty(file))
.collect();
let mut retained = candidates.clone();
if predicates.is_empty() {
return retained;
}
loop {
let count_before = retained.len();
for predicate in predicates {
let mask = match self.file_pruning_mask_for(predicate, &retained, file_statistics) {
Ok(mask) => mask,
Err(e) => {
tracing::debug!(error = %e, "skipping plan-time file pruning");
return candidates;
},
};
debug_assert_eq!(mask.len(), retained.len());
retained = retained
.into_iter()
.zip(mask)
.filter_map(|(file, keep)| keep.then_some(file))
.collect();
}
if retained.len() == count_before {
return retained;
}
}
}
fn physical_conjuncts(
&self,
state: &dyn Session,
filters: &[Expr],
) -> DataFusionResult<Vec<Arc<dyn PhysicalExpr>>> {
let df_schema = DFSchema::try_from(self.physical_schema.as_ref().clone())?;
filters
.iter()
.flat_map(datafusion::logical_expr::utils::split_conjunction)
.map(|expr| state.create_physical_expr(expr.clone(), &df_schema))
.collect()
}
fn pruning_predicates(
&self,
conjuncts: impl IntoIterator<Item = Arc<dyn PhysicalExpr>>,
) -> DataFusionResult<Vec<PruningPredicate>> {
conjuncts
.into_iter()
.map(|conjunct| {
PruningPredicateBuilder::new()
.with_file_schema(Arc::clone(&self.physical_schema))
.try_build(conjunct)
})
.collect()
}
fn file_pruning_mask_for(
&self,
pruning: &PruningPredicate,
table_files: &[&DuckLakeTableFile],
file_statistics: &HashMap<i64, Arc<Statistics>>,
) -> DataFusionResult<Vec<bool>> {
let per_file: Vec<Arc<Statistics>> = table_files
.iter()
.map(|tf| {
file_statistics
.get(&tf.data_file_id)
.map(Arc::clone)
.unwrap_or_else(|| Arc::new(Statistics::new_unknown(&self.physical_schema)))
})
.collect();
let stats = FilePruningStatistics::new(per_file, Arc::clone(&self.physical_schema));
pruning.prune(&stats)
}
fn create_parquet_source(
&self,
state: &dyn Session,
schema: impl Into<TableSchema>,
) -> DataFusionResult<ParquetSource> {
let reader_factory = cached_parquet_reader_factory(state, self.object_store_url.as_ref())?;
#[cfg(feature = "encryption")]
if let Some(factory) = self.encryption_keys.factory.get() {
return Ok(ParquetSource::new(schema)
.with_encryption_factory(Arc::clone(factory))
.with_parquet_file_reader_factory(reader_factory));
}
Ok(ParquetSource::new(schema).with_parquet_file_reader_factory(reader_factory))
}
fn scan_config_builder(&self, source: Arc<dyn FileSource>) -> FileScanConfigBuilder {
FileScanConfigBuilder::new(self.object_store_url.as_ref().clone(), source)
.with_expr_adapter(Some(Arc::new(DuckLakeDefaultExprAdapterFactory)))
}
#[cfg(feature = "encryption")]
fn collect_encryption_keys(
&self,
keys: &mut HashMap<String, String>,
table_files: &[DuckLakeTableFile],
) -> Result<()> {
for table_file in table_files {
let resolved_path = resolve_path(
&self.table_path,
&table_file.file.path,
table_file.file.path_is_relative,
)?;
record_encryption_key(
keys,
&resolved_path,
table_file.file.encryption_key.as_deref(),
);
if let Some(delete_file) = &table_file.delete_file {
let path = resolve_path(
&self.table_path,
&delete_file.path,
delete_file.path_is_relative,
)?;
record_encryption_key(keys, &path, delete_file.encryption_key.as_deref());
}
}
Ok(())
}
#[cfg(feature = "encryption")]
fn install_encryption_keys(&self, keys: HashMap<String, String>) {
if keys.is_empty() {
return;
}
self.encryption_keys.keys.write().unwrap().extend(keys);
self.encryption_keys.factory.get_or_init(|| {
Arc::new(DuckLakeEncryptionFactory::shared(Arc::clone(
&self.encryption_keys.keys,
))) as Arc<dyn EncryptionFactory>
});
}
#[cfg(feature = "encryption")]
fn configure_encryption_factory(&self, table_files: &[DuckLakeTableFile]) -> Result<()> {
let mut keys = HashMap::new();
self.collect_encryption_keys(&mut keys, table_files)?;
self.install_encryption_keys(keys);
Ok(())
}
async fn file_schema_mapping(
&self,
state: &dyn Session,
file: &DuckLakeFileData,
) -> DataFusionResult<SchemaMapping> {
let resolved_path = self.resolve_file_path(file)?;
if let Some(mapping_id) = file.mapping_id {
return self.mapped_schema(mapping_id, &resolved_path);
}
let object_store = state
.runtime_env()
.object_store(self.object_store_url.as_ref())?;
let object_path = ObjectPath::from(resolved_path.as_str());
let object_meta = epoch_object_meta(
object_path,
validated_file_size(file.file_size_bytes, &resolved_path)?,
);
let metadata_size_hint = file
.footer_size
.filter(|&size| size > 0)
.and_then(|size| usize::try_from(size).ok());
#[cfg(feature = "encryption")]
let decryption_properties = match file.encryption_key.as_deref() {
Some(key) if !key.is_empty() => {
let key_bytes = crate::encryption::DuckLakeEncryptionFactory::decode_key(key)?;
let props =
parquet::encryption::decrypt::FileDecryptionProperties::builder(key_bytes)
.build()
.map_err(|e| {
DataFusionError::Execution(format!(
"Failed to create decryption properties: {}",
e
))
})?;
Some(props)
},
_ => None,
};
#[cfg(not(feature = "encryption"))]
let decryption_properties = None;
let metadata = DFParquetMetadata::new(object_store.as_ref(), &object_meta)
.with_metadata_size_hint(metadata_size_hint)
.with_file_metadata_cache(Some(
state.runtime_env().cache_manager.get_file_metadata_cache(),
))
.with_decryption_properties(decryption_properties)
.fetch_metadata()
.await?;
let field_id_map = extract_parquet_field_ids(&metadata);
if field_id_map.is_empty() {
return Ok((self.schema.clone(), HashMap::new(), HashMap::new()));
}
let arrow_schema = parquet::arrow::parquet_to_arrow_schema(
metadata.file_metadata().schema_descr(),
metadata.file_metadata().key_value_metadata(),
)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let (read_schema, name_mapping) = build_read_schema_with_field_id_mapping_from_schema(
&self.columns,
self.physical_schema.as_ref(),
&field_id_map,
Some(&arrow_schema),
)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let read_schema =
apply_initial_default_metadata(&read_schema, &self.columns, &name_mapping);
Ok((Arc::new(read_schema), name_mapping, HashMap::new()))
}
fn present_file_read_config(
&self,
plan: Arc<dyn ExecutionPlan>,
file_cfg: &FileReadConfig,
physical_len: usize,
) -> Arc<dyn ExecutionPlan> {
let mut output_fields: Vec<Arc<Field>> =
self.physical_schema.fields().iter().cloned().collect();
output_fields.extend(plan.schema().fields().iter().skip(physical_len).cloned());
Arc::new(ColumnRenameExec::new_with_constants(
plan,
Arc::new(Schema::new(output_fields)),
file_cfg.name_mapping.clone(),
file_cfg.constants.clone(),
))
}
pub async fn resolve_positions(
&self,
state: &dyn Session,
data_file: &DuckLakeFileData,
predicate: Arc<dyn datafusion::physical_expr::PhysicalExpr>,
) -> DataFusionResult<HashSet<i64>> {
let file_cfg = self.build_file_read_config(state, data_file).await?;
let (table_schema, pos_table_idx, _pos_name) = positional_table_schema_reserving(
file_cfg.read_schema.clone(),
self.physical_schema
.fields()
.iter()
.map(|field| field.name().as_str()),
);
let mut proj: Vec<usize> = (0..self.physical_schema.fields().len()).collect();
proj.push(pos_table_idx);
let pos_idx = proj.len() - 1;
let mut source = self.create_parquet_source(state, table_schema)?;
if self.predicate_is_prunable(&predicate, &file_cfg) {
source = source.with_predicate(Arc::clone(&predicate));
}
let plan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(
self.scan_config_builder(Arc::new(source))
.with_file_group(FileGroup::new(vec![self.partitioned_file(data_file)?]))
.with_projection_indices(Some(proj))?
.build(),
);
let plan = self.split_across_partitions(plan, state, &file_cfg)?;
let plan =
self.present_file_read_config(plan, &file_cfg, self.physical_schema.fields().len());
let batches = datafusion::physical_plan::collect(plan, state.task_ctx()).await?;
let mut positions = HashSet::new();
for batch in &batches {
let mask = predicate.evaluate(batch)?.into_array(batch.num_rows())?;
let mask = mask
.as_any()
.downcast_ref::<BooleanArray>()
.ok_or_else(|| {
DataFusionError::Execution(
"resolve_positions: predicate did not evaluate to a boolean".to_string(),
)
})?;
let pos = batch
.column(pos_idx)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| {
DataFusionError::Internal("physical-position column is not Int64".to_string())
})?;
for i in 0..batch.num_rows() {
if mask.is_valid(i) && mask.value(i) {
positions.insert(pos.value(i));
}
}
}
Ok(positions)
}
async fn classify_delete_file_read_error(
&self,
state: &dyn Session,
resolved_delete_path: &str,
err: DataFusionError,
) -> DataFusionError {
let missing = if is_object_store_not_found(&err) {
true
} else {
match state
.runtime_env()
.object_store(self.object_store_url.as_ref())
{
Ok(store) => matches!(
store.head(&ObjectPath::from(resolved_delete_path)).await,
Err(object_store::Error::NotFound { .. })
),
Err(_) => false,
}
};
if missing {
DataFusionError::Execution(format!(
"Delete file '{resolved_delete_path}' referenced in catalog metadata was not found. This may indicate catalog corruption or that the file was deleted outside of DuckLake."
))
} else {
err
}
}
pub async fn read_delete_file_positions(
&self,
state: &dyn Session,
delete_file: &DuckLakeFileData,
) -> DataFusionResult<HashSet<i64>> {
let delete_schema = delete_file_schema();
let resolved_delete_path = self.resolve_file_path(delete_file)?;
let mut pf = PartitionedFile::new(
&resolved_delete_path,
validated_file_size(delete_file.file_size_bytes, &resolved_delete_path)?,
);
if let Some(footer_size) = delete_file.footer_size
&& footer_size > 0
&& let Ok(hint) = usize::try_from(footer_size)
{
pf = pf.with_metadata_size_hint(hint);
}
let file_scan_config = FileScanConfigBuilder::new(
self.object_store_url.as_ref().clone(),
Arc::new(self.create_parquet_source(state, delete_schema)?),
)
.with_file_group(FileGroup::new(vec![pf]))
.build();
let exec = DataSourceExec::from_data_source(file_scan_config);
let task_ctx = state.task_ctx();
let stream = exec.execute(0, task_ctx)?;
let collected = stream
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<DataFusionResult<Vec<_>>>();
let batches: Vec<RecordBatch> = match collected {
Ok(batches) => batches,
Err(e) => {
return Err(self
.classify_delete_file_read_error(state, &resolved_delete_path, e)
.await);
},
};
let mut positions = HashSet::new();
for batch in batches {
extract_deleted_positions_from_batch(&batch, &mut positions)?;
}
Ok(positions)
}
pub(crate) fn inlined_deletes_by_file(&self) -> DataFusionResult<HashMap<i64, HashSet<i64>>> {
let deletes = self
.provider
.get_inlined_deletes(self.table_id, self.snapshot_id)
.map_err(|error| DataFusionError::External(Box::new(error)))?;
let mut by_file: HashMap<i64, HashSet<i64>> = HashMap::new();
for delete in deletes {
if delete.data_file_id < 0 || delete.row_id < 0 {
return Err(DataFusionError::Execution(format!(
"inlined delete has invalid file_id {} and row_id {}",
delete.data_file_id, delete.row_id
)));
}
by_file
.entry(delete.data_file_id)
.or_default()
.insert(delete.row_id);
}
Ok(by_file)
}
async fn deleted_positions_for_file(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
inlined_positions: Option<&HashSet<i64>>,
) -> DataFusionResult<HashSet<i64>> {
let mut positions = inlined_positions.cloned().unwrap_or_default();
if let Some(delete_file) = &table_file.delete_file {
positions.extend(self.read_delete_file_positions(state, delete_file).await?);
}
Ok(positions)
}
pub async fn file_has_embedded_rowid(
&self,
state: &dyn Session,
file: &DuckLakeFileData,
) -> DataFusionResult<bool> {
let cfg = self.build_file_read_config(state, file).await?;
Ok(cfg.embedded_rowid_parquet_name.is_some())
}
fn roll_up_scan_statistics(
&self,
state: &dyn Session,
file_group: FileGroup,
read_schema: &SchemaRef,
summary_statistics: &[Arc<Statistics>],
) -> (FileGroup, Option<Statistics>) {
if !state.config_options().execution.collect_statistics {
return (file_group, None);
}
if summary_statistics.is_empty() || summary_statistics.len() != file_group.len() {
return (file_group, None);
}
if summary_statistics
.iter()
.any(|statistics| !matches!(statistics.num_rows, Precision::Exact(_)))
{
return (file_group, None);
}
let merged = Statistics::try_merge_iter(
summary_statistics.iter().map(Arc::as_ref),
self.physical_schema.as_ref(),
);
match merged {
Ok(summary) => {
let summary = if read_schema.as_ref() == self.physical_schema.as_ref() {
summary
} else {
let mut unknown = Statistics::new_unknown(read_schema);
unknown.num_rows = summary.num_rows;
unknown.total_byte_size = summary.total_byte_size;
unknown
};
let file_group = file_group.with_statistics(Arc::new(summary.clone()));
(file_group, Some(summary))
},
Err(error) => {
tracing::debug!(%error, "skipping scan statistics rollup");
(file_group, None)
},
}
}
async fn build_exec_for_files_without_deletes(
&self,
state: &dyn Session,
files: &[&DuckLakeTableFile],
file_statistics: &HashMap<i64, Arc<Statistics>>,
summary_statistics: &HashMap<i64, Arc<Statistics>>,
projection: Option<&Vec<usize>>,
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let mut groups: Vec<(SchemaMapping, Vec<PartitionedFile>, Vec<Arc<Statistics>>)> =
Vec::new();
let mut group_index: HashMap<String, usize> = HashMap::new();
for table_file in files {
let mapping = self.file_schema_mapping(state, &table_file.file).await?;
let pf = self.partitioned_data_file(table_file, false, file_statistics)?;
let (read_schema, name_mapping, constants) = &mapping;
let mut key = String::new();
for f in read_schema.fields() {
key.push_str(f.name());
key.push('\u{1}');
key.push_str(&format!("{:?}", f.data_type()));
key.push('\u{2}');
}
let mut pairs: Vec<(&String, &String)> = name_mapping.iter().collect();
pairs.sort();
for (k, v) in pairs {
key.push_str(k);
key.push('\u{3}');
key.push_str(v);
key.push('\u{4}');
}
let mut constants: Vec<_> = constants.iter().collect();
constants.sort_by_key(|(name, _)| *name);
for (name, value) in constants {
key.push_str(name);
key.push('\u{5}');
key.push_str(&format!("{value:?}"));
key.push('\u{6}');
}
let summary = summary_statistics.get(&table_file.data_file_id).cloned();
match group_index.get(&key) {
Some(&gi) => {
groups[gi].1.push(pf);
if let Some(summary) = summary {
groups[gi].2.push(summary);
}
},
None => {
group_index.insert(key, groups.len());
groups.push((mapping, vec![pf], summary.into_iter().collect()));
},
}
}
let output_schema = match projection {
Some(indices) => Arc::new(self.schema.project(indices)?),
None => self.schema.clone(),
};
let nan_unsafe_columns = self.nan_unsafe_float_columns(files, file_statistics);
let mut execs: Vec<Arc<dyn ExecutionPlan>> = Vec::with_capacity(groups.len());
for ((read_schema, name_mapping, constants), partitioned_files, group_summary) in groups {
let (file_group, summary) = self.roll_up_scan_statistics(
state,
FileGroup::new(partitioned_files),
&read_schema,
&group_summary,
);
let mut builder = self
.scan_config_builder(Arc::new(
self.create_parquet_source(state, read_schema.clone())?,
))
.with_limit(limit)
.with_file_group(file_group);
if let Some(summary) = summary {
builder = builder.with_statistics(summary);
}
if let Some(proj) = projection {
builder = builder.with_projection_indices(Some(proj.clone()))?;
}
let parquet_exec: Arc<dyn ExecutionPlan> =
DataSourceExec::from_data_source(builder.build());
let mut exec = if !name_mapping.is_empty() || parquet_exec.schema() != output_schema {
Arc::new(ColumnRenameExec::new_with_constants(
parquet_exec,
output_schema.clone(),
name_mapping,
constants,
)) as Arc<dyn ExecutionPlan>
} else {
parquet_exec
};
if !nan_unsafe_columns.is_empty() {
exec = Arc::new(NanPruningBarrierExec::new(
exec,
Arc::clone(&nan_unsafe_columns),
));
}
execs.push(exec);
}
combine_execution_plans(execs)
}
fn nan_unsafe_float_columns(
&self,
files: &[&DuckLakeTableFile],
file_statistics: &HashMap<i64, Arc<Statistics>>,
) -> Arc<HashSet<String>> {
let mut unsafe_columns = HashSet::new();
for (index, field) in self.physical_schema.fields().iter().enumerate() {
if !matches!(
field.data_type(),
DataType::Float16 | DataType::Float32 | DataType::Float64
) {
continue;
}
let any_unsafe = files.iter().any(|file| {
file_statistics
.get(&file.data_file_id)
.is_none_or(|statistics| {
matches!(
statistics.column_statistics[index].max_value,
Precision::Absent
)
})
});
if any_unsafe {
unsafe_columns.insert(field.name().clone());
}
}
Arc::new(unsafe_columns)
}
fn split_across_partitions(
&self,
plan: Arc<dyn ExecutionPlan>,
state: &dyn Session,
file_cfg: &FileReadConfig,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let options = state.config_options();
if !options.optimizer.repartition_file_scans {
return Ok(plan);
}
let target = state
.config()
.target_partitions()
.min(file_cfg.row_group_count.max(1));
if target <= 1 {
return Ok(plan);
}
Ok(plan.repartitioned(target, options)?.unwrap_or(plan))
}
fn predicate_is_prunable(
&self,
predicate: &Arc<dyn datafusion::physical_expr::PhysicalExpr>,
file_cfg: &FileReadConfig,
) -> bool {
use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion::physical_expr::expressions::Column;
let mut prunable = true;
predicate
.apply(|expr| {
let Some(column) = expr.downcast_ref::<Column>() else {
return Ok(TreeNodeRecursion::Continue);
};
let Some(field) = self.physical_schema.fields().get(column.index()) else {
prunable = false;
return Ok(TreeNodeRecursion::Stop);
};
let is_float = matches!(
field.data_type(),
DataType::Float16 | DataType::Float32 | DataType::Float64
);
if is_float || file_cfg.renamed_column_indices.contains(&column.index()) {
prunable = false;
return Ok(TreeNodeRecursion::Stop);
}
Ok(TreeNodeRecursion::Continue)
})
.expect("scanning a predicate for column references cannot fail");
prunable
}
#[cfg(feature = "write")]
pub fn with_writer(mut self, schema_name: String, writer: Arc<dyn MetadataWriter>) -> Self {
self.schema_name = Some(schema_name);
self.writer = Some(writer);
self
}
#[cfg(feature = "write")]
pub fn with_write_options(
mut self,
options: crate::table_writer::DuckLakeWriteOptions,
) -> Self {
self.write_options = options;
self
}
fn gross_scan_statistics(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
schema: &SchemaRef,
) -> Option<Statistics> {
if !state.config_options().execution.collect_statistics {
return None;
}
let rows = table_file
.max_row_count
.and_then(|value| statistic_usize(value, "record_count"))?;
let mut statistics = Statistics::new_unknown(schema);
statistics.num_rows = Precision::Exact(rows);
Some(statistics)
}
async fn build_exec_for_file_with_deletes(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
inlined_positions: Option<&HashSet<i64>>,
file_statistics: &HashMap<i64, Arc<Statistics>>,
projection: Option<&Vec<usize>>,
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let file_cfg = self.build_file_read_config(state, &table_file.file).await?;
let deleted_positions = self
.deleted_positions_for_file(state, table_file, inlined_positions)
.await?;
let deleted_positions = (!deleted_positions.is_empty()).then_some(deleted_positions);
let output_schema = match projection {
Some(indices) => Arc::new(self.schema.project(indices)?),
None => self.schema.clone(),
};
let proj_indices: Vec<usize> = match projection {
Some(indices) => indices.clone(),
None => (0..self.physical_schema.fields().len()).collect(),
};
let exec_after_delete: Arc<dyn ExecutionPlan> = if let Some(positions) = deleted_positions {
let (table_schema, pos_table_idx, _pos_name) =
positional_table_schema(file_cfg.read_schema.clone());
let gross = self.gross_scan_statistics(state, table_file, table_schema.table_schema());
let mut proj = proj_indices.clone();
proj.push(pos_table_idx);
let pos_index = proj.len() - 1;
let pf = self.partitioned_data_file(
table_file,
file_cfg.embedded_rowid_parquet_name.is_some(),
file_statistics,
)?;
let mut builder = self
.scan_config_builder(Arc::new(self.create_parquet_source(state, table_schema)?))
.with_file_group(FileGroup::new(vec![pf]));
builder = builder.with_projection_indices(Some(proj))?;
if let Some(gross) = gross {
builder = builder.with_statistics(gross);
}
let scan = DataSourceExec::from_data_source(builder.build());
Arc::new(DeleteFilterExec::try_new(
scan,
table_file.file.path.clone(),
Arc::new(positions),
pos_index,
)?)
} else {
let pf = self.partitioned_data_file(
table_file,
file_cfg.embedded_rowid_parquet_name.is_some(),
file_statistics,
)?;
let mut builder = self
.scan_config_builder(Arc::new(
self.create_parquet_source(state, file_cfg.read_schema.clone())?,
))
.with_limit(limit)
.with_file_group(FileGroup::new(vec![pf]));
builder = builder.with_projection_indices(Some(proj_indices.clone()))?;
DataSourceExec::from_data_source(builder.build())
};
let mut exec: Arc<dyn ExecutionPlan> = if !file_cfg.name_mapping.is_empty()
|| !file_cfg.constants.is_empty()
|| exec_after_delete.schema() != output_schema
{
Arc::new(ColumnRenameExec::new_with_constants(
exec_after_delete,
output_schema,
file_cfg.name_mapping.clone(),
file_cfg.constants.clone(),
))
} else {
exec_after_delete
};
let nan_unsafe_columns =
self.nan_unsafe_float_columns(std::slice::from_ref(&table_file), file_statistics);
if !nan_unsafe_columns.is_empty() {
exec = Arc::new(NanPruningBarrierExec::new(exec, nan_unsafe_columns));
}
Ok(exec)
}
async fn build_file_read_config(
&self,
state: &dyn Session,
file: &DuckLakeFileData,
) -> DataFusionResult<Arc<FileReadConfig>> {
let resolved_path = self.resolve_file_path(file)?;
{
let cache = self.file_read_config_cache.lock().unwrap();
if let Some(cfg) = cache.get(&resolved_path) {
return Ok(cfg.clone());
}
}
#[cfg(feature = "encryption")]
let encryption_key = file.encryption_key.as_deref();
#[cfg(not(feature = "encryption"))]
let encryption_key = None;
let layout = read_parquet_file_layout(
state,
self.object_store_url.as_ref(),
&resolved_path,
Some(file.file_size_bytes),
encryption_key,
&self.columns,
&self.physical_schema,
)
.await?;
let (physical_read_schema, mut name_mapping, constants) =
if let Some(mapping_id) = file.mapping_id {
self.mapped_schema(mapping_id, &resolved_path)?
} else {
(
layout.read_schema.clone(),
layout.name_mapping.clone(),
HashMap::new(),
)
};
let read_schema = if let Some(ref parquet_name) = layout.embedded_rowid_parquet_name {
let mut fields: Vec<Arc<Field>> =
physical_read_schema.fields().iter().cloned().collect();
fields.push(Arc::new(Field::new(
parquet_name.clone(),
DataType::Int64,
true,
)));
if parquet_name != ROWID_COLUMN_NAME {
name_mapping.insert(parquet_name.clone(), ROWID_COLUMN_NAME.to_string());
}
Arc::new(Schema::new(fields))
} else {
physical_read_schema.clone()
};
let renamed_column_indices: HashSet<usize> = self
.physical_schema
.fields()
.iter()
.zip(physical_read_schema.fields())
.enumerate()
.filter(|(_, (catalog, physical))| catalog.name() != physical.name())
.map(|(index, _)| index)
.collect();
let cfg = Arc::new(FileReadConfig {
read_schema,
row_group_count: layout.row_group_count,
renamed_column_indices: Arc::new(renamed_column_indices),
name_mapping,
constants,
embedded_rowid_parquet_name: layout.embedded_rowid_parquet_name.clone(),
embedded_snapshot_parquet_name: layout.embedded_snapshot_parquet_name.clone(),
drops_current_columns: layout.drops_current_columns,
});
{
let mut cache = self.file_read_config_cache.lock().unwrap();
cache.entry(resolved_path).or_insert_with(|| cfg.clone());
}
Ok(cfg)
}
#[allow(clippy::too_many_arguments, reason = "per-file row-lineage scan inputs stay explicit")]
async fn build_exec_for_file_with_rowid(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
inlined_positions: Option<&HashSet<i64>>,
file_statistics: &HashMap<i64, Arc<Statistics>>,
user_proj: &[usize],
rowid_idx: usize,
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let file_cfg = self.build_file_read_config(state, &table_file.file).await?;
let has_embedded = file_cfg.embedded_rowid_parquet_name.is_some();
let physical_proj: Vec<usize> = user_proj
.iter()
.filter(|&&i| i != rowid_idx)
.copied()
.collect();
if !has_embedded && table_file.row_id_start.is_none() {
return Err(DataFusionError::Execution(format!(
"File \"{}\" has no embedded `_ducklake_internal_row_id` column and no \
`row_id_start` set in the catalog — row lineage cannot be reconstructed",
table_file.file.path
)));
}
let deleted_positions = self
.deleted_positions_for_file(state, table_file, inlined_positions)
.await?;
let deleted_positions = (!deleted_positions.is_empty()).then_some(deleted_positions);
let has_deletes = deleted_positions.is_some();
let needs_position = !has_embedded || has_deletes;
let parquet_projection: Vec<usize> = if has_embedded {
let rowid_col_in_read_schema = file_cfg.read_schema.fields().len() - 1;
let mut p = physical_proj.clone();
p.push(rowid_col_in_read_schema);
p
} else {
physical_proj.clone()
};
let after_deletes: Arc<dyn ExecutionPlan> = if needs_position {
let (table_schema, pos_table_idx, _pos_name) =
positional_table_schema(file_cfg.read_schema.clone());
let mut proj = parquet_projection;
proj.push(pos_table_idx);
let pos_index = proj.len() - 1;
let pf = self.partitioned_data_file(table_file, has_embedded, file_statistics)?;
let mut builder = self
.scan_config_builder(Arc::new(self.create_parquet_source(state, table_schema)?))
.with_file_group(FileGroup::new(vec![pf]));
builder = builder.with_projection_indices(Some(proj))?;
let mut plan: Arc<dyn ExecutionPlan> =
DataSourceExec::from_data_source(builder.build());
if let Some(p) = deleted_positions {
plan = Arc::new(DeleteFilterExec::try_new(
plan,
table_file.file.path.clone(),
Arc::new(p),
pos_index,
)?);
}
if !has_embedded {
plan = Arc::new(RowIdExec::try_new(
plan,
table_file.row_id_start,
pos_index,
)?);
}
plan
} else {
let pf = self.partitioned_data_file(table_file, true, file_statistics)?;
let mut builder = self
.scan_config_builder(Arc::new(
self.create_parquet_source(state, file_cfg.read_schema.clone())?,
))
.with_limit(limit)
.with_file_group(FileGroup::new(vec![pf]));
builder = builder.with_projection_indices(Some(parquet_projection))?;
DataSourceExec::from_data_source(builder.build())
};
let output_schema = self.output_schema_for_projection(user_proj, rowid_idx);
let mut exec = if !file_cfg.name_mapping.is_empty()
|| !file_cfg.constants.is_empty()
|| after_deletes.schema() != output_schema
{
Arc::new(ColumnRenameExec::new_with_constants(
after_deletes,
output_schema,
file_cfg.name_mapping.clone(),
file_cfg.constants.clone(),
)) as Arc<dyn ExecutionPlan>
} else {
after_deletes
};
let nan_unsafe_columns =
self.nan_unsafe_float_columns(std::slice::from_ref(&table_file), file_statistics);
if !nan_unsafe_columns.is_empty() {
exec = Arc::new(NanPruningBarrierExec::new(exec, nan_unsafe_columns));
}
Ok(exec)
}
fn output_schema_for_projection(&self, user_proj: &[usize], rowid_idx: usize) -> SchemaRef {
let mut fields: Vec<Arc<Field>> = Vec::with_capacity(user_proj.len());
for &i in user_proj {
if i == rowid_idx {
fields.push(Arc::new(rowid_field()));
} else {
fields.push(self.schema.fields()[i].clone());
}
}
Arc::new(Schema::new(fields))
}
fn needs_snapshot_filter(&self, table_file: &DuckLakeTableFile) -> bool {
table_file
.partial_max
.is_some_and(|partial_max| self.snapshot_id < partial_max)
}
async fn build_exec_for_partial_file(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
inlined_positions: Option<&HashSet<i64>>,
output_schema: SchemaRef,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
if table_file.delete_file.is_some() {
return Err(DataFusionError::Internal(format!(
"partial file \"{}\" has a live delete file at read snapshot {}, which the \
snapshot-filtered read path cannot apply",
table_file.file.path, self.snapshot_id
)));
}
let file_cfg = self.build_file_read_config(state, &table_file.file).await?;
let snap_name = file_cfg
.embedded_snapshot_parquet_name
.clone()
.ok_or_else(|| {
DataFusionError::Internal(format!(
"partial file \"{}\" is missing its embedded snapshot-id column",
table_file.file.path
))
})?;
if output_schema.field_with_name(ROWID_COLUMN_NAME).is_ok()
&& file_cfg.embedded_rowid_parquet_name.is_none()
&& table_file.row_id_start.is_none()
{
return Err(DataFusionError::Execution(format!(
"File \"{}\" has no embedded `_ducklake_internal_row_id` column and no \
`row_id_start` set in the catalog — row lineage cannot be reconstructed",
table_file.file.path
)));
}
let mut fields: Vec<Arc<Field>> = file_cfg.read_schema.fields().iter().cloned().collect();
fields.push(Arc::new(Field::new(&snap_name, DataType::Int64, true)));
let read_schema = Arc::new(Schema::new(fields));
let projection: Vec<usize> = (0..read_schema.fields().len()).collect();
let deleted_positions = self
.deleted_positions_for_file(state, table_file, inlined_positions)
.await?;
let scan: Arc<dyn ExecutionPlan> = if deleted_positions.is_empty() {
let resolved_path = self.resolve_file_path(&table_file.file)?;
let mut pf = PartitionedFile::new(
&resolved_path,
validated_file_size(table_file.file.file_size_bytes, &resolved_path)?,
);
if let Some(footer_size) = table_file.file.footer_size
&& footer_size > 0
&& let Ok(hint) = usize::try_from(footer_size)
{
pf = pf.with_metadata_size_hint(hint);
}
let builder = self
.scan_config_builder(Arc::new(
self.create_parquet_source(state, read_schema.clone())?,
))
.with_file_group(FileGroup::new(vec![pf]))
.with_projection_indices(Some(projection))?;
DataSourceExec::from_data_source(builder.build())
} else {
let (table_schema, pos_table_idx, _pos_name) =
positional_table_schema(read_schema.clone());
let mut proj = projection;
proj.push(pos_table_idx);
let pos_index = proj.len() - 1;
let builder = self
.scan_config_builder(Arc::new(self.create_parquet_source(state, table_schema)?))
.with_file_group(FileGroup::new(vec![
self.partitioned_file(&table_file.file)?,
]))
.with_projection_indices(Some(proj))?;
let scan = DataSourceExec::from_data_source(builder.build());
Arc::new(DeleteFilterExec::try_new(
scan,
table_file.file.path.clone(),
Arc::new(deleted_positions),
pos_index,
)?)
};
let filtered: Arc<dyn ExecutionPlan> = Arc::new(SnapshotFilterExec::try_new(
scan,
snap_name,
self.snapshot_id,
)?);
Ok(Arc::new(ColumnRenameExec::new_with_constants(
filtered,
output_schema,
file_cfg.name_mapping.clone(),
file_cfg.constants.clone(),
)))
}
#[cfg(feature = "write")]
fn read_only_clone(&self) -> DuckLakeTable {
DuckLakeTable {
table_id: self.table_id,
table_name: self.table_name.clone(),
provider: Arc::clone(&self.provider),
snapshot_id: self.snapshot_id,
object_store_url: self.object_store_url.clone(),
table_path: self.table_path.clone(),
schema: self.physical_schema.clone(),
physical_schema: self.physical_schema.clone(),
row_lineage: false,
columns: self.columns.clone(),
column_defaults: self.column_defaults.clone(),
fields: self.fields.clone(),
name_mapping_cache: Arc::clone(&self.name_mapping_cache),
table_statistics: self.table_statistics.clone(),
partition_spec: self.partition_spec.clone(),
file_read_config_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
inlined_materialized_row_count: Arc::clone(&self.inlined_materialized_row_count),
#[cfg(feature = "encryption")]
encryption_keys: Arc::clone(&self.encryption_keys),
schema_name: None,
writer: None,
write_options: self.write_options.clone(),
}
}
#[cfg(feature = "write")]
pub(crate) fn physical_schema(&self) -> SchemaRef {
self.physical_schema.clone()
}
#[cfg(feature = "write")]
pub(crate) fn writer(&self) -> Option<&Arc<dyn MetadataWriter>> {
self.writer.as_ref()
}
#[cfg(feature = "write")]
pub(crate) fn write_options(&self) -> &crate::table_writer::DuckLakeWriteOptions {
&self.write_options
}
#[cfg(feature = "write")]
pub(crate) fn schema_name(&self) -> Option<&str> {
self.schema_name.as_deref()
}
#[cfg(feature = "write")]
pub(crate) fn table_name(&self) -> &str {
&self.table_name
}
#[cfg(feature = "write")]
pub(crate) fn table_id(&self) -> i64 {
self.table_id
}
#[cfg(feature = "write")]
pub(crate) fn base_snapshot(&self) -> i64 {
self.snapshot_id
}
#[cfg(feature = "write")]
pub(crate) fn object_store_url(&self) -> &Arc<ObjectStoreUrl> {
&self.object_store_url
}
#[cfg(feature = "write")]
pub(crate) fn live_sort_spec(&self) -> crate::Result<Option<crate::sort::SortSpec>> {
let head = self.provider.get_current_snapshot()?;
self.provider.get_sort_spec(self.table_id, head)
}
#[cfg(feature = "write")]
pub(crate) fn live_partition_spec(&self) -> crate::Result<Option<PartitionSpec>> {
let head = self.provider.get_current_snapshot()?;
self.provider.get_partition_spec(self.table_id, head)
}
#[cfg(feature = "write")]
pub(crate) fn column_ids(&self) -> Vec<i64> {
let mut ids = Vec::new();
for column in &self.columns {
ids.push(column.column_id);
ids.extend_from_slice(&column.nested_column_ids);
}
ids
}
#[cfg(feature = "write")]
pub(crate) fn top_level_column_ids(&self) -> Vec<i64> {
self.columns.iter().map(|column| column.column_id).collect()
}
#[cfg(feature = "write")]
pub(crate) async fn merge_source_facts(
&self,
state: &dyn Session,
file: &DuckLakeFileData,
) -> DataFusionResult<MergeSourceFacts> {
let cfg = self.build_file_read_config(state, file).await?;
Ok(MergeSourceFacts {
drops_current_columns: cfg.drops_current_columns,
has_embedded_snapshot: cfg.embedded_snapshot_parquet_name.is_some(),
})
}
#[cfg(feature = "write")]
pub(crate) async fn build_update_scan(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
inlined_positions: Option<&HashSet<i64>>,
) -> DataFusionResult<UpdateSourceScan> {
self.build_update_scan_with_snapshot(state, table_file, false, inlined_positions)
.await
}
#[cfg(feature = "write")]
pub(crate) async fn build_update_scan_with_snapshot(
&self,
state: &dyn Session,
table_file: &DuckLakeTableFile,
want_snapshot_id: bool,
inlined_positions: Option<&HashSet<i64>>,
) -> DataFusionResult<UpdateSourceScan> {
let file_cfg = self.build_file_read_config(state, &table_file.file).await?;
let has_embedded = file_cfg.embedded_rowid_parquet_name.is_some();
if !has_embedded && table_file.row_id_start.is_none() {
return Err(DataFusionError::Execution(format!(
"File \"{}\" has no embedded `_ducklake_internal_row_id` column and no \
`row_id_start` in the catalog — cannot preserve row lineage through UPDATE",
table_file.file.path
)));
}
let existing_parquet_deleted = match &table_file.delete_file {
Some(delete_file) => self.read_delete_file_positions(state, delete_file).await?,
None => HashSet::new(),
};
let mut deleted_positions = existing_parquet_deleted.clone();
if let Some(inlined) = inlined_positions {
deleted_positions.extend(inlined);
}
let read_schema = if want_snapshot_id {
let snap_name = file_cfg
.embedded_snapshot_parquet_name
.clone()
.ok_or_else(|| {
DataFusionError::Internal(format!(
"File \"{}\" carries no embedded `_ducklake_internal_snapshot_id` \
column, so its rows' origin snapshots cannot be read",
table_file.file.path
))
})?;
let mut fields: Vec<Arc<Field>> =
file_cfg.read_schema.fields().iter().cloned().collect();
fields.push(Arc::new(Field::new(&snap_name, DataType::Int64, true)));
Arc::new(Schema::new(fields))
} else {
Arc::clone(&file_cfg.read_schema)
};
let physical_len = self.physical_schema.fields().len();
let (table_schema, pos_table_idx, _pos_name) = positional_table_schema_reserving(
read_schema.clone(),
self.physical_schema
.fields()
.iter()
.map(|field| field.name().as_str()),
);
let mut proj: Vec<usize> = (0..physical_len).collect();
let embedded_batch_idx = has_embedded.then(|| {
proj.push(file_cfg.read_schema.fields().len() - 1);
proj.len() - 1
});
let embedded_snapshot_batch_idx = want_snapshot_id.then(|| {
proj.push(read_schema.fields().len() - 1);
proj.len() - 1
});
proj.push(pos_table_idx);
let pos_index = proj.len() - 1;
let scan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(
self.scan_config_builder(Arc::new(self.create_parquet_source(state, table_schema)?))
.with_file_group(FileGroup::new(vec![
self.partitioned_file(&table_file.file)?,
]))
.with_projection_indices(Some(proj))?
.build(),
);
let mut plan = self.split_across_partitions(scan, state, &file_cfg)?;
if !deleted_positions.is_empty() {
plan = Arc::new(DeleteFilterExec::try_new(
plan,
table_file.file.path.clone(),
Arc::new(deleted_positions),
pos_index,
)?);
}
plan = self.present_file_read_config(plan, &file_cfg, physical_len);
Ok(UpdateSourceScan {
scan: plan,
physical_len,
embedded_batch_idx,
embedded_snapshot_batch_idx,
pos_index,
row_id_start: table_file.row_id_start,
existing_parquet_deleted,
data_file_id: table_file.data_file_id,
delete_file_id: table_file.delete_file_id,
source_path: table_file.file.path.clone(),
})
}
#[cfg(feature = "write")]
pub(crate) fn apply_update_to_batches(
&self,
scan: &UpdateSourceScan,
batches: &[RecordBatch],
predicate: Option<&Arc<dyn PhysicalExpr>>,
assignments: &[(usize, Arc<dyn PhysicalExpr>)],
) -> DataFusionResult<FileUpdateOutput> {
let out_schema = rewrite_output_schema(&self.physical_schema);
let mut updated_batches: Vec<RecordBatch> = Vec::new();
let mut new_positions: Vec<i64> = Vec::new();
for batch in batches {
let Some(rewritten) = rewrite_scanned_batch(
&self.physical_schema,
&out_schema,
scan,
batch,
predicate,
assignments,
)?
else {
continue;
};
updated_batches.push(rewritten.batch);
new_positions.extend(rewritten.positions);
}
let matched_count = new_positions.len();
let mut cumulative = scan.existing_parquet_deleted.clone();
cumulative.extend(new_positions);
let mut cumulative_positions: Vec<i64> = cumulative.into_iter().collect();
cumulative_positions.sort_unstable();
Ok(FileUpdateOutput {
updated_batches,
matched_count,
cumulative_positions,
})
}
}
#[cfg(feature = "write")]
#[derive(Clone, Copy)]
pub(crate) struct MergeSourceFacts {
pub(crate) drops_current_columns: bool,
pub(crate) has_embedded_snapshot: bool,
}
#[cfg(feature = "write")]
pub(crate) fn rewrite_output_schema(physical_schema: &SchemaRef) -> SchemaRef {
let mut fields: Vec<Arc<Field>> = physical_schema.fields().iter().cloned().collect();
fields.push(Arc::new(rowid_field()));
Arc::new(Schema::new(fields))
}
#[cfg(feature = "write")]
pub(crate) struct RewrittenBatch {
pub(crate) batch: RecordBatch,
pub(crate) positions: Vec<i64>,
pub(crate) origin_snapshots: Option<ArrayRef>,
}
#[cfg(feature = "write")]
pub(crate) fn rewrite_scanned_batch(
physical_schema: &SchemaRef,
out_schema: &SchemaRef,
scan: &UpdateSourceScan,
batch: &RecordBatch,
predicate: Option<&Arc<dyn PhysicalExpr>>,
assignments: &[(usize, Arc<dyn PhysicalExpr>)],
) -> DataFusionResult<Option<RewrittenBatch>> {
let physical_len = scan.physical_len;
let n = batch.num_rows();
if n == 0 {
return Ok(None);
}
let mut phys_cols: Vec<ArrayRef> = Vec::with_capacity(physical_len);
for i in 0..physical_len {
phys_cols.push(crate::column_rename::coerce_column(
batch.column(i),
physical_schema.field(i).data_type(),
)?);
}
let phys_batch = RecordBatch::try_new(Arc::clone(physical_schema), phys_cols.clone())?;
let row_pos = batch
.column(scan.pos_index)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| {
DataFusionError::Internal("physical-position column is not Int64".to_string())
})?;
let mask: BooleanArray = match predicate {
Some(p) => {
let arr = p.evaluate(&phys_batch)?.into_array(n)?;
let b = arr.as_any().downcast_ref::<BooleanArray>().ok_or_else(|| {
DataFusionError::Execution(
"UPDATE predicate did not evaluate to a boolean".to_string(),
)
})?;
BooleanArray::from(
(0..n)
.map(|i| b.is_valid(i) && b.value(i))
.collect::<Vec<bool>>(),
)
},
None => BooleanArray::from(vec![true; n]),
};
if mask.true_count() == 0 {
return Ok(None);
}
let matched_phys: Vec<ArrayRef> = phys_cols
.iter()
.enumerate()
.map(|(i, column)| {
let filtered = arrow::compute::filter(column.as_ref(), &mask)
.map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
crate::column_rename::coerce_column(&filtered, physical_schema.field(i).data_type())
})
.collect::<DataFusionResult<_>>()?;
let matched_batch = RecordBatch::try_new(Arc::clone(physical_schema), matched_phys.clone())?;
let matched_rows = matched_batch.num_rows();
let mut out_cols = matched_phys;
for (col_idx, expr) in assignments {
let val = expr.evaluate(&matched_batch)?.into_array(matched_rows)?;
out_cols[*col_idx] =
crate::column_rename::coerce_column(&val, physical_schema.field(*col_idx).data_type())?;
}
let matched_pos = arrow::compute::filter(row_pos, &mask)
.map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
let matched_pos = matched_pos
.as_any()
.downcast_ref::<Int64Array>()
.expect("filtered Int64Array");
let rowid_col: ArrayRef = if let Some(idx) = scan.embedded_batch_idx {
let embedded = batch
.column(idx)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| {
DataFusionError::Internal("embedded rowid column is not Int64".to_string())
})?;
let embedded: ArrayRef = Arc::new(embedded.clone());
arrow::compute::filter(embedded.as_ref(), &mask)
.map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?
} else {
let start = scan
.row_id_start
.expect("row_id_start checked in build_update_scan");
Arc::new(Int64Array::from(
matched_pos
.values()
.iter()
.map(|p| start + p)
.collect::<Vec<i64>>(),
))
};
out_cols.push(rowid_col);
let origin_snapshots = match scan.embedded_snapshot_batch_idx {
Some(idx) => Some(
arrow::compute::filter(batch.column(idx).as_ref(), &mask)
.map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?,
),
None => None,
};
Ok(Some(RewrittenBatch {
batch: RecordBatch::try_new(Arc::clone(out_schema), out_cols)?,
positions: matched_pos.values().to_vec(),
origin_snapshots,
}))
}
#[cfg(feature = "write")]
#[derive(Clone, Debug)]
pub(crate) struct UpdateSourceScan {
pub(crate) scan: Arc<dyn ExecutionPlan>,
pub(crate) physical_len: usize,
pub(crate) embedded_batch_idx: Option<usize>,
pub(crate) embedded_snapshot_batch_idx: Option<usize>,
pub(crate) pos_index: usize,
pub(crate) row_id_start: Option<i64>,
pub(crate) existing_parquet_deleted: HashSet<i64>,
pub(crate) data_file_id: i64,
pub(crate) delete_file_id: Option<i64>,
pub(crate) source_path: String,
}
#[cfg(feature = "write")]
pub(crate) struct FileUpdateOutput {
pub(crate) updated_batches: Vec<RecordBatch>,
pub(crate) matched_count: usize,
pub(crate) cumulative_positions: Vec<i64>,
}
#[async_trait]
impl TableProvider for DuckLakeTable {
fn schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
fn table_type(&self) -> TableType {
TableType::Base
}
fn get_column_default(&self, column: &str) -> Option<&Expr> {
self.column_defaults.get(column)
}
fn statistics(&self) -> Option<Statistics> {
let mut statistics = self.table_statistics.clone();
if self.row_lineage {
statistics
.column_statistics
.push(ColumnStatistics::new_unknown());
}
Some(statistics)
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
Ok(filters
.iter()
.map(|_| TableProviderFilterPushDown::Inexact)
.collect())
}
async fn scan(
&self,
state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let inlined_filter = translate_inlined_filters(filters);
let rowid_idx = self.rowid_index();
let rowid_in_proj = match (rowid_idx, projection) {
(Some(r), Some(p)) => p.contains(&r),
(Some(_), None) => true,
(None, _) => false,
};
if rowid_in_proj {
let inlined = self.provider.scan_inlined_data(
self.table_id,
self.snapshot_id,
&self.columns,
inlined_filter.as_ref(),
)?;
self.inlined_materialized_row_count
.store(inlined.materialized_row_count, Ordering::Relaxed);
if inlined.batches.iter().any(|batch| batch.num_rows() > 0) {
return Err(crate::DuckLakeError::Unsupported(format!(
"row-lineage (rowid) scan on a table with inlined rows is not supported; \
{INLINED_DATA_REMEDIATION}"
))
.into());
}
}
let mut execs: Vec<Arc<dyn ExecutionPlan>> = Vec::new();
let inlined_deletes = self.inlined_deletes_by_file()?;
let conjuncts = match self.physical_conjuncts(state, filters) {
Ok(conjuncts) => conjuncts,
Err(error) => {
tracing::debug!(%error, "skipping plan-time file pruning");
Vec::new()
},
};
let stats_filter = datafusion::physical_expr::conjunction_opt(conjuncts.iter().cloned())
.and_then(|predicate| {
stats_filter::lower_predicate(&predicate, &self.physical_schema, &self.columns)
});
let pruning = match self.pruning_predicates(conjuncts) {
Ok(pruning) => pruning,
Err(error) => {
tracing::debug!(%error, "skipping plan-time file pruning");
Vec::new()
},
};
for metadata in self.file_metadata_pages("planning", stats_filter.as_ref()) {
let (table_files, file_statistics, summary_statistics) =
self.page_files_with_both_statistics(metadata?);
#[cfg(feature = "encryption")]
self.configure_encryption_factory(&table_files)?;
let table_files =
self.prune_table_files_iteratively(&pruning, &table_files, &file_statistics);
if rowid_in_proj {
let rowid_idx = rowid_idx.unwrap();
let user_proj: Vec<usize> = projection
.cloned()
.unwrap_or_else(|| (0..self.schema.fields().len()).collect());
for table_file in table_files {
let exec = if self.needs_snapshot_filter(table_file) {
let output_schema =
self.output_schema_for_projection(&user_proj, rowid_idx);
self.build_exec_for_partial_file(
state,
table_file,
inlined_deletes.get(&table_file.data_file_id),
output_schema,
)
.await?
} else {
self.build_exec_for_file_with_rowid(
state,
table_file,
inlined_deletes.get(&table_file.data_file_id),
&file_statistics,
&user_proj,
rowid_idx,
limit,
)
.await?
};
execs.push(exec);
}
continue;
}
let (needs_filter, rest): (Vec<_>, Vec<_>) = table_files
.into_iter()
.partition(|table_file| self.needs_snapshot_filter(table_file));
let (files_with_deletes, files_without_deletes): (Vec<_>, Vec<_>) =
rest.into_iter().partition(|table_file| {
table_file.delete_file.is_some()
|| inlined_deletes.contains_key(&table_file.data_file_id)
});
if !files_without_deletes.is_empty() {
execs.push(
self.build_exec_for_files_without_deletes(
state,
&files_without_deletes,
&file_statistics,
&summary_statistics,
projection,
limit,
)
.await?,
);
}
for table_file in files_with_deletes {
execs.push(
self.build_exec_for_file_with_deletes(
state,
table_file,
inlined_deletes.get(&table_file.data_file_id),
&file_statistics,
projection,
limit,
)
.await?,
);
}
for table_file in needs_filter {
let output_schema = match projection {
Some(indices) => Arc::new(self.schema.project(indices)?),
None => self.schema.clone(),
};
execs.push(
self.build_exec_for_partial_file(
state,
table_file,
inlined_deletes.get(&table_file.data_file_id),
output_schema,
)
.await?,
);
}
}
if rowid_in_proj {
if execs.is_empty() {
use datafusion::physical_plan::empty::EmptyExec;
let rowid_idx = rowid_idx.unwrap();
let user_proj: Vec<usize> = projection
.cloned()
.unwrap_or_else(|| (0..self.schema.fields().len()).collect());
let schema = self.output_schema_for_projection(&user_proj, rowid_idx);
return Ok(Arc::new(EmptyExec::new(schema)));
}
return combine_execution_plans(execs);
}
let inlined = self.provider.scan_inlined_data(
self.table_id,
self.snapshot_id,
&self.columns,
inlined_filter.as_ref(),
)?;
self.inlined_materialized_row_count
.store(inlined.materialized_row_count, Ordering::Relaxed);
if inlined.batches.iter().any(|b| b.num_rows() > 0) {
let exec = MemorySourceConfig::try_new_exec(
&[inlined.batches],
self.physical_schema.clone(),
projection.cloned(),
)?;
execs.push(exec);
}
if execs.is_empty() {
use datafusion::physical_plan::empty::EmptyExec;
let projected_schema = match projection {
Some(indices) => Arc::new(self.schema.project(indices)?),
None => self.schema.clone(),
};
return Ok(Arc::new(EmptyExec::new(projected_schema)));
}
combine_execution_plans(execs)
}
#[cfg(feature = "write")]
async fn insert_into(
&self,
_state: &dyn Session,
input: Arc<dyn ExecutionPlan>,
insert_op: InsertOp,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let writer = self.writer.as_ref().ok_or_else(|| {
DataFusionError::Plan(
"Table is read-only. Use DuckLakeCatalog::with_writer() to enable writes."
.to_string(),
)
})?;
let schema_name = self.schema_name.as_ref().ok_or_else(|| {
DataFusionError::Internal("Schema name not set for writable table".to_string())
})?;
self.write_options
.validate()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let write_mode = match insert_op {
InsertOp::Append => WriteMode::Append,
InsertOp::Overwrite | InsertOp::Replace => WriteMode::Replace,
};
let head_snapshot = self
.provider
.get_current_snapshot()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let live_spec = self
.provider
.get_partition_spec(self.table_id, head_snapshot)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let partition = match live_spec.as_ref() {
None => None,
Some(spec) => {
let column_ids: Vec<i64> = self.columns.iter().map(|c| c.column_id).collect();
Some(
crate::partition::PartitionWriteSpec::resolve(
spec,
&column_ids,
self.physical_schema.as_ref(),
)
.map_err(|e| DataFusionError::External(Box::new(e)))?,
)
},
};
let ordering = if self.write_options.sort_on_insert.unwrap_or(true) {
let live_sort = self
.provider
.get_sort_spec(self.table_id, head_snapshot)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
sort_ordering_for(&input.schema(), live_sort.as_ref())?
} else {
None
};
let input = match ordering.clone() {
Some(ordering) => Arc::new(datafusion::physical_plan::sorts::sort::SortExec::new(
ordering, input,
)) as Arc<dyn ExecutionPlan>,
None => input,
};
Ok(Arc::new(DuckLakeInsertExec::new(
input,
Arc::clone(writer),
schema_name.clone(),
self.table_name.clone(),
self.schema(),
write_mode,
self.object_store_url.clone(),
partition,
self.write_options.clone(),
ordering,
)))
}
#[cfg(feature = "write")]
async fn update(
&self,
state: &dyn Session,
assignments: Vec<(String, Expr)>,
filters: Vec<Expr>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let writer = self.writer.as_ref().ok_or_else(|| {
DataFusionError::Plan(
"Table is read-only. Use DuckLakeCatalog::with_writer() to enable writes."
.to_string(),
)
})?;
let schema_name = self.schema_name.as_ref().ok_or_else(|| {
DataFusionError::Internal("Schema name not set for writable table".to_string())
})?;
self.write_options
.validate()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
if !writer.supports_update() {
return Err(DataFusionError::NotImplemented(
"UPDATE not supported on this metadata backend".to_string(),
));
}
let inlined = self.provider.get_inlined_data(
self.table_id,
self.provider.get_current_snapshot()?,
&self.columns,
)?;
if inlined.iter().any(|batch| batch.num_rows() > 0) {
return Err(crate::DuckLakeError::Unsupported(format!(
"UPDATE on a table with inlined rows is not supported; \
{INLINED_DATA_REMEDIATION}"
))
.into());
}
let df_schema = DFSchema::try_from(self.physical_schema.as_ref().clone())
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let mut phys_assignments: Vec<(usize, Arc<dyn PhysicalExpr>)> =
Vec::with_capacity(assignments.len());
for (col_name, expr) in assignments {
let idx = self.physical_schema.index_of(&col_name).map_err(|_| {
DataFusionError::Plan(format!(
"UPDATE assignment targets unknown column '{col_name}'"
))
})?;
let pexpr = state.create_physical_expr(expr, &df_schema)?;
phys_assignments.push((idx, pexpr));
}
let mut predicate: Option<Arc<dyn PhysicalExpr>> = None;
for f in filters {
let pe = state.create_physical_expr(f, &df_schema)?;
predicate = Some(match predicate {
None => pe,
Some(prev) => Arc::new(BinaryExpr::new(prev, Operator::And, pe)),
});
}
let table_files = self
.files()
.map_err(|error| DataFusionError::External(Box::new(error)))?;
let inlined_deletes = self.inlined_deletes_by_file()?;
let mut scans = Vec::with_capacity(table_files.len());
for tf in &table_files {
scans.push(
self.build_update_scan(state, tf, inlined_deletes.get(&tf.data_file_id))
.await?,
);
}
Ok(Arc::new(DuckLakeUpdateExec::new(
Arc::new(self.read_only_clone()),
Arc::clone(writer),
schema_name.clone(),
self.table_name.clone(),
self.snapshot_id,
scans,
phys_assignments,
predicate,
self.object_store_url.clone(),
)))
}
#[cfg(feature = "write")]
async fn delete_from(
&self,
state: &dyn Session,
filters: Vec<Expr>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
use datafusion::logical_expr::utils::conjunction;
let writer = self.writer.as_ref().ok_or_else(|| {
DataFusionError::Plan(
"Table is read-only. Use DuckLakeCatalog::with_writer() to enable writes."
.to_string(),
)
})?;
let schema_name = self.schema_name.as_ref().ok_or_else(|| {
DataFusionError::Internal("Schema name not set for writable table".to_string())
})?;
self.write_options
.validate()
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let predicate = match conjunction(filters) {
None => None,
Some(expr) => {
let df_schema =
datafusion::common::DFSchema::try_from(self.physical_schema.as_ref().clone())?;
Some(state.create_physical_expr(expr, &df_schema)?)
},
};
let session_state = state
.as_any()
.downcast_ref::<datafusion::execution::SessionState>()
.ok_or_else(|| {
DataFusionError::NotImplemented(
"DELETE on a DuckLake table requires a DataFusion SessionState session"
.to_string(),
)
})?
.clone();
Ok(Arc::new(DuckLakeDeleteExec::new(
Arc::new(self.clone()),
session_state,
predicate,
Arc::clone(writer),
schema_name.clone(),
self.table_name.clone(),
self.table_id,
self.snapshot_id,
self.object_store_url.clone(),
)))
}
}
fn combine_execution_plans(
execs: Vec<Arc<dyn ExecutionPlan>>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
if execs.len() == 1 {
Ok(execs.into_iter().next().unwrap())
} else {
use datafusion::physical_plan::union::UnionExec;
UnionExec::try_new(execs)
}
}
fn extract_deleted_positions_from_batch(
batch: &RecordBatch,
positions: &mut HashSet<i64>,
) -> DataFusionResult<()> {
let schema = batch.schema();
let pos_idx = schema.index_of(DELETE_POS_COL)?;
let pos_array = batch
.column(pos_idx)
.as_any()
.downcast_ref::<Int64Array>()
.ok_or_else(|| {
DataFusionError::Internal(format!("{} column not found or wrong type", DELETE_POS_COL))
})?;
for i in 0..batch.num_rows() {
if !pos_array.is_null(i) {
positions.insert(pos_array.value(i));
}
}
Ok(())
}
fn is_object_store_not_found(err: &DataFusionError) -> bool {
if let DataFusionError::ObjectStore(os_err) = err {
return matches!(&**os_err, object_store::Error::NotFound { .. });
}
let mut source = std::error::Error::source(err);
while let Some(e) = source {
if let Some(os_err) = e.downcast_ref::<object_store::Error>() {
return matches!(os_err, object_store::Error::NotFound { .. });
}
source = e.source();
}
false
}
#[cfg(feature = "write")]
fn sort_ordering_for(
schema: &arrow::datatypes::Schema,
sort_spec: Option<&crate::sort::SortSpec>,
) -> datafusion::common::Result<Option<datafusion::physical_expr::LexOrdering>> {
let Some(sort_spec) = sort_spec else {
return Ok(None);
};
let keys = sort_spec.producible_columns().ok_or_else(|| {
DataFusionError::NotImplemented(format!(
"DuckLake sort order {} contains an unsupported expression; \
datafusion-ducklake can write only bare-column sort keys",
sort_spec.sort_id
))
})?;
if keys.is_empty() {
return Ok(None);
}
let mut sort_exprs = Vec::with_capacity(keys.len());
for (name, direction, null_order) in &keys {
let index = schema.index_of(name).map_err(|_| {
DataFusionError::Plan(format!(
"DuckLake sort key '{name}' is not present in the write schema"
))
})?;
let column = Arc::new(datafusion::physical_expr::expressions::Column::new(
name, index,
));
sort_exprs.push(
datafusion::physical_expr_common::sort_expr::PhysicalSortExpr::new(
column,
arrow::compute::SortOptions {
descending: matches!(direction, crate::sort::SortDirection::Desc),
nulls_first: null_order.nulls_first(),
},
),
);
}
Ok(datafusion::physical_expr::LexOrdering::new(sort_exprs))
}
fn restate_in_physical_row_space(
statistics: &mut HashMap<i64, Arc<Statistics>>,
files: &[DuckLakeTableFile],
) {
fn physical<T: Clone + std::fmt::Debug + Eq + PartialOrd>(
value: &Precision<T>,
) -> Precision<T> {
match value {
Precision::Inexact(value) => Precision::Exact(value.clone()),
other => other.clone(),
}
}
for file in files.iter().filter(|f| f.delete_file.is_some()) {
let Some(current) = statistics.get(&file.data_file_id) else {
continue;
};
let mut restated = current.as_ref().clone();
restated.num_rows = Precision::Absent;
for column in &mut restated.column_statistics {
column.min_value = physical(&column.min_value);
column.max_value = physical(&column.max_value);
column.null_count = physical(&column.null_count);
}
statistics.insert(file.data_file_id, Arc::new(restated));
}
}
struct FilePruningStatistics {
base: PrunableStatistics,
statistics: Vec<Arc<Statistics>>,
schema: SchemaRef,
}
impl FilePruningStatistics {
fn new(statistics: Vec<Arc<Statistics>>, schema: SchemaRef) -> Self {
let base = PrunableStatistics::new(statistics.clone(), Arc::clone(&schema));
Self {
base,
statistics,
schema,
}
}
fn exact_values(
&self,
column: &Column,
get_statistic: impl Fn(&ColumnStatistics) -> &Precision<ScalarValue>,
) -> Option<ArrayRef> {
let index = self.schema.index_of(column.name()).ok()?;
let typed_null = ScalarValue::try_new_null(self.schema.field(index).data_type()).ok()?;
let mut has_value = false;
let values = self.statistics.iter().map(|statistics| {
statistics
.column_statistics
.get(index)
.and_then(|statistics| match get_statistic(statistics) {
Precision::Exact(value) => {
has_value = true;
Some(value.clone())
},
_ => None,
})
.unwrap_or_else(|| typed_null.clone())
});
ScalarValue::iter_to_array(values)
.ok()
.filter(|_| has_value)
}
}
impl PruningStatistics for FilePruningStatistics {
fn min_values(&self, column: &Column) -> Option<ArrayRef> {
self.exact_values(column, |statistics| &statistics.min_value)
}
fn max_values(&self, column: &Column) -> Option<ArrayRef> {
self.exact_values(column, |statistics| &statistics.max_value)
}
fn num_containers(&self) -> usize {
self.base.num_containers()
}
fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
self.base.null_counts(column)
}
fn row_counts(&self) -> Option<ArrayRef> {
self.base.row_counts()
}
fn contained(&self, column: &Column, values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
self.base.contained(column, values)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metadata_provider::{
ColumnWithTable, DataFileChange, DeleteFileChange, FileWithTable, SchemaMetadata,
SnapshotMetadata, TableMetadata, TableWithSchema,
};
use crate::partition::{PartitionSpecColumn, PartitionTransform};
use crate::types::build_arrow_schema;
#[cfg(feature = "write")]
use datafusion::datasource::physical_plan::FileScanConfig;
use datafusion::prelude::{SessionContext, col, lit};
#[cfg(feature = "write")]
use object_store::ObjectStore;
use rstest::rstest;
use std::sync::atomic::{AtomicUsize, Ordering};
#[rstest]
#[case("NULL", DataType::Utf8, ScalarValue::Utf8(None))]
#[case("null", DataType::Int32, ScalarValue::Int32(None))]
#[case("", DataType::Int32, ScalarValue::Int32(None))]
#[case("", DataType::Utf8, ScalarValue::Utf8(Some(String::new())))]
#[case("__HIVE_DEFAULT_PARTITION__", DataType::Int32, ScalarValue::Int32(None))]
#[case("0x10", DataType::Int32, ScalarValue::Int32(Some(16)))]
#[case("1e2", DataType::Decimal128(10, 2), ScalarValue::Decimal128(Some(10_000), 10, 2))]
#[case("1.25e-1", DataType::Decimal128(10, 3), ScalarValue::Decimal128(Some(125), 10, 3))]
fn hive_partition_values_match_duckdb(
#[case] encoded: &str,
#[case] data_type: DataType,
#[case] expected: ScalarValue,
) {
let path = format!("file:///tmp/p={encoded}/data.parquet");
let actual = DuckLakeTable::hive_partition_value(&path, "p", &data_type).unwrap();
assert_eq!(actual, expected);
}
#[rstest]
#[case(
"s3://bucket/a=one/../p=7/file.parquet",
"a",
DataType::Utf8,
ScalarValue::Utf8(Some("one".to_string()))
)]
#[case(
"s3://bucket/p=a#b/file.parquet",
"p",
DataType::Utf8,
ScalarValue::Utf8(Some("a#b".to_string()))
)]
#[case(
"s3://bucket/not=a?b/p=7/file.parquet",
"p",
DataType::Int32,
ScalarValue::Int32(Some(7))
)]
#[case(r"s3:\bucket\p=7\file.parquet", "p", DataType::Int32, ScalarValue::Int32(Some(7)))]
#[case(
"s3://bucket/p%20q=7/file.parquet",
"p%20q",
DataType::Int32,
ScalarValue::Int32(Some(7))
)]
#[case(
"s3://bucket/p=a%20b/file.parquet",
"p",
DataType::Utf8,
ScalarValue::Utf8(Some("a b".to_string()))
)]
fn hive_partition_path_parsing_matches_duckdb(
#[case] path: &str,
#[case] source_name: &str,
#[case] data_type: DataType,
#[case] expected: ScalarValue,
) {
let actual = DuckLakeTable::hive_partition_value(path, source_name, &data_type).unwrap();
assert_eq!(actual, expected);
}
#[rstest]
#[case("s3://bucket/p=7=8/file.parquet", "p")]
#[case("s3://bucket/p%20q=7/file.parquet", "p q")]
fn hive_partition_path_rejects_non_candidates(#[case] path: &str, #[case] source_name: &str) {
let error =
DuckLakeTable::hive_partition_value(path, source_name, &DataType::Int32).unwrap_err();
assert!(error.to_string().contains("absent from file path"));
}
#[derive(Debug, Default)]
struct LazyMillionFileProvider {
eager_file_reads: AtomicUsize,
max_page: AtomicUsize,
page_calls: AtomicUsize,
}
impl MetadataProvider for LazyMillionFileProvider {
fn get_current_snapshot(&self) -> Result<i64> {
Ok(1)
}
fn get_data_path(&self) -> Result<String> {
Ok("memory:///".to_string())
}
fn list_snapshots(&self) -> Result<Vec<SnapshotMetadata>> {
unimplemented!()
}
fn list_schemas(&self, _snapshot_id: i64) -> Result<Vec<SchemaMetadata>> {
unimplemented!()
}
fn list_tables(&self, _schema_id: i64, _snapshot_id: i64) -> Result<Vec<TableMetadata>> {
unimplemented!()
}
fn get_table_structure(
&self,
table_id: i64,
_snapshot_id: i64,
) -> Result<Vec<DuckLakeTableColumn>> {
if table_id == 2 {
return Ok(vec![
DuckLakeTableColumn::new(
1,
"range_value".to_string(),
"bigint".to_string(),
false,
),
DuckLakeTableColumn::new(
2,
"partition_key".to_string(),
"bigint".to_string(),
false,
),
]);
}
Ok(vec![DuckLakeTableColumn::new(
1,
"id".to_string(),
"bigint".to_string(),
false,
)])
}
fn get_table_files_for_select(
&self,
_table_id: i64,
_snapshot_id: i64,
) -> Result<Vec<DuckLakeTableFile>> {
self.eager_file_reads.fetch_add(1, Ordering::Relaxed);
panic!("the eager file API must not be used while planning a scan")
}
fn get_table_summary_statistics(
&self,
table_id: i64,
_snapshot_id: i64,
) -> Result<DuckLakeStatistics> {
if table_id == 2 {
return Ok(DuckLakeStatistics::default());
}
Ok(DuckLakeStatistics {
table: Some(crate::metadata_provider::DuckLakeTableStatistics {
record_count: Some(1_000_000),
file_size_bytes: Some(1_000_000),
}),
columns: vec![DuckLakeTableColumnStatistics {
column_id: 1,
contains_null: Some(false),
min_value: Some("1".to_string()),
max_value: Some("1000000".to_string()),
contains_nan: None,
column_size_bytes: Some(1_000_000),
bounds_are_exact: true,
}],
..Default::default()
})
}
fn get_table_file_metadata_page(
&self,
_table_id: i64,
snapshot_id: i64,
after_data_file_id: Option<i64>,
limit: usize,
) -> Result<Vec<DuckLakeFileMetadata>> {
self.page_calls.fetch_add(1, Ordering::Relaxed);
let start = after_data_file_id.unwrap_or(0) + 1;
let end = (start + i64::try_from(limit).unwrap()).min(1_000_001);
let page: Vec<_> = (start..end)
.map(|data_file_id| DuckLakeFileMetadata {
file: DuckLakeTableFile {
data_file_id,
file: DuckLakeFileData::new(
format!("file-{data_file_id}.parquet"),
true,
1,
),
delete_file_id: None,
delete_file: None,
row_id_start: Some(data_file_id - 1),
snapshot_id: Some(snapshot_id),
begin_snapshot: Some(1),
schema_version: Some(0),
partial_max: None,
max_row_count: Some(1),
delete_count: None,
partition_id: None,
partition_values: Vec::new(),
},
column_statistics: vec![DuckLakeFileColumnStatistics {
data_file_id,
column_id: 1,
column_size_bytes: Some(1),
value_count: Some(1),
null_count: Some(0),
min_value: Some(data_file_id.to_string()),
max_value: Some(data_file_id.to_string()),
contains_nan: None,
}],
})
.collect();
self.max_page.fetch_max(page.len(), Ordering::Relaxed);
Ok(page)
}
fn get_schema_by_name(
&self,
_name: &str,
_snapshot_id: i64,
) -> Result<Option<SchemaMetadata>> {
unimplemented!()
}
fn get_table_by_name(
&self,
_schema_id: i64,
_name: &str,
_snapshot_id: i64,
) -> Result<Option<TableMetadata>> {
unimplemented!()
}
fn table_exists(&self, _schema_id: i64, _name: &str, _snapshot_id: i64) -> Result<bool> {
unimplemented!()
}
fn list_all_tables(&self, _snapshot_id: i64) -> Result<Vec<TableWithSchema>> {
unimplemented!()
}
fn list_all_columns(&self, _snapshot_id: i64) -> Result<Vec<ColumnWithTable>> {
unimplemented!()
}
fn list_all_files(&self, _snapshot_id: i64) -> Result<Vec<FileWithTable>> {
unimplemented!()
}
fn get_data_files_added_between_snapshots(
&self,
_table_id: i64,
_start_snapshot: i64,
_end_snapshot: i64,
) -> Result<Vec<DataFileChange>> {
unimplemented!()
}
fn get_delete_files_added_between_snapshots(
&self,
_table_id: i64,
_start_snapshot: i64,
_end_snapshot: i64,
) -> Result<Vec<DeleteFileChange>> {
unimplemented!()
}
}
#[test]
fn planning_prunes_a_lazy_million_file_fixture_in_bounded_pages() -> Result<()> {
let provider = Arc::new(LazyMillionFileProvider::default());
let table = DuckLakeTable::new(
1,
"events",
provider.clone(),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)?;
assert_eq!(provider.eager_file_reads.load(Ordering::Relaxed), 0);
let statistics = table.statistics().unwrap();
assert_eq!(statistics.num_rows, Precision::Exact(1_000_000));
assert_eq!(
statistics.column_statistics[0].min_value,
Precision::Exact(ScalarValue::Int64(Some(1)))
);
assert_eq!(
statistics.column_statistics[0].max_value,
Precision::Exact(ScalarValue::Int64(Some(1_000_000)))
);
assert_eq!(
statistics.column_statistics[0].byte_size,
Precision::Inexact(1_000_000)
);
let state = SessionContext::new().state();
let filters = [col("id").eq(lit(999_999_i64))];
let conjuncts = table.physical_conjuncts(&state, &filters).unwrap();
let pruning = table.pruning_predicates(conjuncts).unwrap();
let mut after = None;
let mut retained = Vec::new();
loop {
let metadata =
provider.get_table_file_metadata_page(1, 1, after, FILE_METADATA_BATCH_SIZE)?;
if metadata.is_empty() {
break;
}
after = metadata.last().map(|entry| entry.file.data_file_id);
let files: Vec<_> = metadata.iter().map(|entry| entry.file.clone()).collect();
let catalog_statistics = metadata
.into_iter()
.flat_map(|entry| entry.column_statistics)
.collect();
let (_, file_statistics) = build_datafusion_statistics(
table.physical_schema.as_ref(),
&table.columns,
&files,
DuckLakeStatistics {
files: catalog_statistics,
..Default::default()
},
false,
true,
);
retained.extend(
table
.prune_table_files_iteratively(&pruning, &files, &file_statistics)
.into_iter()
.map(|file| file.data_file_id),
);
}
assert_eq!(provider.max_page.load(Ordering::Relaxed), 4_096);
assert_eq!(provider.page_calls.load(Ordering::Relaxed), 246);
assert_eq!(retained, vec![999_999]);
assert_eq!(provider.eager_file_reads.load(Ordering::Relaxed), 0);
Ok(())
}
#[test]
fn iterative_pruning_revisits_range_conjunct_after_partition_conjunct() -> Result<()> {
let table = DuckLakeTable::new(
2,
"events",
Arc::new(LazyMillionFileProvider::default()),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)?;
let state = SessionContext::new().state();
let filters = [col("range_value").eq(lit(1_i64)), col("partition_key").eq(lit(1_i64))];
let predicates = table.pruning_predicates(table.physical_conjuncts(&state, &filters)?)?;
let file = |data_file_id| DuckLakeTableFile {
data_file_id,
file: DuckLakeFileData::new(format!("file-{data_file_id}.parquet"), true, 1),
delete_file_id: None,
delete_file: None,
row_id_start: None,
snapshot_id: Some(1),
begin_snapshot: Some(1),
schema_version: Some(0),
partial_max: None,
max_row_count: Some(1),
delete_count: None,
partition_id: None,
partition_values: Vec::new(),
};
let files = vec![file(1), file(2), file(3)];
let statistics = |range_value: Option<i64>, partition_key: i64| {
let mut statistics = Statistics::new_unknown(table.physical_schema.as_ref());
if let Some(value) = range_value {
statistics.column_statistics[0].min_value =
Precision::Exact(ScalarValue::Int64(Some(value)));
statistics.column_statistics[0].max_value =
Precision::Exact(ScalarValue::Int64(Some(value)));
}
statistics.column_statistics[1].min_value =
Precision::Exact(ScalarValue::Int64(Some(partition_key)));
statistics.column_statistics[1].max_value =
Precision::Exact(ScalarValue::Int64(Some(partition_key)));
Arc::new(statistics)
};
let file_statistics = HashMap::from([
(1, statistics(Some(1), 1)),
(2, statistics(Some(2), 1)),
(3, statistics(None, 2)),
]);
let retained = table
.prune_table_files_iteratively(&predicates, &files, &file_statistics)
.into_iter()
.map(|file| file.data_file_id)
.collect::<Vec<_>>();
assert_eq!(retained, vec![1]);
Ok(())
}
const ID_COLUMN_ID: i64 = 1;
const REGION_COLUMN_ID: i64 = 2;
#[derive(Debug)]
struct FixedFileProvider {
files: Vec<DuckLakeFileMetadata>,
partition_spec: Option<PartitionSpec>,
visible_when_filtered: Option<Vec<i64>>,
}
impl MetadataProvider for FixedFileProvider {
fn get_current_snapshot(&self) -> Result<i64> {
Ok(1)
}
fn get_data_path(&self) -> Result<String> {
Ok("memory:///".to_string())
}
fn get_table_structure(
&self,
_table_id: i64,
_snapshot_id: i64,
) -> Result<Vec<DuckLakeTableColumn>> {
Ok(vec![
DuckLakeTableColumn::new(
ID_COLUMN_ID,
"id".to_string(),
"bigint".to_string(),
false,
),
DuckLakeTableColumn::new(
REGION_COLUMN_ID,
"region".to_string(),
"bigint".to_string(),
false,
),
])
}
fn get_partition_spec(
&self,
_table_id: i64,
_snapshot_id: i64,
) -> Result<Option<PartitionSpec>> {
Ok(self.partition_spec.clone())
}
fn get_table_files_for_select(
&self,
_table_id: i64,
_snapshot_id: i64,
) -> Result<Vec<DuckLakeTableFile>> {
Ok(self.files.iter().map(|entry| entry.file.clone()).collect())
}
fn get_table_file_metadata_page(
&self,
_table_id: i64,
_snapshot_id: i64,
after_data_file_id: Option<i64>,
limit: usize,
) -> Result<Vec<DuckLakeFileMetadata>> {
Ok(self
.files
.iter()
.filter(|entry| {
after_data_file_id.is_none_or(|after| entry.file.data_file_id > after)
})
.take(limit)
.cloned()
.collect())
}
fn get_table_file_metadata_page_filtered(
&self,
table_id: i64,
snapshot_id: i64,
after_data_file_id: Option<i64>,
limit: usize,
filter: Option<&crate::stats_filter::StatsFilter>,
) -> Result<Vec<DuckLakeFileMetadata>> {
let Some(visible) = self
.visible_when_filtered
.as_ref()
.filter(|_| filter.is_some())
else {
return self.get_table_file_metadata_page(
table_id,
snapshot_id,
after_data_file_id,
limit,
);
};
Ok(self
.files
.iter()
.filter(|entry| visible.contains(&entry.file.data_file_id))
.filter(|entry| {
after_data_file_id.is_none_or(|after| entry.file.data_file_id > after)
})
.take(limit)
.cloned()
.collect())
}
fn get_table_summary_statistics(
&self,
_table_id: i64,
_snapshot_id: i64,
) -> Result<DuckLakeStatistics> {
Ok(DuckLakeStatistics::default())
}
fn list_snapshots(&self) -> Result<Vec<SnapshotMetadata>> {
unimplemented!()
}
fn list_schemas(&self, _snapshot_id: i64) -> Result<Vec<SchemaMetadata>> {
unimplemented!()
}
fn list_tables(&self, _schema_id: i64, _snapshot_id: i64) -> Result<Vec<TableMetadata>> {
unimplemented!()
}
fn get_schema_by_name(
&self,
_name: &str,
_snapshot_id: i64,
) -> Result<Option<SchemaMetadata>> {
unimplemented!()
}
fn get_table_by_name(
&self,
_schema_id: i64,
_name: &str,
_snapshot_id: i64,
) -> Result<Option<TableMetadata>> {
unimplemented!()
}
fn table_exists(&self, _schema_id: i64, _name: &str, _snapshot_id: i64) -> Result<bool> {
unimplemented!()
}
fn list_all_tables(&self, _snapshot_id: i64) -> Result<Vec<TableWithSchema>> {
unimplemented!()
}
fn list_all_columns(&self, _snapshot_id: i64) -> Result<Vec<ColumnWithTable>> {
unimplemented!()
}
fn list_all_files(&self, _snapshot_id: i64) -> Result<Vec<FileWithTable>> {
unimplemented!()
}
fn get_data_files_added_between_snapshots(
&self,
_table_id: i64,
_start_snapshot: i64,
_end_snapshot: i64,
) -> Result<Vec<DataFileChange>> {
unimplemented!()
}
fn get_delete_files_added_between_snapshots(
&self,
_table_id: i64,
_start_snapshot: i64,
_end_snapshot: i64,
) -> Result<Vec<DeleteFileChange>> {
unimplemented!()
}
}
type RegionValue = Option<Option<&'static str>>;
fn fixture_file(
data_file_id: i64,
region: RegionValue,
id_bounds: Option<(i64, i64)>,
) -> DuckLakeFileMetadata {
let mut file = DuckLakeTableFile::new(DuckLakeFileData::new(
format!("file-{data_file_id}.parquet"),
true,
1,
));
file.data_file_id = data_file_id;
file.snapshot_id = Some(1);
file.begin_snapshot = Some(1);
file.max_row_count = Some(1);
if let Some(value) = region {
file.partition_id = Some(1);
file.partition_values = vec![(0, value.map(str::to_string))];
}
let column_statistics = id_bounds
.map(|(min, max)| DuckLakeFileColumnStatistics {
data_file_id,
column_id: ID_COLUMN_ID,
column_size_bytes: Some(8),
value_count: Some(1),
null_count: Some(0),
min_value: Some(min.to_string()),
max_value: Some(max.to_string()),
contains_nan: None,
})
.into_iter()
.collect();
DuckLakeFileMetadata {
file,
column_statistics,
}
}
fn region_spec(prune_safe: bool) -> PartitionSpec {
PartitionSpec {
partition_id: 1,
columns: vec![PartitionSpecColumn {
partition_key_index: 0,
column_id: REGION_COLUMN_ID,
transform: PartitionTransform::Identity,
}],
prune_safe,
}
}
fn fixed_table(
files: Vec<DuckLakeFileMetadata>,
partition_spec: Option<PartitionSpec>,
) -> Result<DuckLakeTable> {
DuckLakeTable::new(
1,
"events",
Arc::new(FixedFileProvider {
files,
partition_spec,
visible_when_filtered: None,
}),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)
}
fn catalog_filtering_table(
files: Vec<DuckLakeFileMetadata>,
visible_when_filtered: Vec<i64>,
) -> Result<DuckLakeTable> {
DuckLakeTable::new(
1,
"events",
Arc::new(FixedFileProvider {
files,
partition_spec: None,
visible_when_filtered: Some(visible_when_filtered),
}),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)
}
#[test]
fn predicate_is_prunable_ignores_the_synthetic_rowid_rename() -> Result<()> {
let table = DuckLakeTable::new(
1,
"events",
Arc::new(LazyMillionFileProvider::default()),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(5_i64)));
let cfg = |renamed: Vec<usize>, mapping: Vec<(&str, &str)>| FileReadConfig {
read_schema: table.physical_schema.clone(),
name_mapping: mapping
.into_iter()
.map(|(a, b)| (a.to_string(), b.to_string()))
.collect(),
constants: HashMap::new(),
row_group_count: 4,
renamed_column_indices: Arc::new(renamed.into_iter().collect()),
embedded_rowid_parquet_name: None,
embedded_snapshot_parquet_name: None,
drops_current_columns: false,
};
let id_index = table.physical_schema.index_of("id").unwrap();
assert!(
table.predicate_is_prunable(
&predicate,
&cfg(
vec![],
vec![(
crate::row_id::EMBEDDED_ROW_ID_COLUMN_NAME,
ROWID_COLUMN_NAME
)]
)
),
"the synthetic rowid rename must not disable pruning"
);
assert!(
!table.predicate_is_prunable(&predicate, &cfg(vec![id_index], vec![])),
"a predicate on a renamed column must disable pruning"
);
assert!(
table.predicate_is_prunable(&predicate, &cfg(vec![id_index + 1], vec![])),
"an unrelated renamed/absent column must not disable pruning"
);
assert!(table.predicate_is_prunable(&predicate, &cfg(vec![], vec![])));
Ok(())
}
fn physical_predicate(table: &DuckLakeTable, expr: Expr) -> Arc<dyn PhysicalExpr> {
let df_schema = DFSchema::try_from(table.physical_schema.as_ref().clone()).unwrap();
SessionContext::new()
.state()
.create_physical_expr(expr, &df_schema)
.unwrap()
}
fn matching_ids(table: &DuckLakeTable, predicate: &Arc<dyn PhysicalExpr>) -> Result<Vec<i64>> {
Ok(table
.files_matching(predicate)?
.into_iter()
.map(|file| file.data_file_id)
.collect())
}
#[test]
fn files_matching_does_not_judge_a_delete_bearing_file_all_null() -> Result<()> {
let mut entry = fixture_file(1, None, Some((5, 5)));
entry.file.max_row_count = Some(2); entry.file.delete_count = Some(1); entry.file.delete_file_id = Some(1);
entry.file.delete_file = Some(DuckLakeFileData::new(
"delete-1.parquet".to_string(),
true,
1,
));
for stats in &mut entry.column_statistics {
stats.null_count = Some(1);
stats.value_count = Some(2);
}
let table = fixed_table(vec![entry], None)?;
let predicate = physical_predicate(&table, col("id").eq(lit(5_i64)));
assert_eq!(
matching_ids(&table, &predicate)?,
vec![1],
"a physical null_count equal to the LIVE row count must not make the \
file look entirely null -- the matching physical row is still there"
);
Ok(())
}
#[test]
fn files_matching_prunes_a_delete_bearing_file_by_its_partition_bound() -> Result<()> {
let delete_bearing = |data_file_id: i64, value: &'static str| {
let mut entry = fixture_file(data_file_id, Some(Some(value)), None);
entry.file.max_row_count = Some(10);
entry.file.delete_count = Some(1);
entry.file.delete_file_id = Some(data_file_id);
entry.file.delete_file = Some(DuckLakeFileData::new(
format!("delete-{data_file_id}.parquet"),
true,
1,
));
entry
};
let table = fixed_table(
vec![fixture_file(1, Some(Some("10")), None), delete_bearing(2, "9999")],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(&table, col("region").eq(lit(10_i64)));
assert_eq!(
matching_ids(&table, &predicate)?,
vec![1],
"a delete-bearing file whose partition bound excludes the value must \
still prune"
);
Ok(())
}
#[test]
fn files_matching_still_prunes_a_file_that_carries_deletes() -> Result<()> {
let carrying_deletes = |data_file_id: i64, id_bounds: (i64, i64)| {
let mut entry = fixture_file(data_file_id, None, Some(id_bounds));
entry.file.max_row_count = Some(10);
entry.file.delete_count = Some(1);
entry.file.delete_file_id = Some(data_file_id);
entry.file.delete_file = Some(DuckLakeFileData::new(
format!("delete-{data_file_id}.parquet"),
true,
1,
));
entry
};
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
carrying_deletes(2, (100, 110)),
fixture_file(3, None, Some((200, 210))),
],
None,
)?;
let inside = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(
matching_ids(&table, &inside)?,
vec![2],
"a delete-bearing file whose bounds admit the value must be kept"
);
let outside = physical_predicate(&table, col("id").eq(lit(5_i64)));
assert_eq!(
matching_ids(&table, &outside)?,
vec![1],
"a delete-bearing file whose bounds exclude the value must still prune"
);
Ok(())
}
#[test]
fn files_matching_prunes_unpartitioned_files_by_column_statistics() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
fixture_file(2, None, Some((100, 110))),
fixture_file(3, None, Some((200, 210))),
],
None,
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![2]);
Ok(())
}
#[cfg(feature = "encryption")]
const FIXTURE_ENCRYPTION_KEY: &str = "0123456789abcdef0123456789abcdef";
#[cfg(feature = "encryption")]
fn encrypted_fixture_file(data_file_id: i64, id_bounds: (i64, i64)) -> DuckLakeFileMetadata {
let mut entry = fixture_file(data_file_id, None, Some(id_bounds));
entry.file.file.encryption_key = Some(FIXTURE_ENCRYPTION_KEY.to_string());
entry
}
#[cfg(feature = "encryption")]
async fn installed_decryption_properties(
table: &DuckLakeTable,
path: &str,
) -> Option<Arc<parquet::encryption::decrypt::FileDecryptionProperties>> {
let factory = table
.encryption_keys
.factory
.get()
.cloned()
.expect("an encrypted table installs a factory");
factory
.get_file_decryption_properties(&Default::default(), &ObjectPath::from(path))
.await
.expect("key lookup succeeds")
}
#[cfg(feature = "encryption")]
#[tokio::test]
async fn files_matching_keeps_decryption_keys_for_pruned_files() {
let table = fixed_table(
vec![encrypted_fixture_file(1, (1, 10)), encrypted_fixture_file(2, (100, 110))],
None,
)
.expect("table opens");
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
let matching = table.files_matching(&predicate).expect("pruning succeeds");
assert_eq!(
matching.len(),
1,
"file 1's statistics exclude the key, so it is pruned"
);
assert!(
installed_decryption_properties(&table, "file-1.parquet")
.await
.is_some(),
"the pruned file's key must stay installed for other readers",
);
}
#[cfg(feature = "encryption")]
#[tokio::test]
async fn files_matching_keeps_keys_the_catalog_filter_left_unlisted() {
let table = catalog_filtering_table(
vec![encrypted_fixture_file(1, (1, 10)), encrypted_fixture_file(2, (100, 110))],
vec![2],
)
.expect("table opens");
assert_eq!(table.files().expect("listing succeeds").len(), 2);
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
let matching = table.files_matching(&predicate).expect("pruning succeeds");
assert_eq!(
matching
.iter()
.map(|file| file.data_file_id)
.collect::<Vec<_>>(),
vec![2],
"the catalog leaves file 1 out of the listing entirely"
);
assert!(
installed_decryption_properties(&table, "file-1.parquet")
.await
.is_some(),
"a file the catalog never listed must keep the key an earlier listing installed",
);
assert!(
installed_decryption_properties(&table, "file-2.parquet")
.await
.is_some(),
"the returned file's own key must be installed",
);
}
#[cfg(feature = "encryption")]
#[tokio::test]
async fn installing_a_second_key_for_a_path_supersedes_the_first() {
const SUPERSEDED_KEY: &str = "fedcba9876543210fedcba9876543210";
let table = fixed_table(Vec::new(), None).expect("table opens");
let install = |key: &str| {
table.install_encryption_keys(HashMap::from([(
"file-1.parquet".to_string(),
key.to_string(),
)]));
};
install(SUPERSEDED_KEY);
install(FIXTURE_ENCRYPTION_KEY);
let properties = installed_decryption_properties(&table, "file-1.parquet")
.await
.expect("the path has a key installed");
assert_eq!(
properties
.footer_key(None)
.expect("a footer key")
.as_slice(),
DuckLakeEncryptionFactory::decode_key(FIXTURE_ENCRYPTION_KEY)
.expect("the fixture key decodes")
.as_slice(),
"the factory must serve the key installed last, not the superseded one",
);
}
#[cfg(feature = "encryption")]
#[tokio::test]
async fn one_factory_serves_the_keys_of_every_page() {
let table = fixed_table(Vec::new(), None).expect("table opens");
let install = |path: &str| {
table.install_encryption_keys(HashMap::from([(
path.to_string(),
FIXTURE_ENCRYPTION_KEY.to_string(),
)]));
};
let installed_factory = || {
table
.encryption_keys
.factory
.get()
.cloned()
.expect("an encrypted table installs a factory")
};
install("file-1.parquet");
let from_first_page = installed_factory();
install("file-2.parquet");
assert!(
Arc::ptr_eq(&from_first_page, &installed_factory()),
"a later page must reuse the installed factory, not rebuild it",
);
assert!(
from_first_page
.get_file_decryption_properties(
&Default::default(),
&ObjectPath::from("file-2.parquet"),
)
.await
.expect("key lookup succeeds")
.is_some(),
"the factory handed out on page 1 must resolve a key installed on page 2",
);
}
#[test]
fn files_matching_keeps_files_without_statistics() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
fixture_file(2, None, None),
fixture_file(3, None, Some((200, 210))),
],
None,
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![2]);
Ok(())
}
fn empty_file_without_statistics(data_file_id: i64) -> DuckLakeFileMetadata {
let mut entry = fixture_file(data_file_id, None, None);
entry.file.max_row_count = Some(0);
entry
}
fn file_with_unknown_row_count(
data_file_id: i64,
id_bounds: (i64, i64),
) -> DuckLakeFileMetadata {
let mut entry = fixture_file(data_file_id, None, Some(id_bounds));
entry.file.max_row_count = None;
entry
}
#[test]
fn files_matching_prunes_around_an_empty_file_without_statistics() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
fixture_file(2, None, Some((100, 110))),
fixture_file(3, None, Some((200, 210))),
empty_file_without_statistics(4),
],
None,
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![2]);
Ok(())
}
#[test]
fn files_matching_keeps_a_file_whose_row_count_is_unknown() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
file_with_unknown_row_count(2, (100, 110)),
empty_file_without_statistics(3),
],
None,
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![2]);
Ok(())
}
#[test]
fn files_matching_keeps_a_file_with_a_null_partition_value() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("10")), None),
fixture_file(2, Some(Some("9999")), None),
fixture_file(3, Some(None), None),
],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(&table, col("region").eq(lit(10_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![1, 3]);
Ok(())
}
#[test]
fn files_matching_keeps_a_file_with_an_undecodable_partition_value() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("10")), None),
fixture_file(2, Some(Some("9999")), None),
fixture_file(3, Some(Some("not-a-region")), None),
],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(&table, col("region").eq(lit(10_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![1, 3]);
Ok(())
}
#[test]
fn files_matching_never_prunes_a_non_partition_predicate_by_partition_bounds() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("9999")), None),
fixture_file(2, Some(Some("9999")), Some((100, 110))),
],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(&table, col("id").eq(lit(105_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![1, 2]);
Ok(())
}
#[test]
fn files_matching_prunes_a_partition_column_without_column_statistics() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("10")), None),
fixture_file(2, Some(Some("9999")), None),
],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(&table, col("region").eq(lit(10_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![1]);
Ok(())
}
#[test]
fn files_matching_ignores_partition_bounds_when_the_spec_is_not_prune_safe() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("10")), None),
fixture_file(2, Some(Some("9999")), None),
],
Some(region_spec(false)),
)?;
let predicate = physical_predicate(&table, col("region").eq(lit(10_i64)));
assert_eq!(matching_ids(&table, &predicate)?, vec![1, 2]);
Ok(())
}
#[test]
fn files_matching_applies_every_conjunct() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, Some(Some("10")), Some((100, 110))),
fixture_file(2, Some(Some("9999")), Some((100, 110))),
fixture_file(3, Some(Some("10")), Some((1, 10))),
],
Some(region_spec(true)),
)?;
let predicate = physical_predicate(
&table,
col("id")
.eq(lit(105_i64))
.and(col("region").eq(lit(10_i64))),
);
assert_eq!(matching_ids(&table, &predicate)?, vec![1]);
Ok(())
}
#[test]
fn files_matching_returns_every_file_for_a_trivially_true_predicate() -> Result<()> {
let table = fixed_table(
vec![
fixture_file(1, None, Some((1, 10))),
fixture_file(2, None, Some((100, 110))),
fixture_file(3, None, None),
],
None,
)?;
let predicate = physical_predicate(&table, lit(true));
assert_eq!(matching_ids(&table, &predicate)?, vec![1, 2, 3]);
Ok(())
}
#[test]
fn files_matching_prunes_across_every_metadata_page() -> Result<()> {
let provider = Arc::new(LazyMillionFileProvider::default());
let table = DuckLakeTable::new(
1,
"events",
provider.clone(),
1,
Arc::new(ObjectStoreUrl::parse("memory://").unwrap()),
String::new(),
)?;
let predicate = physical_predicate(
&table,
col("id").eq(lit(5_i64)).or(col("id").eq(lit(999_999_i64))),
);
assert_eq!(matching_ids(&table, &predicate)?, vec![5, 999_999]);
assert!(provider.page_calls.load(Ordering::Relaxed) > 1);
assert_eq!(provider.eager_file_reads.load(Ordering::Relaxed), 0);
Ok(())
}
#[test]
fn summary_bounds_with_live_deletes_remain_inexact() -> Result<()> {
let columns =
vec![DuckLakeTableColumn::new(1, "id".to_string(), "integer".to_string(), false)];
let schema = build_arrow_schema(&columns)?;
let (statistics, _) = build_datafusion_statistics(
&schema,
&columns,
&[],
DuckLakeStatistics {
columns: vec![DuckLakeTableColumnStatistics {
column_id: 1,
contains_null: Some(false),
min_value: Some("0".to_string()),
max_value: Some("9".to_string()),
contains_nan: None,
column_size_bytes: Some(40),
bounds_are_exact: false,
}],
..Default::default()
},
true,
false,
);
let column = &statistics.column_statistics[0];
assert_eq!(
column.min_value,
Precision::Inexact(ScalarValue::Int32(Some(0)))
);
assert_eq!(
column.max_value,
Precision::Inexact(ScalarValue::Int32(Some(9)))
);
assert_eq!(column.byte_size, Precision::Inexact(40));
Ok(())
}
#[test]
fn float_file_max_gated_by_contains_nan() -> Result<()> {
let columns =
vec![DuckLakeTableColumn::new(1, "x".to_string(), "double".to_string(), false)];
let schema = build_arrow_schema(&columns)?;
let mut file =
DuckLakeTableFile::new(DuckLakeFileData::new("f.parquet".to_string(), true, 1));
file.data_file_id = 7;
for (contains_nan, max_usable) in [(None, false), (Some(true), false), (Some(false), true)]
{
let (table_stats, file_stats) = build_datafusion_statistics(
&schema,
&columns,
std::slice::from_ref(&file),
DuckLakeStatistics {
files: vec![DuckLakeFileColumnStatistics {
data_file_id: 7,
column_id: 1,
column_size_bytes: Some(16),
value_count: Some(2),
null_count: Some(0),
min_value: Some("1.0".to_string()),
max_value: Some("2.0".to_string()),
contains_nan,
}],
..Default::default()
},
false,
true,
);
let per_file = &file_stats[&7].column_statistics[0];
let table_col = &table_stats.column_statistics[0];
let expected_min = if contains_nan == Some(false) {
Precision::Exact(ScalarValue::Float64(Some(1.0)))
} else {
Precision::Absent
};
assert_eq!(
per_file.min_value, expected_min,
"contains_nan={contains_nan:?}"
);
assert_eq!(
table_col.min_value, expected_min,
"contains_nan={contains_nan:?}"
);
let expected_max = if max_usable {
Precision::Exact(ScalarValue::Float64(Some(2.0)))
} else {
Precision::Absent
};
assert_eq!(
per_file.max_value, expected_max,
"contains_nan={contains_nan:?}"
);
assert_eq!(
table_col.max_value, expected_max,
"contains_nan={contains_nan:?}"
);
}
Ok(())
}
#[test]
fn float_rollup_max_gated_by_contains_nan() -> Result<()> {
let columns =
vec![DuckLakeTableColumn::new(1, "x".to_string(), "double".to_string(), false)];
let schema = build_arrow_schema(&columns)?;
for (contains_nan, max_usable) in [(None, false), (Some(true), false), (Some(false), true)]
{
let (statistics, _) = build_datafusion_statistics(
&schema,
&columns,
&[],
DuckLakeStatistics {
columns: vec![DuckLakeTableColumnStatistics {
column_id: 1,
contains_null: Some(false),
min_value: Some("1.0".to_string()),
max_value: Some("2.0".to_string()),
contains_nan,
column_size_bytes: Some(16),
bounds_are_exact: true,
}],
..Default::default()
},
true,
false,
);
let column = &statistics.column_statistics[0];
let expected_min = if max_usable {
Precision::Exact(ScalarValue::Float64(Some(1.0)))
} else {
Precision::Absent
};
assert_eq!(
column.min_value, expected_min,
"contains_nan={contains_nan:?}"
);
let expected_max = if max_usable {
Precision::Exact(ScalarValue::Float64(Some(2.0)))
} else {
Precision::Absent
};
assert_eq!(
column.max_value, expected_max,
"contains_nan={contains_nan:?}"
);
}
Ok(())
}
#[test]
fn float_table_aggregate_max_needs_all_files_nan_free() -> Result<()> {
let columns =
vec![DuckLakeTableColumn::new(1, "x".to_string(), "double".to_string(), false)];
let schema = build_arrow_schema(&columns)?;
let mut file_a =
DuckLakeTableFile::new(DuckLakeFileData::new("a.parquet".to_string(), true, 1));
file_a.data_file_id = 1;
let mut file_b =
DuckLakeTableFile::new(DuckLakeFileData::new("b.parquet".to_string(), true, 1));
file_b.data_file_id = 2;
let table_files = vec![file_a, file_b];
let stat =
|data_file_id, min: &str, max: &str, contains_nan| DuckLakeFileColumnStatistics {
data_file_id,
column_id: 1,
column_size_bytes: Some(16),
value_count: Some(2),
null_count: Some(0),
min_value: Some(min.to_string()),
max_value: Some(max.to_string()),
contains_nan,
};
let build = |nan_b| {
build_datafusion_statistics(
&schema,
&columns,
&table_files,
DuckLakeStatistics {
files: vec![stat(1, "1.0", "2.0", Some(false)), stat(2, "0.5", "3.0", nan_b)],
..Default::default()
},
false,
true,
)
.0
};
let statistics = build(Some(false));
assert_eq!(
statistics.column_statistics[0].min_value,
Precision::Exact(ScalarValue::Float64(Some(0.5)))
);
assert_eq!(
statistics.column_statistics[0].max_value,
Precision::Exact(ScalarValue::Float64(Some(3.0)))
);
let statistics = build(None);
assert_eq!(statistics.column_statistics[0].min_value, Precision::Absent);
assert_eq!(statistics.column_statistics[0].max_value, Precision::Absent);
Ok(())
}
#[test]
fn test_validated_file_size_positive() {
assert_eq!(validated_file_size(0, "test.parquet").unwrap(), 0);
assert_eq!(validated_file_size(1024, "test.parquet").unwrap(), 1024);
assert_eq!(
validated_file_size(i64::MAX, "test.parquet").unwrap(),
i64::MAX as u64
);
}
#[test]
fn test_validated_file_size_negative() {
let err = validated_file_size(-1, "data/test.parquet").unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("-1"),
"Error should contain the negative value: {}",
msg
);
assert!(
msg.contains("data/test.parquet"),
"Error should contain the file path: {}",
msg
);
}
#[test]
fn test_validated_file_size_large_negative() {
let err = validated_file_size(i64::MIN, "bad.parquet").unwrap_err();
let msg = err.to_string();
assert!(msg.contains("bad.parquet"));
assert!(msg.contains(&i64::MIN.to_string()));
}
#[test]
fn test_validated_record_count_positive() {
assert_eq!(validated_record_count(0, "test.parquet").unwrap(), 0);
assert_eq!(validated_record_count(100, "test.parquet").unwrap(), 100);
assert_eq!(
validated_record_count(i64::MAX, "test.parquet").unwrap(),
i64::MAX as u64
);
}
#[test]
fn test_validated_record_count_negative() {
let err = validated_record_count(-1, "data/test.parquet").unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("-1"),
"Error should contain the negative value: {}",
msg
);
assert!(
msg.contains("data/test.parquet"),
"Error should contain the file path: {}",
msg
);
assert!(
msg.contains("record_count"),
"Error should mention record_count: {}",
msg
);
}
#[test]
fn test_validated_record_count_large_negative() {
let err = validated_record_count(i64::MIN, "bad.parquet").unwrap_err();
let msg = err.to_string();
assert!(msg.contains("bad.parquet"));
assert!(msg.contains(&i64::MIN.to_string()));
}
#[test]
fn test_parse_ducklake_statistic_encodings() {
let boolean = DuckLakeTableColumn::new(1, "flag".to_string(), "boolean".to_string(), true);
assert_eq!(
parse_statistic_scalar("1", &boolean, &DataType::Boolean),
Some(ScalarValue::Boolean(Some(true)))
);
assert_eq!(
parse_statistic_scalar("FALSE", &boolean, &DataType::Boolean),
Some(ScalarValue::Boolean(Some(false)))
);
let blob = DuckLakeTableColumn::new(2, "bytes".to_string(), "blob".to_string(), true);
assert_eq!(
parse_statistic_scalar("68656C6C6F", &blob, &DataType::BinaryView),
Some(ScalarValue::BinaryView(Some(b"hello".to_vec())))
);
assert_eq!(
parse_statistic_scalar("\\x68656C6C6F", &blob, &DataType::BinaryView),
Some(ScalarValue::BinaryView(Some(b"hello".to_vec())))
);
let uuid = DuckLakeTableColumn::new(3, "id".to_string(), "uuid".to_string(), true);
assert_eq!(
parse_statistic_scalar(
"550e8400-e29b-41d4-a716-446655440000",
&uuid,
&DataType::FixedSizeBinary(16),
),
Some(ScalarValue::FixedSizeBinary(
16,
Some(vec![
0x55, 0x0e, 0x84, 0x00, 0xe2, 0x9b, 0x41, 0xd4, 0xa7, 0x16, 0x44, 0x66, 0x55,
0x44, 0x00, 0x00,
]),
))
);
let decimal =
DuckLakeTableColumn::new(4, "amount".to_string(), "decimal(10,2)".to_string(), true);
assert_eq!(
parse_statistic_scalar("123.45", &decimal, &DataType::Decimal128(10, 2)),
Some(ScalarValue::Decimal128(Some(12_345), 10, 2))
);
}
#[test]
fn expression_column_default_does_not_block_table_reads() {
let column =
DuckLakeTableColumn::new(1, "created_at".to_string(), "timestamp".to_string(), false)
.with_defaults(
None,
Some("current_timestamp".to_string()),
Some("expression".to_string()),
Some("duckdb".to_string()),
);
assert!(validate_column_defaults(&[column]).unwrap().is_empty());
}
#[test]
fn null_default_sentinel_does_not_create_an_operational_default() {
let column = DuckLakeTableColumn::new(1, "id".to_string(), "int32".to_string(), false)
.with_defaults(
None,
Some("NULL".to_string()),
Some("literal".to_string()),
Some("duckdb".to_string()),
);
assert!(column.default_value.is_none());
assert!(validate_column_defaults(&[column]).unwrap().is_empty());
}
#[cfg(feature = "write")]
#[test]
fn sort_ordering_rejects_expression_sort_key() {
let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]);
let sort_spec = crate::sort::SortSpec {
sort_id: 9,
fields: vec![crate::sort::SortField {
sort_key_index: 0,
expression: "lower(id)".to_string(),
dialect: crate::sort::DUCKDB_DIALECT.to_string(),
direction: crate::sort::SortDirection::Asc,
null_order: crate::sort::NullOrder::NullsLast,
}],
};
let err = sort_ordering_for(&schema, Some(&sort_spec)).unwrap_err();
assert_eq!(
err.to_string(),
"This feature is not implemented: DuckLake sort order 9 contains an unsupported \
expression; datafusion-ducklake can write only bare-column sort keys",
);
}
fn footer_cache_fixture_bytes() -> Vec<u8> {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("region", DataType::Int64, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int64Array::from((0..10).collect::<Vec<i64>>())) as ArrayRef,
Arc::new(Int64Array::from(vec![1i64; 10])) as ArrayRef,
],
)
.expect("build fixture batch");
let mut buf = Vec::new();
{
let mut writer =
parquet::arrow::ArrowWriter::try_new(&mut buf, schema, None).expect("open writer");
writer.write(&batch).expect("write batch");
writer.close().expect("close writer");
}
buf
}
#[derive(Debug)]
struct FooterCountingStore {
inner: Arc<dyn object_store::ObjectStore>,
file_len: u64,
footer_reads: AtomicUsize,
}
impl FooterCountingStore {
fn new(inner: Arc<dyn object_store::ObjectStore>, file_len: u64) -> Self {
Self {
inner,
file_len,
footer_reads: AtomicUsize::new(0),
}
}
fn footer_reads(&self) -> usize {
self.footer_reads.load(Ordering::SeqCst)
}
fn count_footer_ranges(&self, ranges: &[std::ops::Range<u64>]) {
for range in ranges {
if range.end == self.file_len {
self.footer_reads.fetch_add(1, Ordering::SeqCst);
}
}
}
}
impl std::fmt::Display for FooterCountingStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "FooterCountingStore({})", self.inner)
}
}
#[async_trait::async_trait]
impl object_store::ObjectStore for FooterCountingStore {
async fn put_opts(
&self,
location: &ObjectPath,
payload: object_store::PutPayload,
opts: object_store::PutOptions,
) -> object_store::Result<object_store::PutResult> {
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &ObjectPath,
opts: object_store::PutMultipartOptions,
) -> object_store::Result<Box<dyn object_store::MultipartUpload>> {
self.inner.put_multipart_opts(location, opts).await
}
async fn get_opts(
&self,
location: &ObjectPath,
options: object_store::GetOptions,
) -> object_store::Result<object_store::GetResult> {
if let Some(range) = &options.range
&& let Ok(resolved) = range.as_range(self.file_len)
&& resolved.end == self.file_len
{
self.footer_reads.fetch_add(1, Ordering::SeqCst);
}
self.inner.get_opts(location, options).await
}
async fn get_ranges(
&self,
location: &ObjectPath,
ranges: &[std::ops::Range<u64>],
) -> object_store::Result<Vec<bytes::Bytes>> {
self.count_footer_ranges(ranges);
self.inner.get_ranges(location, ranges).await
}
fn delete_stream(
&self,
locations: futures::stream::BoxStream<'static, object_store::Result<ObjectPath>>,
) -> futures::stream::BoxStream<'static, object_store::Result<ObjectPath>> {
self.inner.delete_stream(locations)
}
fn list(
&self,
prefix: Option<&ObjectPath>,
) -> futures::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>>
{
self.inner.list(prefix)
}
async fn list_with_delimiter(
&self,
prefix: Option<&ObjectPath>,
) -> object_store::Result<object_store::ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy_opts(
&self,
from: &ObjectPath,
to: &ObjectPath,
options: object_store::CopyOptions,
) -> object_store::Result<()> {
self.inner.copy_opts(from, to, options).await
}
}
#[cfg(feature = "write")]
#[tokio::test]
async fn ducklake_scan_attaches_the_cached_parquet_reader_factory() -> Result<()> {
let table = fixed_table(vec![], None)?;
let store: Arc<dyn object_store::ObjectStore> =
Arc::new(object_store::memory::InMemory::new());
let store_url = url::Url::parse(table.object_store_url().as_str()).expect("valid url");
let ctx = SessionContext::new();
ctx.runtime_env().register_object_store(&store_url, store);
let state = ctx.state();
let source = table.create_parquet_source(&state, Arc::clone(&table.physical_schema))?;
let plan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(
table
.scan_config_builder(Arc::new(source))
.with_file_group(FileGroup::new(vec![]))
.build(),
);
let found_source = plan
.downcast_ref::<DataSourceExec>()
.and_then(|exec| exec.data_source().downcast_ref::<FileScanConfig>())
.and_then(|config| config.file_source().downcast_ref::<ParquetSource>())
.expect("plan is a DataSourceExec wrapping a parquet FileScanConfig");
let factory = found_source
.parquet_file_reader_factory()
.expect("a query-time ParquetSource must attach a reader factory (#1323)");
let debug = format!("{factory:?}");
assert!(
debug.contains("CachedParquetFileReaderFactory"),
"expected the metadata-cache-backed reader factory, got: {debug}",
);
Ok(())
}
#[cfg(feature = "write")]
#[tokio::test]
async fn ducklake_parquet_source_reuses_cached_footer_metadata_across_scans() -> Result<()> {
let table = fixed_table(vec![], None)?;
let bytes = footer_cache_fixture_bytes();
let file_len = bytes.len() as u64;
let inner: Arc<dyn object_store::ObjectStore> =
Arc::new(object_store::memory::InMemory::new());
let object_path = ObjectPath::from("footer-cache.parquet");
inner
.put(&object_path, bytes.into())
.await
.expect("seed the fixture file");
let counting = Arc::new(FooterCountingStore::new(inner, file_len));
let store_url = url::Url::parse(table.object_store_url().as_str()).expect("valid url");
let ctx = SessionContext::new();
ctx.runtime_env()
.register_object_store(&store_url, Arc::clone(&counting) as Arc<dyn ObjectStore>);
let state = ctx.state();
let build_plan = |state: &datafusion::execution::SessionState| -> DataFusionResult<Arc<dyn ExecutionPlan>> {
let source = table.create_parquet_source(state, Arc::clone(&table.physical_schema))?;
let pf = PartitionedFile::new("footer-cache.parquet", file_len);
Ok(DataSourceExec::from_data_source(
table
.scan_config_builder(Arc::new(source))
.with_file_group(FileGroup::new(vec![pf]))
.build(),
))
};
let plan1 = build_plan(&state)?;
let batches1 = datafusion::physical_plan::collect(plan1, state.task_ctx())
.await
.expect("first scan executes");
assert_eq!(
batches1.iter().map(|b| b.num_rows()).sum::<usize>(),
10,
"sanity: the fixture file's rows come back"
);
let footer_reads_after_cold_scan = counting.footer_reads();
assert!(
footer_reads_after_cold_scan >= 1,
"the cold scan must read the footer at least once"
);
let plan2 = build_plan(&state)?;
let batches2 = datafusion::physical_plan::collect(plan2, state.task_ctx())
.await
.expect("second scan executes");
assert_eq!(batches2.iter().map(|b| b.num_rows()).sum::<usize>(), 10);
assert_eq!(
counting.footer_reads(),
footer_reads_after_cold_scan,
"a warm scan must issue zero additional footer reads once the reader factory \
shares the session's metadata cache (#1323)",
);
Ok(())
}
}