diff --git a/crates/datafusion/src/read/mod.rs b/crates/datafusion/src/read/mod.rs index 4c0157d..2930e2c 100644 --- a/crates/datafusion/src/read/mod.rs +++ b/crates/datafusion/src/read/mod.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. +pub(crate) mod physical_predicate; pub(crate) mod predicate; pub(crate) mod scan; +pub use physical_predicate::convert_physical_filter_to_predicate; pub use predicate::convert_filters_to_predicate; pub use scan::{IcebergDataSource, IcebergDataSourceBuilder}; diff --git a/crates/datafusion/src/read/physical_predicate.rs b/crates/datafusion/src/read/physical_predicate.rs new file mode 100644 index 0000000..7418839 --- /dev/null +++ b/crates/datafusion/src/read/physical_predicate.rs @@ -0,0 +1,794 @@ +// 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::sync::Arc; + +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{ + BinaryExpr, CaseExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, + Literal, NotExpr, +}; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use datafusion::scalar::ScalarValue; +use iceberg::expr::{Predicate, Reference}; +use iceberg::spec::{Datum, PrimitiveType}; + +/// Turns a DataFusion physical filter into an Iceberg predicate. +/// +/// The predicate matches every row the filter matches. It can match more rows when part of +/// the filter has no Iceberg form, so callers must still apply the filter to the scanned rows. +/// +/// Returns `None` when no part of the filter can be converted. +pub fn convert_physical_filter_to_predicate( + filter: &Arc, +) -> Option { + convert(filter).map(|converted| converted.predicate) +} + +/// An iceberg predicate that matches every row the DataFusion filter matches. +/// It can also match extra rows, so the caller must still apply the original filter. +struct Converted { + predicate: Predicate, + /// True if the predicate is false wherever the filter is false. + /// Only an exact predicate can go under a NOT. + exact: bool, +} + +impl Converted { + fn exact(predicate: Predicate) -> Self { + Self { + predicate, + exact: true, + } + } + + fn wider(predicate: Predicate) -> Self { + Self { + predicate, + exact: false, + } + } + + fn and(self, other: Self) -> Self { + Self { + predicate: self.predicate.and(other.predicate), + exact: self.exact && other.exact, + } + } + + fn or(self, other: Self) -> Self { + Self { + predicate: self.predicate.or(other.predicate), + exact: self.exact && other.exact, + } + } +} + +fn convert(expr: &Arc) -> Option { + if let Some(binary) = expr.downcast_ref::() { + return convert_binary(binary); + } + if let Some(not) = expr.downcast_ref::() { + return convert_not(not); + } + if let Some(case) = expr.downcast_ref::() { + return convert_case(case); + } + if let Some(in_list) = expr.downcast_ref::() { + return convert_in_list(in_list); + } + if let Some(is_null) = expr.downcast_ref::() { + let column = column_name(is_null.arg())?; + return Some(Converted::exact(Reference::new(column).is_null())); + } + if let Some(is_not_null) = expr.downcast_ref::() { + let column = column_name(is_not_null.arg())?; + return Some(Converted::exact(Reference::new(column).is_not_null())); + } + if let Some(like) = expr.downcast_ref::() { + return convert_like(like).map(Converted::exact); + } + if let Some(function) = expr.downcast_ref::() { + return convert_function(function).map(Converted::exact); + } + // A bare boolean column like `WHERE flag` becomes `flag = true`. + if let Some(column) = expr.downcast_ref::() { + let predicate = Reference::new(column.name()).equal_to(Datum::bool(true)); + return Some(Converted::exact(predicate)); + } + // A bare `true`, `false` or NULL. NULL matches nothing. + match literal_value(expr)? { + ScalarValue::Boolean(Some(true)) => Some(Converted::exact(Predicate::AlwaysTrue)), + ScalarValue::Boolean(_) | ScalarValue::Null => { + Some(Converted::exact(Predicate::AlwaysFalse)) + } + _ => None, + } +} + +fn convert_binary(binary: &BinaryExpr) -> Option { + let (left, right) = (binary.left(), binary.right()); + match binary.op() { + Operator::And => both(convert(left), convert(right)), + // An OR needs both sides. Turning `a = 1 OR f(b)` into `a = 1` would skip rows where + // only `f(b)` is true. + Operator::Or => Some(convert(left)?.or(convert(right)?)), + op => compare(left, *op, right).map(Converted::exact), + } +} + +/// ANDs two converted sides. An AND can drop a side it can't convert: +/// `a = 1 AND f(b)` becomes `a = 1`. The result matches extra rows, which is fine. +fn both(left: Option, right: Option) -> Option { + match (left, right) { + (Some(left), Some(right)) => Some(left.and(right)), + (Some(side), None) | (None, Some(side)) => Some(Converted::wider(side.predicate)), + (None, None) => None, + } +} + +/// A NOT needs an exact predicate. In `NOT (a = 1 AND f(b))`, the inside becomes `a = 1`, +/// and `NOT (a = 1)` would skip rows where `a = 1` but `f(b)` is false. +fn convert_not(not: &NotExpr) -> Option { + convert(not.arg()) + .filter(|inner| inner.exact) + .map(|inner| Converted::exact(inner.predicate.negate())) +} + +/// In a CASE statement, each WHEN/THEN pair applies a filter over the rows. +/// +/// `CASE WHEN c1 THEN x WHEN c2 THEN y ELSE z END` becomes `(c1 AND x) OR (c2 AND y) OR z`. +/// It ignores that the second branch only applies when `c1` is false, so it can match +/// extra rows. +/// +/// `CASE a WHEN 1 THEN x ...` works the same way, with `a = 1` as the WHEN. +/// +/// A missing ELSE means NULL, which matches nothing. +fn convert_case(case: &CaseExpr) -> Option { + // This handles the two variants for CASE + // 1. CASE WHEN a = 1 THEN ... + // In this variant, `case.expr()` is `None` because the expression is imbued in the WHEN clause. + // We pass the WHEN clause to our general predicate converter for this. + // + // 2. CASE a WHEN 0 THEN ... + // In this variant, `case.expr()` is `Some(a)`, and each WHEN clause compares its own + // expression with this expression. + // To handle this, we convert this into a comparison predicate. + // Iceberg predicates can't always represent a given expression; if that is the case, we drop + // the predicate for the WHEN clause. + let convert_when = |when| match case.expr() { + Some(value) => compare(value, Operator::Eq, when).map(Converted::exact), + None => convert(when), + }; + + let predicate = case + .when_then_expr() + .iter() + .map(|(when, then)| both(convert_when(when), convert(then))) + .chain(case.else_expr().map(convert)) + .collect::>>()? + .into_iter() + .map(|branch| branch.predicate) + .reduce(Predicate::or)?; + Some(Converted::wider(predicate)) +} + +/// Converts a column compared with a value, like `a < 5`. +/// The value can come first too: `5 > a` becomes `a < 5`. +fn compare( + left: &Arc, + op: Operator, + right: &Arc, +) -> Option { + let (column, op, datum) = match (column_name(left), literal_datum(right)) { + (Some(column), Some(datum)) => (column, op, datum), + _ => (column_name(right)?, op.swap()?, literal_datum(left)?), + }; + let column = Reference::new(column); + match op { + Operator::Eq => Some(column.equal_to(datum)), + Operator::NotEq => Some(column.not_equal_to(datum)), + Operator::Lt => Some(column.less_than(datum)), + Operator::LtEq => Some(column.less_than_or_equal_to(datum)), + Operator::Gt => Some(column.greater_than(datum)), + Operator::GtEq => Some(column.greater_than_or_equal_to(datum)), + _ => None, + } +} + +fn convert_in_list(in_list: &InListExpr) -> Option { + // Something like `struct(a, b) IN (...)`, which checks several columns at once. + let Some(column) = column_name(in_list.expr()) else { + return convert_struct_in_list(in_list).map(Converted::wider); + }; + let column = Reference::new(column); + let values = in_list + .list() + .iter() + .map(literal_value) + .collect::>>()?; + let (nulls, values): (Vec<_>, Vec<_>) = + values.into_iter().partition(|value| value.is_null()); + let datums = values + .into_iter() + .map(scalar_value_to_datum) + .collect::>>()?; + let predicate = match (in_list.negated(), nulls.is_empty()) { + // `x IN (1, NULL)` matches the same rows as `x IN (1)`. + (false, _) => column.is_in(datums), + (true, true) => column.is_not_in(datums), + // `x NOT IN (1, NULL)` never matches anything. It's rare, so we don't convert it. + (true, false) => return None, + }; + Some(Converted::exact(predicate)) +} + +/// Hash joins on two or more columns build filters like `struct(a, b) IN ({1, 3}, {2, 4})`. +/// Iceberg can't check pairs, so each column is checked on its own: +/// `a IN (1, 2) AND b IN (3, 4)`. This also matches pairs like `{1, 4}`, which is fine. +/// A column with a NULL value is left out, because the join decides if NULL matches NULL. +fn convert_struct_in_list(in_list: &InListExpr) -> Option { + if in_list.negated() { + return None; + } + let columns = in_list + .expr() + .downcast_ref::() + .filter(|function| function.name() == "struct")? + .args(); + let items = in_list + .list() + .iter() + .map(|item| match literal_value(item)? { + ScalarValue::Struct(item) => Some(item), + _ => None, + }) + .collect::>>()?; + columns + .iter() + .enumerate() + .filter_map(|(index, column)| { + let datums = items + .iter() + .map(|item| { + let value = + ScalarValue::try_from_array(item.columns().get(index)?, 0) + .ok()?; + scalar_value_to_datum(&value) + }) + .collect::>>()?; + Some(Reference::new(column_name(column)?).is_in(datums)) + }) + .reduce(Predicate::and) +} + +/// Only simple prefix patterns: `a LIKE 'abc%'` becomes `a STARTS WITH 'abc'`. +/// ILIKE can't be converted, because iceberg's STARTS WITH is case-sensitive. +fn convert_like(like: &LikeExpr) -> Option { + if like.case_insensitive() { + return None; + } + + // Iceberg can only represent prefix matches using plain columns. + // E.g. WHERE s LIKE ... + // It cannot represent complex expressions or functions on the left side of the match + // E.g WHERE upper(s) LIKE ... + // `column_name` will return None if the left side is not a plain column and we don't add a predicate for this LikeExpr + let column = Reference::new(column_name(like.expr())?); + + // The pattern being matched *must* be a string literal that ends in a '%', which means that it + // is trying to match the string's prefix. + // + // Iceberg can only support `STARTS WITH` predicates - after stripping the final '%', + // if there are more symbols in the string that could represent a broader match, + // we don't add a predicate for this LikeExpr + let prefix = string_literal(like.pattern())?.strip_suffix('%')?; + if prefix.contains(['%', '_', '\\']) { + return None; + } + + let prefix = prefix_datum(prefix)?; + Some(if like.negated() { + column.not_starts_with(prefix) + } else { + column.starts_with(prefix) + }) +} + +/// Functions are matched by name: `isnan(a)` and `starts_with(a, 'abc')`. +fn convert_function(function: &ScalarFunctionExpr) -> Option { + match (function.name(), function.args()) { + ("isnan", [arg]) => Some(Reference::new(column_name(arg)?).is_nan()), + ("starts_with", [arg, prefix]) => { + let prefix = prefix_datum(string_literal(prefix)?)?; + Some(Reference::new(column_name(arg)?).starts_with(prefix)) + } + _ => None, + } +} + +/// Converts a string used as the prefix in a 'STARTS WITH' clause into a Datum +/// +/// This only works for ASCII strings. This is because iceberg-rust currently has a bug +/// where it can skip rows unintentionally if a partition column truncates the underlying value +/// via truncate(N) and if that underlying value contains non-ascii characters +/// Ref: https://github.com/apache/iceberg-rust/issues/3626 +fn prefix_datum(prefix: &str) -> Option { + prefix.is_ascii().then(|| Datum::string(prefix)) +} + +fn column_name(expr: &Arc) -> Option<&str> { + expr.downcast_ref::().map(Column::name) +} + +fn literal_value(expr: &Arc) -> Option<&ScalarValue> { + expr.downcast_ref::().map(Literal::value) +} + +fn literal_datum(expr: &Arc) -> Option { + scalar_value_to_datum(literal_value(expr)?) +} + +fn string_literal(expr: &Arc) -> Option<&str> { + literal_value(expr)?.try_as_str().flatten() +} + +/// Converts a DataFusion value to an iceberg value, for the types iceberg columns use. +/// +/// Floats aren't converted. DataFusion sorts NaN above every number and -0.0 below 0.0. +/// Iceberg doesn't, so a float filter like `a > 1.0` could skip rows DataFusion keeps. +/// +/// Times aren't converted either, because iceberg-rust's file reader can't compare them yet. +fn scalar_value_to_datum(value: &ScalarValue) -> Option { + match value { + ScalarValue::Boolean(Some(v)) => Some(Datum::bool(*v)), + ScalarValue::Int32(Some(v)) => Some(Datum::int(*v)), + ScalarValue::Int64(Some(v)) => Some(Datum::long(*v)), + ScalarValue::Decimal128(Some(v), precision, scale) => { + let data_type = PrimitiveType::Decimal { + precision: u32::from(*precision), + scale: u32::try_from(*scale).ok()?, + }; + Datum::try_from_bytes(&v.to_be_bytes(), data_type).ok() + } + ScalarValue::Date32(Some(v)) => Some(Datum::date(*v)), + ScalarValue::TimestampMicrosecond(Some(v), None) => { + Some(Datum::timestamp_micros(*v)) + } + ScalarValue::TimestampMicrosecond(Some(v), Some(_)) => { + Some(Datum::timestamptz_micros(*v)) + } + ScalarValue::TimestampNanosecond(Some(v), None) => { + Some(Datum::timestamp_nanos(*v)) + } + ScalarValue::TimestampNanosecond(Some(v), Some(_)) => { + Some(Datum::timestamptz_nanos(*v)) + } + ScalarValue::Utf8(Some(v)) + | ScalarValue::LargeUtf8(Some(v)) + | ScalarValue::Utf8View(Some(v)) => Some(Datum::string(v)), + ScalarValue::Binary(Some(v)) + | ScalarValue::LargeBinary(Some(v)) + | ScalarValue::BinaryView(Some(v)) => Some(Datum::binary(v.clone())), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::{ArrayRef, Int64Array, StructArray}; + use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema}; + use datafusion::common::{DFSchema, Result}; + use datafusion::config::ConfigOptions; + use datafusion::functions::core::r#struct; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, InListExpr, NotExpr, col, in_list, lit, + }; + use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; + use datafusion::prelude::SessionContext; + use datafusion::scalar::ScalarValue; + use insta::assert_snapshot; + + use super::convert_physical_filter_to_predicate; + + // The filters below use the columns from `schema()`: `a` and `b` are integers, `s` is a + // string, `f` is a float and `flag` is a boolean. + + #[test] + fn converts_null_checks() { + assert_snapshot!(convert_sql("a IS NULL"), @"a IS NULL"); + assert_snapshot!(convert_sql("a IS NOT NULL"), @"a IS NOT NULL"); + } + + #[test] + fn converts_comparisons_with_a_literal() { + assert_snapshot!(convert_sql("a = 5"), @"a = 5"); + assert_snapshot!(convert_sql("a != 5"), @"a != 5"); + assert_snapshot!(convert_sql("a < 5"), @"a < 5"); + assert_snapshot!(convert_sql("a <= 5"), @"a <= 5"); + assert_snapshot!(convert_sql("a > 5"), @"a > 5"); + assert_snapshot!(convert_sql("a >= 5"), @"a >= 5"); + } + + #[test] + fn flips_comparisons_with_the_literal_on_the_left() { + assert_snapshot!(convert_sql("5 < a"), @"a > 5"); + assert_snapshot!(convert_sql("5 >= a"), @"a <= 5"); + } + + #[test] + fn leaves_comparisons_between_columns_unconverted() { + assert_snapshot!(convert_sql("a < b"), @"None"); + } + + #[test] + fn leaves_float_comparisons_unconverted() { + assert_snapshot!(convert_sql("f > 1.5"), @"None"); + assert_snapshot!(convert_sql("a = 1 AND f > 1.5"), @"a = 1"); + } + + #[test] + fn converts_in_lists() { + assert_snapshot!(convert_sql("a IN (1, 2)"), @"a IN (1, 2)"); + assert_snapshot!(convert_sql("a NOT IN (1, 2)"), @"a NOT IN (1, 2)"); + } + + #[test] + fn drops_nulls_from_an_in_list() -> Result<()> { + // Built by hand, the way a hash join builds it. SQL would wrap the NULL in a cast. + let schema = schema(); + let null = lit(ScalarValue::Int64(None)); + let values = vec![lit(1i64), Arc::clone(&null)]; + let one_or_null = in_list(col("a", &schema)?, values, &false, &schema)?; + let only_null = in_list(col("a", &schema)?, vec![null], &false, &schema)?; + assert_snapshot!(convert_filter(&one_or_null), @"a IN (1)"); + // `a IN (NULL)` matches nothing. + assert_snapshot!(convert_filter(&only_null), @"a IN ()"); + Ok(()) + } + + #[test] + fn leaves_not_in_with_a_null_unconverted() -> Result<()> { + // Built by hand, the way a hash join builds it. SQL would wrap the NULL in a cast. + let schema = schema(); + let values = vec![lit(1i64), lit(ScalarValue::Int64(None))]; + let filter = in_list(col("a", &schema)?, values, &true, &schema)?; + assert_snapshot!(convert_filter(&filter), @"None"); + Ok(()) + } + + #[test] + fn leaves_in_lists_with_a_column_item_unconverted() { + assert_snapshot!(convert_sql("a IN (b, 1)"), @"None"); + } + + #[test] + fn leaves_an_in_list_over_an_expression_unconverted() { + assert_snapshot!(convert_sql("a + 1 IN (1, 2)"), @"None"); + } + + #[test] + fn converts_like_prefixes() { + assert_snapshot!(convert_sql("s LIKE 'ab%'"), @r#"s STARTS WITH "ab""#); + assert_snapshot!(convert_sql("s NOT LIKE 'ab%'"), @r#"s NOT STARTS WITH "ab""#); + } + + #[test] + fn leaves_other_like_patterns_unconverted() { + assert_snapshot!(convert_sql("s LIKE 'a_b%'"), @"None"); + assert_snapshot!(convert_sql("s LIKE 'a%b%'"), @"None"); + assert_snapshot!(convert_sql(r"s LIKE 'a\b%'"), @"None"); + assert_snapshot!(convert_sql("s LIKE 'ab'"), @"None"); + assert_snapshot!(convert_sql("s ILIKE 'ab%'"), @"None"); + } + + #[test] + fn leaves_non_ascii_prefixes_unconverted() { + assert_snapshot!(convert_sql("s LIKE 'éa%'"), @"None"); + assert_snapshot!(convert_sql("starts_with(s, 'éa')"), @"None"); + } + + #[test] + fn converts_isnan_and_starts_with() { + assert_snapshot!(convert_sql("isnan(f)"), @"f IS NAN"); + assert_snapshot!(convert_sql("starts_with(s, 'ab')"), @r#"s STARTS WITH "ab""#); + } + + #[test] + fn converts_a_boolean_column() { + assert_snapshot!(convert_sql("flag"), @"flag = true"); + } + + #[test] + fn converts_true_false_and_null() { + assert_snapshot!(convert_sql("true"), @"TRUE"); + assert_snapshot!(convert_sql("false"), @"FALSE"); + // A NULL filter keeps no rows. + assert_snapshot!(convert_sql("NULL"), @"FALSE"); + } + + #[test] + fn converts_an_and() { + assert_snapshot!(convert_sql("a = 1 AND b = 2"), @"(a = 1) AND (b = 2)"); + } + + #[test] + fn drops_the_unconverted_side_of_an_and() { + assert_snapshot!(convert_sql("a = 1 AND a < b"), @"a = 1"); + } + + #[test] + fn converts_an_or_only_when_both_sides_convert() { + // Hash joins build this when rows with a NULL key can still match. + assert_snapshot!( + convert_sql("a IS NULL OR a IN (1, 2)"), + @"(a IS NULL) OR (a IN (1, 2))" + ); + assert_snapshot!(convert_sql("a = 1 OR a < b"), @"None"); + } + + #[test] + fn negates_exact_predicates() { + assert_snapshot!(convert_sql("NOT (a < 5)"), @"a >= 5"); + assert_snapshot!(convert_sql("NOT (a IS NULL)"), @"a IS NOT NULL"); + // DataFusion turns `NOT (a IN (1, 2))` into `a NOT IN (1, 2)` itself, and does the + // same for LIKE. Putting them inside an OR keeps the NOT. + assert_snapshot!( + convert_sql("NOT (a IN (1, 2) OR b = 2)"), + @"(a NOT IN (1, 2)) AND (b != 2)" + ); + assert_snapshot!( + convert_sql("NOT (s LIKE 'ab%' OR b = 2)"), + @r#"(s NOT STARTS WITH "ab") AND (b != 2)"# + ); + assert_snapshot!(convert_sql("NOT starts_with(s, 'ab')"), @r#"s NOT STARTS WITH "ab""#); + assert_snapshot!(convert_sql("NOT isnan(f)"), @"f IS NOT NAN"); + assert_snapshot!(convert_sql("NOT flag"), @"flag != true"); + assert_snapshot!(convert_sql("NOT (a = 1 AND b = 2)"), @"(a != 1) OR (b != 2)"); + assert_snapshot!(convert_sql("NOT (a = 1 OR b = 2)"), @"(a != 1) AND (b != 2)"); + } + + #[test] + fn leaves_not_over_a_partly_converted_filter_unconverted() { + assert_snapshot!(convert_sql("NOT (a = 1 AND a < b)"), @"None"); + assert_snapshot!(convert_sql("NOT (a = 1 AND (b = 2 AND a < b))"), @"None"); + assert_snapshot!(convert_sql("NOT (a = 1 OR (b = 2 AND a < b))"), @"None"); + } + + #[test] + fn converts_a_case_to_each_when_and_its_then_ored_together() { + assert_snapshot!( + convert_sql("CASE WHEN a = 1 THEN b = 2 WHEN a = 2 THEN b = 3 ELSE b = 4 END"), + @"(((a = 1) AND (b = 2)) OR ((a = 2) AND (b = 3))) OR (b = 4)" + ); + } + + #[test] + fn converts_case_value_whens_to_equalities() { + assert_snapshot!( + convert_sql("CASE a WHEN 1 THEN b = 2 WHEN 2 THEN b = 3 END"), + @"((a = 1) AND (b = 2)) OR ((a = 2) AND (b = 3))" + ); + } + + #[test] + fn keeps_only_the_thens_when_the_whens_cant_be_converted() { + // `a % 2` stands in for the `hash(a) % N` a partitioned hash join uses. + assert_snapshot!( + convert_sql("CASE a % 2 WHEN 0 THEN b IN (1, 2) WHEN 1 THEN b = 3 ELSE false END"), + @"(b IN (1, 2)) OR (b = 3)" + ); + } + + #[test] + fn keeps_the_when_if_its_then_cant_be_converted() { + assert_snapshot!(convert_sql("CASE WHEN a = 1 THEN a < b ELSE false END"), @"a = 1"); + } + + #[test] + fn converts_a_case_with_an_always_true_branch_to_true() { + // Hash joins use `true` for a partition whose build side was canceled. + assert_snapshot!(convert_sql("CASE a % 2 WHEN 0 THEN b = 1 ELSE true END"), @"TRUE"); + } + + #[test] + fn leaves_a_case_with_an_unconverted_branch_unconverted() { + assert_snapshot!( + convert_sql("CASE a % 2 WHEN 0 THEN b IN (1, 2) WHEN 1 THEN a < b ELSE false END"), + @"None" + ); + assert_snapshot!(convert_sql("CASE WHEN a = 1 THEN b = 2 ELSE a < b END"), @"None"); + } + + #[test] + fn leaves_not_over_a_case_unconverted() { + assert_snapshot!(convert_sql("NOT (CASE WHEN a = 1 THEN b = 2 ELSE false END)"), @"None"); + } + + #[test] + fn converts_a_multi_column_join_filter_to_an_in_list_per_column() -> Result<()> { + let filter = struct_in(vec![Some(1), Some(2)], vec![Some(5), Some(6)])?; + assert_snapshot!(convert_filter(&filter), @"(a IN (1, 2)) AND (b IN (5, 6))"); + Ok(()) + } + + #[test] + fn drops_join_filter_columns_with_a_null() -> Result<()> { + let filter = struct_in(vec![Some(1), Some(2)], vec![Some(3), None])?; + assert_snapshot!(convert_filter(&filter), @"a IN (1, 2)"); + Ok(()) + } + + #[test] + fn leaves_a_multi_column_not_in_unconverted() -> Result<()> { + let filter = struct_not_in(vec![Some(1), Some(2)], vec![Some(3), Some(4)])?; + assert_snapshot!(convert_filter(&filter), @"None"); + Ok(()) + } + + #[test] + fn leaves_not_over_a_multi_column_join_filter_unconverted() -> Result<()> { + let filter = struct_in(vec![Some(1), Some(2)], vec![Some(3), Some(4)])?; + let filter: Arc = Arc::new(NotExpr::new(filter)); + assert_snapshot!(convert_filter(&filter), @"None"); + Ok(()) + } + + #[test] + fn converts_literals_to_iceberg_values() { + // Each value is compared with `a`. The translator doesn't check that the types match. + let values = [ + ScalarValue::Boolean(Some(true)), + ScalarValue::Int32(Some(1)), + ScalarValue::Int64(Some(1)), + ScalarValue::Decimal128(Some(12345), 10, 2), + ScalarValue::Decimal128(Some(-12345), 10, 2), + ScalarValue::Date32(Some(19723)), + ScalarValue::TimestampMicrosecond(Some(1_704_067_200_000_000), None), + ScalarValue::TimestampMicrosecond( + Some(1_704_067_200_000_000), + Some("+00:00".into()), + ), + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_000_000), None), + ScalarValue::TimestampNanosecond( + Some(1_704_067_200_000_000_000), + Some("+00:00".into()), + ), + ScalarValue::Utf8(Some("ab".into())), + ScalarValue::LargeUtf8(Some("ab".into())), + ScalarValue::Utf8View(Some("ab".into())), + ScalarValue::Binary(Some(vec![1, 2])), + ScalarValue::LargeBinary(Some(vec![1, 2])), + ScalarValue::BinaryView(Some(vec![1, 2])), + ScalarValue::Float32(Some(1.5)), + ScalarValue::Float64(Some(1.5)), + ScalarValue::Time64Microsecond(Some(3_600_000_000)), + ScalarValue::Int64(None), + ]; + let converted = values + .into_iter() + .map(|value| { + let column = Arc::new(Column::new("a", 0)); + let filter: Arc = + Arc::new(BinaryExpr::new(column, Operator::Eq, lit(value.clone()))); + format!("{value:?} => {}", convert_filter(&filter)) + }) + .collect::>() + .join("\n"); + assert_snapshot!(converted, @r#" + Boolean(true) => a = true + Int32(1) => a = 1 + Int64(1) => a = 1 + Decimal128(123.45,10,2) => a = 123.45 + Decimal128(-123.45,10,2) => a = -123.45 + Date32("2024-01-01") => a = 2024-01-01 + TimestampMicrosecond(1704067200000000, None) => a = 2024-01-01 00:00:00 + TimestampMicrosecond(1704067200000000, Some("+00:00")) => a = 2024-01-01 00:00:00 UTC + TimestampNanosecond(1704067200000000000, None) => a = 2024-01-01 00:00:00 + TimestampNanosecond(1704067200000000000, Some("+00:00")) => a = 2024-01-01 00:00:00 UTC + Utf8("ab") => a = "ab" + LargeUtf8("ab") => a = "ab" + Utf8View("ab") => a = "ab" + Binary("1,2") => a = 0102 + LargeBinary("1,2") => a = 0102 + BinaryView("1,2") => a = 0102 + Float32(1.5) => None + Float64(1.5) => None + Time64Microsecond("3600000000") => None + Int64(NULL) => None + "#); + } + + fn schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + Field::new("s", DataType::Utf8, true), + Field::new("f", DataType::Float64, true), + Field::new("flag", DataType::Boolean, true), + ]) + } + + /// Turns SQL like `a < 5` into a DataFusion filter over [`schema`], then converts it. + fn convert_sql(filter: &str) -> String { + let schema = DFSchema::try_from(schema()).unwrap(); + let ctx = SessionContext::new(); + let expr = ctx.parse_sql_expr(filter, &schema).unwrap(); + let expr = ctx.state().create_physical_expr(expr, &schema).unwrap(); + convert_filter(&expr) + } + + /// Converts the filter and prints the predicate, or `None`. + fn convert_filter(filter: &Arc) -> String { + match convert_physical_filter_to_predicate(filter) { + Some(predicate) => predicate.to_string(), + None => "None".to_string(), + } + } + + /// Builds the filter a hash join on `a` and `b` makes: `struct(a, b) IN (...)`. + /// Each pair takes the values at the same position in `a` and `b`. + fn struct_in( + a: Vec>, + b: Vec>, + ) -> Result> { + struct_in_list(a, b, false) + } + + /// Like [`struct_in`], but `NOT IN`. Hash joins never build this. + fn struct_not_in( + a: Vec>, + b: Vec>, + ) -> Result> { + struct_in_list(a, b, true) + } + + fn struct_in_list( + a: Vec>, + b: Vec>, + negated: bool, + ) -> Result> { + let schema = schema(); + let fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + ]); + let columns = ScalarFunctionExpr::new( + "struct", + r#struct(), + vec![col("a", &schema)?, col("b", &schema)?], + Arc::new(Field::new("struct", DataType::Struct(fields.clone()), true)), + Arc::new(ConfigOptions::default()), + ); + let values = StructArray::new( + fields, + vec![ + Arc::new(Int64Array::from(a)) as ArrayRef, + Arc::new(Int64Array::from(b)) as ArrayRef, + ], + None, + ); + Ok(Arc::new(InListExpr::try_new_from_array( + Arc::new(columns), + Arc::new(values), + negated, + &schema, + )?)) + } +} diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index a3259fd..1f9c497 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -22,11 +22,13 @@ use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::common::Statistics; use datafusion::common::tree_node::TreeNodeRecursion; +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_plan::filter_pushdown::{FilterPushdownPropagation, PushedDown}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayFormatType, Partitioning}; use datafusion::prelude::Expr; @@ -34,6 +36,7 @@ use futures::{Stream, TryStreamExt}; use iceberg::expr::Predicate; use iceberg::table::Table; +use super::physical_predicate::convert_physical_filter_to_predicate; use super::predicate::convert_filters_to_predicate; use crate::to_datafusion_error; @@ -51,6 +54,9 @@ pub struct IcebergDataSource { projection: Vec, /// Filter to apply to the table scan predicate: Option, + /// Physical filters pushed down by DataFusion. The scan ANDs the parts that convert + /// to an Iceberg predicate onto `predicate`. + filters: Vec>, /// Optional limit on the number of rows to return fetch: Option, } @@ -90,6 +96,11 @@ impl IcebergDataSource { self.predicate.as_ref() } + /// The physical filters DataFusion pushed down into this source. + pub fn filters(&self) -> &[Arc] { + &self.filters + } + pub fn fetch(&self) -> Option { self.fetch } @@ -166,6 +177,7 @@ impl IcebergDataSourceBuilder { schema, projection, predicate: self.predicate, + filters: Vec::new(), fetch: self.fetch, }) } @@ -184,11 +196,21 @@ impl DataSource for IcebergDataSource { _partition: usize, context: Arc, ) -> Result { + let predicate = self + .predicate + .clone() + .into_iter() + .chain( + self.filters + .iter() + .filter_map(convert_physical_filter_to_predicate), + ) + .reduce(Predicate::and); let fut = get_batch_stream( self.table.clone(), self.snapshot_id, Some(self.projection.clone()), - self.predicate.clone(), + predicate, context .session_config() .options() @@ -269,6 +291,24 @@ impl DataSource for IcebergDataSource { ) -> Result>> { Ok(None) } + + fn try_pushdown_filters( + &self, + filters: Vec>, + _config: &ConfigOptions, + ) -> Result>> { + // The scan can match extra rows, so DataFusion still applies every filter. + let result = FilterPushdownPropagation::with_parent_pushdown_result(vec![ + PushedDown::No; + filters.len() + ]); + if filters.is_empty() { + return Ok(result); + } + let mut source = self.clone(); + source.filters.extend(filters); + Ok(result.with_updated_node(Arc::new(source))) + } } /// Asynchronously retrieves a stream of [`RecordBatch`] instances diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 42d350a..ceb5798 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -31,10 +31,13 @@ use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::datasource::TableProvider; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::{SessionConfig, SessionContext}; +use datafusion::logical_expr::Operator; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; +use datafusion::physical_expr::expressions::{binary, col, lit}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::common::collect; +use datafusion::physical_plan::filter_pushdown::PushedDown; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ @@ -1205,6 +1208,55 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { Ok(()) } +/// A scan keeps the filters DataFusion pushes into it and reads only the rows that +/// match them. It reports each filter as not pushed down, so DataFusion still +/// applies it. +#[tokio::test] +async fn test_scan_applies_pushed_filters() -> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_pushed_filters".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + let creation = get_table_creation(temp_path(), "my_table", None)?; + iceberg_catalog.create_table(&namespace, creation).await?; + + let ctx = SessionContext::new(); + let catalog = IcebergCatalogProvider::try_new(Arc::new(iceberg_catalog)).await?; + ctx.register_catalog("catalog", Arc::new(catalog)); + ctx.sql("INSERT INTO catalog.test_pushed_filters.my_table VALUES (1, 'alan'), (2, 'turing')") + .await? + .collect() + .await?; + + let provider = ctx + .table_provider("catalog.test_pushed_filters.my_table") + .await?; + let plan = provider.scan(&ctx.state(), None, &[], None).await?; + let scan = plan + .downcast_ref::() + .expect("a data source") + .data_source(); + let schema = provider.schema(); + let filter = binary(col("foo1", &schema)?, Operator::Eq, lit(1), &schema)?; + let pushed = scan.try_pushdown_filters(vec![filter], ctx.state().config_options())?; + assert!(matches!(pushed.filters[..], [PushedDown::No])); + + let source = pushed.updated_node.expect("a scan holding the filter"); + let filters = source + .downcast_ref::() + .expect("an Iceberg data source") + .filters(); + assert_eq!(filters.len(), 1); + assert_eq!(filters[0].to_string(), "foo1@0 = 1"); + insta::assert_snapshot!(&run(&DataSourceExec::new(source), &ctx).await?, @r#" + +------+------+ + | foo1 | foo2 | + +------+------+ + | 1 | alan | + +------+------+"#); + + Ok(()) +} + /// A catalog-backed provider keeps the schema it was built with, but scans the /// table as it is now. After a column is added to the table and written to, a /// scan with no projection still returns only the provider's columns, matching