From d5cccf7ff09253647da627f87d097a8854bab461 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Sun, 9 Aug 2026 15:38:47 +0200 Subject: [PATCH 1/6] Add basic Iceberg integration --- iceberg/AGENTS.md | 12 + iceberg/README.md | 34 + iceberg/src/codec.rs | 65 ++ iceberg/src/common/error.rs | 105 ++ iceberg/src/common/expr_to_predicate.rs | 931 ++++++++++++++++++ iceberg/src/common/mod.rs | 5 + iceberg/src/config.rs | 72 ++ iceberg/src/data_source.rs | 299 ++++++ .../distributed_desired_task_count_handler.rs | 21 + iceberg/src/iceberg_ext.rs | 214 ++++ iceberg/src/lib.rs | 34 +- iceberg/src/table_provider/catalog.rs | 95 ++ iceberg/src/table_provider/factory.rs | 161 +++ iceberg/src/table_provider/mod.rs | 7 + iceberg/src/table_provider/static.rs | 96 ++ iceberg/src/test_utils/harness.rs | 155 +++ iceberg/src/test_utils/mod.rs | 3 + iceberg/src/work_unit_feed.rs | 223 +++++ iceberg/tests/external_table.rs | 96 ++ iceberg/tests/filter_pushdown.rs | 156 +++ iceberg/tests/limit_pushdown.rs | 137 +++ iceberg/tests/projection_pushdown.rs | 136 +++ 22 files changed, 3056 insertions(+), 1 deletion(-) create mode 100644 iceberg/AGENTS.md create mode 100644 iceberg/README.md create mode 100644 iceberg/src/codec.rs create mode 100644 iceberg/src/common/error.rs create mode 100644 iceberg/src/common/expr_to_predicate.rs create mode 100644 iceberg/src/common/mod.rs create mode 100644 iceberg/src/config.rs create mode 100644 iceberg/src/data_source.rs create mode 100644 iceberg/src/distributed_desired_task_count_handler.rs create mode 100644 iceberg/src/iceberg_ext.rs create mode 100644 iceberg/src/table_provider/catalog.rs create mode 100644 iceberg/src/table_provider/factory.rs create mode 100644 iceberg/src/table_provider/mod.rs create mode 100644 iceberg/src/table_provider/static.rs create mode 100644 iceberg/src/test_utils/harness.rs create mode 100644 iceberg/src/test_utils/mod.rs create mode 100644 iceberg/src/work_unit_feed.rs create mode 100644 iceberg/tests/external_table.rs create mode 100644 iceberg/tests/filter_pushdown.rs create mode 100644 iceberg/tests/limit_pushdown.rs create mode 100644 iceberg/tests/projection_pushdown.rs diff --git a/iceberg/AGENTS.md b/iceberg/AGENTS.md new file mode 100644 index 000000000..5d68d69c3 --- /dev/null +++ b/iceberg/AGENTS.md @@ -0,0 +1,12 @@ +# Iceberg crate guide + +## Tests + +- Reuse `src/test_utils/harness.rs` and `testdata/iceberg/taxi` as much as + possible. +- Prefer integration tests over unit tests when possible. Each isolated feature + should have its own dedicated file in `tests/`. +- Never create long tests with a lot of setup, the pattern should be several + tests with a very small body, and have any necessary helper at the bottom + of the `mod tests {};` block, below the actual tests. +- Prefer snapshot testing with `insta` when possible. diff --git a/iceberg/README.md b/iceberg/README.md new file mode 100644 index 000000000..c035069b3 --- /dev/null +++ b/iceberg/README.md @@ -0,0 +1,34 @@ +# DataFusion Distributed Iceberg + +Read-only Apache Iceberg tables for DataFusion Distributed. + +```rust +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionContext; +use datafusion_distributed_iceberg::{IcebergExt, IcebergIntegrationOptions}; + +async fn example() -> datafusion::error::Result<()> { + let state = SessionStateBuilder::new() + .with_default_features() + .with_iceberg_integration(IcebergIntegrationOptions::default()) + .build(); + let ctx = SessionContext::new_with_state(state); + + ctx.sql( + "CREATE EXTERNAL TABLE taxi STORED AS ICEBERG \ + LOCATION 's3://warehouse/taxi/metadata/v1.metadata.json'", + ) + .await? + .collect() + .await?; + Ok(()) +} +``` + +The default storage factory resolves `file://`, S3 (`s3://`, `s3a://`, +`s3n://`), and GCS (`gs://`, `gcs://`) URIs. Use +`IcebergIntegrationOptions` to supply custom storage or an Iceberg runtime. + +```bash +cargo test -p datafusion-distributed-iceberg +``` diff --git a/iceberg/src/codec.rs b/iceberg/src/codec.rs new file mode 100644 index 000000000..2f5104a0f --- /dev/null +++ b/iceberg/src/codec.rs @@ -0,0 +1,65 @@ +use crate::work_unit_feed::FileScanTaskMessage; +use bytes::{Buf, BufMut}; +use datafusion::common::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use prost::encoding::{DecodeContext, WireType}; +use prost::{DecodeError, Message}; +use std::sync::Arc; + +#[derive(Debug)] +pub struct IcebergCodec; + +// TODO: Implement protobuf codecs for the IcebergDataSource. +impl PhysicalExtensionCodec for IcebergCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + ) -> Result> { + unimplemented!() + } + + fn try_encode(&self, _node: Arc, _buf: &mut Vec) -> Result<()> { + unimplemented!() + } +} + +// TODO: Implement serde for FileScanTaskMessage +// This message is an individual WorkUnit, but it cannot be serialized yet. During distributed +// execution, this will be streamed over the wire from coordinator to workers, but for that to +// happen, it will need to be represented as a prost::Message. +// +// WARNING: for the ones who end up implementing this. I have no idea if implementing Message here +// is really the best option for serialization, it might not be. +impl Message for FileScanTaskMessage { + fn encode_raw(&self, _buf: &mut impl BufMut) + where + Self: Sized, + { + unimplemented!() + } + + fn merge_field( + &mut self, + _tag: u32, + _wire_type: WireType, + _buf: &mut impl Buf, + _ctx: DecodeContext, + ) -> std::result::Result<(), DecodeError> + where + Self: Sized, + { + unimplemented!() + } + + fn encoded_len(&self) -> usize { + unimplemented!() + } + + fn clear(&mut self) { + unimplemented!() + } +} diff --git a/iceberg/src/common/error.rs b/iceberg/src/common/error.rs new file mode 100644 index 000000000..0a2badd04 --- /dev/null +++ b/iceberg/src/common/error.rs @@ -0,0 +1,105 @@ +use datafusion::error::DataFusionError; +use iceberg::{Error, ErrorKind}; + +/// Converts a datafusion error into an iceberg error. +pub fn iceberg_err(error: DataFusionError) -> Error { + let fallback_message = error.to_string(); + let DataFusionError::Context(ctx, err) = error else { + return Error::new( + ErrorKind::Unexpected, + format!("DataFusion execution failed: {fallback_message}"), + ); + }; + + let Some(kind) = parse_iceberg_error_kind(&ctx) else { + return Error::new( + ErrorKind::Unexpected, + format!("DataFusion execution failed: {fallback_message}"), + ); + }; + let DataFusionError::Execution(message) = err.as_ref() else { + return Error::new( + ErrorKind::Unexpected, + format!("DataFusion execution failed: {fallback_message}"), + ); + }; + + Error::new(kind, strip_error_kind(kind, message)) +} + +/// Converts an Iceberg error into a DataFusion error. +pub fn df_err(error: Error) -> DataFusionError { + DataFusionError::Context( + format!("IcebergError({})", error.kind()), + Box::new(DataFusionError::Execution(error.to_string())), + ) +} + +fn parse_iceberg_error_kind(context: &str) -> Option { + let kind = context.strip_prefix("IcebergError(")?.strip_suffix(')')?; + + match kind { + "PreconditionFailed" => Some(ErrorKind::PreconditionFailed), + "Unexpected" => Some(ErrorKind::Unexpected), + "DataInvalid" => Some(ErrorKind::DataInvalid), + "NamespaceAlreadyExists" => Some(ErrorKind::NamespaceAlreadyExists), + "TableAlreadyExists" => Some(ErrorKind::TableAlreadyExists), + "NamespaceNotFound" => Some(ErrorKind::NamespaceNotFound), + "TableNotFound" => Some(ErrorKind::TableNotFound), + "FeatureUnsupported" => Some(ErrorKind::FeatureUnsupported), + "CatalogCommitConflicts" => Some(ErrorKind::CatalogCommitConflicts), + _ => None, + } +} + +fn strip_error_kind(kind: ErrorKind, message: &str) -> String { + let kind = kind.into_static(); + if message == kind { + String::new() + } else { + message + .strip_prefix(&format!("{kind} => ")) + .unwrap_or(message) + .to_string() + } +} + +#[cfg(test)] +mod tests { + use datafusion::error::DataFusionError; + use iceberg::{Error, ErrorKind}; + + use super::{df_err, iceberg_err}; + + #[test] + fn roundtrips_iceberg_error_kind_and_message() { + let error = Error::new(ErrorKind::DataInvalid, "invalid manifest"); + let roundtripped = iceberg_err(df_err(error)); + + assert_eq!(roundtripped.kind(), ErrorKind::DataInvalid); + assert_eq!(roundtripped.to_string(), "DataInvalid => invalid manifest"); + } + + #[test] + fn encodes_iceberg_errors_with_native_datafusion_variants() { + let error = df_err(Error::new(ErrorKind::DataInvalid, "invalid manifest")); + + assert!(matches!( + error, + DataFusionError::Context(context, inner) + if context == "IcebergError(DataInvalid)" + && matches!(inner.as_ref(), DataFusionError::Execution(message) if message == "DataInvalid => invalid manifest") + )); + } + + #[test] + fn maps_non_iceberg_datafusion_errors_to_unexpected() { + let error = iceberg_err(DataFusionError::Execution("worker failed".to_string())); + + assert_eq!(error.kind(), ErrorKind::Unexpected); + assert_eq!( + error.to_string(), + "Unexpected => DataFusion execution failed: Execution error: worker failed" + ); + } +} diff --git a/iceberg/src/common/expr_to_predicate.rs b/iceberg/src/common/expr_to_predicate.rs new file mode 100644 index 000000000..e71662d87 --- /dev/null +++ b/iceberg/src/common/expr_to_predicate.rs @@ -0,0 +1,931 @@ +// This file was copied from https://github.com/apache/iceberg-rust/blob/86d9d7dc4297302496e6a121934de719c3ba9bea/crates/integrations/datafusion/src/physical_plan/expr_to_predicate.rs, +// and therefore, it has the same license header. + +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::vec; + +use datafusion::arrow::datatypes::DataType; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{BinaryExpr, Expr, Like, Operator}; +use datafusion::scalar::ScalarValue; +use iceberg::expr::{BinaryExpression, Predicate, PredicateOperator, Reference, UnaryExpression}; +use iceberg::spec::{Datum, PrimitiveLiteral}; + +// A datafusion expression could be an Iceberg predicate, column, or literal. +enum TransformedResult { + Predicate(Predicate), + Column(Reference), + Literal(Datum), + NotTransformed, +} + +enum OpTransformedResult { + Operator(PredicateOperator), + And, + Or, + NotTransformed, +} + +/// Converts DataFusion filters ([`Expr`]) to an iceberg [`Predicate`]. +/// If none of the filters could be converted, return `None` which adds no predicates to the scan operation. +/// If the conversion was successful, return the converted predicates combined with an AND operator. +pub fn convert_filters_to_predicate(filters: &[Expr]) -> Option { + filters + .iter() + .filter_map(convert_filter_to_predicate) + .reduce(Predicate::and) +} + +fn convert_filter_to_predicate(expr: &Expr) -> Option { + match to_iceberg_predicate(expr) { + TransformedResult::Predicate(predicate) => Some(predicate), + TransformedResult::Column(column) => { + // A bare column in a filter context represents a boolean column check + // Convert it to: column = true + Some(Predicate::Binary(BinaryExpression::new( + PredicateOperator::Eq, + column, + Datum::bool(true), + ))) + } + TransformedResult::Literal(_) => { + // Literal values in filter context cannot be pushed down + None + } + _ => None, + } +} + +fn to_iceberg_predicate(expr: &Expr) -> TransformedResult { + match expr { + Expr::BinaryExpr(binary) => { + let left = to_iceberg_predicate(&binary.left); + let right = to_iceberg_predicate(&binary.right); + let op = to_iceberg_operation(binary.op); + match op { + OpTransformedResult::Operator(op) => to_iceberg_binary_predicate(left, right, op), + OpTransformedResult::And => to_iceberg_and_predicate(left, right), + OpTransformedResult::Or => to_iceberg_or_predicate(left, right), + OpTransformedResult::NotTransformed => TransformedResult::NotTransformed, + } + } + Expr::Not(exp) => { + let expr = to_iceberg_predicate(exp); + match expr { + TransformedResult::Predicate(p) => TransformedResult::Predicate(!p), + TransformedResult::Column(column) => { + // NOT of a bare boolean column: NOT col => col = false + TransformedResult::Predicate(Predicate::Binary(BinaryExpression::new( + PredicateOperator::Eq, + column, + Datum::bool(false), + ))) + } + _ => TransformedResult::NotTransformed, + } + } + Expr::Column(column) => TransformedResult::Column(Reference::new(column.name())), + Expr::Literal(literal, _) => match scalar_value_to_datum(literal) { + Some(data) => TransformedResult::Literal(data), + None => TransformedResult::NotTransformed, + }, + Expr::InList(inlist) => { + let mut datums = vec![]; + for expr in &inlist.list { + let p = to_iceberg_predicate(expr); + match p { + TransformedResult::Literal(l) => datums.push(l), + _ => return TransformedResult::NotTransformed, + } + } + + let expr = to_iceberg_predicate(&inlist.expr); + match expr { + TransformedResult::Column(r) => match inlist.negated { + false => TransformedResult::Predicate(r.is_in(datums)), + true => TransformedResult::Predicate(r.is_not_in(datums)), + }, + _ => TransformedResult::NotTransformed, + } + } + Expr::IsNull(expr) => { + let p = to_iceberg_predicate(expr); + match p { + TransformedResult::Column(r) => TransformedResult::Predicate(Predicate::Unary( + UnaryExpression::new(PredicateOperator::IsNull, r), + )), + _ => TransformedResult::NotTransformed, + } + } + Expr::IsNotNull(expr) => { + let p = to_iceberg_predicate(expr); + match p { + TransformedResult::Column(r) => TransformedResult::Predicate(Predicate::Unary( + UnaryExpression::new(PredicateOperator::NotNull, r), + )), + _ => TransformedResult::NotTransformed, + } + } + Expr::Cast(c) => { + if *c.field.data_type() == DataType::Date32 || *c.field.data_type() == DataType::Date64 + { + // Casts to date truncate the expression, we cannot simply extract it as it + // can create erroneous predicates. + return TransformedResult::NotTransformed; + } + to_iceberg_predicate(&c.expr) + } + Expr::Like(Like { + negated, + expr, + pattern, + escape_char, + case_insensitive, + }) => { + // Only support simple prefix patterns (e.g., 'prefix%') + // Note: Iceberg's StartsWith operator is case-sensitive, so we cannot + // push down case-insensitive LIKE (ILIKE) patterns + // Escape characters are also not supported for pushdown + if escape_char.is_some() || *case_insensitive { + return TransformedResult::NotTransformed; + } + + // Extract the pattern string + let pattern_str = match to_iceberg_predicate(pattern) { + TransformedResult::Literal(d) => match d.literal() { + PrimitiveLiteral::String(s) => s.clone(), + _ => return TransformedResult::NotTransformed, + }, + _ => return TransformedResult::NotTransformed, + }; + + // Check if it's a simple prefix pattern (ends with % and no other wildcards) + if pattern_str.ends_with('%') + && !pattern_str[..pattern_str.len() - 1].contains(['%', '_']) + { + // Extract the prefix (remove trailing %) + let prefix = pattern_str[..pattern_str.len() - 1].to_string(); + + // Get the column reference + let column = match to_iceberg_predicate(expr) { + TransformedResult::Column(r) => r, + _ => return TransformedResult::NotTransformed, + }; + + // Create the appropriate predicate + let predicate = if *negated { + column.not_starts_with(Datum::string(prefix)) + } else { + column.starts_with(Datum::string(prefix)) + }; + + TransformedResult::Predicate(predicate) + } else { + // Complex LIKE patterns cannot be pushed down + TransformedResult::NotTransformed + } + } + Expr::ScalarFunction(ScalarFunction { func, args }) => { + scalar_function_to_iceberg_predicate(func.name(), args) + } + _ => TransformedResult::NotTransformed, + } +} + +fn to_iceberg_operation(op: Operator) -> OpTransformedResult { + match op { + Operator::Eq => OpTransformedResult::Operator(PredicateOperator::Eq), + Operator::NotEq => OpTransformedResult::Operator(PredicateOperator::NotEq), + Operator::Lt => OpTransformedResult::Operator(PredicateOperator::LessThan), + Operator::LtEq => OpTransformedResult::Operator(PredicateOperator::LessThanOrEq), + Operator::Gt => OpTransformedResult::Operator(PredicateOperator::GreaterThan), + Operator::GtEq => OpTransformedResult::Operator(PredicateOperator::GreaterThanOrEq), + // AND OR + Operator::And => OpTransformedResult::And, + Operator::Or => OpTransformedResult::Or, + // Others not supported + _ => OpTransformedResult::NotTransformed, + } +} + +/// Translates a DataFusion scalar function into an Iceberg predicate. +/// Unlike dedicated Expr variants (e.g. `Expr::IsNull`), scalar functions are +/// identified by name at runtime, so we need to handle them here. +fn scalar_function_to_iceberg_predicate(func_name: &str, args: &[Expr]) -> TransformedResult { + match func_name { + "isnan" if args.len() == 1 => match resolve_nan_preserving_reference(&args[0]) { + Some(r) => TransformedResult::Predicate(r.is_nan()), + None => TransformedResult::NotTransformed, + }, + _ => TransformedResult::NotTransformed, + } +} + +/// Attempts to resolve a numeric expression argument down to a single column +/// [`Reference`] such that `isnan(arg)` is logically equivalent to +/// `isnan(reference)`. +/// +/// Filter pushdown is reported as `Inexact` (see +/// [`IcebergTableProvider::supports_filters_pushdown`]), so DataFusion +/// re-applies the original predicate after scanning. We therefore only need the +/// pushed-down predicate to be implied by the original filter (it may match +/// extra rows, but must never drop a matching one). Every transformation handled +/// here preserves NaN-ness *exactly* — the result is NaN if and only if the +/// wrapped column is NaN — so both `isnan(arg)` and `NOT isnan(arg)` are sound: +/// +/// * negation: `-x` is NaN iff `x` is NaN +/// * `abs(x)`: `abs(x)` is NaN iff `x` is NaN +/// * casts between numeric types preserve NaN +/// * `x + c`, `c + x`, `x - c`, `c - x` for a finite literal `c` +/// * `x * c`, `c * x`, `x / c` for a finite, non-zero literal `c` +/// +/// Multiplication/division by zero and `c / x` are intentionally rejected: e.g. +/// `x * 0` is NaN when `x` is `±inf`, so it does not imply `x` is NaN. +/// +/// [`IcebergTableProvider::supports_filters_pushdown`]: crate::CatalogIcebergTableProvider +fn resolve_nan_preserving_reference(expr: &Expr) -> Option { + match expr { + Expr::Column(column) => Some(Reference::new(column.name())), + Expr::Negative(inner) => resolve_nan_preserving_reference(inner), + Expr::Cast(cast) => { + // Casts to date truncate the value and are not numeric, so they + // cannot be treated as NaN-preserving. + if *cast.field.data_type() == DataType::Date32 + || *cast.field.data_type() == DataType::Date64 + { + return None; + } + resolve_nan_preserving_reference(&cast.expr) + } + Expr::ScalarFunction(ScalarFunction { func, args }) + if func.name() == "abs" && args.len() == 1 => + { + resolve_nan_preserving_reference(&args[0]) + } + Expr::BinaryExpr(binary) => resolve_nan_preserving_binary(binary), + _ => None, + } +} + +/// Resolves the column reference from an arithmetic expression that combines a +/// single column with a finite literal while preserving NaN-ness. See +/// [`resolve_nan_preserving_reference`] for the soundness argument. +/// +/// Expressions with column references on both sides (e.g. `(x + 1) * (x - 2)`) +/// are not supported. Handling them safely would require both operands to +/// resolve to the *same* column (`x + y` cannot be expressed as a single +/// `col IS NAN`) and the operator combination itself to be NaN-preserving: +/// `(x + 1) * (x - 2)` is NaN iff `x` is NaN, but `(x + 1) - (x - 2)` is NaN +/// for `x = inf` (`inf - inf`) even though `x` is not. +/// +/// TODO: support NaN-preserving expressions with column references on both +/// sides, see . +fn resolve_nan_preserving_binary(binary: &BinaryExpr) -> Option { + let (left, right) = (&binary.left, &binary.right); + match binary.op { + // `x + c`, `c + x`, `x - c` and `c - x` are NaN iff `x` is NaN, for any + // finite literal `c`. The column may be on either side. + Operator::Plus | Operator::Minus => { + if finite_literal(right).is_some() { + resolve_nan_preserving_reference(left) + } else if finite_literal(left).is_some() { + resolve_nan_preserving_reference(right) + } else { + None + } + } + + // `x * c` and `c * x` are NaN iff `x` is NaN, but only when `c` is + // non-zero. Per IEEE-754: + // - inf is not NaN + // - inf * 0 is NaN + // so multiplying by zero is rejected. The column may be on either side. + Operator::Multiply => { + if matches!(finite_literal(right), Some(c) if c != 0.0) { + resolve_nan_preserving_reference(left) + } else if matches!(finite_literal(left), Some(c) if c != 0.0) { + resolve_nan_preserving_reference(right) + } else { + None + } + } + + // `x / c` is NaN iff `x` is NaN, for a finite non-zero literal `c`. + // `c / x` is rejected and the column must be the dividend (left side). + // Per IEEE-754: + // - 0 is not NaN + // - 0 / 0 is NaN + // so `c / x` is not NaN-preserving. + Operator::Divide => { + if matches!(finite_literal(right), Some(c) if c != 0.0) { + resolve_nan_preserving_reference(left) + } else { + None + } + } + + _ => None, + } +} + +/// Returns the value of `expr` as an `f64` if it is a finite numeric literal +/// (i.e. not a non-literal, non-numeric, or infinite/NaN value). The numeric +/// conversion is delegated to DataFusion's [`ScalarValue::cast_to`]; the value +/// is only used to inspect finiteness and sign (precision loss is irrelevant). +fn finite_literal(expr: &Expr) -> Option { + let Expr::Literal(value, _) = expr else { + return None; + }; + match value.cast_to(&DataType::Float64).ok()? { + ScalarValue::Float64(Some(v)) if v.is_finite() => Some(v), + _ => None, + } +} + +fn to_iceberg_and_predicate( + left: TransformedResult, + right: TransformedResult, +) -> TransformedResult { + match (left, right) { + (TransformedResult::Predicate(left), TransformedResult::Predicate(right)) => { + TransformedResult::Predicate(left.and(right)) + } + (TransformedResult::Predicate(left), _) => TransformedResult::Predicate(left), + (_, TransformedResult::Predicate(right)) => TransformedResult::Predicate(right), + _ => TransformedResult::NotTransformed, + } +} + +fn to_iceberg_or_predicate(left: TransformedResult, right: TransformedResult) -> TransformedResult { + match (left, right) { + (TransformedResult::Predicate(left), TransformedResult::Predicate(right)) => { + TransformedResult::Predicate(left.or(right)) + } + _ => TransformedResult::NotTransformed, + } +} + +fn to_iceberg_binary_predicate( + left: TransformedResult, + right: TransformedResult, + op: PredicateOperator, +) -> TransformedResult { + let (r, d, op) = match (left, right) { + (TransformedResult::NotTransformed, _) => return TransformedResult::NotTransformed, + (_, TransformedResult::NotTransformed) => return TransformedResult::NotTransformed, + (TransformedResult::Column(r), TransformedResult::Literal(d)) => (r, d, op), + (TransformedResult::Literal(d), TransformedResult::Column(r)) => { + (r, d, reverse_predicate_operator(op)) + } + _ => return TransformedResult::NotTransformed, + }; + TransformedResult::Predicate(Predicate::Binary(BinaryExpression::new(op, r, d))) +} + +fn reverse_predicate_operator(op: PredicateOperator) -> PredicateOperator { + match op { + PredicateOperator::Eq => PredicateOperator::Eq, + PredicateOperator::NotEq => PredicateOperator::NotEq, + PredicateOperator::GreaterThan => PredicateOperator::LessThan, + PredicateOperator::GreaterThanOrEq => PredicateOperator::LessThanOrEq, + PredicateOperator::LessThan => PredicateOperator::GreaterThan, + PredicateOperator::LessThanOrEq => PredicateOperator::GreaterThanOrEq, + _ => unreachable!("Reverse {}", op), + } +} + +const MILLIS_PER_DAY: i64 = 24 * 60 * 60 * 1000; + +/// Convert a scalar value to an iceberg datum. +fn scalar_value_to_datum(value: &ScalarValue) -> Option { + match value { + ScalarValue::Boolean(Some(v)) => Some(Datum::bool(*v)), + ScalarValue::Int8(Some(v)) => Some(Datum::int(*v as i32)), + ScalarValue::Int16(Some(v)) => Some(Datum::int(*v as i32)), + ScalarValue::Int32(Some(v)) => Some(Datum::int(*v)), + ScalarValue::Int64(Some(v)) => Some(Datum::long(*v)), + ScalarValue::Float32(Some(v)) => Some(Datum::double(*v as f64)), + ScalarValue::Float64(Some(v)) => Some(Datum::double(*v)), + ScalarValue::Utf8(Some(v)) => Some(Datum::string(v.clone())), + ScalarValue::LargeUtf8(Some(v)) => Some(Datum::string(v.clone())), + ScalarValue::Binary(Some(v)) => Some(Datum::binary(v.clone())), + ScalarValue::LargeBinary(Some(v)) => Some(Datum::binary(v.clone())), + ScalarValue::Date32(Some(v)) => Some(Datum::date(*v)), + ScalarValue::Date64(Some(v)) => Some(Datum::date((*v / MILLIS_PER_DAY) as i32)), + // Timestamp conversions + // Note: TimestampSecond and TimestampMillisecond are not handled here because + // DataFusion's type coercion always converts them to match the column type + // (either TimestampMicrosecond or TimestampNanosecond) before predicate pushdown. + // See unit tests for how those conversions would work if needed. + ScalarValue::TimestampMicrosecond(Some(v), _) => Some(Datum::timestamp_micros(*v)), + ScalarValue::TimestampNanosecond(Some(v), _) => Some(Datum::timestamp_nanos(*v)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + #[allow(clippy::disallowed_types)] + use std::collections::HashMap; + + use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion::common::DFSchema; + use datafusion::logical_expr::utils::split_conjunction; + use datafusion::prelude::{Expr, SessionContext}; + use iceberg::expr::{Predicate, Reference}; + use iceberg::spec::Datum; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + + use super::convert_filters_to_predicate; + + fn create_test_schema() -> DFSchema { + let arrow_schema = Schema::new(vec![ + Field::new("foo", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + Field::new("bar", DataType::Utf8, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + Field::new("ts", DataType::Timestamp(TimeUnit::Second, None), true).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())]), + ), + Field::new("qux", DataType::Float64, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "4".to_string(), + )])), + ]); + DFSchema::try_from_qualified_schema("my_table", &arrow_schema).unwrap() + } + + fn convert_to_iceberg_predicate(sql: &str) -> Option { + let df_schema = create_test_schema(); + let expr = SessionContext::new() + .parse_sql_expr(sql, &df_schema) + .unwrap(); + let exprs: Vec = split_conjunction(&expr).into_iter().cloned().collect(); + convert_filters_to_predicate(&exprs[..]) + } + + #[test] + fn test_predicate_conversion_with_single_condition() { + let predicate = convert_to_iceberg_predicate("foo = 1").unwrap(); + assert_eq!(predicate, Reference::new("foo").equal_to(Datum::long(1))); + + let predicate = convert_to_iceberg_predicate("foo != 1").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").not_equal_to(Datum::long(1)) + ); + + let predicate = convert_to_iceberg_predicate("foo > 1").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").greater_than(Datum::long(1)) + ); + + let predicate = convert_to_iceberg_predicate("foo >= 1").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").greater_than_or_equal_to(Datum::long(1)) + ); + + let predicate = convert_to_iceberg_predicate("foo < 1").unwrap(); + assert_eq!(predicate, Reference::new("foo").less_than(Datum::long(1))); + + let predicate = convert_to_iceberg_predicate("foo <= 1").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").less_than_or_equal_to(Datum::long(1)) + ); + + let predicate = convert_to_iceberg_predicate("foo is null").unwrap(); + assert_eq!(predicate, Reference::new("foo").is_null()); + + let predicate = convert_to_iceberg_predicate("foo is not null").unwrap(); + assert_eq!(predicate, Reference::new("foo").is_not_null()); + + let predicate = convert_to_iceberg_predicate("foo in (5, 6)").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").is_in([Datum::long(5), Datum::long(6)]) + ); + + let predicate = convert_to_iceberg_predicate("foo not in (5, 6)").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").is_not_in([Datum::long(5), Datum::long(6)]) + ); + + let predicate = convert_to_iceberg_predicate("not foo = 1").unwrap(); + assert_eq!(predicate, !Reference::new("foo").equal_to(Datum::long(1))); + } + + #[test] + fn test_predicate_conversion_with_single_unsupported_condition() { + let predicate = convert_to_iceberg_predicate("foo + 1 = 1"); + assert_eq!(predicate, None); + + let predicate = convert_to_iceberg_predicate("length(bar) = 1"); + assert_eq!(predicate, None); + + let predicate = convert_to_iceberg_predicate("foo in (1, 2, foo)"); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_single_condition_rev() { + let predicate = convert_to_iceberg_predicate("1 < foo").unwrap(); + assert_eq!( + predicate, + Reference::new("foo").greater_than(Datum::long(1)) + ); + } + + #[test] + fn test_predicate_conversion_with_and_condition() { + let sql = "foo > 1 and bar = 'test'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Predicate::and( + Reference::new("foo").greater_than(Datum::long(1)), + Reference::new("bar").equal_to(Datum::string("test")), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_and_condition_unsupported() { + let sql = "foo > 1 and length(bar) = 1"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Reference::new("foo").greater_than(Datum::long(1)); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_and_condition_both_unsupported() { + let sql = "foo in (1, 2, foo) and length(bar) = 1"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_or_condition_unsupported() { + let sql = "foo > 1 or length(bar) = 1"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_or_condition_supported() { + let sql = "foo > 1 or bar = 'test'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Predicate::or( + Reference::new("foo").greater_than(Datum::long(1)), + Reference::new("bar").equal_to(Datum::string("test")), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_complex_binary_expr() { + let sql = "(foo > 1 and bar = 'test') or foo < 0 "; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + + let inner_predicate = Predicate::and( + Reference::new("foo").greater_than(Datum::long(1)), + Reference::new("bar").equal_to(Datum::string("test")), + ); + let expected_predicate = Predicate::or( + inner_predicate, + Reference::new("foo").less_than(Datum::long(0)), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_one_and_expr_supported() { + let sql = "(foo > 1 and length(bar) = 1 ) or foo < 0 "; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + + let inner_predicate = Reference::new("foo").greater_than(Datum::long(1)); + let expected_predicate = Predicate::or( + inner_predicate, + Reference::new("foo").less_than(Datum::long(0)), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_complex_binary_expr_unsupported() { + let sql = "(foo > 1 or length(bar) = 1 ) and foo < 0 "; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Reference::new("foo").less_than(Datum::long(0)); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_cast() { + let sql = "ts >= timestamp '2023-01-05T00:00:00'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = + Reference::new("ts").greater_than_or_equal_to(Datum::string("2023-01-05T00:00:00")); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_date_cast() { + let sql = "ts >= date '2023-01-05T11:00:00'"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_scalar_value_to_datum_timestamp() { + use datafusion::common::ScalarValue; + + // Test TimestampMicrosecond - maps directly to Datum::timestamp_micros + let ts_micros = 1672876800000000i64; // 2023-01-05 00:00:00 UTC in microseconds + let datum = + super::scalar_value_to_datum(&ScalarValue::TimestampMicrosecond(Some(ts_micros), None)); + assert_eq!(datum, Some(Datum::timestamp_micros(ts_micros))); + + // Test TimestampNanosecond - maps to Datum::timestamp_nanos to preserve precision + let ts_nanos = 1672876800000000500i64; // 2023-01-05 00:00:00.000000500 UTC in nanoseconds + let datum = + super::scalar_value_to_datum(&ScalarValue::TimestampNanosecond(Some(ts_nanos), None)); + assert_eq!(datum, Some(Datum::timestamp_nanos(ts_nanos))); + + // Test None timestamp + let datum = super::scalar_value_to_datum(&ScalarValue::TimestampMicrosecond(None, None)); + assert_eq!(datum, None); + + // Note: TimestampSecond and TimestampMillisecond are not supported because + // DataFusion's type coercion converts them to TimestampMicrosecond or TimestampNanosecond + // before they reach scalar_value_to_datum in SQL queries. + // + // These return None (not pushed down): + let ts_seconds = 1672876800i64; // 2023-01-05 00:00:00 UTC in seconds + let datum = + super::scalar_value_to_datum(&ScalarValue::TimestampSecond(Some(ts_seconds), None)); + assert_eq!(datum, None); + + let ts_millis = 1672876800000i64; // 2023-01-05 00:00:00 UTC in milliseconds + let datum = + super::scalar_value_to_datum(&ScalarValue::TimestampMillisecond(Some(ts_millis), None)); + assert_eq!(datum, None); + } + + #[test] + fn test_scalar_value_to_datum_binary() { + use datafusion::common::ScalarValue; + + let bytes = vec![1u8, 2u8, 3u8]; + let datum = super::scalar_value_to_datum(&ScalarValue::Binary(Some(bytes.clone()))); + assert_eq!(datum, Some(Datum::binary(bytes.clone()))); + + let datum = super::scalar_value_to_datum(&ScalarValue::LargeBinary(Some(bytes.clone()))); + assert_eq!(datum, Some(Datum::binary(bytes))); + + let datum = super::scalar_value_to_datum(&ScalarValue::Binary(None)); + assert_eq!(datum, None); + } + + #[test] + fn test_predicate_conversion_with_binary() { + let sql = "foo = 1 and bar = X'0102'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + // Binary literals are converted to Datum::binary + // Note: SQL literal 1 is converted to Long by DataFusion + let expected_predicate = Reference::new("foo") + .equal_to(Datum::long(1)) + .and(Reference::new("bar").equal_to(Datum::binary(vec![1u8, 2u8]))); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_scalar_value_to_datum_boolean() { + use datafusion::common::ScalarValue; + + // Test boolean true + let datum = super::scalar_value_to_datum(&ScalarValue::Boolean(Some(true))); + assert_eq!(datum, Some(Datum::bool(true))); + + // Test boolean false + let datum = super::scalar_value_to_datum(&ScalarValue::Boolean(Some(false))); + assert_eq!(datum, Some(Datum::bool(false))); + + // Test None boolean + let datum = super::scalar_value_to_datum(&ScalarValue::Boolean(None)); + assert_eq!(datum, None); + } + + #[test] + fn test_predicate_conversion_with_like_starts_with() { + let sql = "bar LIKE 'test%'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!( + predicate, + Reference::new("bar").starts_with(Datum::string("test")) + ); + } + + #[test] + fn test_predicate_conversion_with_not_like_starts_with() { + let sql = "bar NOT LIKE 'test%'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!( + predicate, + Reference::new("bar").not_starts_with(Datum::string("test")) + ); + } + + #[test] + fn test_predicate_conversion_with_like_empty_prefix() { + let sql = "bar LIKE '%'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!( + predicate, + Reference::new("bar").starts_with(Datum::string("")) + ); + } + + #[test] + fn test_predicate_conversion_with_like_complex_pattern() { + // Patterns with wildcards in the middle cannot be pushed down + let sql = "bar LIKE 'te%st'"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_like_underscore_wildcard() { + // Patterns with underscore wildcard cannot be pushed down + let sql = "bar LIKE 'test_'"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_like_no_wildcard() { + // Patterns without trailing % cannot be pushed down as StartsWith + let sql = "bar LIKE 'test'"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_ilike() { + // Case-insensitive LIKE (ILIKE) is not supported + let sql = "bar ILIKE 'test%'"; + let predicate = convert_to_iceberg_predicate(sql); + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_like_and_other_conditions() { + let sql = "bar LIKE 'test%' AND foo > 1"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Predicate::and( + Reference::new("bar").starts_with(Datum::string("test")), + Reference::new("foo").greater_than(Datum::long(1)), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_like_special_characters() { + // Test LIKE with special characters in prefix + let sql = "bar LIKE 'test-abc_123%'"; + let predicate = convert_to_iceberg_predicate(sql); + // This should not be pushed down because it contains underscore + assert_eq!(predicate, None); + } + + #[test] + fn test_predicate_conversion_with_like_unicode() { + // Test LIKE with unicode characters in prefix + let sql = "bar LIKE '测试%'"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!( + predicate, + Reference::new("bar").starts_with(Datum::string("测试")) + ); + } + + #[test] + fn test_predicate_conversion_with_isnan() { + let predicate = convert_to_iceberg_predicate("isnan(qux)").unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan()); + } + + #[test] + fn test_predicate_conversion_with_not_isnan() { + let predicate = convert_to_iceberg_predicate("NOT isnan(qux)").unwrap(); + assert_eq!(predicate, !Reference::new("qux").is_nan()); + } + + #[test] + fn test_predicate_conversion_with_isnan_and_other_condition() { + let sql = "isnan(qux) AND foo > 1"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Predicate::and( + Reference::new("qux").is_nan(), + Reference::new("foo").greater_than(Datum::long(1)), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_isnan_negation() { + // -x is NaN iff x is NaN + let predicate = convert_to_iceberg_predicate("isnan(-qux)").unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan()); + + let predicate = convert_to_iceberg_predicate("NOT isnan(-qux)").unwrap(); + assert_eq!(predicate, !Reference::new("qux").is_nan()); + } + + #[test] + fn test_predicate_conversion_with_isnan_abs() { + // abs(x) is NaN iff x is NaN + let predicate = convert_to_iceberg_predicate("isnan(abs(qux))").unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan()); + } + + #[test] + fn test_predicate_conversion_with_isnan_additive() { + // x + c, c + x, x - c, c - x are NaN iff x is NaN (for finite c) + for sql in [ + "isnan(qux + 1)", + "isnan(1 + qux)", + "isnan(qux - 1)", + "isnan(1 - qux)", + "isnan(qux + 1.5)", + ] { + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan(), "sql: {sql}"); + } + } + + #[test] + fn test_predicate_conversion_with_isnan_multiplicative() { + // x * c, c * x, x / c are NaN iff x is NaN (for finite non-zero c) + for sql in ["isnan(qux * 2)", "isnan(2 * qux)", "isnan(qux / 2)"] { + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan(), "sql: {sql}"); + } + } + + #[test] + fn test_predicate_conversion_with_isnan_nested_expr() { + // Nested NaN-preserving transformations resolve to the inner column + let predicate = convert_to_iceberg_predicate("isnan(-(abs(qux) + 1) * 3)").unwrap(); + assert_eq!(predicate, Reference::new("qux").is_nan()); + } + + #[test] + fn test_predicate_conversion_with_isnan_and_other_complex_condition() { + let sql = "isnan(qux + 1) AND foo > 1"; + let predicate = convert_to_iceberg_predicate(sql).unwrap(); + let expected_predicate = Predicate::and( + Reference::new("qux").is_nan(), + Reference::new("foo").greater_than(Datum::long(1)), + ); + assert_eq!(predicate, expected_predicate); + } + + #[test] + fn test_predicate_conversion_with_isnan_unsupported_arg() { + // Multiplying/dividing by zero does not preserve NaN-ness: `x * 0` is NaN + // when `x` is ±inf, so it cannot be pushed down. + assert_eq!(convert_to_iceberg_predicate("isnan(qux * 0)"), None); + assert_eq!(convert_to_iceberg_predicate("isnan(qux / 0)"), None); + + // `c / x` is not NaN-preserving (e.g. `0 / 0` is NaN while `0` is not). + assert_eq!(convert_to_iceberg_predicate("isnan(1 / qux)"), None); + + // Expressions referencing more than one column cannot be reduced to a + // single column reference. + assert_eq!(convert_to_iceberg_predicate("isnan(qux + foo)"), None); + + // Unknown scalar functions are not pushed down. + assert_eq!(convert_to_iceberg_predicate("isnan(sqrt(qux))"), None); + } +} diff --git a/iceberg/src/common/mod.rs b/iceberg/src/common/mod.rs new file mode 100644 index 000000000..b5b574bee --- /dev/null +++ b/iceberg/src/common/mod.rs @@ -0,0 +1,5 @@ +mod error; +mod expr_to_predicate; + +pub(crate) use error::{df_err, iceberg_err}; +pub(crate) use expr_to_predicate::convert_filters_to_predicate; diff --git a/iceberg/src/config.rs b/iceberg/src/config.rs new file mode 100644 index 000000000..4bf74271c --- /dev/null +++ b/iceberg/src/config.rs @@ -0,0 +1,72 @@ +use std::sync::Arc; + +use datafusion::common::extensions_options; +use datafusion::config::{ConfigExtension, ConfigOptions}; +use datafusion::execution::TaskContext; +use datafusion::prelude::SessionConfig; + +extensions_options! { + /// Configuration for Iceberg table reads. + pub struct IcebergConfig { + /// Maximum number of data files to read concurrently. Must be greater than zero. + pub data_file_concurrency_limit: usize, default = 8 + /// Whether to prune Parquet row groups using their statistics. + pub row_group_filtering_enabled: bool, default = true + /// Whether to apply row-level selections while reading Parquet files. + pub row_selection_enabled: bool, default = false + } +} + +impl ConfigExtension for IcebergConfig { + const PREFIX: &'static str = "iceberg"; +} + +impl IcebergConfig { + /// Returns the registered Iceberg configuration, or its defaults when it + /// has not been registered. + /// + /// Register [`IcebergConfig::default`] through + /// [`SessionConfig::with_option_extension`] to configure these settings + /// with `iceberg.*` DataFusion options. + pub fn from_config_options(cfg: &ConfigOptions) -> Self { + cfg.extensions.get::().cloned().unwrap_or_default() + } + + /// Returns the registered Iceberg configuration, or its defaults. + pub fn from_session_config(session_cfg: &SessionConfig) -> Self { + Self::from_config_options(session_cfg.options()) + } + + /// Returns the registered Iceberg configuration, or its defaults. + pub fn from_task_context(ctx: &Arc) -> Self { + Self::from_session_config(ctx.session_config()) + } +} + +#[cfg(test)] +mod tests { + use super::IcebergConfig; + use datafusion::prelude::SessionConfig; + + #[test] + fn config_options_can_be_set_through_session_config() { + let config = SessionConfig::new() + .with_option_extension(IcebergConfig::default()) + .set_usize("iceberg.data_file_concurrency_limit", 8) + .set_bool("iceberg.row_group_filtering_enabled", false) + .set_bool("iceberg.row_selection_enabled", true); + + let iceberg_config = IcebergConfig::from_session_config(&config); + assert_eq!(iceberg_config.data_file_concurrency_limit, 8); + assert!(!iceberg_config.row_group_filtering_enabled); + assert!(iceberg_config.row_selection_enabled); + } + + #[test] + fn unregistered_config_uses_defaults() { + assert_eq!( + IcebergConfig::from_session_config(&SessionConfig::new()).data_file_concurrency_limit, + IcebergConfig::default().data_file_concurrency_limit, + ); + } +} diff --git a/iceberg/src/data_source.rs b/iceberg/src/data_source.rs new file mode 100644 index 000000000..ed15343dc --- /dev/null +++ b/iceberg/src/data_source.rs @@ -0,0 +1,299 @@ +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Statistics, exec_datafusion_err}; +use datafusion::config::ConfigOptions; +use datafusion::datasource::source::DataSource; +use datafusion::error::Result; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_expr::{Partitioning, PhysicalSortExpr}; +use datafusion::physical_plan::filter_pushdown::{FilterPushdownPropagation, PushedDown}; +use datafusion::physical_plan::limit::LimitStream; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayFormatType, SortOrderPushdownResult}; +use datafusion::prelude::Expr; +use datafusion_distributed::WorkUnitFeed; +use futures::{StreamExt, TryStreamExt}; +use iceberg::arrow::ArrowReaderBuilder; + +use crate::common::{convert_filters_to_predicate, df_err, iceberg_err}; +use crate::{IcebergConfig, IcebergWorkUnitFeed}; + +/// Consumes a stream of [iceberg::scan::FileScanTask]s per partition and reads the underlying +/// files into an Arrow stream. +/// +/// [iceberg::scan::FileScanTask] are discovered progressively during execution by the +/// [IcebergWorkUnitFeed], and this [DataSource] executes those tasks as they come, also in +/// a streaming fashion. This works seamlessly in both single-node and distributed execution: +/// +/// ## Single Node +/// +/// [iceberg::scan::FileScanTask] are streamed in-memory, with as many parallel streams as +/// partitions this [IcebergDataSource] exposes: +/// +/// ```text +/// ┌────────────────────────────────────────────┐ +/// │ IcebergDataSource │ +/// │ │ +/// │┌──────────────────────────────────────────┐│ +/// ││ IcebergWorkUnitFeed ││ +/// ││┌────────────┐┌────────────┐┌────────────┐││ +/// │││ Feed 0 ││ Feed 1 ││ Feed 2 │││ +/// ││└──────┬─────┘└──────┬─────┘└──────┬─────┘││ +/// │└───────┼─────────────┼─────────────┼──────┘│ +/// │ .─────▼─────. .─────▼─────. .─────▼─────. │ +/// │ (FileScanTask (FileScanTask (FileScanTask )│ +/// │ .───────────. `─────┬─────' .───────────. │ +/// │ (FileScanTask ) │ (FileScanTask )│ +/// │ `─────┬─────' │ .───────────. │ +/// │ │ │ (FileScanTask )│ +/// │ │ │ `─────┬─────' │ +/// │ │ │ │ │ +/// │ ┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ │ +/// │ │Partition 0 ││Partition 1 ││Partition 2 │ │ +/// │ │ArrowReader ││ArrowReader ││ArrowReader │ │ +/// │ └──────┬─────┘└──────┬─────┘└──────┬─────┘ │ +/// │ │ │ │ │ +/// │ .─────▼─────. │ .─────▼─────. │ +/// │ ( RecordBatch ).─────▼─────.( RecordBatch )│ +/// │ `─────┬─────'( RecordBatch ).───────────. │ +/// │ │ `─────┬─────'( RecordBatch )│ +/// │ │ │ `───────────' │ +/// └────────┼─────────────┼─────────────┼───────┘ +/// ▼ ▼ ▼ +/// ``` +/// +/// ## Distributed +/// +/// [iceberg::scan::FileScanTask] are streamed over the network, with as many parallel streams as +/// partitions * distributed tasks: +/// +/// ```text +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// Coordinating Context │ +/// │ +/// ┌────────────────────────────────────────────────────────────────────────────────────────┐│ +/// ││ IcebergWorkUnitFeed │ +/// │┌─────────────┐┌─────────────┐┌────────────┐┌────────────┐┌─────────────┐┌─────────────┐││ +/// │││ Feed 0 ││ Feed 1 ││ Feed 2 ││ Feed 3 ││ Feed 4 ││ Feed 5 ││ +/// │└──────┬──────┘└─────┬───────┘└────┬───────┘└───────┬────┘└───────┬─────┘└──────┬──────┘││ +/// └└───────┼─────────────┼─────────────┼────────────────┼─────────────┼─────────────┼───────┴ +/// .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. +/// (FileScanTask (FileScanTask (FileScanTask ) (FileScanTask (FileScanTask (FileScanTask ) +/// .───────────. `─────┬─────' .───────────. `─────┬─────' .───────────. `─────┬─────' +/// (FileScanTask ) │ (FileScanTask ) │ (FileScanTask ) │ +/// `─────┬─────' │ .───────────. │ `───────────' │ +/// │ │ (FileScanTask ) │ │ │ +/// Worker 0 │ │ `─────┬─────' │ │ │ Worker 1 +/// ┌ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ┐┌ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ┐ +/// ┌───────┼─────────────┼─────────────┼───────┐┌───────┼─────────────┼─────────────┼───────┐ +/// │ │ │ IcebergD│taSource │ ││ │ IcebergD│taSource │ │ │ +/// │ │ │ │ ││ │ │ │ │ +/// │ │┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ ││┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ │ │ +/// ││Partition 0 ││Partition 1 ││Partition 2 │ │││Partition 0 ││Partition 1 ││Partition 2 │ │ +/// │ ││ArrowReader ││ArrowReader ││ArrowReader │ │││ArrowReader ││ArrowReader ││ArrowReader │ │ │ +/// │└──────┬─────┘└──────┬─────┘└──────┬─────┘ ││└──────┬─────┘└──────┬─────┘└──────┬─────┘ │ +/// │ │ │ │ │ ││ │ │ │ │ │ +/// │ .─────▼─────. │ .─────▼─────. ││ │ ▼ ▼ │ +/// │ │( RecordBatch ).─────▼─────.( RecordBatch )││ .─────▼─────. .───────────. .───────────. │ │ +/// │ `─────┬─────'( RecordBatch ).───────────. ││( RecordBatch ( RecordBatch ) RecordBatch )│ +/// │ │ │ `─────┬─────'( RecordBatch )││ `─────┬─────' `───────────' `─────┬─────' │ │ +/// │ │ │ `───────────' ││ │ ( RecordBatch ) │ │ +/// │ │ │ │ │ ││ │ `─────┬─────' │ │ │ +/// └───────┼─────────────┼─────────────┼───────┘└───────┼─────────────┼─────────────┼───────┘ +/// │ ▼ ▼ ▼ ││ ▼ ▼ ▼ │ +/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ``` +/// +/// This distributed mechanism is transparent to this [DataSource]. +#[derive(Debug, Clone)] +pub struct IcebergDataSource { + schema: SchemaRef, + partitioning: Partitioning, + fetch: Option, + metrics: ExecutionPlanMetricsSet, + iceberg_file_io: iceberg::io::FileIO, + iceberg_runtime: iceberg::Runtime, + feed: WorkUnitFeed, +} + +/// Optional fields for building an [IcebergDataSource]. +#[derive(Default, Clone)] +pub(crate) struct IcebergDataSourceOptions<'a> { + pub(crate) snapshot_id: Option, + pub(crate) projection: Option<&'a Vec>, + pub(crate) fetch: Option, + pub(crate) filters: &'a [Expr], + pub(crate) iceberg_runtime: Option, +} + +impl IcebergDataSource { + /// Creates a new [`IcebergDataSource`] object. + pub(crate) fn new( + table: iceberg::table::Table, + schema: SchemaRef, + partitioning: Partitioning, + opts: IcebergDataSourceOptions, + ) -> Self { + let output_schema = match opts.projection { + None => schema.clone(), + Some(projection) => Arc::new(schema.project(projection).unwrap()), + }; + let projection = opts.projection.map(|v| { + v.iter() + .map(|p| schema.field(*p).name().clone()) + .collect::>() + }); + + let predicates = convert_filters_to_predicate(opts.filters); + + Self { + schema: output_schema, + iceberg_file_io: table.file_io().clone(), + partitioning: partitioning.clone(), + fetch: opts.fetch, + metrics: ExecutionPlanMetricsSet::new(), + iceberg_runtime: opts + .iceberg_runtime + .unwrap_or_else(iceberg::Runtime::current), + feed: WorkUnitFeed::new(IcebergWorkUnitFeed { + iceberg_table: table, + snapshot_id: opts.snapshot_id, + projection, + predicates, + partitioning, + sync_manager: Default::default(), + }), + } + } +} + +impl IcebergDataSource { + /// Returns the [WorkUnitFeed] implementation that feeds this + /// DataSource with [iceberg::scan::FileScanTask] messages. + pub fn feed(&self) -> &WorkUnitFeed { + &self.feed + } +} + +impl DataSource for IcebergDataSource { + fn open( + &self, + partition: usize, + context: Arc, + ) -> Result { + let config = IcebergConfig::from_task_context(&context); + + let reader = + ArrowReaderBuilder::new(self.iceberg_file_io.clone(), self.iceberg_runtime.clone()) + .with_batch_size(context.session_config().batch_size()) + .with_data_file_concurrency_limit(config.data_file_concurrency_limit) + .with_row_group_filtering_enabled(config.row_group_filtering_enabled) + .with_row_selection_enabled(config.row_selection_enabled) + .build(); + + let feed = self + .feed + .feed(partition, context)? + .map(|msg_or_err| match msg_or_err { + Ok(msg) => match msg.inner { + Some(msg) => Ok(msg), + None => Err(iceberg_err(exec_datafusion_err!("Missing inner"))), + }, + Err(err) => Err(iceberg_err(err)), + }) + .boxed(); + + let stream = reader + .read(feed) + .map(|result| result.stream()) + .map_err(df_err)? + .map_err(df_err); + + let stream = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + )) as SendableRecordBatchStream; + + let metrics = BaselineMetrics::new(&self.metrics, partition); + + Ok(Box::pin(LimitStream::new(stream, 0, self.fetch, metrics))) + } + + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "format=iceberg")?; + let Some(feed) = self.feed.inner() else { + return Ok(()); + }; + if let Some(projection) = &feed.projection { + write!(f, ", projection=[{}]", projection.join(", "))?; + } + if let Some(predicate) = &feed.predicates { + write!(f, ", predicate={predicate}")?; + } + if let Some(fetch) = self.fetch { + write!(f, ", fetch={fetch}")?; + } + Ok(()) + } + + fn output_partitioning(&self) -> Partitioning { + self.partitioning.clone() + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new(Arc::clone(&self.schema)) + } + + fn partition_statistics(&self, _partition: Option) -> Result> { + // TODO: Implement planning time statistics for this DataSource. + // At this point, we have information about the iceberg::table::Table which we are about + // to read, so maybe there's something we can get from there. + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + + fn with_fetch(&self, fetch: Option) -> Option> { + let mut self_clone = self.clone(); + self_clone.fetch = fetch; + Some(Arc::new(self_clone)) + } + + fn fetch(&self) -> Option { + self.fetch + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + Ok(None) + } + + fn metrics(&self) -> ExecutionPlanMetricsSet { + self.metrics.clone() + } + + fn try_pushdown_filters( + &self, + filters: Vec>, + _config: &ConfigOptions, + ) -> Result>> { + // TODO: Allow this DataSource to be pushed down filters. Some filters might be more + // straight forward to accept, like simple predicates, but some others might require + // a bit more work, like dynamic filters. + Ok(FilterPushdownPropagation::with_parent_pushdown_result( + vec![PushedDown::No; filters.len()], + )) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + // TODO: Allow this DataSource to be pushed down sort expressions. + Ok(SortOrderPushdownResult::Unsupported) + } +} diff --git a/iceberg/src/distributed_desired_task_count_handler.rs b/iceberg/src/distributed_desired_task_count_handler.rs new file mode 100644 index 000000000..77cd9bad3 --- /dev/null +++ b/iceberg/src/distributed_desired_task_count_handler.rs @@ -0,0 +1,21 @@ +use datafusion::common::Result; +use datafusion::datasource::source::DataSourceExec; +use datafusion_distributed::{DesiredTaskCountEvent, DesiredTaskCountEventResponse}; + +use crate::IcebergDataSource; + +// TODO: read the inner iceberg::table::Table and based on the contents attempt +// to estimate a good desired task count. +pub fn iceberg_desired_task_count( + ev: DesiredTaskCountEvent, +) -> Option> { + let _iceberg_data_source = ev + .plan + .downcast_ref::()? + .data_source() + .downcast_ref::()?; + + // TODO: not implemented. + + None +} diff --git a/iceberg/src/iceberg_ext.rs b/iceberg/src/iceberg_ext.rs new file mode 100644 index 000000000..a318282f8 --- /dev/null +++ b/iceberg/src/iceberg_ext.rs @@ -0,0 +1,214 @@ +use std::sync::Arc; + +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_distributed::DistributedExt; +use delegate::delegate; +use iceberg::io::StorageFactory; +use iceberg_storage_opendal::OpenDalResolvingStorageFactory; + +use crate::codec::IcebergCodec; +use crate::distributed_desired_task_count_handler::iceberg_desired_task_count; +use crate::{IcebergConfig, IcebergTableProviderFactory}; + +/// Configuration required to register the Iceberg SQL integration. +pub struct IcebergIntegrationOptions { + /// Builds storage implementations for table metadata and data files. + pub storage_factory: Arc, + /// Executes Iceberg IO-bound and CPU-bound work. + /// + /// Construct this from the application's Tokio runtime with + /// [`iceberg::Runtime::new`] or [`iceberg::Runtime::new_with_split`]. + pub iceberg_runtime: iceberg::Runtime, +} + +impl Default for IcebergIntegrationOptions { + fn default() -> Self { + Self { + storage_factory: Arc::new(OpenDalResolvingStorageFactory::new()), + iceberg_runtime: iceberg::Runtime::current(), + } + } +} + +/// Extends DataFusion session types with the Iceberg integration and its +/// configuration. +pub trait IcebergExt: Sized { + /// Registers the `ICEBERG` table provider factory and the default Iceberg + /// configuration. + fn set_iceberg_integration(&mut self, options: IcebergIntegrationOptions); + + /// Registers the `ICEBERG` table provider factory and the default Iceberg + /// configuration. + fn with_iceberg_integration(self, options: IcebergIntegrationOptions) -> Self; + + /// Sets the maximum number of Iceberg data files read concurrently. + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize); + + /// Sets the maximum number of Iceberg data files read concurrently. + fn with_iceberg_data_file_concurrency_limit(self, limit: usize) -> Self; + + /// Enables or disables Parquet row-group filtering for Iceberg reads. + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool); + + /// Enables or disables Parquet row-group filtering for Iceberg reads. + fn with_iceberg_row_group_filtering_enabled(self, enabled: bool) -> Self; + + /// Enables or disables row-level selection for Iceberg reads. + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool); + + /// Enables or disables row-level selection for Iceberg reads. + fn with_iceberg_row_selection_enabled(self, enabled: bool) -> Self; +} + +trait IcebergConfigExt { + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize); + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool); + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool); +} + +impl IcebergConfigExt for SessionConfig { + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize) { + iceberg_config_mut(self).data_file_concurrency_limit = limit; + } + + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool) { + iceberg_config_mut(self).row_group_filtering_enabled = enabled; + } + + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool) { + iceberg_config_mut(self).row_selection_enabled = enabled; + } +} + +fn iceberg_config_mut(config: &mut SessionConfig) -> &mut IcebergConfig { + if config.options().extensions.get::().is_none() { + config + .options_mut() + .extensions + .insert(IcebergConfig::default()); + } + + config + .options_mut() + .extensions + .get_mut::() + .expect("IcebergConfig was inserted above") +} + +const TABLE_FACTORY_IDENTIFIER: &str = "ICEBERG"; + +fn set_iceberg_integration(state: &mut SessionState, options: IcebergIntegrationOptions) { + iceberg_config_mut(state.config_mut()); + state.table_factories_mut().insert( + TABLE_FACTORY_IDENTIFIER.to_string(), + Arc::new(IcebergTableProviderFactory::new_with_runtime( + options.storage_factory, + options.iceberg_runtime, + )), + ); + state.set_distributed_desired_task_count_handler(iceberg_desired_task_count); + state.set_distributed_user_codec(IcebergCodec); +} + +impl IcebergExt for SessionStateBuilder { + fn set_iceberg_integration(&mut self, options: IcebergIntegrationOptions) { + iceberg_config_mut(self.config().get_or_insert_default()); + self.table_factories().get_or_insert_default().insert( + TABLE_FACTORY_IDENTIFIER.to_string(), + Arc::new(IcebergTableProviderFactory::new_with_runtime( + options.storage_factory, + options.iceberg_runtime, + )), + ); + self.set_distributed_desired_task_count_handler(iceberg_desired_task_count); + self.set_distributed_user_codec(IcebergCodec); + } + + delegate! { + to self.config().get_or_insert_default() { + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize); + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool); + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool); + } + + to self { + #[call(set_iceberg_integration)] + #[expr($;self)] + fn with_iceberg_integration(mut self, options: IcebergIntegrationOptions) -> Self; + + #[call(set_iceberg_data_file_concurrency_limit)] + #[expr($;self)] + fn with_iceberg_data_file_concurrency_limit(mut self, limit: usize) -> Self; + + #[call(set_iceberg_row_group_filtering_enabled)] + #[expr($;self)] + fn with_iceberg_row_group_filtering_enabled(mut self, enabled: bool) -> Self; + + #[call(set_iceberg_row_selection_enabled)] + #[expr($;self)] + fn with_iceberg_row_selection_enabled(mut self, enabled: bool) -> Self; + } + } +} + +impl IcebergExt for SessionState { + fn set_iceberg_integration(&mut self, options: IcebergIntegrationOptions) { + set_iceberg_integration(self, options); + } + + delegate! { + to self.config_mut() { + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize); + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool); + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool); + } + + to self { + #[call(set_iceberg_integration)] + #[expr($;self)] + fn with_iceberg_integration(mut self, options: IcebergIntegrationOptions) -> Self; + + #[call(set_iceberg_data_file_concurrency_limit)] + #[expr($;self)] + fn with_iceberg_data_file_concurrency_limit(mut self, limit: usize) -> Self; + + #[call(set_iceberg_row_group_filtering_enabled)] + #[expr($;self)] + fn with_iceberg_row_group_filtering_enabled(mut self, enabled: bool) -> Self; + + #[call(set_iceberg_row_selection_enabled)] + #[expr($;self)] + fn with_iceberg_row_selection_enabled(mut self, enabled: bool) -> Self; + } + } +} + +impl IcebergExt for SessionContext { + delegate! { + to self.state_ref().write() { + fn set_iceberg_integration(&mut self, options: IcebergIntegrationOptions); + fn set_iceberg_data_file_concurrency_limit(&mut self, limit: usize); + fn set_iceberg_row_group_filtering_enabled(&mut self, enabled: bool); + fn set_iceberg_row_selection_enabled(&mut self, enabled: bool); + } + + to self { + #[call(set_iceberg_integration)] + #[expr($;self)] + fn with_iceberg_integration(mut self, options: IcebergIntegrationOptions) -> Self; + + #[call(set_iceberg_data_file_concurrency_limit)] + #[expr($;self)] + fn with_iceberg_data_file_concurrency_limit(mut self, limit: usize) -> Self; + + #[call(set_iceberg_row_group_filtering_enabled)] + #[expr($;self)] + fn with_iceberg_row_group_filtering_enabled(mut self, enabled: bool) -> Self; + + #[call(set_iceberg_row_selection_enabled)] + #[expr($;self)] + fn with_iceberg_row_selection_enabled(mut self, enabled: bool) -> Self; + } + } +} diff --git a/iceberg/src/lib.rs b/iceberg/src/lib.rs index 4440687be..cb90c5e8a 100644 --- a/iceberg/src/lib.rs +++ b/iceberg/src/lib.rs @@ -1 +1,33 @@ -//! Apache Iceberg integration for DataFusion Distributed. +#![allow(clippy::disallowed_types)] + +//! Read-only Apache Iceberg integration for DataFusion Distributed. +//! +//! This crate ports the read path from Apache Iceberg Rust's DataFusion +//! integration. It deliberately contains no distributed execution adaptation +//! and no Iceberg write or commit implementation. + +mod common; +mod config; +mod data_source; +mod distributed_desired_task_count_handler; +mod iceberg_ext; +mod table_provider; +mod work_unit_feed; + +mod codec; +#[doc(hidden)] +pub mod test_utils; + +pub use codec::IcebergCodec; +pub use config::IcebergConfig; +pub use data_source::IcebergDataSource; +pub use distributed_desired_task_count_handler::iceberg_desired_task_count; +pub use iceberg_ext::IcebergExt; +pub use iceberg_ext::IcebergIntegrationOptions; +pub use table_provider::IcebergCatalogTableProvider; +pub use table_provider::IcebergStaticTableProvider; +pub use table_provider::IcebergTableProviderFactory; +pub use work_unit_feed::IcebergWorkUnitFeed; + +// re-export of iceberg-rust. +pub use iceberg; diff --git a/iceberg/src/table_provider/catalog.rs b/iceberg/src/table_provider/catalog.rs new file mode 100644 index 000000000..776b9724f --- /dev/null +++ b/iceberg/src/table_provider/catalog.rs @@ -0,0 +1,95 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::Session; +use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::Result; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_expr::Partitioning; +use datafusion::physical_plan::ExecutionPlan; +use iceberg::arrow::schema_to_arrow_schema; +use iceberg::{Catalog, NamespaceIdent, TableIdent}; + +use crate::IcebergDataSource; +use crate::common::df_err; +use crate::data_source::IcebergDataSourceOptions; + +/// Catalog-backed, read-only table provider with automatic metadata refresh. +/// +/// The provider loads fresh table metadata from the catalog on every scan. For +/// a fixed snapshot, use IcebergStaticTableProvider instead. +#[derive(Debug, Clone)] +pub struct IcebergCatalogTableProvider { + catalog: Arc, + table_ident: TableIdent, + schema: SchemaRef, + iceberg_runtime: iceberg::Runtime, +} + +impl IcebergCatalogTableProvider { + /// Creates a read-only catalog-backed provider. + pub async fn try_new( + catalog: Arc, + namespace: NamespaceIdent, + name: impl Into, + iceberg_runtime: iceberg::Runtime, + ) -> Result { + let table_ident = TableIdent::new(namespace, name.into()); + let table = catalog.load_table(&table_ident).await.map_err(df_err)?; + let schema = schema_to_arrow_schema(table.metadata().current_schema()).map_err(df_err)?; + + Ok(Self { + catalog, + table_ident, + schema: Arc::new(schema), + iceberg_runtime, + }) + } +} + +#[async_trait] +impl TableProvider for IcebergCatalogTableProvider { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + let table = self + .catalog + .load_table(&self.table_ident) + .await + .map_err(df_err)?; + + Ok(DataSourceExec::from_data_source(IcebergDataSource::new( + table, + self.schema.clone(), + Partitioning::UnknownPartitioning(state.config().target_partitions()), + IcebergDataSourceOptions { + snapshot_id: None, + projection, + filters, + fetch: limit, + iceberg_runtime: Some(self.iceberg_runtime.clone()), + }, + ))) + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } +} diff --git a/iceberg/src/table_provider/factory.rs b/iceberg/src/table_provider/factory.rs new file mode 100644 index 000000000..72356e76e --- /dev/null +++ b/iceberg/src/table_provider/factory.rs @@ -0,0 +1,161 @@ +use std::borrow::Cow; +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider, TableProviderFactory}; +use datafusion::common::{plan_datafusion_err, plan_err}; +use datafusion::error::Result; +use datafusion::logical_expr::CreateExternalTable; +use datafusion::sql::TableReference; +use iceberg::TableIdent; +use iceberg::io::{FileIOBuilder, StorageFactory}; +use iceberg::table::StaticTable; +use iceberg_storage_opendal::OpenDalResolvingStorageFactory; + +use crate::IcebergStaticTableProvider; +use crate::common::df_err; + +const SNAPSHOT_ID_OPTION: &str = "iceberg.snapshot_id"; + +/// Creates table providers for `CREATE EXTERNAL TABLE ... STORED AS ICEBERG`. +/// +/// # Example +/// +/// ```no_run +/// use datafusion::execution::SessionStateBuilder; +/// use datafusion::prelude::SessionContext; +/// use datafusion_distributed_iceberg::{IcebergExt, IcebergIntegrationOptions}; +/// +/// # async fn example() -> datafusion::error::Result<()> { +/// let state = SessionStateBuilder::new() +/// .with_default_features() +/// .with_iceberg_integration(IcebergIntegrationOptions::default()) +/// .build(); +/// let ctx = SessionContext::new_with_state(state); +/// +/// ctx.sql( +/// "CREATE EXTERNAL TABLE taxi STORED AS ICEBERG \ +/// LOCATION 's3://warehouse/taxi/metadata/v1.metadata.json'", +/// ) +/// .await? +/// .collect() +/// .await?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug)] +pub struct IcebergTableProviderFactory { + storage_factory: Arc, + iceberg_runtime: iceberg::Runtime, +} + +impl IcebergTableProviderFactory { + pub fn new() -> Self { + Self { + storage_factory: Arc::new(OpenDalResolvingStorageFactory::new()), + iceberg_runtime: iceberg::Runtime::current(), + } + } + + /// Create a new factory with a custom storage factory for creating FileIO instances. + pub fn new_with_storage_factory(storage_factory: Arc) -> Self { + Self::new_with_runtime(storage_factory, iceberg::Runtime::current()) + } + + /// Creates a factory with custom storage and Tokio runtime handles. + pub fn new_with_runtime( + storage_factory: Arc, + iceberg_runtime: iceberg::Runtime, + ) -> Self { + Self { + storage_factory, + iceberg_runtime, + } + } +} + +impl Default for IcebergTableProviderFactory { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl TableProviderFactory for IcebergTableProviderFactory { + async fn create( + &self, + _state: &dyn Session, + cmd: &CreateExternalTable, + ) -> Result> { + check_cmd(cmd)?; + + let table_name = &cmd.name; + let metadata_file_path = &cmd.location; + let options = &cmd.options; + let snapshot_id = parse_snapshot_id(options)?; + let mut storage_props = options.clone(); + storage_props.remove(SNAPSHOT_ID_OPTION); + + let table_name_with_ns = match table_name { + TableReference::Bare { table } => { + Cow::Owned(TableReference::partial("default", table.as_ref())) + } + other => Cow::Borrowed(other), + }; + + let table_ident = TableIdent::from_strs(table_name_with_ns.to_vec()).map_err(df_err)?; + let file_io = FileIOBuilder::new(self.storage_factory.clone()) + .with_props(&storage_props) + .build(); + let table = StaticTable::from_metadata_file(metadata_file_path, table_ident, file_io) + .await + .map_err(df_err)? + .into_table(); + + Ok(Arc::new(IcebergStaticTableProvider::try_new( + table, + snapshot_id, + self.iceberg_runtime.clone(), + )?)) + } +} + +fn parse_snapshot_id(options: &HashMap) -> Result> { + options + .get(SNAPSHOT_ID_OPTION) + .map(|snapshot_id| { + snapshot_id.parse::().map_err(|error| { + plan_datafusion_err!( + "{SNAPSHOT_ID_OPTION} must be a valid Iceberg snapshot ID: {error}" + ) + }) + }) + .transpose() +} + +fn check_cmd(cmd: &CreateExternalTable) -> Result<()> { + let CreateExternalTable { + schema, + table_partition_cols, + order_exprs, + constraints, + column_defaults, + .. + } = cmd; + + // Check if any of the fields violate the constraints in a single condition + let is_invalid = !schema.fields().is_empty() + || !table_partition_cols.is_empty() + || !order_exprs.is_empty() + || !constraints.is_empty() + || !column_defaults.is_empty(); + + if is_invalid { + return plan_err!( + "Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider." + ); + } + + Ok(()) +} diff --git a/iceberg/src/table_provider/mod.rs b/iceberg/src/table_provider/mod.rs new file mode 100644 index 000000000..b72eb2de7 --- /dev/null +++ b/iceberg/src/table_provider/mod.rs @@ -0,0 +1,7 @@ +mod catalog; +mod factory; +mod r#static; + +pub use catalog::IcebergCatalogTableProvider; +pub use factory::IcebergTableProviderFactory; +pub use r#static::IcebergStaticTableProvider; diff --git a/iceberg/src/table_provider/static.rs b/iceberg/src/table_provider/static.rs new file mode 100644 index 000000000..03d3bd968 --- /dev/null +++ b/iceberg/src/table_provider/static.rs @@ -0,0 +1,96 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::memory::DataSourceExec; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::{Result, plan_datafusion_err}; +use datafusion::datasource::TableType; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_expr::Partitioning; +use datafusion::physical_plan::ExecutionPlan; +use iceberg::arrow::schema_to_arrow_schema; + +use crate::IcebergDataSource; +use crate::common::df_err; +use crate::data_source::IcebergDataSourceOptions; + +/// Static, read-only provider for a table or a specific snapshot. +#[derive(Debug, Clone)] +pub struct IcebergStaticTableProvider { + table: iceberg::table::Table, + snapshot_id: Option, + schema: SchemaRef, + iceberg_runtime: iceberg::Runtime, +} + +impl IcebergStaticTableProvider { + /// Creates a provider that reads the provided table snapshot, or the current snapshot + /// if none provided. + pub fn try_new( + table: iceberg::table::Table, + snapshot_id: Option, + iceberg_runtime: iceberg::Runtime, + ) -> Result { + let table_schema = if let Some(snapshot_id) = snapshot_id { + let snapshot = table + .metadata() + .snapshot_by_id(snapshot_id) + .ok_or_else(|| { + plan_datafusion_err!( + "snapshot id {snapshot_id} not found in table {}", + table.identifier().name() + ) + })?; + snapshot.schema(table.metadata()).map_err(df_err)? + } else { + Arc::clone(table.metadata().current_schema()) + }; + + Ok(Self { + table, + snapshot_id, + schema: Arc::new(schema_to_arrow_schema(&table_schema).map_err(df_err)?), + iceberg_runtime, + }) + } +} + +#[async_trait] +impl TableProvider for IcebergStaticTableProvider { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + Ok(DataSourceExec::from_data_source(IcebergDataSource::new( + self.table.clone(), + self.schema.clone(), + Partitioning::UnknownPartitioning(state.config().target_partitions()), + IcebergDataSourceOptions { + snapshot_id: self.snapshot_id, + projection, + filters, + fetch: limit, + iceberg_runtime: Some(self.iceberg_runtime.clone()), + }, + ))) + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } +} diff --git a/iceberg/src/test_utils/harness.rs b/iceberg/src/test_utils/harness.rs new file mode 100644 index 000000000..a29a907e9 --- /dev/null +++ b/iceberg/src/test_utils/harness.rs @@ -0,0 +1,155 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_plan::displayable; +use datafusion::prelude::SessionContext; +use futures::StreamExt; +use futures::stream::BoxStream; +use iceberg::io::{ + FileMetadata, FileRead, FileWrite, InputFile, LocalFsStorage, OutputFile, Storage, + StorageConfig, StorageFactory, +}; +use iceberg::{Error, ErrorKind, Result as IcebergResult}; +use serde::{Deserialize, Serialize}; + +use crate::{IcebergExt, IcebergIntegrationOptions}; + +pub const FIXTURE_URI: &str = "s3://iceberg-test/warehouse/taxi"; +const WAREHOUSE_URI: &str = "s3://iceberg-test/warehouse/"; + +pub struct IcebergTestHarness { + ctx: SessionContext, +} + +impl IcebergTestHarness { + pub async fn new() -> Result { + let state = SessionStateBuilder::new() + .with_default_features() + .with_iceberg_integration(IcebergIntegrationOptions { + storage_factory: Arc::new(FixtureStorageFactory::default()), + iceberg_runtime: iceberg::Runtime::current(), + }) + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.sql(&format!( + "CREATE EXTERNAL TABLE taxi STORED AS ICEBERG \ + LOCATION '{FIXTURE_URI}/metadata/v1.metadata.json'" + )) + .await? + .collect() + .await?; + Ok(Self { ctx }) + } + + pub async fn query(&self, sql: &str) -> Result<(String, String)> { + let dataframe: DataFrame = self.ctx.sql(sql).await?; + let plan = dataframe.create_physical_plan().await?; + let batches = dataframe.collect().await?; + + Ok(( + displayable(plan.as_ref()).indent(true).to_string(), + pretty_format_batches(&batches)?.to_string(), + )) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct FixtureStorageFactory { + root: PathBuf, +} + +impl Default for FixtureStorageFactory { + fn default() -> Self { + Self { + root: PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../testdata/iceberg"), + } + } +} + +#[typetag::serde] +impl StorageFactory for FixtureStorageFactory { + fn build(&self, _config: &StorageConfig) -> IcebergResult> { + Ok(Arc::new(FixtureStorage { + root: self.root.clone(), + })) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct FixtureStorage { + root: PathBuf, +} + +impl FixtureStorage { + fn local_path(&self, path: &str) -> IcebergResult { + let relative = path.strip_prefix(WAREHOUSE_URI).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("unsupported fixture URI: {path}"), + ) + })?; + Ok(self.root.join(relative).display().to_string()) + } + + fn local(&self) -> LocalFsStorage { + LocalFsStorage::new() + } +} + +#[async_trait] +#[typetag::serde] +impl Storage for FixtureStorage { + async fn exists(&self, path: &str) -> IcebergResult { + self.local().exists(&self.local_path(path)?).await + } + + async fn metadata(&self, path: &str) -> IcebergResult { + self.local().metadata(&self.local_path(path)?).await + } + + async fn read(&self, path: &str) -> IcebergResult { + self.local().read(&self.local_path(path)?).await + } + + async fn reader(&self, path: &str) -> IcebergResult> { + self.local().reader(&self.local_path(path)?).await + } + + async fn write(&self, path: &str, bytes: Bytes) -> IcebergResult<()> { + self.local().write(&self.local_path(path)?, bytes).await + } + + async fn writer(&self, path: &str) -> IcebergResult> { + self.local().writer(&self.local_path(path)?).await + } + + async fn delete(&self, path: &str) -> IcebergResult<()> { + self.local().delete(&self.local_path(path)?).await + } + + async fn delete_prefix(&self, path: &str) -> IcebergResult<()> { + self.local().delete_prefix(&self.local_path(path)?).await + } + + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> IcebergResult<()> { + let mut paths = paths; + while let Some(path) = paths.next().await { + self.delete(&path).await?; + } + Ok(()) + } + + fn new_input(&self, path: &str) -> IcebergResult { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> IcebergResult { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } +} diff --git a/iceberg/src/test_utils/mod.rs b/iceberg/src/test_utils/mod.rs new file mode 100644 index 000000000..d3fc3550c --- /dev/null +++ b/iceberg/src/test_utils/mod.rs @@ -0,0 +1,3 @@ +mod harness; + +pub use harness::*; diff --git a/iceberg/src/work_unit_feed.rs b/iceberg/src/work_unit_feed.rs new file mode 100644 index 000000000..e153c4905 --- /dev/null +++ b/iceberg/src/work_unit_feed.rs @@ -0,0 +1,223 @@ +use std::sync::{Arc, Mutex, OnceLock}; + +use bytes::{Buf, BufMut}; +use datafusion::common::runtime::SpawnedTask; +use datafusion::common::{Result, exec_err, internal_err}; +use datafusion::error::DataFusionError; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::Partitioning; +use datafusion_distributed::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider}; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +use iceberg::expr::Predicate; +use iceberg::scan::FileScanTask; +use prost::encoding::{DecodeContext, WireType}; +use prost::{DecodeError, Message}; +use tokio::sync::mpsc::UnboundedReceiver; +use tokio_stream::wrappers::UnboundedReceiverStream; + +use crate::common::df_err; + +/// Work unit feed implementation that yields [FileScanTask] messages at execution time. +/// +/// It lazily spawns a task that scans an Iceberg table, and places each resolved [FileScanTask] +/// in P * T output channels, where: +/// - P is the number of partition in each distributed task. +/// - T is the number of distributed tasks +/// +/// ```text +/// ┌───────────────────────────┐ +/// │ Lazily spawned task │ +/// │ ┌───────────────────────┐ │ +/// │ │ Iceberg Table Scan │ │ +/// │ └───.─────────────.─────┘ │ +/// │ ( FileScanTask ) │ +/// │ .─────────────. │ +/// │ ( FileScanTask ) │ +/// │ `─────────────' │ +/// │ ... │ +/// │ .─────────────. │ +/// │ ( FileScanTask ) │ +/// │ `─────┬┬┬┬────' │ +/// ┌──────────┼───────────┘││└────────────┼───────────┐ +/// │ └────────────┼┼─────────────┘ │ +/// │ ┌──────┘└────────┐ │ +/// │ │ │ │ +/// ▼ ▼ ▼ ▼ +/// ┌───────────────┐┌───────────────┐┌───────────────┐┌───────────────┐ +/// │ Output mpsc 0 ││ Output mpsc 1 ││ Output mpsc 2 ││ Output mpsc 3 │ +/// │.─────────────.││.─────────────.││.─────────────.││.─────────────.│ +/// ( FileScanTask )( FileScanTask )( FileScanTask )( FileScanTask ) +/// │`─────────────'││.─────────────.││`─────────────'││.─────────────.│ +/// │ │( FileScanTask )│ │( FileScanTask ) +/// │ ││`─────────────'││ ││.─────────────.│ +/// │ ││ ││ │( FileScanTask ) +/// │ ││ ││ ││`─────────────'│ +/// │ ││ ││ ││ │ +/// └───────────────┘└───────────────┘└───────────────┘└───────────────┘ +/// ``` +/// +/// Each individual output channel ends up being a [datafusion_distributed::WorkUnit] stream that +/// goes to one partition of one distributed task: +/// +/// ```text +/// ┌───────────────┐┌───────────────┐┌───────────────┐┌───────────────┐ +/// │ Output mpsc 0 ││ Output mpsc 1 ││ Output mpsc 2 ││ Output mpsc 3 │ +/// └───────────────┘└───────────────┘└───────────────┘└───────────────┘ +/// │ │ │ │ +/// │ │ │ │ +/// ┌───────┼────────────────┼───────┐┌───────┼────────────────┼───────┐ +/// │ ▼ Task 0 ▼ ││ ▼ Task 1 ▼ │ +/// │┌──────────────┐┌──────────────┐││┌──────────────┐┌──────────────┐│ +/// ││ Partition 0 ││ Partition 1 ││││ Partition 2 ││ Partition 3 ││ +/// │└──────────────┘└──────────────┘││└──────────────┘└──────────────┘│ +/// └────────────────────────────────┘└────────────────────────────────┘ +/// ``` +/// +/// This works seamlessly in single-node and distributed mode using +/// [datafusion_distributed::WorkUnitFeed] machinery: +/// - If the query was not distributed, the [FileScanTask]s will be streamed in-memory. +/// - If the query was distributed, the [FileScanTask]s will be streamed over the network from +/// coordinator to workers. +#[derive(Debug)] +pub struct IcebergWorkUnitFeed { + /// A table in the catalog. + pub(crate) iceberg_table: iceberg::table::Table, + /// Snapshot of the table to scan. + pub(crate) snapshot_id: Option, + /// Projection column names, None means all columns. + pub(crate) projection: Option>, + /// Filters to apply to the table scan. + pub(crate) predicates: Option, + /// Partitioning scheme to which the feeds should adhere. + /// TODO: Today, only Partitioning::UnknownPartitioning partitioning is supported. + /// Ideally, both Range partitioning and hash partitioning should be supported. + pub(crate) partitioning: Partitioning, + /// Container for the lazily initialized task that scans the Iceberg table. + /// It will start as soon as the first [IcebergWorkUnitFeed::feed] is called. + pub(crate) sync_manager: OnceLock>>, +} + +impl Clone for IcebergWorkUnitFeed { + fn clone(&self) -> Self { + Self { + iceberg_table: self.iceberg_table.clone(), + snapshot_id: self.snapshot_id, + projection: self.projection.clone(), + predicates: self.predicates.clone(), + partitioning: self.partitioning.clone(), + sync_manager: Default::default(), + } + } +} + +type TakeableVec = Vec>>; + +#[derive(Debug)] +pub(crate) struct SyncManager { + task: Arc>, + feeds: TakeableVec>>, +} + +#[derive(Debug, Clone, Default)] +pub struct FileScanTaskMessage { + pub(crate) inner: Option, +} + +impl FileScanTaskMessage { + fn new(inner: FileScanTask) -> Self { + Self { inner: Some(inner) } + } +} + +impl WorkUnitFeedProvider for IcebergWorkUnitFeed { + type WorkUnit = FileScanTaskMessage; + + fn feed( + &self, + partition: usize, + ctx: Arc, + ) -> Result>> { + let wuf_ctx = DistributedWorkUnitFeedContext::from_ctx(&ctx); + + // This lazily spawns the tokio task that scans the Iceberg table. + // Only the first IcebergWorkUnitFeed::feed call will get to execute it, and the + // rest will just observe the already initialized result. + let sync_manager_or_err = self.sync_manager.get_or_init(|| { + // Start the table scan only once for all the .feed() calls. + let scan_builder = match self.snapshot_id { + Some(snapshot_id) => self.iceberg_table.scan().snapshot_id(snapshot_id), + None => self.iceberg_table.scan(), + }; + + let mut scan_builder = match &self.projection { + Some(column_names) => scan_builder.select(column_names), + None => scan_builder.select_all(), + }; + if let Some(pred) = &self.predicates { + scan_builder = scan_builder.with_filter(pred.clone()); + } + let table_scan = scan_builder.build().map_err(df_err)?; + + // Fanout the FileScanTask stream across P * T output channels where: + // - P is the number of output partitions per distributed task (`partition_count`) + // - T is the number of distributed tasks (`fan_out_tasks`) + let out_partitions = wuf_ctx.fan_out_tasks * self.partitioning.partition_count(); + let mut rxs = Vec::with_capacity(out_partitions); + let mut txs = Vec::with_capacity(out_partitions); + for _ in 0..out_partitions { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + rxs.push(Mutex::new(Some(rx))); + txs.push(tx); + } + + // A reference to this spawned task needs to be held, otherwise it will automatically + // be canceled. The lifetime of this `task` variable needs to leave as long as any of + // the return streams of the `feed()` method. + let task = SpawnedTask::spawn(async move { + let mut stream = match table_scan.plan_files().await { + Ok(stream) => stream.map_ok(FileScanTaskMessage::new), + Err(err) => { + let _ = txs[0].send(Err(df_err(err))); + return; + } + }; + + // Round robing across output partitions. + // TODO: this is fine for Partitioning::UnknownPartitioning, but any other + // partitioning will require smarter routing across output channels. + let mut i = 0; + while let Some(scan_task_or_err) = stream.next().await { + let _ = txs[i % txs.len()].send(scan_task_or_err.map_err(df_err)); + i += 1; + } + }); + + Ok(SyncManager { + task: Arc::new(task), + feeds: rxs, + }) + }); + + let sync_manager = match sync_manager_or_err { + Ok(sync_manager) => sync_manager, + Err(err) => return Err(DataFusionError::Shared(Arc::clone(err))), + }; + + let Some(feed) = sync_manager.feeds.get(partition) else { + return internal_err!("Invalid feed index {partition}"); + }; + + let Some(feed) = feed.lock().unwrap().take() else { + return exec_err!("Feed with index {partition} already taken"); + }; + + let task_ref = Arc::clone(&sync_manager.task); + + Ok(UnboundedReceiverStream::new(feed) + .inspect(move |_| { + let _ = &task_ref; // Keep the task alive as long as one feed is alive. + }) + .boxed()) + } +} diff --git a/iceberg/tests/external_table.rs b/iceberg/tests/external_table.rs new file mode 100644 index 000000000..c2b226689 --- /dev/null +++ b/iceberg/tests/external_table.rs @@ -0,0 +1,96 @@ +#[cfg(test)] +mod tests { + use datafusion::error::Result; + use datafusion_distributed_iceberg::test_utils::{FIXTURE_URI, IcebergTestHarness}; + + #[tokio::test] + async fn registers_the_fixture_with_the_iceberg_schema() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (_, batches) = harness.query("DESCRIBE taxi").await?; + + insta::assert_snapshot!(batches, @r" + +---------------------+---------------+-------------+ + | column_name | data_type | is_nullable | + +---------------------+---------------+-------------+ + | vendor_id | Int32 | YES | + | pickup_at | Timestamp(µs) | YES | + | dropoff_at | Timestamp(µs) | YES | + | passenger_count | Int64 | YES | + | trip_distance | Float64 | YES | + | pickup_location_id | Int32 | YES | + | dropoff_location_id | Int32 | YES | + | payment_type | Int64 | YES | + | fare_amount | Float64 | YES | + | tip_amount | Float64 | YES | + | tolls_amount | Float64 | YES | + | total_amount | Float64 | YES | + | pickup_date | Date32 | YES | + +---------------------+---------------+-------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn rejects_schema_definitions_for_existing_iceberg_tables() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let error = harness + .query(&format!( + "CREATE EXTERNAL TABLE invalid (id INT) STORED AS ICEBERG \ + LOCATION '{FIXTURE_URI}/metadata/v1.metadata.json'" + )) + .await + .unwrap_err(); + + insta::assert_snapshot!(error.to_string(), @r" + Error during planning: Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider. + "); + + Ok(()) + } + + #[tokio::test] + async fn registers_a_table_at_a_specific_snapshot() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + harness + .query(&format!( + "CREATE EXTERNAL TABLE taxi_snapshot STORED AS ICEBERG \ + LOCATION '{FIXTURE_URI}/metadata/v1.metadata.json' \ + OPTIONS ('iceberg.snapshot_id' '3167948105555765929')" + )) + .await?; + + let (_, batches) = harness + .query("SELECT COUNT(*) AS trips FROM taxi_snapshot") + .await?; + + insta::assert_snapshot!(batches, @r" + +--------+ + | trips | + +--------+ + | 175000 | + +--------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn rejects_an_invalid_snapshot_id() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let error = harness + .query(&format!( + "CREATE EXTERNAL TABLE invalid_snapshot STORED AS ICEBERG \ + LOCATION '{FIXTURE_URI}/metadata/v1.metadata.json' \ + OPTIONS ('iceberg.snapshot_id' 'not-a-snapshot-id')" + )) + .await + .unwrap_err(); + + insta::assert_snapshot!(error.to_string(), @r" + Error during planning: iceberg.snapshot_id must be a valid Iceberg snapshot ID: invalid digit found in string + "); + + Ok(()) + } +} diff --git a/iceberg/tests/filter_pushdown.rs b/iceberg/tests/filter_pushdown.rs new file mode 100644 index 000000000..6b70cfe4c --- /dev/null +++ b/iceberg/tests/filter_pushdown.rs @@ -0,0 +1,156 @@ +#[cfg(test)] +mod tests { + use datafusion::error::Result; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + + #[tokio::test] + async fn pushes_down_compound_predicates() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT COUNT(*) AS trips + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + AND payment_type IN (1, 2) + AND trip_distance >= 2.0", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: pickup_date@2 = 2024-01-10 AND (payment_type@1 = 1 OR payment_type@1 = 2) AND trip_distance@0 >= 2, projection=[] + DataSourceExec: format=iceberg, projection=[trip_distance, payment_type, pickup_date], predicate=((pickup_date = 2024-01-10) AND ((payment_type = 1) OR (payment_type = 2))) AND (trip_distance >= 2) + "); + insta::assert_snapshot!(batches, @r" + +-------+ + | trips | + +-------+ + | 9891 | + +-------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn pushes_down_disjunctions() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT pickup_date, COUNT(*) AS trips + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + OR pickup_date = DATE '2024-01-11' + GROUP BY pickup_date + ORDER BY pickup_date", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] + SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] + AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([pickup_date@0], 16), input_partitions=16 + AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + FilterExec: pickup_date@0 = 2024-01-10 OR pickup_date@0 = 2024-01-11 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=(pickup_date = 2024-01-10) OR (pickup_date = 2024-01-11) + "); + insta::assert_snapshot!(batches, @r" + +-------------+-------+ + | pickup_date | trips | + +-------------+-------+ + | 2024-01-10 | 25000 | + | 2024-01-11 | 25000 | + +-------------+-------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn pushes_down_null_predicates() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query("SELECT COUNT(*) AS trips FROM taxi WHERE passenger_count IS NULL") + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: passenger_count@0 IS NULL, projection=[] + DataSourceExec: format=iceberg, projection=[passenger_count], predicate=passenger_count IS NULL + "); + insta::assert_snapshot!(batches, @r" + +-------+ + | trips | + +-------+ + | 8829 | + +-------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn retains_unsupported_filter_after_safe_pushdown() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT COUNT(*) AS trips + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + AND trip_distance + 1.0 > 3.0", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: pickup_date@1 = 2024-01-10 AND trip_distance@0 + 1 > 3, projection=[] + DataSourceExec: format=iceberg, projection=[trip_distance, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-------+ + | trips | + +-------+ + | 10516 | + +-------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn executes_wholly_unsupported_filters_without_pushdown() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query("SELECT COUNT(*) AS trips FROM taxi WHERE trip_distance + 1.0 > 20.0") + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: trip_distance@0 + 1 > 20, projection=[] + DataSourceExec: format=iceberg, projection=[trip_distance] + "); + insta::assert_snapshot!(batches, @r" + +-------+ + | trips | + +-------+ + | 2757 | + +-------+ + "); + + Ok(()) + } +} diff --git a/iceberg/tests/limit_pushdown.rs b/iceberg/tests/limit_pushdown.rs new file mode 100644 index 000000000..532845a98 --- /dev/null +++ b/iceberg/tests/limit_pushdown.rs @@ -0,0 +1,137 @@ +#[cfg(test)] +mod tests { + use datafusion::error::Result; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + + #[tokio::test] + async fn applies_sql_limit_to_the_iceberg_scan_output() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT vendor_id, pickup_at, passenger_count, trip_distance, pickup_date + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + ORDER BY pickup_at + LIMIT 3", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortPreservingMergeExec: [pickup_at@1 ASC NULLS LAST], fetch=3 + SortExec: TopK(fetch=3), expr=[pickup_at@1 ASC NULLS LAST], preserve_partitioning=[true] + FilterExec: pickup_date@4 = 2024-01-10 + DataSourceExec: format=iceberg, projection=[vendor_id, pickup_at, passenger_count, trip_distance, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-----------+---------------------+-----------------+---------------+-------------+ + | vendor_id | pickup_at | passenger_count | trip_distance | pickup_date | + +-----------+---------------------+-----------------+---------------+-------------+ + | 2 | 2024-01-10T00:00:09 | | 0.78 | 2024-01-10 | + | 2 | 2024-01-10T00:00:10 | 1 | 0.88 | 2024-01-10 | + | 1 | 2024-01-10T00:00:10 | 1 | 3.4 | 2024-01-10 | + +-----------+---------------------+-----------------+---------------+-------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn keeps_unordered_limits_above_the_iceberg_source() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query("SELECT pickup_date FROM taxi WHERE pickup_date = DATE '2024-01-10' LIMIT 3") + .await?; + + insta::assert_snapshot!(plan, @r" + CoalescePartitionsExec: fetch=3 + FilterExec: pickup_date@0 = 2024-01-10, fetch=3 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-------------+ + | pickup_date | + +-------------+ + | 2024-01-10 | + | 2024-01-10 | + | 2024-01-10 | + +-------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn passes_the_scan_limit_to_the_iceberg_source() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query("SELECT pickup_date FROM taxi LIMIT 3") + .await?; + + insta::assert_snapshot!(plan, @r" + CoalescePartitionsExec: fetch=3 + DataSourceExec: format=iceberg, projection=[pickup_date], fetch=3 + "); + insta::assert_snapshot!(batches, @r" + +-------------+ + | pickup_date | + +-------------+ + | 2024-01-08 | + | 2024-01-08 | + | 2024-01-08 | + +-------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn avoids_scanning_for_limit_zero() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query("SELECT vendor_id FROM taxi WHERE pickup_date = DATE '2024-01-10' LIMIT 0") + .await?; + + insta::assert_snapshot!(plan, @r" + EmptyExec + "); + insta::assert_snapshot!(batches, @r" + ++ + ++ + "); + + Ok(()) + } + + #[tokio::test] + async fn applies_offset_before_limit() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT vendor_id, pickup_at + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + ORDER BY pickup_at, vendor_id, pickup_location_id + LIMIT 2 OFFSET 3", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[vendor_id@0 as vendor_id, pickup_at@1 as pickup_at] + GlobalLimitExec: skip=3, fetch=2 + SortPreservingMergeExec: [pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@2 ASC NULLS LAST], fetch=5 + SortExec: TopK(fetch=5), expr=[pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@2 ASC NULLS LAST], preserve_partitioning=[true] + FilterExec: pickup_date@3 = 2024-01-10, projection=[vendor_id@0, pickup_at@1, pickup_location_id@2] + DataSourceExec: format=iceberg, projection=[vendor_id, pickup_at, pickup_location_id, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-----------+---------------------+ + | vendor_id | pickup_at | + +-----------+---------------------+ + | 2 | 2024-01-10T00:00:12 | + | 2 | 2024-01-10T00:00:14 | + +-----------+---------------------+ + "); + + Ok(()) + } +} diff --git a/iceberg/tests/projection_pushdown.rs b/iceberg/tests/projection_pushdown.rs new file mode 100644 index 000000000..a428fdacc --- /dev/null +++ b/iceberg/tests/projection_pushdown.rs @@ -0,0 +1,136 @@ +#[cfg(test)] +mod tests { + use datafusion::error::Result; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + + #[tokio::test] + async fn projects_only_columns_required_by_the_query() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT pickup_date, COUNT(*) AS trips + FROM taxi + WHERE pickup_date >= DATE '2024-01-10' + GROUP BY pickup_date + ORDER BY pickup_date", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] + SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] + AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([pickup_date@0], 16), input_partitions=16 + AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + FilterExec: pickup_date@0 >= 2024-01-10 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date >= 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-------------+-------+ + | pickup_date | trips | + +-------------+-------+ + | 2024-01-10 | 25000 | + | 2024-01-11 | 25000 | + | 2024-01-12 | 25000 | + | 2024-01-13 | 25000 | + | 2024-01-14 | 25000 | + +-------------+-------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn includes_filter_columns_that_are_not_in_the_output() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT vendor_id, pickup_location_id + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + ORDER BY pickup_at, vendor_id + LIMIT 3", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[vendor_id@0 as vendor_id, pickup_location_id@1 as pickup_location_id] + SortPreservingMergeExec: [pickup_at@2 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST], fetch=3 + SortExec: TopK(fetch=3), expr=[pickup_at@2 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST], preserve_partitioning=[true] + FilterExec: pickup_date@3 = 2024-01-10, projection=[vendor_id@0, pickup_location_id@2, pickup_at@1] + DataSourceExec: format=iceberg, projection=[vendor_id, pickup_at, pickup_location_id, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-----------+--------------------+ + | vendor_id | pickup_location_id | + +-----------+--------------------+ + | 2 | 75 | + | 1 | 161 | + | 2 | 162 | + +-----------+--------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn projects_source_columns_for_computed_expressions() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT MAX(trip_distance * fare_amount) AS max_weighted_fare + FROM taxi + WHERE pickup_date = DATE '2024-01-10'", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + ProjectionExec: expr=[max(taxi.trip_distance * taxi.fare_amount)@0 as max_weighted_fare] + AggregateExec: mode=Final, gby=[], aggr=[max(taxi.trip_distance * taxi.fare_amount)] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[max(taxi.trip_distance * taxi.fare_amount)] + FilterExec: pickup_date@2 = 2024-01-10, projection=[trip_distance@0, fare_amount@1] + DataSourceExec: format=iceberg, projection=[trip_distance, fare_amount, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-------------------+ + | max_weighted_fare | + +-------------------+ + | 207508.9584 | + +-------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn reads_all_columns_when_selecting_star() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness + .query( + r"SELECT * + FROM taxi + WHERE pickup_date = DATE '2024-01-10' + ORDER BY pickup_at, vendor_id, pickup_location_id + LIMIT 1", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortPreservingMergeExec: [pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@5 ASC NULLS LAST], fetch=1 + SortExec: TopK(fetch=1), expr=[pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@5 ASC NULLS LAST], preserve_partitioning=[true] + FilterExec: pickup_date@12 = 2024-01-10 + DataSourceExec: format=iceberg, projection=[vendor_id, pickup_at, dropoff_at, passenger_count, trip_distance, pickup_location_id, dropoff_location_id, payment_type, fare_amount, tip_amount, tolls_amount, total_amount, pickup_date], predicate=pickup_date = 2024-01-10 + "); + insta::assert_snapshot!(batches, @r" + +-----------+---------------------+---------------------+-----------------+---------------+--------------------+---------------------+--------------+-------------+------------+--------------+--------------+-------------+ + | vendor_id | pickup_at | dropoff_at | passenger_count | trip_distance | pickup_location_id | dropoff_location_id | payment_type | fare_amount | tip_amount | tolls_amount | total_amount | pickup_date | + +-----------+---------------------+---------------------+-----------------+---------------+--------------------+---------------------+--------------+-------------+------------+--------------+--------------+-------------+ + | 2 | 2024-01-10T00:00:09 | 2024-01-10T00:03:30 | | 0.78 | 75 | 236 | 0 | 1.74 | 3.15 | 0.0 | 8.89 | 2024-01-10 | + +-----------+---------------------+---------------------+-----------------+---------------+--------------------+---------------------+--------------+-------------+------------+--------------+--------------+-------------+ + "); + + Ok(()) + } +} From 2abc5a32944918715835adb2f3d5a483c8adb057 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 24 Aug 2026 21:37:16 +0200 Subject: [PATCH 2/6] Add CI for iceberg --- .github/workflows/ci.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3c5da3fde..43390cedd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -118,6 +118,13 @@ jobs: key: "data" - run: cargo test --features clickbench --test clickbench_plans_test + iceberg-test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/setup + - run: cargo test -p datafusion-distributed-iceberg + format-check: runs-on: ubuntu-latest steps: From a6cb6af15558f3a82b01a209802bd0ab38db5418 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 25 Aug 2026 07:50:09 +0200 Subject: [PATCH 3/6] Fix clippy errors --- iceberg/src/work_unit_feed.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/iceberg/src/work_unit_feed.rs b/iceberg/src/work_unit_feed.rs index e153c4905..6e73eb3d1 100644 --- a/iceberg/src/work_unit_feed.rs +++ b/iceberg/src/work_unit_feed.rs @@ -1,6 +1,5 @@ use std::sync::{Arc, Mutex, OnceLock}; -use bytes::{Buf, BufMut}; use datafusion::common::runtime::SpawnedTask; use datafusion::common::{Result, exec_err, internal_err}; use datafusion::error::DataFusionError; @@ -11,8 +10,6 @@ use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use iceberg::expr::Predicate; use iceberg::scan::FileScanTask; -use prost::encoding::{DecodeContext, WireType}; -use prost::{DecodeError, Message}; use tokio::sync::mpsc::UnboundedReceiver; use tokio_stream::wrappers::UnboundedReceiverStream; From 8f7d284ff64e61e88bd1f0b6b1c279007fc95063 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 25 Aug 2026 08:47:29 +0200 Subject: [PATCH 4/6] Ignore failing test in the CI --- iceberg/tests/external_table.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/iceberg/tests/external_table.rs b/iceberg/tests/external_table.rs index c2b226689..c6c227619 100644 --- a/iceberg/tests/external_table.rs +++ b/iceberg/tests/external_table.rs @@ -50,6 +50,7 @@ mod tests { } #[tokio::test] + #[ignore = "Fails just on Linux with: Failed to load Parquet metadata, source: External: DataInvalid => Failed to read 524288 bytes: failed to fill whole buffer"] async fn registers_a_table_at_a_specific_snapshot() -> Result<()> { let harness = IcebergTestHarness::new().await?; harness From 6f340372514de02c7736c2e3b298509b1dd6ffb6 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 25 Aug 2026 09:18:15 +0200 Subject: [PATCH 5/6] Add lfx to CI and remove ignored test --- .github/workflows/ci.yml | 2 ++ iceberg/tests/external_table.rs | 1 - 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 43390cedd..3e8ee4d47 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -122,6 +122,8 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 + with: + lfs: true - uses: ./.github/actions/setup - run: cargo test -p datafusion-distributed-iceberg diff --git a/iceberg/tests/external_table.rs b/iceberg/tests/external_table.rs index c6c227619..c2b226689 100644 --- a/iceberg/tests/external_table.rs +++ b/iceberg/tests/external_table.rs @@ -50,7 +50,6 @@ mod tests { } #[tokio::test] - #[ignore = "Fails just on Linux with: Failed to load Parquet metadata, source: External: DataInvalid => Failed to read 524288 bytes: failed to fill whole buffer"] async fn registers_a_table_at_a_specific_snapshot() -> Result<()> { let harness = IcebergTestHarness::new().await?; harness From 58eaa415d081ba58808734874187ef8ad6bb31d2 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 25 Aug 2026 09:42:51 +0200 Subject: [PATCH 6/6] Fix target_partitions to 4 --- iceberg/src/test_utils/harness.rs | 3 ++- iceberg/tests/filter_pushdown.rs | 18 +++++++++--------- iceberg/tests/projection_pushdown.rs | 18 +++++++++--------- 3 files changed, 20 insertions(+), 19 deletions(-) diff --git a/iceberg/src/test_utils/harness.rs b/iceberg/src/test_utils/harness.rs index a29a907e9..c8be724c3 100644 --- a/iceberg/src/test_utils/harness.rs +++ b/iceberg/src/test_utils/harness.rs @@ -8,7 +8,7 @@ use datafusion::dataframe::DataFrame; use datafusion::error::Result; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::displayable; -use datafusion::prelude::SessionContext; +use datafusion::prelude::{SessionConfig, SessionContext}; use futures::StreamExt; use futures::stream::BoxStream; use iceberg::io::{ @@ -31,6 +31,7 @@ impl IcebergTestHarness { pub async fn new() -> Result { let state = SessionStateBuilder::new() .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(4)) .with_iceberg_integration(IcebergIntegrationOptions { storage_factory: Arc::new(FixtureStorageFactory::default()), iceberg_runtime: iceberg::Runtime::current(), diff --git a/iceberg/tests/filter_pushdown.rs b/iceberg/tests/filter_pushdown.rs index 6b70cfe4c..1daf717d5 100644 --- a/iceberg/tests/filter_pushdown.rs +++ b/iceberg/tests/filter_pushdown.rs @@ -50,15 +50,15 @@ mod tests { .await?; insta::assert_snapshot!(plan, @r" - SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] - SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] - ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] - AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] - RepartitionExec: partitioning=Hash([pickup_date@0], 16), input_partitions=16 - AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] - FilterExec: pickup_date@0 = 2024-01-10 OR pickup_date@0 = 2024-01-11 - DataSourceExec: format=iceberg, projection=[pickup_date], predicate=(pickup_date = 2024-01-10) OR (pickup_date = 2024-01-11) - "); + SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] + SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] + AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([pickup_date@0], 4), input_partitions=4 + AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + FilterExec: pickup_date@0 = 2024-01-10 OR pickup_date@0 = 2024-01-11 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=(pickup_date = 2024-01-10) OR (pickup_date = 2024-01-11) + "); insta::assert_snapshot!(batches, @r" +-------------+-------+ | pickup_date | trips | diff --git a/iceberg/tests/projection_pushdown.rs b/iceberg/tests/projection_pushdown.rs index a428fdacc..1e4078fb9 100644 --- a/iceberg/tests/projection_pushdown.rs +++ b/iceberg/tests/projection_pushdown.rs @@ -17,15 +17,15 @@ mod tests { .await?; insta::assert_snapshot!(plan, @r" - SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] - SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] - ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] - AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] - RepartitionExec: partitioning=Hash([pickup_date@0], 16), input_partitions=16 - AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] - FilterExec: pickup_date@0 >= 2024-01-10 - DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date >= 2024-01-10 - "); + SortPreservingMergeExec: [pickup_date@0 ASC NULLS LAST] + SortExec: expr=[pickup_date@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] + AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + RepartitionExec: partitioning=Hash([pickup_date@0], 4), input_partitions=4 + AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))] + FilterExec: pickup_date@0 >= 2024-01-10 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date >= 2024-01-10 + "); insta::assert_snapshot!(batches, @r" +-------------+-------+ | pickup_date | trips |