diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index e047db39a5740..2da9c5e6d40f2 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -1407,9 +1407,9 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result Ok(PrimitiveArray::new(rows_number.into(), None)) } -/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array. -/// For non-float arrays returns the input unchanged. NaN payloads are -/// preserved. +/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array, +/// including dictionary-wrapped floats. For other arrays, returns the input +/// unchanged. NaN payloads are preserved. /// /// Arrow's comparison kernels (`arrow::compute::kernels::cmp::eq` etc.) and /// row-encoding (`arrow::row::RowConverter`) use IEEE 754 totalOrder @@ -1430,6 +1430,17 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { const NEG_ZERO_F32_BITS: u32 = (-0.0_f32).to_bits(); const NEG_ZERO_F64_BITS: u64 = (-0.0_f64).to_bits(); match array.data_type() { + DataType::Dictionary(_, value_type) + if is_float_or_dictionary_float(value_type) => + { + let dictionary = array.as_any_dictionary(); + let values = normalize_float_zero(dictionary.values()); + if Arc::ptr_eq(&values, dictionary.values()) { + Arc::clone(array) + } else { + dictionary.with_values(values) + } + } DataType::Float32 => { let arr: &Float32Array = array.as_primitive::(); if !arr @@ -1439,8 +1450,13 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { { return Arc::clone(array); } - let normalized: Float32Array = - arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f32 } else { v }); + let normalized: Float32Array = arr.unary(|v| { + if v.to_bits() == NEG_ZERO_F32_BITS { + 0.0_f32 + } else { + v + } + }); Arc::new(normalized) } DataType::Float64 => { @@ -1452,8 +1468,13 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { { return Arc::clone(array); } - let normalized: Float64Array = - arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f64 } else { v }); + let normalized: Float64Array = arr.unary(|v| { + if v.to_bits() == NEG_ZERO_F64_BITS { + 0.0_f64 + } else { + v + } + }); Arc::new(normalized) } DataType::Float16 => { @@ -1466,8 +1487,8 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { return Arc::clone(array); } let normalized: Float16Array = arr.unary(|v| { - if v.to_bits() << 1 == 0 { - half::f16::from_bits(0) + if v.to_bits() == NEG_ZERO_F16_BITS { + half::f16::ZERO } else { v } @@ -1478,22 +1499,38 @@ pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { } } +fn is_float_or_dictionary_float(mut data_type: &DataType) -> bool { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type.is_floating() +} + /// Replace `-0.0` with `+0.0` in `Float16`, `Float32`, or `Float64` scalar -/// values. Other variants are returned unchanged. See [`normalize_float_zero`] -/// for context. -pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue { - match scalar { - ScalarValue::Float32(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float32(Some(0.0)) +/// values, including dictionary-wrapped floats. Other variants are returned +/// unchanged. See [`normalize_float_zero`] for context. +pub fn normalize_float_zero_scalar(mut scalar: ScalarValue) -> ScalarValue { + let mut value = &mut scalar; + while let ScalarValue::Dictionary(_, dictionary_value) = value { + value = dictionary_value.as_mut(); + } + + match value { + ScalarValue::Float32(Some(value)) if value.to_bits() == (-0.0_f32).to_bits() => { + *value = 0.0 } - ScalarValue::Float64(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float64(Some(0.0)) + ScalarValue::Float64(Some(value)) if value.to_bits() == (-0.0_f64).to_bits() => { + *value = 0.0 } - ScalarValue::Float16(Some(v)) if v.to_bits() << 1 == 0 => { - ScalarValue::Float16(Some(half::f16::from_bits(0))) + ScalarValue::Float16(Some(value)) + if value.to_bits() == half::f16::NEG_ZERO.to_bits() => + { + *value = half::f16::ZERO; } - other => other, + _ => {} } + + scalar } #[cfg(test)] @@ -1503,9 +1540,9 @@ mod tests { use super::*; use crate::ScalarValue::Null; use arrow::{ - array::{Float64Array, Int32Array}, + array::{DictionaryArray, Float64Array, Int8Array, Int32Array}, buffer::NullBuffer, - datatypes::Int32Type, + datatypes::{Float64Type, Int8Type, Int32Type}, }; #[cfg(feature = "sql")] use sqlparser::ast::Ident; @@ -1534,6 +1571,49 @@ mod tests { } } + #[test] + fn normalize_float_zero_in_dictionary_arrays_and_scalars() -> Result<()> { + let nan = f64::from_bits(0x7ff8_0000_0000_0001); + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let array: ArrayRef = Arc::new(DictionaryArray::try_new( + keys.clone(), + Arc::new(Float64Array::from(vec![-0.0, nan, 1.0])), + )?); + + let normalized = normalize_float_zero(&array); + assert!(!Arc::ptr_eq(&normalized, &array)); + let dictionary = normalized.as_dictionary::(); + assert_eq!(dictionary.keys(), &keys); + let values = dictionary.values().as_primitive::(); + assert_eq!(values.value(0).to_bits(), 0.0_f64.to_bits()); + assert_eq!(values.value(1).to_bits(), nan.to_bits()); + assert_eq!(values.value(2), 1.0); + + let without_negative_zero: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1]), + Arc::new(Float64Array::from(vec![0.0, nan])), + )?); + assert!(Arc::ptr_eq( + &normalize_float_zero(&without_negative_zero), + &without_negative_zero + )); + + let scalar = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Float64(Some(-0.0))), + ); + let ScalarValue::Dictionary(_, value) = normalize_float_zero_scalar(scalar) + else { + unreachable!() + }; + let ScalarValue::Float64(Some(value)) = *value else { + unreachable!() + }; + assert_eq!(value.to_bits(), 0.0_f64.to_bits()); + + Ok(()) + } + #[test] fn test_bisect_linear_left_and_right() -> Result<()> { let arrays: Vec = vec![ diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..b76af09183618 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -245,6 +245,10 @@ name = "math_query_sql" harness = false name = "filter_query_sql" +[[bench]] +harness = false +name = "in_list_rewrite" + [[bench]] harness = false name = "struct_query_sql" diff --git a/datafusion/core/benches/in_list_rewrite.rs b/datafusion/core/benches/in_list_rewrite.rs new file mode 100644 index 0000000000000..bd20f5e6cc9ea --- /dev/null +++ b/datafusion/core/benches/in_list_rewrite.rs @@ -0,0 +1,567 @@ +// 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. + +//! Compares short `IN` lists with equivalent `OR`/`AND` chains. The explicit +//! chain SQL uses the same left-deep shape produced when the logical short-list +//! rewrite is selected; planning is deliberately unoptimized so both candidate +//! forms remain available for comparison. +//! +//! Both forms start as SQL and are planned once. Criterion measures only +//! physical-expression evaluation so table scans, scheduling, and planning do +//! not hide the predicate cost this benchmark is intended to compare. +//! [`BinaryExpr`] evaluates its vectorized children eagerly, so a first-value +//! hit does not provide scalar-style short-circuiting. + +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ + ArrayRef, AsArray, FixedSizeBinaryArray, Float64Array, Int32Array, StringArray, + StringViewArray, +}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion::datasource::MemTable; +use datafusion::prelude::SessionContext; +use datafusion_common::ScalarValue; +use datafusion_expr::{Expr, LogicalPlan, Operator}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, InListExpr}; +use rand::prelude::*; +use tokio::runtime::Runtime; + +const ALL_LIST_LENGTHS: &[usize] = &[1, 2, 3, 4]; +const NULL_LIST_LENGTHS: &[usize] = &[2, 3]; +const MISS_VALUE_BASE: usize = 10_000; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ValueKind { + Int32, + Float64, + Utf8, + Utf8View, + FixedSizeBinary8, +} + +impl ValueKind { + const ALL: [Self; 5] = [ + Self::Int32, + Self::Float64, + Self::Utf8, + Self::Utf8View, + Self::FixedSizeBinary8, + ]; + + fn name(self) -> &'static str { + match self { + Self::Int32 => "int32", + Self::Float64 => "float64", + Self::Utf8 => "utf8", + Self::Utf8View => "utf8view_inline", + Self::FixedSizeBinary8 => "fixed_size_binary_8", + } + } + + fn table_name(self) -> &'static str { + match self { + Self::Int32 => "in_list_i32", + Self::Float64 => "in_list_f64", + Self::Utf8 => "in_list_utf8", + Self::Utf8View => "in_list_utf8view", + Self::FixedSizeBinary8 => "in_list_fsb8", + } + } + + fn data_type(self) -> DataType { + match self { + Self::Int32 => DataType::Int32, + Self::Float64 => DataType::Float64, + Self::Utf8 => DataType::Utf8, + Self::Utf8View => DataType::Utf8View, + Self::FixedSizeBinary8 => DataType::FixedSizeBinary(8), + } + } + + fn seed_tag(self) -> u64 { + match self { + Self::Int32 => 1, + Self::Float64 => 2, + Self::Utf8 => 3, + Self::Utf8View => 4, + Self::FixedSizeBinary8 => 5, + } + } + + fn sql_literal(self, value: usize) -> String { + match self { + Self::Int32 => value.to_string(), + Self::Float64 => format!("{value}.0"), + Self::Utf8 | Self::Utf8View => format!("'{}'", string_value(value)), + Self::FixedSizeBinary8 => format!("${}", value + 1), + } + } + + // The cast only checks alignment; the pointer is never dereferenced. + #[expect(clippy::cast_ptr_alignment)] + fn make_array(self, values: &[Option]) -> ArrayRef { + match self { + Self::Int32 => Arc::new(Int32Array::from_iter( + values + .iter() + .copied() + .map(|value| value.map(|value| value as i32)), + )), + Self::Float64 => Arc::new(Float64Array::from_iter( + values + .iter() + .copied() + .map(|value| value.map(|value| value as f64)), + )), + Self::Utf8 => Arc::new(StringArray::from_iter( + values.iter().copied().map(|value| value.map(string_value)), + )), + Self::Utf8View => Arc::new(StringViewArray::from_iter( + values.iter().copied().map(|value| value.map(string_value)), + )), + Self::FixedSizeBinary8 => { + let array = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values + .iter() + .copied() + .map(|value| value.map(|value| (value as u64).to_be_bytes())), + 8, + ) + .unwrap(); + assert!( + array.values().as_ptr().cast::().is_aligned(), + "builder-created FixedSizeBinary(8) buffer is unexpectedly unaligned" + ); + Arc::new(array) + } + } + } + + fn param_values( + self, + list_len: usize, + list_has_null: bool, + ) -> Option> { + (self == Self::FixedSizeBinary8).then(|| { + let value_count = list_len - usize::from(list_has_null); + (0..value_count) + .map(|value| { + ScalarValue::FixedSizeBinary( + 8, + Some((value as u64).to_be_bytes().to_vec()), + ) + }) + .collect() + }) + } +} + +/// All generated strings are eight bytes and therefore use Arrow's inline +/// byte-view representation, which is the specialized path under test. +fn string_value(value: usize) -> String { + let value = format!("v{value:07}"); + assert_eq!(value.len(), 8); + value +} + +#[derive(Debug, Clone, Copy)] +struct Profile { + name: &'static str, + batch_size: usize, + null_percent: usize, + match_percent: usize, + first_value_hits: bool, + list_has_null: bool, +} + +const PROFILES: [Profile; 8] = [ + Profile { + name: "miss", + batch_size: 8192, + null_percent: 0, + match_percent: 0, + first_value_hits: false, + list_has_null: false, + }, + Profile { + name: "balanced", + batch_size: 8192, + null_percent: 0, + match_percent: 50, + first_value_hits: false, + list_has_null: false, + }, + Profile { + name: "skewed_hit", + batch_size: 8192, + null_percent: 0, + match_percent: 90, + first_value_hits: true, + list_has_null: false, + }, + Profile { + name: "balanced_nullable", + batch_size: 8192, + null_percent: 20, + match_percent: 50, + first_value_hits: false, + list_has_null: false, + }, + Profile { + name: "balanced_small_batch", + batch_size: 64, + null_percent: 0, + match_percent: 50, + first_value_hits: false, + list_has_null: false, + }, + Profile { + name: "single_row_miss", + batch_size: 1, + null_percent: 0, + match_percent: 0, + first_value_hits: false, + list_has_null: false, + }, + Profile { + name: "single_row_hit", + batch_size: 1, + null_percent: 0, + match_percent: 100, + first_value_hits: true, + list_has_null: false, + }, + Profile { + name: "list_null_balanced", + batch_size: 8192, + null_percent: 0, + match_percent: 50, + first_value_hits: false, + list_has_null: true, + }, +]; + +struct PlannedPair { + list_len: usize, + negated: bool, + list_has_null: bool, + chain: Arc, + in_list: Arc, +} + +fn register_table(ctx: &SessionContext, kind: ValueKind) { + // A real row prevents an optimizer from replacing the scan with an empty + // relation. Benchmark batches are evaluated directly and are not stored in + // this table. + let schema = Arc::new(Schema::new(vec![Field::new("a", kind.data_type(), true)])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![kind.make_array(&[Some(0)])]) + .unwrap(); + let table = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_table(kind.table_name(), Arc::new(table)) + .unwrap(); +} + +fn list_literals(kind: ValueKind, list_len: usize, list_has_null: bool) -> Vec { + let value_count = list_len - usize::from(list_has_null); + let mut literals = (0..value_count) + .map(|value| kind.sql_literal(value)) + .collect::>(); + if list_has_null { + literals.push("NULL".to_owned()); + } + literals +} + +fn in_list_sql( + kind: ValueKind, + list_len: usize, + negated: bool, + list_has_null: bool, +) -> String { + let literals = list_literals(kind, list_len, list_has_null).join(", "); + let not = if negated { "NOT " } else { "" }; + format!( + "SELECT * FROM {} WHERE a {not}IN ({literals})", + kind.table_name() + ) +} + +fn chain_sql( + kind: ValueKind, + list_len: usize, + negated: bool, + list_has_null: bool, +) -> String { + let comparison = if negated { "<>" } else { "=" }; + let conjunction = if negated { "AND" } else { "OR" }; + let mut literals = list_literals(kind, list_len, list_has_null).into_iter(); + let first = literals.next().expect("an IN list is never empty"); + let mut predicate = format!("a {comparison} {first}"); + for literal in literals { + predicate = format!("({predicate}) {conjunction} (a {comparison} {literal})"); + } + format!("SELECT * FROM {} WHERE {predicate}", kind.table_name()) +} + +fn find_filter(plan: &LogicalPlan) -> Option<(&Expr, &LogicalPlan)> { + if let LogicalPlan::Filter(filter) = plan { + return Some((&filter.predicate, filter.input.as_ref())); + } + + for input in plan.inputs() { + if let Some(filter) = find_filter(input) { + return Some(filter); + } + } + None +} + +/// Plan a SQL filter predicate through SQL parsing and analysis before +/// converting it to a physical expression. +/// +/// Logical optimization is deliberately skipped: the one-row planning table +/// has exact statistics that can simplify a comparison chain for those stored +/// values. The benchmark evaluates different batches, so that table-specific +/// simplification would not be a valid counterfactual. +fn plan_filter_expr( + ctx: &SessionContext, + runtime: &Runtime, + sql: &str, + param_values: Option<&[ScalarValue]>, +) -> Arc { + let dataframe = runtime + .block_on(ctx.sql(sql)) + .unwrap_or_else(|error| panic!("failed to plan SQL `{sql}`: {error}")); + let dataframe = if let Some(param_values) = param_values { + dataframe + .with_param_values(param_values.to_vec()) + .unwrap_or_else(|error| panic!("failed to bind SQL `{sql}`: {error}")) + } else { + dataframe + }; + let plan = dataframe.into_unoptimized_plan(); + let (predicate, input) = find_filter(&plan) + .unwrap_or_else(|| panic!("planned SQL contains no Filter: `{sql}`\n{plan}")); + ctx.create_physical_expr(predicate.clone(), input.schema().as_ref()) + .unwrap_or_else(|error| { + panic!("failed to create physical predicate for `{sql}`: {error}") + }) +} + +fn assert_static_in_list(expr: &Arc, list_len: usize, negated: bool) { + let in_list = expr + .downcast_ref::() + .unwrap_or_else(|| panic!("expected retained InListExpr, got `{expr}`")); + assert_eq!(in_list.len(), list_len); + assert_eq!(in_list.negated(), negated); + assert!( + expr.to_string().contains("IN (SET)"), + "literal list did not produce a static filter: `{expr}`" + ); +} + +fn assert_comparison(expr: &dyn PhysicalExpr, comparison: Operator) { + let binary = expr + .downcast_ref::() + .unwrap_or_else(|| panic!("expected comparison BinaryExpr, got `{expr}`")); + assert_eq!(binary.op(), &comparison, "unexpected leaf `{expr}`"); +} + +fn assert_left_deep_chain(expr: &dyn PhysicalExpr, list_len: usize, negated: bool) { + let comparison = if negated { + Operator::NotEq + } else { + Operator::Eq + }; + if list_len == 1 { + assert_comparison(expr, comparison); + return; + } + + let conjunction = if negated { Operator::And } else { Operator::Or }; + let binary = expr + .downcast_ref::() + .unwrap_or_else(|| panic!("expected logical BinaryExpr, got `{expr}`")); + assert_eq!(binary.op(), &conjunction, "unexpected chain `{expr}`"); + assert_left_deep_chain(binary.left().as_ref(), list_len - 1, negated); + assert_comparison(binary.right().as_ref(), comparison); +} + +fn plan_pairs( + ctx: &SessionContext, + runtime: &Runtime, + kind: ValueKind, +) -> Vec { + let mut pairs = Vec::with_capacity(12); + for list_has_null in [false, true] { + let list_lengths = if list_has_null { + NULL_LIST_LENGTHS + } else { + ALL_LIST_LENGTHS + }; + for &list_len in list_lengths { + for negated in [false, true] { + let in_sql = in_list_sql(kind, list_len, negated, list_has_null); + let param_values = kind.param_values(list_len, list_has_null); + let in_list = + plan_filter_expr(ctx, runtime, &in_sql, param_values.as_deref()); + assert_static_in_list(&in_list, list_len, negated); + + // SQL-plan an explicit left-deep comparison-chain baseline. + // Optimizing it would merge it back into an InList or use + // planning-table statistics to remove comparisons. + let chain = plan_filter_expr( + ctx, + runtime, + &chain_sql(kind, list_len, negated, list_has_null), + param_values.as_deref(), + ); + assert_left_deep_chain(chain.as_ref(), list_len, negated); + + pairs.push(PlannedPair { + list_len, + negated, + list_has_null, + chain, + in_list, + }); + } + } + } + pairs +} + +fn make_batch( + kind: ValueKind, + profile: Profile, + profile_index: usize, + list_len: usize, + list_has_null: bool, +) -> RecordBatch { + let null_count = profile.batch_size * profile.null_percent / 100; + let non_null_count = profile.batch_size - null_count; + let match_count = non_null_count * profile.match_percent / 100; + let miss_count = non_null_count - match_count; + + let mut values = Vec::with_capacity(profile.batch_size); + values.extend(std::iter::repeat_n(None, null_count)); + values.extend((0..match_count).map(|index| { + Some(if profile.first_value_hits { + 0 + } else { + index % (list_len - usize::from(list_has_null)) + }) + })); + values.extend((0..miss_count).map(|index| Some(MISS_VALUE_BASE + index))); + + let seed = 0x1A11_1575_5EED_u64 + ^ (kind.seed_tag() << 48) + ^ ((profile_index as u64) << 32) + ^ list_len as u64; + values.shuffle(&mut StdRng::seed_from_u64(seed)); + + let schema = Arc::new(Schema::new(vec![Field::new("a", kind.data_type(), true)])); + RecordBatch::try_new(schema, vec![kind.make_array(&values)]).unwrap() +} + +fn assert_same_output(pair: &PlannedPair, batch: &RecordBatch) { + let chain = pair + .chain + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let in_list = pair + .in_list + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + assert_eq!( + chain.as_boolean(), + in_list.as_boolean(), + "chain and InList disagree for list={}, negated={}, list_has_null={}", + pair.list_len, + pair.negated, + pair.list_has_null + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let runtime = Runtime::new().unwrap(); + let ctx = SessionContext::new(); + for kind in ValueKind::ALL { + register_table(&ctx, kind); + } + + for kind in ValueKind::ALL { + let pairs = plan_pairs(&ctx, &runtime, kind); + for (profile_index, profile) in PROFILES.into_iter().enumerate() { + let mut group = c.benchmark_group(format!( + "in_list_rewrite/{}/{}/batch={}/nulls={}%/match={}%", + kind.name(), + profile.name, + profile.batch_size, + profile.null_percent, + profile.match_percent + )); + group.throughput(Throughput::Elements(profile.batch_size as u64)); + + for pair in pairs + .iter() + .filter(|pair| pair.list_has_null == profile.list_has_null) + { + let batch = make_batch( + kind, + profile, + profile_index, + pair.list_len, + pair.list_has_null, + ); + assert_same_output(pair, &batch); + + let operation = if pair.negated { "not_in" } else { "in" }; + let case = format!("{operation}/list={}", pair.list_len); + group.bench_function(BenchmarkId::new(&case, "chain"), |b| { + b.iter(|| black_box(pair.chain.evaluate(black_box(&batch)).unwrap())) + }); + group.bench_function(BenchmarkId::new(&case, "in_list"), |b| { + b.iter(|| { + black_box(pair.in_list.evaluate(black_box(&batch)).unwrap()) + }) + }); + } + group.finish(); + } + } +} + +criterion_group! { + name = benches; + config = Criterion::default() + .warm_up_time(Duration::from_millis(100)) + .measurement_time(Duration::from_millis(500)); + targets = criterion_benchmark +} +criterion_main!(benches); diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index 2ca04ad6b3fac..7e26fb94c1292 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/datafusion/core/tests/sql/explain_analyze.rs @@ -519,7 +519,7 @@ async fn csv_explain_verbose() { async fn csv_explain_inlist_verbose() { let ctx = SessionContext::new(); register_aggregate_csv_by_sql(&ctx).await; - // Inlist len <=3 case will be transformed to OR List so we test with len=4 + // Use len=4 because specialized Int8 filters retain lists of lengths 2 and 3. let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4,5)"; let actual = execute(&ctx, sql).await; diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index bb4a04bb7b6fc..6d0b34de4e769 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -117,9 +117,16 @@ pub struct ExprSimplifier { max_simplifier_cycles: u32, } +/// Largest `IN` list considered for expansion into binary comparisons. +/// +/// The final decision also accounts for whether the list can use a specialized +/// physical filter. pub const THRESHOLD_INLINE_INLIST: usize = 3; pub const DEFAULT_MAX_SIMPLIFIER_CYCLES: u32 = 3; +/// Avoid unbounded expression growth when rewriting `IN` lists with preimages. +const MAX_PREIMAGE_IN_LIST_VALUES: usize = 3; + impl ExprSimplifier { /// Create a new `ExprSimplifier` with the given [`SimplifyContext`]. /// See [`simplify`](Self::simplify) for an example. @@ -199,7 +206,7 @@ impl ExprSimplifier { let mut simplifier = Simplifier::new(&self.info); let config_options = Some(Arc::clone(self.info.config_options())); let mut const_evaluator = ConstEvaluator::try_new(config_options)?; - let mut shorten_in_list_simplifier = ShortenInListSimplifier::new(); + let mut shorten_in_list_simplifier = ShortenInListSimplifier::new(&self.info); let guarantees_map: HashMap<&Expr, &NullableInterval> = self.guarantees.iter().map(|(k, v)| (k, v)).collect(); @@ -2156,7 +2163,7 @@ impl TreeNodeRewriter for Simplifier<'_> { list, negated, }) => { - if list.len() > THRESHOLD_INLINE_INLIST || list.iter().any(is_null) { + if list.len() > MAX_PREIMAGE_IN_LIST_VALUES || list.iter().any(is_null) { return Ok(Transformed::no(Expr::InList(InList { expr, list, @@ -2507,7 +2514,9 @@ mod tests { use super::*; use crate::test::test_table_scan_with_name; use arrow::{ - array::{BooleanArray, Float64Array, Int32Array, StructArray}, + array::{ + BooleanArray, Float64Array, Int32Array, MAX_INLINE_VIEW_LEN, StructArray, + }, datatypes::{FieldRef, Fields}, }; use datafusion_common::{DFSchemaRef, ToDFSchema, assert_contains}; @@ -4709,6 +4718,25 @@ mod tests { (col("c1") * lit(10)).eq(lit(2)) ); + let expr = in_list(col("c3"), vec![lit(1_i64), lit(2_i64)], false); + assert_eq!(simplify(expr.clone()), expr); + + let expr = in_list(col("c3"), vec![lit(1_i64), lit(2_i64), lit(3_i64)], true); + assert_eq!(simplify(expr.clone()), expr); + + // Dynamic lists cannot use a static filter and retain the existing + // comparison expansion even for a specialized primitive type. + assert_eq!( + simplify(in_list( + col("c3"), + vec![lit(1_i64), col("c3_non_null")], + false, + )), + col("c3") + .eq(lit(1_i64)) + .or(col("c3").eq(col("c3_non_null"))) + ); + assert_eq!( simplify(in_list(col("c1"), vec![lit(1), lit(2)], false)), col("c1").eq(lit(1)).or(col("c1").eq(lit(2))) @@ -5107,6 +5135,117 @@ mod tests { } } + #[test] + fn simplify_short_inlist_specialized_filter_boundaries() { + fn simplify_typed(expr: Expr, data_type: DataType) -> Expr { + let schema = Schema::new(vec![Field::new("value", data_type, true)]) + .to_dfschema_ref() + .unwrap(); + ExprSimplifier::new(SimplifyContext::builder().with_schema(schema).build()) + .simplify(expr) + .unwrap() + } + + let expr = in_list(col("value"), vec![lit(0.0_f64), lit(-0.0_f64)], false); + assert_eq!(simplify_typed(expr.clone(), DataType::Float64), expr); + + let expr = in_list( + col("value"), + vec![lit(1_i64), lit(ScalarValue::Int64(None))], + false, + ); + assert_eq!(simplify_typed(expr.clone(), DataType::Int64), expr); + + let expr = in_list( + col("value"), + vec![lit(1_i64), lit(2_i64), lit(ScalarValue::Int64(None))], + true, + ); + assert_eq!(simplify_typed(expr.clone(), DataType::Int64), expr); + + let dictionary_type = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int64), + )), + ); + let dictionary_literal = |value| { + lit(ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Dictionary( + Box::new(DataType::Int16), + Box::new(ScalarValue::Int64(Some(value))), + )), + )) + }; + let expr = in_list( + col("value"), + vec![dictionary_literal(1), dictionary_literal(2)], + false, + ); + assert_eq!(simplify_typed(expr.clone(), dictionary_type), expr); + + let fixed_size_literal = |width, byte| { + lit(ScalarValue::FixedSizeBinary( + width, + Some(vec![byte; width as usize]), + )) + }; + let expr = in_list( + col("value"), + vec![fixed_size_literal(16, 1), fixed_size_literal(16, 2)], + false, + ); + assert_eq!( + simplify_typed(expr.clone(), DataType::FixedSizeBinary(16)), + expr + ); + + let list = vec![fixed_size_literal(3, 1), fixed_size_literal(3, 2)]; + let expr = in_list(col("value"), list.clone(), false); + let expected = col("value") + .eq(list[0].clone()) + .or(col("value").eq(list[1].clone())); + assert_eq!(simplify_typed(expr, DataType::FixedSizeBinary(3)), expected); + + let inline = vec![ + lit(ScalarValue::Utf8View(Some("abcdefghijkl".into()))), + lit(ScalarValue::Utf8View(Some("short".into()))), + ]; + let expr = in_list(col("value"), inline, false); + assert_eq!(simplify_typed(expr.clone(), DataType::Utf8View), expr); + + let mixed = vec![ + lit(ScalarValue::Utf8View(Some("abcdefghijklm".into()))), + lit(ScalarValue::Utf8View(Some("short".into()))), + ]; + let expr = in_list(col("value"), mixed.clone(), false); + let expected = col("value") + .eq(mixed[0].clone()) + .or(col("value").eq(mixed[1].clone())); + assert_eq!(simplify_typed(expr, DataType::Utf8View), expected); + + let binary_view_literal = + |len, byte| lit(ScalarValue::BinaryView(Some(vec![byte; len]))); + let inline = vec![ + binary_view_literal(MAX_INLINE_VIEW_LEN as usize, 1), + binary_view_literal(1, 2), + ]; + let expr = in_list(col("value"), inline, false); + assert_eq!(simplify_typed(expr.clone(), DataType::BinaryView), expr); + + let mixed = vec![ + binary_view_literal(MAX_INLINE_VIEW_LEN as usize + 1, 1), + binary_view_literal(1, 2), + ]; + let expr = in_list(col("value"), mixed.clone(), false); + let expected = col("value") + .eq(mixed[0].clone()) + .or(col("value").eq(mixed[1].clone())); + assert_eq!(simplify_typed(expr, DataType::BinaryView), expected); + } + #[test] fn simplify_inlist_set_operation_propagates_nullability_errors() { let info = SimplifyContext::builder() diff --git a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs index 17112d4f0ae24..04049094f1921 100644 --- a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs @@ -19,24 +19,46 @@ use super::THRESHOLD_INLINE_INLIST; -use datafusion_common::Result; +use arrow::array::MAX_INLINE_VIEW_LEN; +use arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion_common::{Result, ScalarValue}; use datafusion_expr::Expr; use datafusion_expr::expr::InList; +use datafusion_expr::simplify::SimplifyContext; -pub(super) struct ShortenInListSimplifier {} +pub(super) struct ShortenInListSimplifier<'a> { + info: &'a SimplifyContext, +} + +impl<'a> ShortenInListSimplifier<'a> { + pub(super) fn new(info: &'a SimplifyContext) -> Self { + Self { info } + } + + /// Returns true when the physical `IN` expression has a specialized fast + /// path that is preferable to expanding a short static list into ORs. + fn has_specialized_static_filter(&self, expr: &Expr, list: &[Expr]) -> bool { + if !list.iter().all(|expr| matches!(expr, Expr::Literal(_, _))) { + return false; + } -impl ShortenInListSimplifier { - pub(super) fn new() -> Self { - Self {} + // Valid optimizer inputs have already been type-coerced, so the tested + // expression determines the physical list representation. + let Ok(data_type) = self.info.get_data_type(expr) else { + // Type errors are reported elsewhere. Preserve the existing + // shortening behavior instead of making simplification fail. + return false; + }; + supports_specialized_static_filter(&data_type, list) } } -impl TreeNodeRewriter for ShortenInListSimplifier { +impl TreeNodeRewriter for ShortenInListSimplifier<'_> { type Node = Expr; fn f_up(&mut self, expr: Expr) -> Result> { - // if expr is a single column reference: + // Rewrite eligible short lists to left-deep comparison chains: // expr IN (A, B, ...) --> (expr = A) OR (expr = B) OR (expr = C) if let Expr::InList(InList { ref expr, @@ -52,6 +74,7 @@ impl TreeNodeRewriter for ShortenInListSimplifier { list.len() == 1 || list.len() <= THRESHOLD_INLINE_INLIST && expr.try_as_col().is_some() + && !self.has_specialized_static_filter(expr, list) ) { let first_val = list[0].clone(); @@ -93,3 +116,69 @@ impl TreeNodeRewriter for ShortenInListSimplifier { Ok(Transformed::no(expr)) } } + +/// Mirrors the specialized physical filters used by `InListExpr`. +/// +/// Keep this type and representation match synchronized with the primitive, +/// fixed-size-binary, and byte-view selectors under +/// `datafusion/physical-expr/src/expressions/in_list/`. +fn supports_specialized_static_filter(data_type: &DataType, list: &[Expr]) -> bool { + let data_type = dictionary_value_type(data_type); + match data_type { + DataType::Int8 + | DataType::UInt8 + | DataType::Int16 + | DataType::UInt16 + | DataType::Float16 + | DataType::Int32 + | DataType::UInt32 + | DataType::Float32 + | DataType::Date32 + | DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::Date64 + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Decimal128(_, _) + | DataType::Interval(IntervalUnit::MonthDayNano) => true, + DataType::Time32(TimeUnit::Second | TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond | TimeUnit::Nanosecond) => true, + DataType::FixedSizeBinary(width) => matches!(*width, 1 | 2 | 4 | 8 | 16), + DataType::Utf8View | DataType::BinaryView => { + list.iter().all(|expr| inline_view_literal(expr, data_type)) + } + _ => false, + } +} + +fn dictionary_value_type(mut data_type: &DataType) -> &DataType { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type +} + +fn inline_view_literal(expr: &Expr, data_type: &DataType) -> bool { + let Expr::Literal(value, _) = expr else { + return false; + }; + let value = dictionary_scalar_value(value); + if value.is_null() { + return true; + } + + let len = match (data_type, value) { + (DataType::Utf8View, ScalarValue::Utf8View(Some(value))) => value.len(), + (DataType::BinaryView, ScalarValue::BinaryView(Some(value))) => value.len(), + _ => return false, + }; + len <= MAX_INLINE_VIEW_LEN as usize +} + +fn dictionary_scalar_value(mut value: &ScalarValue) -> &ScalarValue { + while let ScalarValue::Dictionary(_, dictionary_value) = value { + value = dictionary_value; + } + value +} diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 1c5a4a1869ddb..1e1563a79dfaa 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -767,7 +767,7 @@ mod tests { assert_optimized_plan_equal!( plan, @ r" - Filter: test.d != Int32(1) AND test.d != Int32(2) AND test.d != Int32(3) + Filter: test.d NOT IN ([Int32(1), Int32(2), Int32(3)]) TableScan: test " ) @@ -784,7 +784,7 @@ mod tests { assert_optimized_plan_equal!( plan, @ r" - Filter: test.d = Int32(1) OR test.d = Int32(2) OR test.d = Int32(3) + Filter: test.d IN ([Int32(1), Int32(2), Int32(3)]) TableScan: test " ) diff --git a/datafusion/optimizer/src/simplify_expressions/udf_preimage.rs b/datafusion/optimizer/src/simplify_expressions/udf_preimage.rs index d888a54d56574..142165d1f18ff 100644 --- a/datafusion/optimizer/src/simplify_expressions/udf_preimage.rs +++ b/datafusion/optimizer/src/simplify_expressions/udf_preimage.rs @@ -169,6 +169,24 @@ mod test { )?), }) } + Expr::Literal(ScalarValue::Int32(Some(700)), _) => { + Ok(PreimageResult::Range { + expr, + interval: Box::new(Interval::try_new( + ScalarValue::Int32(Some(500)), + ScalarValue::Int32(Some(600)), + )?), + }) + } + Expr::Literal(ScalarValue::Int32(Some(800)), _) => { + Ok(PreimageResult::Range { + expr, + interval: Box::new(Interval::try_new( + ScalarValue::Int32(Some(700)), + ScalarValue::Int32(Some(800)), + )?), + }) + } _ => Ok(PreimageResult::None), } } @@ -321,10 +339,13 @@ mod test { #[test] fn test_preimage_in_list_rewrite() { let schema = test_schema(); - let expr = preimage_udf_expr().in_list(vec![lit(500), lit(600)], false); + let expr = preimage_udf_expr().in_list(vec![lit(500), lit(600), lit(700)], false); let expected = or( - and(col("x").gt_eq(lit(100)), col("x").lt(lit(200))), - and(col("x").gt_eq(lit(300)), col("x").lt(lit(400))), + or( + and(col("x").gt_eq(lit(100)), col("x").lt(lit(200))), + and(col("x").gt_eq(lit(300)), col("x").lt(lit(400))), + ), + and(col("x").gt_eq(lit(500)), col("x").lt(lit(600))), ); assert_eq!(optimize_test(expr, &schema), expected); @@ -343,9 +364,10 @@ mod test { } #[test] - fn test_preimage_in_list_long_list_no_rewrite() { + fn test_preimage_in_list_above_limit_no_rewrite() { let schema = test_schema(); - let expr = preimage_udf_expr().in_list((1..100).map(lit).collect(), false); + let expr = preimage_udf_expr() + .in_list(vec![lit(500), lit(600), lit(700), lit(800)], false); assert_eq!(optimize_test(expr.clone(), &schema), expr); } diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 7a7cb25317c1d..6e554e1e33ef4 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -31,6 +31,7 @@ use arrow::compute::kernels::boolean::{not, or_kleene}; use arrow::compute::kernels::cmp::eq as arrow_eq; use arrow::datatypes::*; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; use datafusion_common::{ DFSchema, Result, ScalarValue, assert_or_internal_err, exec_err, }; @@ -82,6 +83,27 @@ fn supports_arrow_eq(dt: &DataType) -> bool { } } +fn normalize_in_list_float_zero_value(value: ColumnarValue) -> ColumnarValue { + match value { + ColumnarValue::Array(array) + if is_float_or_dictionary_float(array.data_type()) => + { + ColumnarValue::Array(normalize_float_zero(&array)) + } + ColumnarValue::Scalar(scalar) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(scalar)) + } + value => value, + } +} + +fn is_float_or_dictionary_float(mut data_type: &DataType) -> bool { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type.is_floating() +} + /// Evaluates the list of expressions into an array, flattening any dictionaries fn evaluate_list( list: &[Arc], @@ -370,12 +392,15 @@ impl PhysicalExpr for InListExpr { // Use Arrow's vectorized eq kernel for types it supports (primitive, // boolean, string, binary, dictionary), falling back to row-by-row // comparator for unsupported types (nested, RunEndEncoded, etc.). - let value = value.into_array(num_rows)?; + // Normalize the left side once for the whole list. Doing this + // outside `compare_one` avoids rescanning it for every item. + let value = + normalize_in_list_float_zero_value(value).into_array(num_rows)?; let lhs_supports_arrow_eq = supports_arrow_eq(value.data_type()); // Helper: compare value against a single list expression let compare_one = |expr: &Arc| -> Result { - match expr.evaluate(batch)? { + match normalize_in_list_float_zero_value(expr.evaluate(batch)?) { ColumnarValue::Array(array) => { if lhs_supports_arrow_eq && supports_arrow_eq(array.data_type()) @@ -3364,6 +3389,111 @@ mod tests { Ok(()) } + #[test] + fn test_in_list_with_columns_float_signed_zero() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + ]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0])), + Arc::new(Float64Array::from(vec![-0.0, 0.0, 2.0])), + ], + )?; + + let expr = make_in_list_with_columns( + col("a", &schema)?, + vec![col("b", &schema)?], + false, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + Ok(()) + } + + #[test] + fn test_in_list_with_columns_float_scalar_signed_zero() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Float32Array::from(vec![0.0, -0.0, 1.0]))], + )?; + let list = vec![lit(ScalarValue::Float32(Some(-0.0)))]; + + for (negated, expected) in [ + (false, BooleanArray::from(vec![true, true, false])), + (true, BooleanArray::from(vec![false, false, true])), + ] { + let expr = + make_in_list_with_columns(col("a", &schema)?, list.clone(), negated); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(as_boolean_array(&result), &expected); + } + + // A scalar left-hand side is normalized before it is broadcast. + let expr = make_in_list_with_columns( + lit(ScalarValue::Float32(Some(-0.0))), + vec![col("a", &schema)?], + false, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + + Ok(()) + } + + #[test] + fn test_in_list_with_columns_dictionary_float_signed_zero() -> Result<()> { + let left: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1, 2]), + Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0])), + )?); + let right: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1, 2]), + Arc::new(Float64Array::from(vec![-0.0, 0.0, 2.0])), + )?); + let data_type = left.data_type().clone(); + let schema = Schema::new(vec![ + Field::new("a", data_type.clone(), false), + Field::new("b", data_type, false), + ]); + let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![left, right])?; + + for (negated, expected) in [ + (false, BooleanArray::from(vec![true, true, false])), + (true, BooleanArray::from(vec![false, false, true])), + ] { + let expr = make_in_list_with_columns( + col("a", &schema)?, + vec![col("b", &schema)?], + negated, + ); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(as_boolean_array(&result), &expected); + } + + let scalar = lit(ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Float64(Some(-0.0))), + )); + let expr = make_in_list_with_columns(col("a", &schema)?, vec![scalar], false); + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, false]) + ); + + Ok(()) + } + /// Tests that short-circuit evaluation produces correct results. /// When all rows match after the first list item, remaining items /// should be skipped without affecting correctness. @@ -3889,6 +4019,29 @@ mod tests { Ok(()) } + #[test] + fn test_try_new_from_array_dict_haystack_float64_signed_zero() -> Result<()> { + // One value beyond the branchless limit selects the hash-set strategy. + let list_len = + ::MAX_LIST_LEN + 1; + let mut list_values = vec![Some(-0.0)]; + list_values.extend((1..list_len).map(|value| Some(value as f64))); + let haystack = make_f64_dict_array(list_values); + let needles: ArrayRef = Arc::new(Float64Array::from(vec![0.0, -0.0, -1.0])); + let expected = BooleanArray::from(vec![true, true, false]); + + assert_eq!( + eval_in_list_from_array(Arc::clone(&needles), Arc::clone(&haystack))?, + expected + ); + assert_eq!( + eval_in_list_from_array(wrap_in_dict(needles), haystack)?, + expected + ); + + Ok(()) + } + #[test] fn test_try_new_from_array_type_mismatch_rejects() -> Result<()> { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); diff --git a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs index 86cea37b3a98d..aa19bc44dadc8 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -29,16 +29,17 @@ //! //! # How does it work? //! -//! When the filter is built, it stores the non-null list values and chooses a -//! comparison function for that list length. Only this small function is -//! specialized for each length. The rest of [`BranchlessFilter`] is shared, -//! which keeps the generated code small. +//! When the filter is built, it stores the values needed by the comparison +//! chain and chooses a function for that chain's length. Only this small +//! function is specialized for each length. The rest of [`BranchlessFilter`] +//! is shared, which keeps the generated code small. //! //! Some Arrow types share the same in-memory representation. For example, a //! `Float32` and a `UInt32` both use four bytes per value. The filter compares //! those stored bits through an unsigned type of the same size, without copying //! the value buffer. A bit pattern is simply the bytes Arrow uses to store a -//! value. Comparing it preserves details such as `0.0` versus `-0.0` and +//! value. Float comparisons treat `0.0` and `-0.0` as equal under SQL +//! semantics, while preserving distinctions between //! different NaN values. [`BranchlessFilterType`] defines these safe, //! same-sized mappings and checks their sizes at compile time. //! @@ -75,6 +76,7 @@ use arrow::buffer::{BooleanBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; +use half::f16; use super::result::build_result_from_contains; use super::static_filter::StaticFilter; @@ -117,16 +119,20 @@ const BRANCHLESS_MAX_16B: usize = 4; /// /// `T` is the logical Arrow type accepted by the filter. `CompareType` is the /// same-width type used for the fixed comparison chain. Signed integers, -/// floats, and temporal values use an unsigned comparison type so they compare -/// by their raw bit pattern. +/// floats, and temporal values use an unsigned comparison type. Most values +/// compare by raw bit pattern; floats additionally treat both signed-zero +/// patterns as equal. pub(super) trait BranchlessFilterType: - ArrowPrimitiveType + Send + Sync + 'static + ArrowPrimitiveType + Send + Sync + Sized + 'static { type CompareType: ArrowPrimitiveType + Send + Sync + 'static; /// Maximum number of non-null IN-list values to handle with /// [`BranchlessFilter`] for this primitive type. const MAX_LIST_LEN: usize; + + /// The two signed-zero encodings for float types. + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = None; } macro_rules! branchless_filter_type { @@ -151,18 +157,33 @@ branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); -branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); +impl BranchlessFilterType for Float16Type { + type CompareType = UInt16Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_2B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([f16::ZERO.to_bits(), f16::NEG_ZERO.to_bits()]); +} branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); +impl BranchlessFilterType for Float32Type { + type CompareType = UInt32Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_4B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([0.0_f32.to_bits(), (-0.0_f32).to_bits()]); +} branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); +impl BranchlessFilterType for Float64Type { + type CompareType = UInt64Type; + const MAX_LIST_LEN: usize = BRANCHLESS_MAX_8B; + const SIGNED_ZERO_BITS: Option<[BranchlessNative; 2]> = + Some([0.0_f64.to_bits(), (-0.0_f64).to_bits()]); +} branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); @@ -189,9 +210,9 @@ type MembershipCheck = fn(in_list_values: &[C], input_values: &[C]) -> Boolea /// `T::MAX_LIST_LEN` values. /// /// The filter stores the non-null `IN`-list values in a slice and chooses a -/// comparison function for that length. Keeping the length out of -/// `BranchlessFilter` avoids generating a full copy of the filter for every -/// supported length. +/// comparison function for that length. If a float zero occurs, both signed +/// encodings are stored once. Keeping the length out of `BranchlessFilter` +/// avoids generating a full copy of the filter for every supported length. pub(super) struct BranchlessFilter { expected_data_type: DataType, null_count: usize, @@ -218,8 +239,10 @@ where } let all_values = branchless_values::(in_array); - let mut in_list_values = Vec::with_capacity(non_null_count); - + // Float zero may add its other signed encoding. + let mut in_list_values = Vec::with_capacity( + non_null_count + usize::from(T::SIGNED_ZERO_BITS.is_some()), + ); match in_array.nulls() { None => { in_list_values.extend(all_values.iter().copied()); @@ -232,8 +255,9 @@ where } } } + materialize_signed_zero_encodings::(&mut in_list_values); - debug_assert_eq!(in_list_values.len(), non_null_count); + debug_assert!(in_list_values.len() <= non_null_count + 1); let in_list_values = in_list_values.into_boxed_slice(); let check_values = membership_check_for_len::(in_list_values.len()); @@ -246,6 +270,27 @@ where } } +fn materialize_signed_zero_encodings(values: &mut Vec>) +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq, +{ + let Some([positive_zero, negative_zero]) = T::SIGNED_ZERO_BITS else { + return; + }; + + // This list has at most 32 entries, so checking both SQL-equivalent + // encodings once is cheap and keeps per-row lookup branch-free. + match ( + values.contains(&positive_zero), + values.contains(&negative_zero), + ) { + (true, false) => values.push(negative_zero), + (false, true) => values.push(positive_zero), + _ => {} + } +} + impl StaticFilter for BranchlessFilter where T: BranchlessFilterType, @@ -299,6 +344,18 @@ where }; } + // A single float zero expands to its two physical encodings. Only these + // three extra lengths can therefore exceed the logical list-size limit. + if T::SIGNED_ZERO_BITS.is_some() && len > T::MAX_LIST_LEN { + debug_assert_eq!(len, T::MAX_LIST_LEN + 1); + return match T::MAX_LIST_LEN { + 8 => check_values::, 9>, + 16 => check_values::, 17>, + 32 => check_values::, 33>, + _ => unreachable!("signed-zero expansion exceeded an unsupported limit"), + }; + } + // Avoid creating checks for lengths a type does not support. match T::MAX_LIST_LEN { 4 => choose!(0, 1, 2, 3, 4), @@ -451,11 +508,11 @@ mod tests { assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![None, Some(true), Some(true), None, None]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), None, None]) ); assert_eq!( filter.contains(&needles, true)?, - BooleanArray::from(vec![None, Some(false), Some(false), None, None]) + BooleanArray::from(vec![Some(false), Some(false), Some(false), None, None]) ); let wrong_type = UInt16Array::from(vec![Some(0x8000), Some(0x7e01)]); @@ -466,20 +523,26 @@ mod tests { } #[test] - fn branchless_filter_floats_use_bit_equality() -> Result<()> { + fn branchless_filter_floats_use_sql_zero_equality() -> Result<()> { let nan_a = f32::from_bits(0x7fc0_0001); let nan_b = f32::from_bits(0x7fc0_0002); let haystack: ArrayRef = - Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + Arc::new(Float32Array::from(vec![Some(0.0), Some(nan_a)])); let filter = BranchlessFilter::::try_new(&haystack)?; let needles = Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]) ); + // A list containing both encodings is not expanded. + let zero_only: ArrayRef = + Arc::new(Float32Array::from(vec![Some(0.0), Some(-0.0)])); + let filter = BranchlessFilter::::try_new(&zero_only)?; + assert_eq!(filter.in_list_values.len(), 2); + let nan_a = f64::from_bits(0x7ff8_0000_0000_0001); let nan_b = f64::from_bits(0x7ff8_0000_0000_0002); let haystack: ArrayRef = @@ -490,7 +553,7 @@ mod tests { assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]) ); Ok(()) diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index ceb9bd525b965..564b96ae36028 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -27,6 +27,7 @@ use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; +use half::f16; use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; use super::result::build_in_list_result; @@ -199,6 +200,12 @@ trait BitmapFilterType: ArrowPrimitiveType + Send + Sync + 'static { /// Returns the index in the bitmap to check for this value. fn index(value: Self::Native) -> usize; + + /// Adds one value to the bitmap. + #[inline] + fn insert(bitmap: &mut Self::Storage, value: Self::Native) { + bitmap.set_bit(Self::index(value)); + } } /// `Int8` has 256 possible bit patterns, so four `u64` words cover the full domain. @@ -254,6 +261,17 @@ impl BitmapFilterType for Float16Type { fn index(value: Self::Native) -> usize { value.to_bits() as usize } + + fn insert(bitmap: &mut Self::Storage, value: Self::Native) { + if value == f16::ZERO { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + bitmap.set_bit(f16::ZERO.to_bits() as usize); + bitmap.set_bit(f16::NEG_ZERO.to_bits() as usize); + } else { + bitmap.set_bit(Self::index(value)); + } + } } /// `IN` filter backed by one bit per possible value. @@ -280,14 +298,14 @@ where match prim_array.nulls() { None => { for &v in values { - bits.set_bit(T::index(v)); + T::insert(&mut bits, v); } } Some(nulls) => { for i in BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) { - bits.set_bit(T::index(values[i])); + T::insert(&mut bits, values[i]); } } } @@ -332,6 +350,16 @@ where } } +/// Converts native values to hash keys and inserts any additional physical +/// encodings that belong to the same SQL equality class. +trait HashSetKey: From + Eq + Hash + Sized { + fn insert(values: &mut HashSet, value: V) { + values.insert(Self::from(value)); + } +} + +impl HashSetKey for T where T: Copy + Eq + Hash {} + /// Wrapper for f32 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -357,6 +385,19 @@ impl From for OrderedFloat32 { } } +impl HashSetKey for OrderedFloat32 { + fn insert(values: &mut HashSet, value: f32) { + if value == 0.0 { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + values.insert(Self(0.0)); + values.insert(Self(-0.0)); + } else { + values.insert(Self(value)); + } + } +} + /// Wrapper for f64 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -382,6 +423,19 @@ impl From for OrderedFloat64 { } } +impl HashSetKey for OrderedFloat64 { + fn insert(values: &mut HashSet, value: f64) { + if value == 0.0 { + // Keep lookup branch-free by materializing both SQL-equivalent + // signed-zero encodings when either appears in the list. + values.insert(Self(0.0)); + values.insert(Self(-0.0)); + } else { + values.insert(Self(value)); + } + } +} + /// Hash-set membership for primitive types. /// /// `K` defaults to the Arrow type's native value. Floats use an ordered wrapper @@ -399,7 +453,7 @@ impl PrimitiveHashSetFilter where T: ArrowPrimitiveType, T::Native: Copy, - K: From + Eq + Hash, + K: HashSetKey, { fn try_new(in_array: &ArrayRef) -> Result { let in_array = in_array.as_primitive_opt::().ok_or_else(|| { @@ -410,7 +464,7 @@ where })?; let mut values = HashSet::with_capacity(in_array.len() - in_array.null_count()); for value in in_array.iter().flatten() { - values.insert(K::from(value)); + K::insert(&mut values, value); } Ok(Self { @@ -425,7 +479,7 @@ impl StaticFilter for PrimitiveHashSetFilter where T: ArrowPrimitiveType + Send + Sync + 'static, T::Native: Copy + Send + Sync, - K: From + Eq + Hash + Send + Sync + 'static, + K: HashSetKey + Send + Sync + 'static, { fn null_count(&self) -> usize { self.null_count @@ -461,7 +515,6 @@ mod tests { DictionaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, UInt8Array, UInt16Array, UInt32Array, }; - use half::f16; use super::super::dictionary_filter::DictionaryFilter; @@ -507,6 +560,45 @@ mod tests { Ok(()) } + #[test] + fn branchless_float_zero_expansion_handles_max_list_len() -> Result<()> { + fn assert_routed_filter(haystack: ArrayRef, needles: &dyn Array) -> Result<()> { + let filter = instantiate_primitive_filter(&haystack)? + .expect("top-level float arrays always use a primitive filter"); + assert_contains( + filter.as_ref(), + needles, + vec![Some(true), Some(true), Some(false)], + ) + } + + // Adding the mirror zero to a full logical list exercises the extra + // generated comparison length for each branchless size class. + let len = ::MAX_LIST_LEN; + let mut values = vec![f16::NEG_ZERO]; + values.extend((1..len).map(|value| f16::from_f32(value as f32))); + let haystack: ArrayRef = Arc::new(Float16Array::from(values)); + let needles = + Float16Array::from(vec![f16::ZERO, f16::NEG_ZERO, f16::from_f32(100.0)]); + assert_routed_filter(haystack, &needles)?; + + let len = ::MAX_LIST_LEN; + let mut values = vec![-0.0_f32]; + values.extend((1..len).map(|value| value as f32)); + let haystack: ArrayRef = Arc::new(Float32Array::from(values)); + let needles = Float32Array::from(vec![0.0, -0.0, 100.0]); + assert_routed_filter(haystack, &needles)?; + + let len = ::MAX_LIST_LEN; + let mut values = vec![-0.0_f64]; + values.extend((1..len).map(|value| value as f64)); + let haystack: ArrayRef = Arc::new(Float64Array::from(values)); + let needles = Float64Array::from(vec![0.0, -0.0, 100.0]); + assert_routed_filter(haystack, &needles)?; + + Ok(()) + } + #[test] fn primitive_hash_filter_handles_float_keys() -> Result<()> { let nan32 = f32::NAN; @@ -524,14 +616,14 @@ mod tests { assert_contains( &filter, &needles, - vec![Some(true), Some(false), Some(true), Some(false), None], + vec![Some(true), Some(true), Some(true), Some(false), None], )?; let nan64 = f64::NAN; - let haystack: ArrayRef = Arc::new(Float64Array::from(vec![1.0, nan64])); + let haystack: ArrayRef = Arc::new(Float64Array::from(vec![-0.0, nan64])); let filter = PrimitiveHashSetFilter::::try_new(&haystack)?; - let needles = Float64Array::from(vec![Some(1.0), Some(nan64), Some(2.0)]); + let needles = Float64Array::from(vec![Some(0.0), Some(nan64), Some(2.0)]); assert_contains(&filter, &needles, vec![Some(true), Some(true), Some(false)]) } @@ -660,21 +752,22 @@ mod tests { ); let filter = BitmapFilter::::try_new(&haystack)?; let needles = Float16Array::from(vec![ + Some(f16::from_f32(9.0)), Some(f16::from_f32(0.0)), Some(f16::from_f32(-0.0)), Some(nan_a), Some(nan_b), None, ]) - .slice(1, 4); + .slice(1, 5); assert_eq!( filter.contains(&needles, false)?, - BooleanArray::from(vec![Some(true), Some(true), None, None]) + BooleanArray::from(vec![Some(true), Some(true), Some(true), None, None]) ); assert_eq!( filter.contains(&needles, true)?, - BooleanArray::from(vec![Some(false), Some(false), None, None]) + BooleanArray::from(vec![Some(false), Some(false), Some(false), None, None]) ); Ok(()) diff --git a/datafusion/sqllogictest/test_files/array/array_has.slt b/datafusion/sqllogictest/test_files/array/array_has.slt index d7b6680fab062..6cf6d2fcee896 100644 --- a/datafusion/sqllogictest/test_files/array/array_has.slt +++ b/datafusion/sqllogictest/test_files/array/array_has.slt @@ -491,7 +491,8 @@ select array_has_all(arrow_cast(make_array(1,2,3), 'FixedSizeList(3, Int64)'), a true false true false false false true true false false true false true # rewrite various array_has operations to InList where the haystack is a literal list -# NB that `col in (a, b, c)` is simplified to OR if there are <= 3 elements, so we make 4-element haystack lists +# Short lists on a column are simplified to OR when they cannot use a specialized +# static filter. This list has an out-of-line Utf8View value, so use four items. query I with test AS (SELECT substr(md5(i::text)::text, 1, 32) as needle FROM generate_series(1, 100000) t(i)) diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index 7cb5547383c38..cd3526c718b9d 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -1100,8 +1100,8 @@ logical_plan 04)------Aggregate: groupBy=[[hits.URLHash, hits.EventDate]], aggr=[[count(Int64(1))]] 05)--------SubqueryAlias: hits 06)----------Projection: hits_raw.URLHash, CAST(CAST(hits_raw.EventDate AS Int32) AS Date32) AS EventDate -07)------------Filter: hits_raw.CounterID = Int32(62) AND hits_raw.EventDate >= UInt16(15887) AND hits_raw.EventDate <= UInt16(15917) AND hits_raw.IsRefresh = Int16(0) AND (hits_raw.TraficSourceID = Int16(-1) OR hits_raw.TraficSourceID = Int16(6)) AND hits_raw.RefererHash = Int64(3594120000172545465) -08)--------------TableScan: hits_raw projection=[EventDate, CounterID, IsRefresh, TraficSourceID, RefererHash, URLHash], partial_filters=[hits_raw.CounterID = Int32(62), hits_raw.EventDate >= UInt16(15887), hits_raw.EventDate <= UInt16(15917), hits_raw.IsRefresh = Int16(0), hits_raw.TraficSourceID = Int16(-1) OR hits_raw.TraficSourceID = Int16(6), hits_raw.RefererHash = Int64(3594120000172545465)] +07)------------Filter: hits_raw.CounterID = Int32(62) AND hits_raw.EventDate >= UInt16(15887) AND hits_raw.EventDate <= UInt16(15917) AND hits_raw.IsRefresh = Int16(0) AND hits_raw.TraficSourceID IN ([Int16(-1), Int16(6)]) AND hits_raw.RefererHash = Int64(3594120000172545465) +08)--------------TableScan: hits_raw projection=[EventDate, CounterID, IsRefresh, TraficSourceID, RefererHash, URLHash], partial_filters=[hits_raw.CounterID = Int32(62), hits_raw.EventDate >= UInt16(15887), hits_raw.EventDate <= UInt16(15917), hits_raw.IsRefresh = Int16(0), hits_raw.TraficSourceID IN ([Int16(-1), Int16(6)]), hits_raw.RefererHash = Int64(3594120000172545465)] physical_plan 01)GlobalLimitExec: skip=100, fetch=10 02)--SortPreservingMergeExec: [pageviews@2 DESC], fetch=110 @@ -1111,9 +1111,9 @@ physical_plan 06)----------RepartitionExec: partitioning=Hash([URLHash@0, EventDate@1], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[URLHash@0 as URLHash, EventDate@1 as EventDate], aggr=[count(Int64(1))] 08)--------------ProjectionExec: expr=[URLHash@0 as URLHash, CAST(CAST(EventDate@1 AS Int32) AS Date32) as EventDate] -09)----------------FilterExec: CounterID@1 = 62 AND EventDate@0 >= 15887 AND EventDate@0 <= 15917 AND IsRefresh@2 = 0 AND (TraficSourceID@3 = -1 OR TraficSourceID@3 = 6) AND RefererHash@4 = 3594120000172545465, projection=[URLHash@5, EventDate@0] +09)----------------FilterExec: CounterID@1 = 62 AND EventDate@0 >= 15887 AND EventDate@0 <= 15917 AND IsRefresh@2 = 0 AND TraficSourceID@3 IN (SET) ([-1, 6]) AND RefererHash@4 = 3594120000172545465, projection=[URLHash@5, EventDate@0] 10)------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventDate, CounterID, IsRefresh, TraficSourceID, RefererHash, URLHash], file_type=parquet, predicate=CounterID@6 = 62 AND EventDate@5 >= 15887 AND EventDate@5 <= 15917 AND IsRefresh@15 = 0 AND (TraficSourceID@37 = -1 OR TraficSourceID@37 = 6) AND RefererHash@102 = 3594120000172545465, pruning_predicate=CounterID_null_count@2 != row_count@3 AND CounterID_min@0 <= 62 AND 62 <= CounterID_max@1 AND EventDate_null_count@5 != row_count@3 AND EventDate_max@4 >= 15887 AND EventDate_null_count@5 != row_count@3 AND EventDate_min@6 <= 15917 AND IsRefresh_null_count@9 != row_count@3 AND IsRefresh_min@7 <= 0 AND 0 <= IsRefresh_max@8 AND (TraficSourceID_null_count@12 != row_count@3 AND TraficSourceID_min@10 <= -1 AND -1 <= TraficSourceID_max@11 OR TraficSourceID_null_count@12 != row_count@3 AND TraficSourceID_min@10 <= 6 AND 6 <= TraficSourceID_max@11) AND RefererHash_null_count@15 != row_count@3 AND RefererHash_min@13 <= 3594120000172545465 AND 3594120000172545465 <= RefererHash_max@14, required_guarantees=[CounterID in (62), IsRefresh in (0), RefererHash in (3594120000172545465), TraficSourceID in (-1, 6)] +11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventDate, CounterID, IsRefresh, TraficSourceID, RefererHash, URLHash], file_type=parquet, predicate=CounterID@6 = 62 AND EventDate@5 >= 15887 AND EventDate@5 <= 15917 AND IsRefresh@15 = 0 AND TraficSourceID@37 IN (SET) ([-1, 6]) AND RefererHash@102 = 3594120000172545465, pruning_predicate=CounterID_null_count@2 != row_count@3 AND CounterID_min@0 <= 62 AND 62 <= CounterID_max@1 AND EventDate_null_count@5 != row_count@3 AND EventDate_max@4 >= 15887 AND EventDate_null_count@5 != row_count@3 AND EventDate_min@6 <= 15917 AND IsRefresh_null_count@9 != row_count@3 AND IsRefresh_min@7 <= 0 AND 0 <= IsRefresh_max@8 AND (TraficSourceID_null_count@12 != row_count@3 AND TraficSourceID_min@10 <= -1 AND -1 <= TraficSourceID_max@11 OR TraficSourceID_null_count@12 != row_count@3 AND TraficSourceID_min@10 <= 6 AND 6 <= TraficSourceID_max@11) AND RefererHash_null_count@15 != row_count@3 AND RefererHash_min@13 <= 3594120000172545465 AND 3594120000172545465 <= RefererHash_max@14, required_guarantees=[CounterID in (62), IsRefresh in (0), RefererHash in (3594120000172545465), TraficSourceID in (-1, 6)] query IDI SELECT "URLHash", "EventDate", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "TraficSourceID" IN (-1, 6) AND "RefererHash" = 3594120000172545465 GROUP BY "URLHash", "EventDate" ORDER BY PageViews DESC LIMIT 10 OFFSET 100; diff --git a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt index a17ce11a0bff0..7e4d100b74237 100644 --- a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt +++ b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt @@ -45,9 +45,9 @@ explain select * from t1 left join t2 on t1.a = t2.x where t2.x in (1, 2, 3); ---- logical_plan 01)Inner Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) OR t1.a = Int32(3) +02)--Filter: t1.a IN ([Int32(1), Int32(2), Int32(3)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) OR t2.x = Int32(3) +04)--Filter: t2.x IN ([Int32(1), Int32(2), Int32(3)]) 05)----TableScan: t2 projection=[x, y, z] # Verify result correctness @@ -63,9 +63,9 @@ explain select * from t1 right join t2 on t1.a = t2.x where t1.a in (1, 2); ---- logical_plan 01)Inner Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) +02)--Filter: t1.a IN ([Int32(1), Int32(2)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) +04)--Filter: t2.x IN ([Int32(1), Int32(2)]) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort @@ -80,9 +80,9 @@ explain select * from t1 full join t2 on t1.a = t2.x where t1.a in (1, 2) and t2 ---- logical_plan 01)Inner Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) +02)--Filter: t1.a IN ([Int32(1), Int32(2)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) +04)--Filter: t2.x IN ([Int32(1), Int32(2)]) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort @@ -97,9 +97,9 @@ explain select * from t1 left join t2 on t1.a = t2.x where t2.x in (1, 2, null); ---- logical_plan 01)Inner Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) +02)--Filter: t1.a IN ([Int32(1), Int32(2), Int32(NULL)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) +04)--Filter: t2.x IN ([Int32(1), Int32(2), Int32(NULL)]) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort @@ -168,9 +168,9 @@ explain select * from t1 left join t2 on t1.a = t2.x where t1.a in (1, 2, 3); ---- logical_plan 01)Left Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) OR t1.a = Int32(3) +02)--Filter: t1.a IN ([Int32(1), Int32(2), Int32(3)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) OR t2.x = Int32(3) +04)--Filter: t2.x IN ([Int32(1), Int32(2), Int32(3)]) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort @@ -185,7 +185,7 @@ query TT explain select * from t1 left join t2 on t1.a = t2.x where t2.x in (1, 2) or t2.x is null; ---- logical_plan -01)Filter: t2.x = Int32(1) OR t2.x = Int32(2) OR t2.x IS NULL +01)Filter: t2.x IN ([Int32(1), Int32(2)]) OR t2.x IS NULL 02)--Left Join: t1.a = t2.x 03)----TableScan: t1 projection=[a, b, c] 04)----TableScan: t2 projection=[x, y, z] @@ -459,9 +459,9 @@ explain select * from t1 full join t2 on t1.a = t2.x where t2.x in (1, 2); ---- logical_plan 01)Right Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) +02)--Filter: t1.a IN ([Int32(1), Int32(2)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: t2.x = Int32(1) OR t2.x = Int32(2) +04)--Filter: t2.x IN ([Int32(1), Int32(2)]) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort @@ -512,9 +512,9 @@ explain select * from t1 left join t2 on t1.a = t2.x where t2.x in (1, 2) and t2 ---- logical_plan 01)Inner Join: t1.a = t2.x -02)--Filter: t1.a = Int32(1) OR t1.a = Int32(2) +02)--Filter: t1.a IN ([Int32(1), Int32(2)]) 03)----TableScan: t1 projection=[a, b, c] -04)--Filter: (t2.x = Int32(1) OR t2.x = Int32(2)) AND t2.y >= Int32(50) AND t2.y <= Int32(250) +04)--Filter: t2.x IN ([Int32(1), Int32(2)]) AND t2.y >= Int32(50) AND t2.y <= Int32(250) 05)----TableScan: t2 projection=[x, y, z] query IITIIT rowsort diff --git a/datafusion/sqllogictest/test_files/in_list.slt b/datafusion/sqllogictest/test_files/in_list.slt index 3ff47afb643fa..c0b9dfd45d6dd 100644 --- a/datafusion/sqllogictest/test_files/in_list.slt +++ b/datafusion/sqllogictest/test_files/in_list.slt @@ -20,9 +20,10 @@ # # This file focuses on the IN operator and its various specializations # -# Note that "short" IN LISTS do not go through the InList implementation at all, -# instead they are rewritten into a series of OR expressions. See: -# https://github.com/apache/datafusion/blob/ed37b6c9555bc278130dc774ed833b8c0bd29bfa/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs#L39-L88 +# One-item IN lists are rewritten to a comparison. Two- and three-item lists on +# a column are rewritten to OR/AND expressions unless a specialized static +# InList filter is available. See: +# https://github.com/apache/datafusion/blob/main/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs ########## diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt index 8ea1122880e14..c2c7cbb16e323 100644 --- a/datafusion/sqllogictest/test_files/negative_zero.slt +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -47,6 +47,72 @@ SELECT 0.0 IS DISTINCT FROM -0.0 AS is_distinct; ---- false +##### +## IN / NOT IN predicates +##### + +statement ok +CREATE TABLE negative_zero_in_list AS +SELECT arrow_cast(0.0, 'Float16') AS positive_f16, + arrow_cast(0.0, 'Float32') AS positive_f32, + arrow_cast(0.0, 'Float64') AS positive_f64, + arrow_cast(-0.0, 'Float16') AS negative_f16, + arrow_cast(-0.0, 'Float32') AS negative_f32, + arrow_cast(-0.0, 'Float64') AS negative_f64, + arrow_cast(-0.0, 'Dictionary(Int32, Float64)') AS negative_dict_f64; + +# A three-item list is rewritten to comparisons, while a four-item list remains +# an InList. All paths must implement the same signed-zero equality. +query BBBB +SELECT + positive_f64 IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64')), + positive_f16 IN (arrow_cast(-0.0, 'Float16'), arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'), arrow_cast(3.0, 'Float16')), + positive_f32 IN (arrow_cast(-0.0, 'Float32'), arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float32'), arrow_cast(3.0, 'Float32')), + positive_f64 IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true true true + +# Check the mirror direction for every floating-point width. +query BBB +SELECT + negative_f16 IN (arrow_cast(0.0, 'Float16'), arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'), arrow_cast(3.0, 'Float16')), + negative_f32 IN (arrow_cast(0.0, 'Float32'), arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float32'), arrow_cast(3.0, 'Float32')), + negative_f64 IN (arrow_cast(0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true true + +# NOT IN uses the same membership result before negation. +query B +SELECT + positive_f64 NOT IN (arrow_cast(-0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +false + +# Dictionary encoding must not change the result on either side of the rewrite +# threshold. A one-item list becomes equality; a four-item list stays InList. +query BB +SELECT + negative_dict_f64 IN (arrow_cast(0.0, 'Float64')), + negative_dict_f64 IN (arrow_cast(0.0, 'Float64'), arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true + +# A column-valued list item forces non-static evaluation. Check both directions. +query BB +SELECT + positive_f64 IN (negative_f64, arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')), + negative_f64 IN (positive_f64, arrow_cast(1.0, 'Float64'), arrow_cast(2.0, 'Float64'), arrow_cast(3.0, 'Float64')) +FROM negative_zero_in_list; +---- +true true + +statement ok +DROP TABLE negative_zero_in_list; + ##### ## SELECT DISTINCT with +0.0 / -0.0 (Float64) ##### diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index 6f3865c089458..afce47fe578a7 100644 --- a/datafusion/sqllogictest/test_files/predicates.slt +++ b/datafusion/sqllogictest/test_files/predicates.slt @@ -242,7 +242,7 @@ query TT EXPLAIN SELECT * FROM test_regex_utf8view WHERE s ~ '^(foo|Bazzz)$' ---- logical_plan -01)Filter: test_regex_utf8view.s = Utf8View("foo") OR test_regex_utf8view.s = Utf8View("Bazzz") +01)Filter: test_regex_utf8view.s IN ([Utf8View("foo"), Utf8View("Bazzz")]) 02)--TableScan: test_regex_utf8view projection=[s] # `~*` anchored alternation -> NOT simplified: it falls back to a regex match, @@ -275,7 +275,7 @@ query TT EXPLAIN SELECT * FROM test_regex_utf8view WHERE s !~ '^(foo|Bazzz)$' ---- logical_plan -01)Filter: test_regex_utf8view.s != Utf8View("foo") AND test_regex_utf8view.s != Utf8View("Bazzz") +01)Filter: test_regex_utf8view.s NOT IN ([Utf8View("foo"), Utf8View("Bazzz")]) 02)--TableScan: test_regex_utf8view projection=[s] # `!~*` anchored alternation -> NOT simplified: it falls back to a regex match, @@ -580,8 +580,8 @@ EXPLAIN SELECT * FROM aggregate_test_100 WHERE (c2 = 1 OR c3 = 100) OR (c2 = 2 OR c2 = 3 OR c2 = 4) ---- logical_plan -01)Filter: aggregate_test_100.c2 = Int8(1) OR aggregate_test_100.c3 = Int16(100) OR aggregate_test_100.c2 = Int8(2) OR aggregate_test_100.c2 = Int8(3) OR aggregate_test_100.c2 = Int8(4) -02)--TableScan: aggregate_test_100 projection=[c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13], partial_filters=[aggregate_test_100.c2 = Int8(1) OR aggregate_test_100.c3 = Int16(100) OR aggregate_test_100.c2 = Int8(2) OR aggregate_test_100.c2 = Int8(3) OR aggregate_test_100.c2 = Int8(4)] +01)Filter: aggregate_test_100.c2 = Int8(1) OR aggregate_test_100.c3 = Int16(100) OR aggregate_test_100.c2 IN ([Int8(2), Int8(3), Int8(4)]) +02)--TableScan: aggregate_test_100 projection=[c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13], partial_filters=[aggregate_test_100.c2 = Int8(1) OR aggregate_test_100.c3 = Int16(100) OR aggregate_test_100.c2 IN ([Int8(2), Int8(3), Int8(4)])] # Partially simplifiable, mixed column query TT @@ -909,8 +909,8 @@ logical_plan 05)--------Inner Join: lineitem.l_partkey = part.p_partkey 06)----------TableScan: lineitem projection=[l_partkey, l_extendedprice, l_discount] 07)----------Projection: part.p_partkey -08)------------Filter: part.p_brand = Utf8View("Brand#12") OR part.p_brand = Utf8View("Brand#23") -09)--------------TableScan: part projection=[p_partkey, p_brand], partial_filters=[part.p_brand = Utf8View("Brand#12") OR part.p_brand = Utf8View("Brand#23")] +08)------------Filter: part.p_brand IN ([Utf8View("Brand#12"), Utf8View("Brand#23")]) +09)--------------TableScan: part projection=[p_partkey, p_brand], partial_filters=[part.p_brand IN ([Utf8View("Brand#12"), Utf8View("Brand#23")])] 10)------TableScan: partsupp projection=[ps_partkey, ps_suppkey] physical_plan 01)AggregateExec: mode=SinglePartitioned, gby=[p_partkey@2 as p_partkey], aggr=[sum(lineitem.l_extendedprice), avg(lineitem.l_discount), count(DISTINCT partsupp.ps_suppkey)] @@ -920,7 +920,7 @@ physical_plan 05)------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=1 06)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/tpch-csv/lineitem.csv]]}, projection=[l_partkey, l_extendedprice, l_discount], file_type=csv, has_header=true 07)------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 -08)--------FilterExec: p_brand@1 = Brand#12 OR p_brand@1 = Brand#23, projection=[p_partkey@0] +08)--------FilterExec: p_brand@1 IN (SET) ([Brand#12, Brand#23]), projection=[p_partkey@0] 09)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 10)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/tpch-csv/part.csv]]}, projection=[p_partkey, p_brand], file_type=csv, has_header=true diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 03670c06fd4cd..d2286dbf693f4 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -257,7 +257,7 @@ logical_plan 01)Sort: t1.t1_id ASC NULLS LAST, t2.t2_id ASC NULLS LAST 02)--Inner Join: Filter: t1.t1_id <= t2.t2_id 03)----SubqueryAlias: t1 -04)------Filter: join_t1.t1_id = Int32(11) OR join_t1.t1_id = Int32(44) +04)------Filter: join_t1.t1_id IN ([Int32(11), Int32(44)]) 05)--------TableScan: join_t1 projection=[t1_id] 06)----SubqueryAlias: t2 07)------Projection: join_t2.t2_id @@ -268,7 +268,7 @@ physical_plan 02)--SortExec: expr=[t1_id@0 ASC NULLS LAST, t2_id@1 ASC NULLS LAST], preserve_partitioning=[true] 03)----PiecewiseMergeJoin: operator=LtEq, join_type=Inner, on=(t1_id <= t2_id) 04)------SortExec: expr=[t1_id@0 DESC], preserve_partitioning=[false] -05)--------FilterExec: t1_id@0 = 11 OR t1_id@0 = 44 +05)--------FilterExec: t1_id@0 IN (SET) ([11, 44]) 06)----------DataSourceExec: partitions=1, partition_sizes=[1] 07)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 08)--------FilterExec: t2_name@1 != y, projection=[t2_id@0] diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index ec374b3d62a28..dd63a5c58d0a2 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1863,7 +1863,7 @@ JOIN range_partitioned p ON b.range_key = p.range_key; ---- physical_plan 01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] -02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20, pruning_predicate=range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 5 AND 5 <= range_key_max@1 OR range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 20 AND 20 <= range_key_max@1, required_guarantees=[range_key in (20, 5)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=range_key@0 IN (SET) ([5, 20]), pruning_predicate=range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 5 AND 5 <= range_key_max@1 OR range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 20 AND 20 <= range_key_max@1, required_guarantees=[range_key in (20, 5)] 03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible query TT @@ -1877,7 +1877,7 @@ JOIN range_partitioned p ON b.range_key = p.range_key; ---- Plan with Metrics 01)HashJoinExec: mode=Partitionedmetrics=[output_rows=2,] -02)--DataSourceExec: file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20metrics=[output_rows=2,] +02)--DataSourceExec: file_type=parquet, predicate=range_key@0 IN (SET) ([5, 20])metrics=[output_rows=2,] 03)--DataSourceExec: file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN range_key@0 >= 5 AND range_key@0 <= 5 AND range_key@0 IN (SET) ([5]) WHEN 2 THEN range_key@0 >= 20 AND range_key@0 <= 20 AND range_key@0 IN (SET) ([20]) ELSE false END ]metrics=[output_rows=2,] query III diff --git a/datafusion/sqllogictest/test_files/simplify_expr.slt b/datafusion/sqllogictest/test_files/simplify_expr.slt index c70ff2b955e19..75fbdc8eefef4 100644 --- a/datafusion/sqllogictest/test_files/simplify_expr.slt +++ b/datafusion/sqllogictest/test_files/simplify_expr.slt @@ -381,10 +381,10 @@ explain select id from date_unwrap where arrow_cast(d32, 'Date64') in (arrow_cas ---- logical_plan 01)Projection: date_unwrap.id -02)--Filter: date_unwrap.d32 = Date32("2025-01-01") OR date_unwrap.d32 = Date32("1969-12-31") +02)--Filter: date_unwrap.d32 IN ([Date32("2025-01-01"), Date32("1969-12-31")]) 03)----TableScan: date_unwrap projection=[id, d32] physical_plan -01)FilterExec: d32@1 = 2025-01-01 OR d32@1 = 1969-12-31, projection=[id@0] +01)FilterExec: d32@1 IN (SET) ([2025-01-01, 1969-12-31]), projection=[id@0] 02)--DataSourceExec: partitions=1, partition_sizes=[1] query I diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index a173e76d6c262..91b690e0c41ec 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -168,11 +168,11 @@ ORDER BY id DESC LIMIT 5; ---- logical_plan 01)Sort: multi_rg_sorted.id DESC NULLS FIRST, fetch=5 -02)--Filter: multi_rg_sorted.category = Utf8View("alpha") OR multi_rg_sorted.category = Utf8View("gamma") -03)----TableScan: multi_rg_sorted projection=[id, category, value], partial_filters=[multi_rg_sorted.category = Utf8View("alpha") OR multi_rg_sorted.category = Utf8View("gamma")] +02)--Filter: multi_rg_sorted.category IN ([Utf8View("alpha"), Utf8View("gamma")]) +03)----TableScan: multi_rg_sorted projection=[id, category, value], partial_filters=[multi_rg_sorted.category IN ([Utf8View("alpha"), Utf8View("gamma")])] physical_plan 01)SortExec: TopK(fetch=5), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_rg_sorted.parquet]]}, projection=[id, category, value], file_type=parquet, predicate=(category@1 = alpha OR category@1 = gamma) AND DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1 OR category_null_count@2 != row_count@3 AND category_min@0 <= gamma AND gamma <= category_max@1, required_guarantees=[category in (alpha, gamma)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_rg_sorted.parquet]]}, projection=[id, category, value], file_type=parquet, predicate=category@1 IN (SET) ([alpha, gamma]) AND DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1 OR category_null_count@2 != row_count@3 AND category_min@0 <= gamma AND gamma <= category_max@1, required_guarantees=[category in (alpha, gamma)] # Verify the results are correct despite reverse scanning with row selection # Expected: gamma values (6, 5) then alpha values (2, 1), in DESC order by id @@ -504,11 +504,11 @@ LIMIT 3; ---- logical_plan 01)Sort: timeseries_parquet.period_end DESC NULLS FIRST, fetch=3 -02)--Filter: timeseries_parquet.timeframe = Utf8View("daily") OR timeseries_parquet.timeframe = Utf8View("weekly") -03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe = Utf8View("daily") OR timeseries_parquet.timeframe = Utf8View("weekly")] +02)--Filter: timeseries_parquet.timeframe IN ([Utf8View("daily"), Utf8View("weekly")]) +03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe IN ([Utf8View("daily"), Utf8View("weekly")])] physical_plan 01)SortExec: TopK(fetch=3), expr=[period_end@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=(timeframe@0 = daily OR timeframe@0 = weekly) AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= daily AND daily <= timeframe_max@1 OR timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= weekly AND weekly <= timeframe_max@1, required_guarantees=[timeframe in (daily, weekly)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=timeframe@0 IN (SET) ([daily, weekly]) AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= daily AND daily <= timeframe_max@1 OR timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= weekly AND weekly <= timeframe_max@1, required_guarantees=[timeframe in (daily, weekly)] # Test 2.9: Complex case - literal constant in sort expression itself # The literal 'constant' is ignored in sort analysis diff --git a/datafusion/sqllogictest/test_files/string_numeric_coercion.slt b/datafusion/sqllogictest/test_files/string_numeric_coercion.slt index 196da416a2037..c09f2c986e7c0 100644 --- a/datafusion/sqllogictest/test_files/string_numeric_coercion.slt +++ b/datafusion/sqllogictest/test_files/string_numeric_coercion.slt @@ -229,10 +229,10 @@ query TT EXPLAIN SELECT * FROM t_int WHERE column1 IN ('5', '325'); ---- logical_plan -01)Filter: t_int.column1 = Int64(5) OR t_int.column1 = Int64(325) +01)Filter: t_int.column1 IN ([Int64(5), Int64(325)]) 02)--TableScan: t_int projection=[column1] physical_plan -01)FilterExec: column1@0 = 5 OR column1@0 = 325 +01)FilterExec: column1@0 IN (SET) ([5, 325]) 02)--DataSourceExec: partitions=1, partition_sizes=[1] # Error on invalid string in IN list diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 9d953f5075425..e88ac5478f2ed 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -365,14 +365,14 @@ EXPLAIN SELECT id, name FROM t1 WHERE id = 1 UNION SELECT id, name FROM t1 WHERE ---- logical_plan 01)Aggregate: groupBy=[[id, name]], aggr=[[]] -02)--Filter: t1.id = Int32(1) OR t1.id = Int32(2) +02)--Filter: t1.id IN ([Int32(1), Int32(2)]) 03)----TableScan: t1 projection=[id, name] physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id, name@1 as name], aggr=[] 02)--RepartitionExec: partitioning=Hash([id@0, name@1], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[id@0 as id, name@1 as name], aggr=[] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)--------FilterExec: id@0 = 1 OR id@0 = 2 +05)--------FilterExec: id@0 IN (SET) ([1, 2]) 06)----------DataSourceExec: partitions=1, partition_sizes=[1] # Regression: schema recomputation must preserve the unqualified UNION @@ -415,14 +415,14 @@ SELECT x.id FROM t1 AS x WHERE x.id = 2 logical_plan 01)Aggregate: groupBy=[[id]], aggr=[[]] 02)--SubqueryAlias: x -03)----Filter: t1.id = Int32(1) OR t1.id = Int32(2) +03)----Filter: t1.id IN ([Int32(1), Int32(2)]) 04)------TableScan: t1 projection=[id] physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[] 02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)--------FilterExec: id@0 = 1 OR id@0 = 2 +05)--------FilterExec: id@0 IN (SET) ([1, 2]) 06)----------DataSourceExec: partitions=1, partition_sizes=[1] # Matching computed projections remain eligible for the rewrite and must diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 6ffe1b4cc087c..9aa6cb4629e71 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -5340,7 +5340,7 @@ logical_plan 01)Sort: t1.c1 ASC NULLS LAST, t1.c2 ASC NULLS LAST, rank ASC NULLS LAST 02)--Projection: t1.c1, t1.c2, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rank 03)----WindowAggr: windowExpr=[[rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] -04)------Filter: t1.c1 = Int32(2) OR t1.c1 = Int32(3) +04)------Filter: t1.c1 IN ([Int32(2), Int32(3)]) 05)--------TableScan: t1 projection=[c1, c2] physical_plan 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank@2 ASC NULLS LAST] @@ -5348,7 +5348,7 @@ physical_plan 03)----BoundedWindowAggExec: wdw=[rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 04)------SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST], preserve_partitioning=[true] 05)--------RepartitionExec: partitioning=Hash([c1@0], 2), input_partitions=2 -06)----------FilterExec: c1@0 = 2 OR c1@0 = 3 +06)----------FilterExec: c1@0 IN (SET) ([2, 3]) 07)------------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 08)--------------DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 813d0ed6c3489..d989926adf786 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -707,18 +707,22 @@ async fn aggregate_case() -> Result<()> { #[tokio::test] async fn roundtrip_inlist_1() -> Result<()> { + // Int64 has a specialized static filter, so this short list remains an + // InList through expression simplification. roundtrip("SELECT * FROM data WHERE a IN (1, 2, 3)").await } #[tokio::test] -// Test with length <= datafusion_optimizer::simplify_expressions::expr_simplifier::THRESHOLD_INLINE_INLIST async fn roundtrip_inlist_2() -> Result<()> { + // Utf8 uses the generic physical filter, so a list at the shortening + // threshold is expanded into OR comparisons. roundtrip("SELECT * FROM data WHERE f IN ('a', 'b', 'c')").await } #[tokio::test] -// Test with length > datafusion_optimizer::simplify_expressions::expr_simplifier::THRESHOLD_INLINE_INLIST async fn roundtrip_inlist_3() -> Result<()> { + // Even for generic Utf8, a list above the shortening threshold remains an + // InList to avoid growing the expression tree. let inlist = (0..=THRESHOLD_INLINE_INLIST) .map(|i| format!("'{i}'")) .collect::>()