diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 2de3be4c10fa4..2c27defb5f213 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -317,6 +317,14 @@ impl ScalarUDF { self.inner.short_circuits() } + /// Returns true if each output row of this function depends only on the + /// corresponding input row. + /// + /// See [ScalarUDFImpl::evaluates_elementwise] for more information. + pub fn evaluates_elementwise(&self) -> bool { + self.inner.evaluates_elementwise() + } + /// Computes the output interval for a [`ScalarUDF`], given the input /// intervals. /// @@ -885,6 +893,22 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { false } + /// Returns true if each output row depends only on the corresponding input + /// row, with no cross-row state. + /// + /// When true, a dictionary-encoded argument is unwrapped before the call: + /// the function is evaluated over the distinct values and the result + /// re-mapped through the keys, or, where that does not pay off, over the + /// expanded column. [`ScalarFunctionArgs::number_rows`] and the argument + /// fields then describe the array actually passed, not the planned batch. + /// + /// Values that no key references may still be evaluated, so a function that + /// can error on valid input should not opt in. Volatile functions are never + /// unwrapped, nor are arguments whose field carries metadata. + fn evaluates_elementwise(&self) -> bool { + false + } + /// Determines which of the arguments passed to this function are evaluated eagerly /// and which may be evaluated lazily. /// @@ -1182,6 +1206,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.short_circuits() } + fn evaluates_elementwise(&self) -> bool { + self.inner.evaluates_elementwise() + } + fn evaluate_bounds(&self, input: &[&Interval]) -> Result { self.inner.evaluate_bounds(input) } diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index a170e9f07c39f..d22dd5665ee6b 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -403,8 +403,3 @@ required-features = ["math_expressions"] harness = false name = "round" required-features = ["math_expressions"] - -[[bench]] -harness = false -name = "dictionary_encoding" -required-features = ["string_expressions", "unicode_expressions"] diff --git a/datafusion/functions/benches/dictionary_encoding.rs b/datafusion/functions/benches/dictionary_encoding.rs deleted file mode 100644 index 4ba04a4940e61..0000000000000 --- a/datafusion/functions/benches/dictionary_encoding.rs +++ /dev/null @@ -1,104 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::hint::black_box; -use std::sync::Arc; - -use arrow::array::{ArrayRef, DictionaryArray}; -use arrow::compute::cast; -use arrow::datatypes::{Field, Int32Type}; -use criterion::{Criterion, criterion_group, criterion_main}; -use datafusion_common::config::ConfigOptions; -use datafusion_expr::type_coercion::functions::fields_with_udf; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; - -const NUM_ROWS: usize = 8_192; -const DICTIONARY_CARDINALITIES: [usize; 4] = [10, 100, 1_000, 8_192]; - -fn create_string_dictionary(cardinality: usize) -> ArrayRef { - let values = (0..NUM_ROWS) - .map(|index| Some(format!("value_{:04}", index % cardinality))) - .collect::>(); - Arc::new( - values - .iter() - .map(|value| value.as_deref()) - .collect::>(), - ) -} - -fn benchmark_dictionary_string_udfs(c: &mut Criterion) { - let udfs = [ - ("ascii", datafusion_functions::string::ascii()), - ("bit_length", datafusion_functions::string::bit_length()), - ("btrim", datafusion_functions::string::btrim()), - ( - "character_length", - datafusion_functions::unicode::character_length(), - ), - ("initcap", datafusion_functions::unicode::initcap()), - ("ltrim", datafusion_functions::string::ltrim()), - ("octet_length", datafusion_functions::string::octet_length()), - ("reverse", datafusion_functions::unicode::reverse()), - ("rtrim", datafusion_functions::string::rtrim()), - ]; - let config_options = Arc::new(ConfigOptions::default()); - - for cardinality in DICTIONARY_CARDINALITIES { - let dictionary = create_string_dictionary(cardinality); - let mut group = c.benchmark_group(format!( - "dictionary_encoding/string/cardinality_{cardinality}" - )); - for (name, udf) in &udfs { - let input_field = - Field::new("a", dictionary.data_type().clone(), false).into(); - let coerced_field = fields_with_udf(&[input_field], udf.as_ref()) - .unwrap() - .into_iter() - .next() - .unwrap(); - let coerced_type = coerced_field.data_type(); - let return_type = - udf.return_type(std::slice::from_ref(coerced_type)).unwrap(); - let return_field = Field::new("f", return_type, false).into(); - let input = if dictionary.data_type() == coerced_type { - Arc::clone(&dictionary) - } else { - cast(dictionary.as_ref(), coerced_type).unwrap() - }; - - group.bench_function(*name, |b| { - b.iter(|| { - black_box( - udf.invoke_with_args(ScalarFunctionArgs { - args: vec![ColumnarValue::Array(Arc::clone(&input))], - arg_fields: vec![Arc::clone(&coerced_field)], - number_rows: NUM_ROWS, - return_field: Arc::clone(&return_field), - config_options: Arc::clone(&config_options), - }) - .unwrap(), - ) - }) - }); - } - group.finish(); - } -} - -criterion_group!(benches, benchmark_dictionary_string_udfs); -criterion_main!(benches); diff --git a/datafusion/functions/src/encoding/inner.rs b/datafusion/functions/src/encoding/inner.rs index 8b57033fa0de6..d9c38b7f6788b 100644 --- a/datafusion/functions/src/encoding/inner.rs +++ b/datafusion/functions/src/encoding/inner.rs @@ -39,8 +39,8 @@ use datafusion_common::{ }, }; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::fmt; @@ -94,7 +94,8 @@ impl EncodeFunc { TypeSignatureClass::Binary, vec![TypeSignatureClass::Native(logical_string())], NativeType::Binary, - ), + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ], Volatility::Immutable, @@ -119,6 +120,10 @@ impl ScalarUDFImpl for EncodeFunc { } } + fn evaluates_elementwise(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [expression, encoding] = take_function_args("encode", &args.args)?; let encoding = Encoding::try_from(encoding)?; diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs index 0332ab5d4427f..10ffbcf8a4387 100644 --- a/datafusion/functions/src/unicode/initcap.rs +++ b/datafusion/functions/src/unicode/initcap.rs @@ -86,6 +86,10 @@ impl ScalarUDFImpl for InitcapFunc { Ok(arg_types[0].clone()) } + fn evaluates_elementwise(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { match &args.args[0] { ColumnarValue::Scalar(scalar) => { @@ -140,6 +144,8 @@ fn initcap_array(array: &ArrayRef) -> Result { DataType::Utf8 => initcap::(&[Arc::clone(array)]), DataType::LargeUtf8 => initcap::(&[Arc::clone(array)]), DataType::Utf8View => initcap_utf8view(&[Arc::clone(array)]), + // Serves the calls the physical layer leaves encoded, such as fields + // carrying extension metadata. DataType::Dictionary(_, _) => { let dictionary = array.as_any_dictionary(); let converted = initcap_array(dictionary.values())?; diff --git a/datafusion/functions/src/unicode/reverse.rs b/datafusion/functions/src/unicode/reverse.rs index e42240ba6697b..d3e409531c1b9 100644 --- a/datafusion/functions/src/unicode/reverse.rs +++ b/datafusion/functions/src/unicode/reverse.rs @@ -87,6 +87,10 @@ impl ScalarUDFImpl for ReverseFunc { Ok(arg_types[0].clone()) } + fn evaluates_elementwise(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { make_scalar_function(reverse, vec![])(&args.args) } @@ -114,6 +118,8 @@ fn reverse(args: &[ArrayRef]) -> Result { &args[0].as_string_view(), StringViewArrayBuilder::with_capacity(len), ), + // Serves the calls the physical layer leaves encoded, such as fields + // carrying extension metadata. DataType::Dictionary(_, _) => { let dictionary = args[0].as_any_dictionary(); let converted = reverse(&[Arc::clone(dictionary.values())])?; diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..0afc69f375dc8 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -102,5 +102,9 @@ name = "simplify" harness = false name = "string_concat" +[[bench]] +harness = false +name = "dictionary_encoding" + [package.metadata.cargo-machete] ignored = ["half"] diff --git a/datafusion/physical-expr/benches/dictionary_encoding.rs b/datafusion/physical-expr/benches/dictionary_encoding.rs new file mode 100644 index 0000000000000..f6460e864ce68 --- /dev/null +++ b/datafusion/physical-expr/benches/dictionary_encoding.rs @@ -0,0 +1,356 @@ +// 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. + +//! Scalar functions over dictionary-encoded columns, measured at two layers. +//! +//! `string/*` calls the function directly, so it measures whatever +//! dictionary handling the function itself has: `reverse` and the trim family +//! peel the dictionary by hand, others see it materialized. These groups +//! moved here from `datafusion-functions` (#23930) unchanged, so their +//! numbers stay comparable, and now sit beside the expression-layer groups +//! they are the baseline for. +//! +//! `expression/*` evaluates through [`ScalarFunctionExpr`], the layer that +//! decides what the function receives. `cold` gives every batch its own +//! dictionary, as a projection building one per batch does; `warm` shares one +//! across batches, as a Parquet scan does within a column chunk; `flat` is the +//! same rows with the encoding materialized; `cast_away` is what a function +//! without encoding preservation pays, since coercion casts the dictionary +//! away before the call. + +use std::cell::Cell; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryArray, DictionaryArray, Int32Array, StringArray, +}; +use arrow::compute::{cast, take}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::type_coercion::functions::fields_with_udf; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + +const ROWS: usize = 8192; + +/// Distinct values per batch. 10, 100, 1000 and 8192 are the grid #23930 used, +/// kept so its numbers remain comparable; 256 is added because that is where +/// real dictionaries sit — a Parquet page for a categorical column typically +/// holds a few hundred values. +const CARDINALITIES: [usize; 5] = [10, 100, 256, 1000, 8192]; + +/// The distinct values a batch draws from, in the type the function receives +/// once coercion has run: `encode` takes binary, the string functions take +/// strings. The textual form is the one #23930 used, so its numbers stay +/// comparable. +fn values_of(distinct: usize, binary: bool) -> ArrayRef { + let values: Vec = (0..distinct).map(|i| format!("value_{i:04}")).collect(); + if binary { + Arc::new(BinaryArray::from( + values.iter().map(|v| v.as_bytes()).collect::>(), + )) + } else { + Arc::new(StringArray::from(values)) + } +} + +/// `ROWS` rows cycling through `cardinality` distinct values. +fn create_string_dictionary(cardinality: usize) -> ArrayRef { + let keys = Int32Array::from( + (0..ROWS) + .map(|index| (index % cardinality) as i32) + .collect::>(), + ); + Arc::new( + DictionaryArray::::try_new(keys, values_of(cardinality, false)) + .expect("dictionary array"), + ) +} + +/// The function's own dictionary handling: the call receives the dictionary +/// exactly as coercion would deliver it. +fn benchmark_function_layer(c: &mut Criterion) { + let udfs = [ + ("ascii", datafusion_functions::string::ascii()), + ("bit_length", datafusion_functions::string::bit_length()), + ("btrim", datafusion_functions::string::btrim()), + ( + "character_length", + datafusion_functions::unicode::character_length(), + ), + ("initcap", datafusion_functions::unicode::initcap()), + ("ltrim", datafusion_functions::string::ltrim()), + ("octet_length", datafusion_functions::string::octet_length()), + ("reverse", datafusion_functions::unicode::reverse()), + ("rtrim", datafusion_functions::string::rtrim()), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + for cardinality in CARDINALITIES { + let dictionary = create_string_dictionary(cardinality); + let mut group = c.benchmark_group(format!( + "dictionary_encoding/string/cardinality_{cardinality}" + )); + for (name, udf) in &udfs { + let input_field = + Field::new("a", dictionary.data_type().clone(), false).into(); + let coerced_field = fields_with_udf(&[input_field], udf.as_ref()) + .unwrap() + .into_iter() + .next() + .unwrap(); + let coerced_type = coerced_field.data_type(); + let return_type = + udf.return_type(std::slice::from_ref(coerced_type)).unwrap(); + let return_field = Field::new("f", return_type, false).into(); + let input = if dictionary.data_type() == coerced_type { + Arc::clone(&dictionary) + } else { + cast(dictionary.as_ref(), coerced_type).unwrap() + }; + + group.bench_function(*name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::clone(&input))], + arg_fields: vec![Arc::clone(&coerced_field)], + number_rows: ROWS, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } + group.finish(); + } +} + +/// `ROWS` rows drawn from `distinct` values, dictionary-encoded. `shift` moves +/// which value each row lands on, so batches built separately differ in their +/// keys as well as in the memory their values occupy. +fn dictionary_batch( + distinct: usize, + shift: usize, + binary: bool, +) -> (Schema, RecordBatch) { + let keys = Int32Array::from( + (0..ROWS) + .map(|i| ((i + shift) % distinct) as i32) + .collect::>(), + ); + let dict = DictionaryArray::::try_new(keys, values_of(distinct, binary)) + .expect("dictionary array"); + let schema = Schema::new(vec![Field::new("c", dict.data_type().clone(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(dict)]) + .expect("batch"); + (schema, batch) +} + +/// The same rows without the encoding — what the function receives when the +/// dictionary is materialized before evaluation. +fn flat_batch(distinct: usize, binary: bool) -> (Schema, RecordBatch) { + let keys = + Int32Array::from((0..ROWS).map(|i| (i % distinct) as i32).collect::>()); + let array = take(values_of(distinct, binary).as_ref(), &keys, None) + .expect("materialized column"); + let data_type = array.data_type().clone(); + let schema = Schema::new(vec![Field::new("c", data_type, true)]); + let batch = + RecordBatch::try_new(Arc::new(schema.clone()), vec![array]).expect("batch"); + (schema, batch) +} + +/// Consecutive batches of one column chunk: their own keys, one dictionary. +fn chunk(distinct: usize, batches: usize, binary: bool) -> (Schema, Vec) { + let values = values_of(distinct, binary); + let schema = Schema::new(vec![Field::new( + "c", + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(values.data_type().clone()), + ), + true, + )]); + let batches = (0..batches) + .map(|b| { + let keys = Int32Array::from( + (0..ROWS) + .map(|i| ((i + b * 7) % distinct) as i32) + .collect::>(), + ); + let dict = DictionaryArray::::try_new(keys, Arc::clone(&values)) + .expect("dictionary array"); + RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(dict)]) + .expect("batch") + }) + .collect(); + (schema, batches) +} + +/// `batches` batches, each with a dictionary of its own. +fn separate(distinct: usize, batches: usize, binary: bool) -> (Schema, Vec) { + let mut schema = None; + let batches = (0..batches) + .map(|b| { + let (built, batch) = dictionary_batch(distinct, b, binary); + schema.get_or_insert(built); + batch + }) + .collect(); + (schema.expect("at least one batch"), batches) +} + +fn expr_over(udf: Arc, schema: &Schema, base64: bool) -> ScalarFunctionExpr { + let mut args: Vec> = vec![Arc::new(Column::new("c", 0))]; + if base64 { + args.push(Arc::new(Literal::new(ScalarValue::from("base64")))); + } + ScalarFunctionExpr::try_new(udf, args, schema, Arc::new(ConfigOptions::new())) + .expect("scalar function expr") +} + +/// The expression layer, which decides what the function receives. +fn benchmark_expression_layer(c: &mut Criterion) { + // (name, function, takes a base64 argument) + let functions: Vec<(&str, Arc, bool)> = vec![ + ("encode", datafusion_functions::encoding::encode(), true), + ("reverse", datafusion_functions::unicode::reverse(), false), + ]; + + for (name, udf, binary) in &functions { + let mut group = + c.benchmark_group(format!("dictionary_encoding/expression/{name}")); + + // A dictionary of its own per batch: nothing carries over. + // + // The cursor lives outside the routine, which criterion calls + // afresh for every sample: restarted per sample it would revisit + // the first batches often enough for a result to still be + // remembered, and the group would quietly measure a warm + // dictionary under a cold name. + for distinct in CARDINALITIES { + let (schema, batches) = separate(distinct, 16, *binary); + let expr = expr_over(Arc::clone(udf), &schema, *binary); + let cursor = Cell::new(0usize); + group.bench_with_input( + BenchmarkId::new("cold", distinct), + &batches, + |b, batches| { + b.iter(|| { + cursor.set((cursor.get() + 1) % batches.len()); + black_box( + expr.evaluate(black_box(&batches[cursor.get()])).unwrap(), + ) + }) + }, + ); + } + + // One dictionary across the batches of a column chunk. + for distinct in CARDINALITIES { + let (schema, batches) = chunk(distinct, 8, *binary); + let expr = expr_over(Arc::clone(udf), &schema, *binary); + for batch in &batches { + expr.evaluate(batch).unwrap(); + } + let cursor = Cell::new(0usize); + group.bench_with_input( + BenchmarkId::new("warm", distinct), + &batches, + |b, batches| { + b.iter(|| { + cursor.set((cursor.get() + 1) % batches.len()); + black_box( + expr.evaluate(black_box(&batches[cursor.get()])).unwrap(), + ) + }) + }, + ); + } + + // The same rows with the encoding materialized: one call per row. + for distinct in [CARDINALITIES[0], ROWS] { + let (schema, batch) = flat_batch(distinct, *binary); + let expr = expr_over(Arc::clone(udf), &schema, *binary); + group.bench_with_input( + BenchmarkId::new("flat", distinct), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + + // What a dictionary column costs without encoding preservation, which + // is what `encode` did before this change and what every function + // without it still does: coercion casts the dictionary away, and the + // call sees one row per row. The `flat` rows above are not this — they + // never were a dictionary and so never pay for materializing one. + for distinct in CARDINALITIES { + let (schema, batches) = separate(distinct, 16, *binary); + let values_type = match schema.field(0).data_type() { + DataType::Dictionary(_, values) => values.as_ref().clone(), + other => other.clone(), + }; + let cast_expr: Arc = Arc::new(CastExpr::new( + Arc::new(Column::new("c", 0)), + values_type, + None, + )); + let mut args: Vec> = vec![cast_expr]; + if *binary { + args.push(Arc::new(Literal::new(ScalarValue::from("base64")))); + } + let expr = ScalarFunctionExpr::try_new( + Arc::clone(udf), + args, + &schema, + Arc::new(ConfigOptions::new()), + ) + .expect("scalar function expr"); + let cursor = Cell::new(0usize); + group.bench_with_input( + BenchmarkId::new("cast_away", distinct), + &batches, + |b, batches| { + b.iter(|| { + cursor.set((cursor.get() + 1) % batches.len()); + black_box( + expr.evaluate(black_box(&batches[cursor.get()])).unwrap(), + ) + }) + }, + ); + } + + group.finish(); + } +} + +criterion_group!( + benches, + benchmark_function_layer, + benchmark_expression_layer +); +criterion_main!(benches); diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 6a5ab219aa8dd..e1fcd0b1ea563 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -31,13 +31,14 @@ use std::fmt::{self, Debug, Formatter}; use std::hash::{Hash, Hasher}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock, RwLock}; use crate::PhysicalExpr; use crate::expressions::Literal; -use arrow::array::{Array, RecordBatch}; -use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::array::{AnyDictionaryArray, Array, ArrayRef, AsArray, RecordBatch}; +use arrow::compute::take; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use datafusion_common::config::{ConfigEntry, ConfigOptions}; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::interval_arithmetic::Interval; @@ -48,6 +49,12 @@ use datafusion_expr::{ ScalarUDFImpl, Volatility, expr_vec_fmt, }; +mod dictionary; +use dictionary::{ + Lookup, Memo, PeeledFields, Recollection, ValuesIdentity, compact_dictionary, + scalar_arguments, unwrap_scalar_dictionaries, +}; + /// Physical expression of a scalar function pub struct ScalarFunctionExpr { fun: Arc, @@ -55,6 +62,11 @@ pub struct ScalarFunctionExpr { args: Vec>, return_field: FieldRef, config_options: Arc, + /// Fields for the peeled call, built once and reused across batches. + peeled: OnceLock, + /// Results computed for a dictionary's values: a batch carries its own + /// keys but its column chunk's whole dictionary, so the same values recur. + memoized: RwLock, } impl Debug for ScalarFunctionExpr { @@ -83,6 +95,8 @@ impl ScalarFunctionExpr { args, return_field, config_options, + peeled: OnceLock::new(), + memoized: RwLock::new(Memo::default()), } } @@ -117,6 +131,8 @@ impl ScalarFunctionExpr { args, return_field, config_options, + peeled: OnceLock::new(), + memoized: RwLock::new(Memo::default()), }) } @@ -147,6 +163,8 @@ impl ScalarFunctionExpr { .clone() .with_nullable(nullable) .into(); + // Derived from the field just replaced. + self.peeled = OnceLock::new(); self } @@ -171,6 +189,319 @@ impl ScalarFunctionExpr { _ => None, } } + + /// Evaluates an elementwise function over the distinct values of a + /// dictionary-encoded argument and re-maps the result through its keys, so + /// the function runs once per distinct value instead of once per row. + /// + /// Returns `None` when peeling does not apply and the caller must evaluate + /// the function over the batch as-is. + fn try_invoke_peeled( + &self, + args: &[ColumnarValue], + arg_fields: &[FieldRef], + num_rows: usize, + ) -> Result> { + if num_rows == 0 + || !self.fun.evaluates_elementwise() + || self.fun.signature().volatility == Volatility::Volatile + { + return Ok(None); + } + + let returns_dictionary = + matches!(self.return_field.data_type(), DataType::Dictionary(_, _)); + + // Exactly one dictionary-encoded array; peeling several at once is only + // sound when their keys line up. Dictionary-encoded scalars carry a + // single value and are unwrapped rather than peeled. + let mut dictionary_index = None; + let mut scalar_dictionaries = false; + for (index, arg) in args.iter().enumerate() { + match arg { + ColumnarValue::Array(array) + if matches!(array.data_type(), DataType::Dictionary(_, _)) => + { + if dictionary_index.replace(index).is_some() { + return Ok(None); + } + } + ColumnarValue::Scalar(ScalarValue::Dictionary(_, _)) => { + scalar_dictionaries = true + } + ColumnarValue::Scalar(_) => {} + _ => return Ok(None), + } + } + let Some(dictionary_index) = dictionary_index else { + if !scalar_dictionaries || returns_dictionary { + return Ok(None); + } + let (args, arg_fields) = unwrap_scalar_dictionaries(args, arg_fields); + return self + .fun + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: num_rows, + return_field: Arc::clone(&self.return_field), + config_options: Arc::clone(&self.config_options), + }) + .map(Some); + }; + let ColumnarValue::Array(array) = &args[dictionary_index] else { + return internal_err!("dictionary argument is not an array"); + }; + // An extension type is bound to its storage type, so fields carrying + // metadata keep the dictionary rather than have their type rewritten. + if !arg_fields[dictionary_index].metadata().is_empty() + || !self.return_field.metadata().is_empty() + { + return Ok(None); + } + let raw = array.as_any_dictionary(); + // A dictionary return keyed differently cannot be rebuilt from these keys. + if let DataType::Dictionary(key_type, _) = self.return_field.data_type() + && key_type.as_ref() != raw.keys().data_type() + { + return Ok(None); + } + let fields = self.peeled_fields(dictionary_index, arg_fields, raw.values()); + + // Both tiers below re-use the keys as they are, which only a strict + // `f` (NULL in -> NULL out) can answer for when some are null. + let keys_are_answerable = raw.keys().null_count() == 0 || self.fun.is_strict(); + + // Memoized tier: already evaluated for this expression — checked + // before profitability, since work already done is worth re-using + // however large the dictionary is. + if keys_are_answerable { + let scalars = scalar_arguments(args, fields.index); + let recollection = self.memoized(raw.values(), &scalars); + if let Recollection::Evaluated(output) = recollection { + return self.remap(raw, output, &fields).map(Some); + } + + // Fast tier: invoke over the values as-is — the cost profile of + // the hand-rolled dictionary arms. Values no key references are + // worth evaluating once a second sighting has made a repeat + // likely; an error is not returned, because it may come from such + // a value, and the tiers below settle that. + let values_len = raw.values().len(); + let repeats = matches!(recollection, Recollection::SeenBefore); + let profitable = if repeats { + true + } else if returns_dictionary { + values_len <= num_rows + } else { + values_len * 2 <= num_rows + }; + if profitable + && let Ok(output) = self.invoke_on_array( + args, + arg_fields, + &fields, + Arc::clone(raw.values()), + &fields.output, + ) + { + // Kept only for a dictionary that has proved it comes back; + // otherwise the entry would never be read. + if repeats { + self.memoize(raw.values(), &scalars, &output); + } + return self.remap(raw, output, &fields).map(Some); + } + } + + // Guarded tier: `f` sees exactly the values this batch references, + // null keys redirected to one appended NULL slot it evaluates like + // any other value — correct even where `f(NULL)` is not NULL. A + // dictionary return re-wraps in O(1); only flat returns carry a + // row budget. + let row_budget = (!returns_dictionary).then_some(num_rows); + if let Some(dictionary) = compact_dictionary(array, row_budget)? { + let peeled = self.invoke_on_values( + args, + arg_fields, + &fields, + dictionary.as_any_dictionary(), + ); + // A dictionary return cannot be rebuilt from a flat call, so its + // errors surface here; a flat one retries over the whole column. + if returns_dictionary || peeled.is_ok() { + return peeled.map(Some); + } + } + + // A dictionary return cannot be rebuilt from a call over expanded rows, + // so the column is left as it arrived and the function sees the + // dictionary — exactly as it does wherever peeling declines. + if returns_dictionary { + return Ok(None); + } + + // Peeling does not pay off for this batch, but `f` declared that it + // evaluates elementwise, so it gets the expanded array rather than the + // dictionary: exactly the input type coercion produces for functions + // that do not preserve the encoding. + let flattened = take(raw.values().as_ref(), raw.keys(), None)?; + self.invoke_on_array(args, arg_fields, &fields, flattened, &self.return_field) + .map(ColumnarValue::Array) + .map(Some) + } + + /// The fields for a peeled call, built once and reused across batches: they + /// depend on the plan, not on the data. Callers have already established + /// that neither field carries metadata. + fn peeled_fields( + &self, + index: usize, + arg_fields: &[FieldRef], + values: &ArrayRef, + ) -> PeeledFields { + let source = &arg_fields[index]; + if let Some(cached) = self.peeled.get() + && cached.index == index + && cached.argument.data_type() == values.data_type() + && cached.source.as_ref() == source.as_ref() + { + return cached.clone(); + } + // Nullable: a peeled call can be handed a NULL value slot even where + // the planned output was known not to be null. + let output_type = match self.return_field.data_type() { + DataType::Dictionary(_, value_type) => value_type.as_ref().clone(), + flat => flat.clone(), + }; + let built = PeeledFields { + index, + source: Arc::clone(source), + argument: Arc::new(Field::new( + source.name(), + values.data_type().clone(), + true, + )), + output: Arc::new(Field::new(self.return_field.name(), output_type, true)), + }; + let _ = self.peeled.set(built.clone()); + built + } + + /// Invokes the function over `dictionary.values()` and re-maps the result + /// to batch length through the keys (`with_values` when the planned type is + /// a dictionary, `take` otherwise). + fn invoke_on_values( + &self, + args: &[ColumnarValue], + arg_fields: &[FieldRef], + fields: &PeeledFields, + dictionary: &dyn AnyDictionaryArray, + ) -> Result { + let values = dictionary.values(); + let output = self.invoke_on_array( + args, + arg_fields, + fields, + Arc::clone(values), + &fields.output, + )?; + self.remap(dictionary, output, fields) + } + + /// Spreads a per-value result back over the rows the keys address. + fn remap( + &self, + dictionary: &dyn AnyDictionaryArray, + output: ArrayRef, + fields: &PeeledFields, + ) -> Result { + // Checked before the re-map, which is not defined for every type a + // misbehaving UDF could return. + if output.data_type() != fields.output.data_type() { + return internal_err!( + "UDF {} returned type {} under dictionary peeling, expected {}", + self.name, + output.data_type(), + fields.output.data_type() + ); + } + let remapped = + if matches!(self.return_field.data_type(), DataType::Dictionary(_, _)) { + dictionary.with_values(output) + } else { + take(output.as_ref(), dictionary.keys(), None)? + }; + Ok(ColumnarValue::Array(remapped)) + } + + /// What this expression already knows about `values`: its result if it + /// has one, or whether it has seen these values before. Hits take only + /// the read lock, so partitions sharing this expression do not serialize + /// on each other's warm batches; between the two locks another partition + /// may record the same dictionary, which costs at most one extra full + /// evaluation — what an unlucky first sighting costs anyway. + fn memoized(&self, values: &ArrayRef, scalars: &[ScalarValue]) -> Recollection { + let identity = { + let Ok(memo) = self.memoized.read() else { + return Recollection::Unknown; + }; + match memo.find(values, scalars) { + Lookup::Evaluated(output) => return Recollection::Evaluated(output), + Lookup::Absent(identity) => identity, + } + }; + let hash = match identity { + Some(identity) => identity.hash(), + None => ValuesIdentity::hash_of(values), + }; + let Ok(mut memo) = self.memoized.write() else { + return Recollection::Unknown; + }; + memo.note(hash) + } + + /// Keeps `output` for the next batch that arrives with the same values. + fn memoize(&self, values: &ArrayRef, scalars: &[ScalarValue], output: &ArrayRef) { + if let Ok(mut memo) = self.memoized.write() { + memo.keep(values, scalars, output); + } + } + + /// Invokes the function with `array` in place of the dictionary argument, + /// whose field is rewritten to match. Returns `array.len()` rows. + fn invoke_on_array( + &self, + args: &[ColumnarValue], + arg_fields: &[FieldRef], + fields: &PeeledFields, + array: ArrayRef, + return_field: &FieldRef, + ) -> Result { + let rows = array.len(); + let (mut peeled_args, mut peeled_fields) = + unwrap_scalar_dictionaries(args, arg_fields); + peeled_args[fields.index] = ColumnarValue::Array(array); + peeled_fields[fields.index] = Arc::clone(&fields.argument); + let output = self.fun.invoke_with_args(ScalarFunctionArgs { + args: peeled_args, + arg_fields: peeled_fields, + number_rows: rows, + return_field: Arc::clone(return_field), + config_options: Arc::clone(&self.config_options), + })?; + + let output = output.into_array(rows)?; + if output.len() != rows { + return internal_err!( + "UDF {} returned {} rows for {} input rows", + self.name, + output.len(), + rows + ); + } + Ok(output) + } } impl fmt::Display for ScalarFunctionExpr { @@ -191,6 +522,8 @@ impl PartialEq for ScalarFunctionExpr { args, return_field, config_options, + peeled: _, // derived from the fields above + memoized: _, } = self; fun.eq(&o.fun) && name.eq(&o.name) @@ -210,6 +543,8 @@ impl Hash for ScalarFunctionExpr { args, return_field, config_options: _, // expensive to hash, and often equal + peeled: _, + memoized: _, } = self; fun.hash(state); name.hash(state); @@ -251,14 +586,17 @@ impl PhysicalExpr for ScalarFunctionExpr { .iter() .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); - // evaluate the function - let output = self.fun.invoke_with_args(ScalarFunctionArgs { - args, - arg_fields, - number_rows: batch.num_rows(), - return_field: Arc::clone(&self.return_field), - config_options: Arc::clone(&self.config_options), - })?; + // evaluate the function, over the distinct dictionary values when possible + let output = match self.try_invoke_peeled(&args, &arg_fields, batch.num_rows())? { + Some(output) => output, + None => self.fun.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: batch.num_rows(), + return_field: Arc::clone(&self.return_field), + config_options: Arc::clone(&self.config_options), + })?, + }; if let ColumnarValue::Array(array) = &output && array.len() != batch.num_rows() @@ -355,11 +693,16 @@ impl PhysicalExpr for ScalarFunctionExpr { #[cfg(test)] mod tests { + use super::dictionary::MEMOIZED_BYTES; use super::*; use crate::expressions::Column; - use arrow::datatypes::Field; + use arrow::array::DictionaryArray; + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{Field, Int32Type}; use datafusion_expr::{ScalarUDFImpl, Signature}; use datafusion_physical_expr_common::physical_expr::is_volatile; + use std::hash::{Hash, Hasher}; + use std::sync::Mutex; /// Test helper to create a mock UDF with a specific volatility #[derive(Debug, PartialEq, Eq, Hash)] @@ -385,6 +728,893 @@ mod tests { } } + /// Records the length of `args[0]` on every invocation, so tests can assert + /// how many values the function actually saw. + #[derive(Debug)] + struct ObservingUdf { + signature: Signature, + seen: Arc>>, + saw_dictionary: Arc, + elementwise: bool, + return_type: DataType, + strict: bool, + fail_if_contains: Option<&'static str>, + } + + impl PartialEq for ObservingUdf { + fn eq(&self, other: &Self) -> bool { + self.signature == other.signature + } + } + impl Eq for ObservingUdf {} + impl Hash for ObservingUdf { + fn hash(&self, state: &mut H) { + self.signature.hash(state); + } + } + + impl ScalarUDFImpl for ObservingUdf { + fn name(&self) -> &str { + "observing_udf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.return_type.clone()) + } + fn evaluates_elementwise(&self) -> bool { + self.elementwise + } + fn is_strict(&self) -> bool { + self.strict + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let len = match &args.args[0] { + ColumnarValue::Array(array) => { + if matches!(array.data_type(), DataType::Dictionary(_, _)) { + self.saw_dictionary + .store(true, std::sync::atomic::Ordering::Relaxed); + } + array.len() + } + ColumnarValue::Scalar(scalar) => { + if matches!(scalar, ScalarValue::Dictionary(_, _)) { + self.saw_dictionary + .store(true, std::sync::atomic::Ordering::Relaxed); + } + 1 + } + }; + self.seen.lock().unwrap().push(len); + if let (Some(needle), ColumnarValue::Array(array)) = + (self.fail_if_contains, &args.args[0]) + && let Some(strings) = array.as_any().downcast_ref::() + && strings.iter().flatten().any(|s| s == needle) + { + return datafusion_common::exec_err!("poisoned value {needle:?}"); + } + // The hand-rolled arm a dictionary-returning function keeps: when a + // declined peel hands over the dictionary itself, evaluate its + // values and re-wrap, so the planned type still comes back. + if let (ColumnarValue::Array(array), DataType::Dictionary(_, _)) = + (&args.args[0], &self.return_type) + && matches!(array.data_type(), DataType::Dictionary(_, _)) + { + let dictionary = array.as_any_dictionary(); + let converted: Int32Array = (0..dictionary.values().len()) + .map(|i| Some(i as i32)) + .collect(); + return Ok(ColumnarValue::Array( + dictionary.with_values(Arc::new(converted)), + )); + } + let values: Int32Array = (0..len).map(|i| Some(i as i32)).collect(); + Ok(ColumnarValue::Array(Arc::new(values))) + } + } + + struct PeelFixture { + expr: ScalarFunctionExpr, + batch: RecordBatch, + values: ArrayRef, + seen: Arc>>, + saw_dictionary: Arc, + } + + impl PeelFixture { + /// Another batch of the same column chunk: new keys, the very same + /// values array, which is what the memo keys on. + fn batch_over(&self, keys: Vec>) -> RecordBatch { + dictionary_batch(Int32Array::from(keys), &self.values) + } + } + + /// What a peeling test varies. Everything not named takes the default: + /// an immutable elementwise function returning `Int32` over the three + /// string values `ab`, `cd`, `ef`. + struct PeelSetup { + keys: Int32Array, + values: ArrayRef, + volatility: Volatility, + elementwise: bool, + return_type: DataType, + strict: bool, + fail_if_contains: Option<&'static str>, + } + + impl Default for PeelSetup { + fn default() -> Self { + Self { + keys: Int32Array::from(vec![0, 1, 0, 2, 1, 0, 2, 0]), + values: Arc::new(StringArray::from(vec!["ab", "cd", "ef"])), + volatility: Volatility::Immutable, + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + } + } + } + + impl PeelSetup { + /// The keys wrapped in a `ScalarFunctionExpr` over the values. + fn build(self) -> PeelFixture { + let seen = Arc::new(Mutex::new(Vec::new())); + let saw_dictionary = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, self.volatility), + seen: Arc::clone(&seen), + saw_dictionary: Arc::clone(&saw_dictionary), + elementwise: self.elementwise, + return_type: self.return_type.clone(), + strict: self.strict, + fail_if_contains: self.fail_if_contains, + })); + + let values = self.values; + let batch = dictionary_batch(self.keys, &values); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", self.return_type, true)), + Arc::new(ConfigOptions::new()), + ); + PeelFixture { + expr, + batch, + values, + seen, + saw_dictionary, + } + } + } + + /// A one-column batch holding `keys` over `values`. + fn dictionary_batch(keys: Int32Array, values: &ArrayRef) -> RecordBatch { + let dict = + DictionaryArray::::try_new(keys, Arc::clone(values)).unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap() + } + + fn keys_8_over_3() -> Int32Array { + Int32Array::from(vec![0, 1, 0, 2, 1, 0, 2, 0]) + } + + #[test] + fn peel_evaluates_once_per_distinct_value() { + let f = PeelSetup::default().build(); + let out = f.expr.evaluate(&f.batch).unwrap(); + + // Three distinct values seen instead of eight rows... + assert_eq!(*f.seen.lock().unwrap(), vec![3]); + // ...and the result is still one value per row, re-mapped through the keys. + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!( + array.as_primitive::().values(), + &[0, 1, 0, 2, 1, 0, 2, 0] + ); + } + + #[test] + fn peel_skipped_when_flag_is_off() { + let f = PeelSetup { + elementwise: false, + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![8]); + } + + #[test] + fn peel_skipped_for_volatile_functions() { + let f = PeelSetup { + volatility: Volatility::Volatile, + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![8]); + } + + #[test] + fn peel_normalizes_null_keys_to_a_null_slot() { + // Null keys are redirected to one appended NULL value: `f` sees the + // three referenced values plus that slot, and rows with null keys get + // `f(NULL)` (here: the slot's positional result), not an assumed NULL. + let keys = Int32Array::from(vec![ + Some(0), + None, + Some(1), + Some(2), + Some(0), + Some(1), + Some(2), + Some(0), + ]); + let f = PeelSetup { + keys, + ..Default::default() + } + .build(); + let out = f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![4]); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!( + array.as_primitive::().values(), + &[0, 3, 1, 2, 0, 1, 2, 0] + ); + } + + #[test] + fn peel_ignores_garbage_under_null_key_slots() { + // Arrow permits arbitrary key values under null slots; they must not + // count as references (a fallible `f` would fail on data no row holds) + // nor be read as indices. + use arrow::buffer::NullBuffer; + let keys = Int32Array::new( + vec![0, 1, 0, 0, 0, 0, 0, 0].into(), // garbage `1` under the null slot + Some(NullBuffer::from(vec![ + true, false, true, true, true, true, true, true, + ])), + ); + let f = PeelSetup { + keys, + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + // Only value slot 0 is live; +1 appended NULL slot. Slot 1 ("cd") is + // never evaluated despite the garbage key pointing at it. + assert_eq!(*f.seen.lock().unwrap(), vec![2]); + } + + #[test] + fn peel_reads_no_out_of_bounds_key_under_a_null_slot() { + // Same as above on the strict fast path, where the keys are re-mapped + // as they are: the garbage index points past the values array. + use arrow::buffer::NullBuffer; + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: true, + fail_if_contains: None, + })); + let keys = Int32Array::new( + vec![0, 9999, 1, 2, 0, 1, 2, 0].into(), + Some(NullBuffer::from(vec![ + true, false, true, true, true, true, true, true, + ])), + ); + let values = Arc::new(StringArray::from(vec!["ab", "cd", "ef"])); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + + let ColumnarValue::Array(array) = expr.evaluate(&batch).unwrap() else { + panic!("expected an array"); + }; + assert_eq!(*seen.lock().unwrap(), vec![3]); + assert!(array.is_null(1)); + } + + #[test] + fn peel_declines_when_the_null_slot_overflows_the_key_type() { + // Int8 keys address 128 values; a batch referencing all of them plus a + // null key has nowhere to put the appended NULL slot. + use arrow::datatypes::Int8Type; + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + })); + let values = Arc::new(StringArray::from( + (0..128).map(|i| format!("v{i}")).collect::>(), + )); + let mut keys: Vec> = (0..128).map(|i| Some(i as i8)).collect(); + keys.push(None); + keys.extend((0..128).map(|i| Some(i as i8))); + let rows = keys.len(); + let dict = + DictionaryArray::::try_new(keys.into_iter().collect(), values) + .unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + + // The batch is expanded rather than failed. + expr.evaluate(&batch).unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![rows]); + } + + #[test] + fn a_dictionary_return_declines_rather_than_expands_on_overflow() { + // The same overflow, but the planned type is a dictionary: a call over + // expanded rows would come back flat, so the column must be handed over + // as it arrived — exactly as if the function had not opted in. + use arrow::datatypes::Int8Type; + let return_type = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)); + let seen = Arc::new(Mutex::new(Vec::new())); + let saw_dictionary = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::clone(&saw_dictionary), + elementwise: true, + return_type: return_type.clone(), + strict: false, + fail_if_contains: None, + })); + let values = Arc::new(StringArray::from( + (0..128).map(|i| format!("v{i}")).collect::>(), + )); + let mut keys: Vec> = (0..128).map(|i| Some(i as i8)).collect(); + keys.push(None); + let rows = keys.len(); + let dict = + DictionaryArray::::try_new(keys.into_iter().collect(), values) + .unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", return_type.clone(), true)), + Arc::new(ConfigOptions::new()), + ); + + let out = expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!(array.data_type(), &return_type); + assert_eq!(array.len(), rows); + assert!(array.is_null(rows - 1)); + assert!(saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + assert_eq!(*seen.lock().unwrap(), vec![rows]); + } + + #[test] + fn peel_skipped_for_empty_batches() { + let f = PeelSetup { + keys: Int32Array::from(Vec::::new()), + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![0]); + } + + #[test] + fn peel_skipped_for_fields_with_metadata() { + // Metadata (e.g. extension types) binds to the field's storage type, + // which peeling would rewrite: such fields take the unpeeled path. + let mut f = PeelSetup::default().build(); + let metadata = std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "myorg.uuid".to_string(), + )]); + f.expr = ScalarFunctionExpr::new( + "observing_udf", + Arc::new(f.expr.fun().clone()), + f.expr.args().to_vec(), + Arc::new(Field::new("f", DataType::Int32, true).with_metadata(metadata)), + Arc::new(ConfigOptions::new()), + ); + f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![8]); + } + + #[test] + fn peel_fast_path_evaluates_raw_values() { + // No null keys: the fast tier invokes over the values as-is (three, + // including one unreferenced) with no discovery scan — the hand-rolled + // cost profile. Harmless for an infallible `f`. + let keys = Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 0]); + let f = PeelSetup { + keys, + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![3]); + } + + #[test] + fn peel_falls_back_to_compaction_when_an_unreferenced_value_errors() { + // The fast attempt sees the poisoned unreferenced value and errors; the + // guarded fallback compacts to the two referenced values and succeeds. + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: Some("ef"), + })); + let keys = Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 0]); + let values = Arc::new(StringArray::from(vec!["ab", "cd", "ef"])); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = + RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + let out = expr.evaluate(&batch).unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![3, 2]); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!(array.len(), 8); + } + + #[test] + fn peel_surfaces_an_error_from_a_referenced_value() { + // The counterpart of the test above: an error the query can actually + // observe is not swallowed by the retry. + let f = PeelSetup { + keys: Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 0]), + fail_if_contains: Some("cd"), + ..Default::default() + } + .build(); + + let error = f.expr.evaluate(&f.batch).unwrap_err().to_string(); + assert!(error.contains("poisoned value"), "{error}"); + } + + #[test] + fn peel_fast_path_keeps_null_keys_for_strict_functions() { + // `is_strict` promises NULL in -> NULL out, so null keys pass through + // untouched: the fast tier runs (three raw values, no null slot) and + // the null-key row stays NULL in the output. + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: true, + fail_if_contains: None, + })); + use arrow::buffer::NullBuffer; + let keys = Int32Array::new( + vec![0, 0, 1, 2, 0, 1, 2, 0].into(), + Some(NullBuffer::from(vec![ + true, false, true, true, true, true, true, true, + ])), + ); + let values = Arc::new(StringArray::from(vec!["ab", "cd", "ef"])); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = + RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + let out = expr.evaluate(&batch).unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![3]); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert!(array.is_null(1), "null-key row must stay NULL"); + } + + #[test] + fn unprofitable_batch_is_expanded_instead_of_peeled() { + // Two rows referencing two values: the gather would cost more than the + // two saved invocations, so the guard declines — and `f`, which opted + // in to elementwise evaluation, gets the expanded array rather than the + // dictionary it cannot be assumed to handle. + let f = PeelSetup { + keys: Int32Array::from(vec![0, 1]), + ..Default::default() + } + .build(); + let out = f.expr.evaluate(&f.batch).unwrap(); + + assert_eq!(*f.seen.lock().unwrap(), vec![2]); + assert!(!f.saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!(array.as_primitive::().values(), &[0, 1]); + } + + #[test] + fn values_are_evaluated_once_across_batches() { + // A batch carries its own keys but the dictionary of a whole column + // chunk: the second batch over the same values evaluates nothing. + let f = PeelSetup::default().build(); + + let keys = [ + vec![0, 1, 0, 2, 1, 0, 2, 0], + vec![2, 2, 1, 0, 1, 1, 0, 2], + vec![1, 0, 2, 2, 0, 1, 1, 0], + ]; + let outputs: Vec<_> = keys + .iter() + .map(|keys| { + let batch = f.batch_over(keys.iter().map(|k| Some(*k)).collect()); + match f.expr.evaluate(&batch).unwrap() { + ColumnarValue::Array(array) => array, + _ => panic!("expected an array"), + } + }) + .collect(); + + // The first batch cannot know the dictionary will come back; the second + // proves it and is what the result is kept from. The third is free. + assert_eq!(*f.seen.lock().unwrap(), vec![3, 3]); + for (output, keys) in outputs.iter().zip(&keys) { + assert_eq!(output.as_primitive::().values(), keys.as_slice()); + } + } + + #[test] + fn concurrent_hits_share_the_memo() { + // One expression is shared by every partition of a plan. Once a result + // is remembered, readers take it concurrently without evaluating — + // whatever the interleaving, the counter must not move. + let f = PeelSetup::default().build(); + let expr = Arc::new(f.expr); + let batch_over = + |keys: Vec| dictionary_batch(Int32Array::from(keys), &f.values); + + // Prove the dictionary repeats so its result is remembered. + expr.evaluate(&batch_over(vec![0, 1, 2])).unwrap(); + expr.evaluate(&batch_over(vec![2, 1, 0])).unwrap(); + assert_eq!(*f.seen.lock().unwrap(), vec![3, 3]); + + let handles: Vec<_> = (0..8) + .map(|t| { + let expr = Arc::clone(&expr); + let batch = batch_over(vec![t % 3, (t + 1) % 3, 0]); + std::thread::spawn(move || { + for _ in 0..200 { + let out = expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!( + array.as_primitive::().values(), + &[t % 3, (t + 1) % 3, 0] + ); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + // Every concurrent pass was a hit. + assert_eq!(*f.seen.lock().unwrap(), vec![3, 3]); + } + + #[test] + fn a_dictionary_too_large_to_keep_is_evaluated_each_time() { + // The memo holds its entries alive, so it is bounded by bytes as well + // as by count: values larger than the whole budget are never admitted, + // and every batch over them is evaluated — correctly, just not freely. + let big = "x".repeat(5 * 1024 * 1024); + let values = Arc::new(StringArray::from(vec![big.as_str(), "cd", "ef"])); + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + })); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + + let batch_over = |keys: Vec, values: &Arc| { + let dict = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(values) as ArrayRef, + ) + .unwrap(); + let schema = + Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap() + }; + + // The values hold five megabytes against a four megabyte budget: the + // repeat is recognised, but its result is never kept. + assert!(values.get_array_memory_size() > MEMOIZED_BYTES); + for keys in [ + vec![0, 1, 0, 2, 1, 0, 2, 0], + vec![2, 2, 1, 0, 1, 1, 0, 2], + vec![1, 0, 2, 2, 0, 1, 1, 0], + ] { + expr.evaluate(&batch_over(keys, &values)).unwrap(); + } + assert_eq!(*seen.lock().unwrap(), vec![3, 3, 3]); + } + + #[test] + fn different_values_are_evaluated_again() { + let seen = Arc::new(Mutex::new(Vec::new())); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::new(std::sync::atomic::AtomicBool::new(false)), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + })); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + + for values in [vec!["ab", "cd", "ef"], vec!["gh", "ij", "kl"]] { + let dict = DictionaryArray::::try_new( + keys_8_over_3(), + Arc::new(StringArray::from(values)), + ) + .unwrap(); + let schema = + Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap(); + expr.evaluate(&batch).unwrap(); + } + + // Same shape, different contents: the second dictionary is its own work. + assert_eq!(*seen.lock().unwrap(), vec![3, 3]); + } + + #[test] + fn peel_boundary_is_two_rows_per_referenced_value() { + // The flat-return bound is `values * 2 <= rows`: three values are + // peeled over six rows, and expanded over five. + let peeled = PeelSetup { + keys: Int32Array::from(vec![0, 1, 2, 0, 1, 2]), + ..Default::default() + } + .build(); + peeled.expr.evaluate(&peeled.batch).unwrap(); + assert_eq!(*peeled.seen.lock().unwrap(), vec![3]); + + let expanded = PeelSetup { + keys: Int32Array::from(vec![0, 1, 2, 0, 1]), + ..Default::default() + } + .build(); + expanded.expr.evaluate(&expanded.batch).unwrap(); + assert_eq!(*expanded.seen.lock().unwrap(), vec![5]); + } + + #[test] + fn functions_that_did_not_opt_in_still_receive_the_dictionary() { + let f = PeelSetup { + elementwise: false, + ..Default::default() + } + .build(); + f.expr.evaluate(&f.batch).unwrap(); + assert!(f.saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + } + + #[test] + fn peel_rewraps_when_the_planned_type_is_a_dictionary() { + let return_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)); + let f = PeelSetup { + return_type, + ..Default::default() + } + .build(); + let out = f.expr.evaluate(&f.batch).unwrap(); + + assert_eq!(*f.seen.lock().unwrap(), vec![3]); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + let dictionary = array.as_any_dictionary(); + assert_eq!(dictionary.len(), 8); + assert_eq!(dictionary.values().len(), 3); + } + + #[test] + fn two_dictionary_arguments_are_left_alone() { + // Peeling one argument holds the other rows fixed, which is only sound + // when there is exactly one; two dictionaries go through as they are. + let seen = Arc::new(Mutex::new(Vec::new())); + let saw_dictionary = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(2, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::clone(&saw_dictionary), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + })); + let values: ArrayRef = Arc::new(StringArray::from(vec!["ab", "cd", "ef"])); + let left = + DictionaryArray::::try_new(keys_8_over_3(), Arc::clone(&values)) + .unwrap(); + let right = + DictionaryArray::::try_new(keys_8_over_3(), values).unwrap(); + let schema = Schema::new(vec![ + Field::new("a", left.data_type().clone(), true), + Field::new("b", right.data_type().clone(), true), + ]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(left), Arc::new(right)]) + .unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![ + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("b", 1)) as Arc, + ], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + expr.evaluate(&batch).unwrap(); + assert!(saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + assert_eq!(*seen.lock().unwrap(), vec![8]); + } + + #[test] + fn a_dictionary_return_compacts_when_larger_than_the_batch() { + // A dictionary result re-wraps in O(1), so a batch referencing a slice + // of a large dictionary is still peeled — through compaction, which + // hands the function only the values the batch uses. + let return_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)); + let seen = Arc::new(Mutex::new(Vec::new())); + let saw_dictionary = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::clone(&saw_dictionary), + elementwise: true, + return_type: return_type.clone(), + strict: false, + fail_if_contains: None, + })); + let values = Arc::new(StringArray::from( + (0..100).map(|i| format!("v{i}")).collect::>(), + )); + let keys = Int32Array::from(vec![7, 42, 7]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let schema = Schema::new(vec![Field::new("d", dict.data_type().clone(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dict)]).unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Column::new("d", 0)) as Arc], + Arc::new(Field::new("f", return_type.clone(), true)), + Arc::new(ConfigOptions::new()), + ); + let out = expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(array) = out else { + panic!("expected an array"); + }; + assert_eq!(array.data_type(), &return_type); + assert_eq!(array.len(), 3); + // Only the two referenced values reached the function. + assert_eq!(*seen.lock().unwrap(), vec![2]); + assert!(!saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + } + + #[test] + fn dictionary_encoded_scalars_are_unwrapped() { + // There is nothing to peel in a scalar, but its encoding is still one + // the function did not ask to see. + let seen = Arc::new(Mutex::new(Vec::new())); + let saw_dictionary = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let udf = Arc::new(ScalarUDF::from(ObservingUdf { + signature: Signature::any(1, Volatility::Immutable), + seen: Arc::clone(&seen), + saw_dictionary: Arc::clone(&saw_dictionary), + elementwise: true, + return_type: DataType::Int32, + strict: false, + fail_if_contains: None, + })); + let scalar = ScalarValue::Dictionary( + Box::new(DataType::Int32), + Box::new(ScalarValue::from("ab")), + ); + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let expr = ScalarFunctionExpr::new( + "observing_udf", + udf, + vec![Arc::new(Literal::new(scalar)) as Arc], + Arc::new(Field::new("f", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + ); + expr.evaluate(&batch).unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![1]); + assert!(!saw_dictionary.load(std::sync::atomic::Ordering::Relaxed)); + } + #[test] fn test_scalar_function_volatile_node() { // Create a volatile UDF diff --git a/datafusion/physical-expr/src/scalar_function/dictionary.rs b/datafusion/physical-expr/src/scalar_function/dictionary.rs new file mode 100644 index 0000000000000..965a0aba4b251 --- /dev/null +++ b/datafusion/physical-expr/src/scalar_function/dictionary.rs @@ -0,0 +1,391 @@ +// 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. + +//! The dictionary side of a peeled scalar-function call: the fields the call +//! is made with, the compaction that narrows a dictionary to what a batch +//! references, and the memo that carries results across the batches of a +//! column chunk. [`super::ScalarFunctionExpr`] owns the locks and the tier +//! decisions; this module owns the policy under them. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, DictionaryArray, PrimitiveArray, UInt64Array, new_null_array, +}; +use arrow::compute::{concat, take}; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, FieldRef}; +use arrow::downcast_dictionary_array; +use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion_expr::ColumnarValue; + +/// The two fields a peeled call needs, which depend only on the plan: the +/// dictionary argument re-typed to its values, and `return_field` with any +/// planned dictionary wrapper stripped. +#[derive(Debug, Clone)] +pub(super) struct PeeledFields { + pub(super) index: usize, + /// The argument field these were built from; a batch presenting a + /// different one rebuilds them rather than reusing these. + pub(super) source: FieldRef, + pub(super) argument: FieldRef, + pub(super) output: FieldRef, +} + +/// What a lookup found: the stored result, or the identity built on the way, +/// so the caller can hash it without reading the array again. +pub(super) enum Lookup { + Evaluated(ArrayRef), + Absent(Option), +} + +/// What an expression remembers about a dictionary it is handed. +pub(super) enum Recollection { + /// Never seen; now recorded, so a repeat is recognisable. + Unknown, + /// Seen before, so evaluating all of its values will pay for itself. + SeenBefore, + /// Already evaluated. + Evaluated(ArrayRef), +} + +/// How many dictionaries results are kept for. One expression is shared by +/// every partition of a plan, so a single slot would be evicted by whichever +/// partition ran last. +pub(super) const MEMOIZED_DICTIONARIES: usize = 8; + +/// How much one memo may hold alive, values and results together — sized to +/// admit a worst-case Parquet page dictionary (at most 1 MiB by default) with +/// its result. +pub(super) const MEMOIZED_BYTES: usize = 4 * 1024 * 1024; + +/// What an expression keeps between batches: results for the dictionaries it +/// has evaluated, and a cheap note of the ones it has only glimpsed. Both are +/// evicted in insertion order and only ever hold dictionaries that proved +/// they repeat, so an eviction costs one re-evaluation, not a wrong result. +#[derive(Debug, Default)] +pub(super) struct Memo { + evaluated: Vec, + /// What the entries hold alive, maintained so admission is O(1). + bytes: usize, + /// Hashes of dictionaries seen once. A false match only costs evaluating + /// a dictionary in full one batch early; results are matched exactly. + glimpsed: Vec, +} + +impl Memo { + /// The stored result for `values`, if any. + pub(super) fn find(&self, values: &ArrayRef, scalars: &[ScalarValue]) -> Lookup { + if self.evaluated.is_empty() { + return Lookup::Absent(None); + } + let known = ValuesIdentity::of(values); + let found = self.evaluated.iter().find(|entry| { + entry.scalars == scalars + && (Arc::ptr_eq(&entry.values, values) || entry.identity == known) + }); + match found { + Some(entry) => Lookup::Evaluated(Arc::clone(&entry.output)), + None => Lookup::Absent(Some(known)), + } + } + + /// Records a sighting; answers whether this dictionary was seen before. + pub(super) fn note(&mut self, hash: u64) -> Recollection { + if self.glimpsed.contains(&hash) { + return Recollection::SeenBefore; + } + if self.glimpsed.len() == MEMOIZED_DICTIONARIES { + self.glimpsed.remove(0); + } + self.glimpsed.push(hash); + Recollection::Unknown + } + + /// Keeps `output` for the next batch that arrives with the same values. + /// The memo holds its entries alive, so admission is bounded by bytes as + /// well as by count; what cannot fit is recomputed per batch instead. + pub(super) fn keep( + &mut self, + values: &ArrayRef, + scalars: &[ScalarValue], + output: &ArrayRef, + ) { + let bytes = values.get_array_memory_size() + output.get_array_memory_size(); + if bytes > MEMOIZED_BYTES { + return; + } + while self.evaluated.len() >= MEMOIZED_DICTIONARIES + || self.bytes + bytes > MEMOIZED_BYTES + { + let evicted = self.evaluated.remove(0); + self.bytes -= evicted.bytes; + } + self.bytes += bytes; + self.evaluated.push(Memoized { + values: Arc::clone(values), + identity: ValuesIdentity::of(values), + scalars: scalars.to_vec(), + output: Arc::clone(output), + bytes, + }); + } +} + +/// The result of evaluating a function over one dictionary's values. +#[derive(Debug)] +struct Memoized { + /// The values these results came from, kept alive so the addresses + /// [`ValuesIdentity`] compares cannot be reused by an unrelated array. + values: ArrayRef, + identity: ValuesIdentity, + /// The other arguments at the time; a different format string is a + /// different result. + scalars: Vec, + output: ArrayRef, + /// What this entry holds alive, counted once at admission. + bytes: usize, +} + +/// The memory an array views: two arrays that agree on all of this hold the +/// same elements. +#[derive(Debug, PartialEq)] +pub(super) struct ValuesIdentity { + len: usize, + offset: usize, + buffers: Vec<(usize, usize)>, + nulls: Option<(usize, usize, usize)>, +} + +impl ValuesIdentity { + fn of(array: &ArrayRef) -> Self { + let data = array.to_data(); + Self { + len: data.len(), + offset: data.offset(), + buffers: data + .buffers() + .iter() + .map(|buffer| (buffer.as_ptr() as usize, buffer.len())) + .collect(), + nulls: data.nulls().map(|nulls| { + ( + nulls.buffer().as_ptr() as usize, + nulls.offset(), + nulls.len(), + ) + }), + } + } + + /// This identity as a hash, once it has been built anyway. + pub(super) fn hash(&self) -> u64 { + use std::hash::{DefaultHasher, Hasher}; + let mut hasher = DefaultHasher::new(); + hasher.write_usize(self.len); + hasher.write_usize(self.offset); + for (address, len) in &self.buffers { + hasher.write_usize(*address); + hasher.write_usize(*len); + } + if let Some((address, offset, _)) = self.nulls { + hasher.write_usize(address); + hasher.write_usize(offset); + } + hasher.finish() + } + + /// The same hash without building the identity, which allocates. + pub(super) fn hash_of(array: &ArrayRef) -> u64 { + use std::hash::{DefaultHasher, Hasher}; + let data = array.to_data(); + let mut hasher = DefaultHasher::new(); + hasher.write_usize(data.len()); + hasher.write_usize(data.offset()); + for buffer in data.buffers() { + hasher.write_usize(buffer.as_ptr() as usize); + hasher.write_usize(buffer.len()); + } + if let Some(nulls) = data.nulls() { + hasher.write_usize(nulls.buffer().as_ptr() as usize); + hasher.write_usize(nulls.offset()); + } + hasher.finish() + } +} + +/// Rewrites a dictionary so its values are exactly the ones referenced by this +/// batch: unreferenced values are dropped, and null keys are redirected to one +/// appended NULL value slot. Returns `None` when the batch references more +/// than half of `row_budget` distinct values, i.e. peeling would not pay off; +/// `None` as a budget disables that check. +/// +/// [`DictionaryArray::occupancy`] answers the same question, but always scans +/// every key; this pass abandons high-cardinality batches part-way. +pub(super) fn compact_dictionary( + array: &ArrayRef, + row_budget: Option, +) -> Result> { + fn rebuild( + dictionary: &DictionaryArray, + row_budget: Option, + ) -> Result> { + let values = dictionary.values(); + let keys = dictionary.keys(); + + // Discovery over a bitmap: everything here is sized by the dictionary, + // and nothing is built until the batch is known to be worth compacting. + let mut referenced = vec![0u64; values.len().div_ceil(64)]; + let mut referenced_count = 0usize; + let mut null_slots = 0usize; + let mark = |key: usize, referenced: &mut [u64], count: &mut usize| { + let bit = 1u64 << (key % 64); + let word = &mut referenced[key / 64]; + let fresh = *word & bit == 0; + *word |= bit; + *count += usize::from(fresh); + fresh + }; + if keys.null_count() == 0 { + // Hot path: raw key slice, no per-key validity checks. + for key in keys.values() { + if mark(key.as_usize(), &mut referenced, &mut referenced_count) + && row_budget.is_some_and(|rows| referenced_count * 2 > rows) + { + return Ok(None); + } + } + } else { + for key in keys.iter() { + let grew = match key { + None if null_slots == 0 => { + null_slots = 1; + true + } + None => false, + Some(key) => { + mark(key.as_usize(), &mut referenced, &mut referenced_count) + } + }; + if grew + && row_budget + .is_some_and(|rows| (referenced_count + null_slots) * 2 > rows) + { + return Ok(None); + } + } + } + if referenced_count == values.len() && null_slots == 0 { + return Ok(Some(Arc::new(dictionary.clone()))); + } + + // A value's new position is the number of referenced values before it, + // counted from the bitmap. Null keys go to the appended NULL slot, + // which sits one past the referenced values and can overflow a narrow + // key type — such a batch is left unpeeled rather than failed. + let null_slot = K::Native::from_usize(referenced_count); + if null_slots > 0 && null_slot.is_none() { + return Ok(None); + } + let null_slot = null_slot.unwrap_or_default(); + let mut preceding = Vec::with_capacity(referenced.len()); + let mut compacted_indices: Vec = Vec::with_capacity(referenced_count); + for (index, word) in referenced.iter().enumerate() { + preceding.push(compacted_indices.len()); + let mut bits = *word; + while bits != 0 { + let bit = bits.trailing_zeros() as usize; + compacted_indices.push((index * 64 + bit) as u64); + bits &= bits - 1; + } + } + let position = |key: usize| { + let before = referenced[key / 64] & ((1u64 << (key % 64)) - 1); + // In range: compacted positions only ever shrink. + K::Native::from_usize(preceding[key / 64] + before.count_ones() as usize) + .unwrap_or_default() + }; + // Garbage under a null key is not a valid position, so those batches + // take the checked path. + let new_keys: PrimitiveArray = if keys.null_count() == 0 { + keys.unary(|key| position(key.as_usize())) + } else { + PrimitiveArray::from_iter_values( + keys.iter() + .map(|key| key.map_or(null_slot, |key| position(key.as_usize()))), + ) + }; + + let compacted = + take(values.as_ref(), &UInt64Array::from(compacted_indices), None)?; + let new_values = if null_slots > 0 { + concat(&[ + compacted.as_ref(), + new_null_array(values.data_type(), 1).as_ref(), + ])? + } else { + compacted + }; + Ok(Some(Arc::new(DictionaryArray::::try_new( + new_keys, new_values, + )?))) + } + + downcast_dictionary_array!( + array => rebuild(array, row_budget), + other => internal_err!("expected a dictionary array, got {other:?}") + ) +} + +/// The scalar arguments a memoized result was computed with; a different trim +/// set or format string makes it a different result. +pub(super) fn scalar_arguments( + args: &[ColumnarValue], + dictionary_index: usize, +) -> Vec { + args.iter() + .enumerate() + .filter_map(|(index, arg)| match arg { + ColumnarValue::Scalar(scalar) if index != dictionary_index => { + Some(scalar.clone()) + } + _ => None, + }) + .collect() +} + +/// Replaces dictionary-encoded scalars, and their fields' data types, with the +/// single value each of them wraps. +pub(super) fn unwrap_scalar_dictionaries( + args: &[ColumnarValue], + arg_fields: &[FieldRef], +) -> (Vec, Vec) { + let mut args = args.to_vec(); + let mut arg_fields = arg_fields.to_vec(); + for (index, arg) in args.iter_mut().enumerate() { + if let ColumnarValue::Scalar(ScalarValue::Dictionary(_, value)) = arg { + let value = value.as_ref().clone(); + arg_fields[index] = Arc::new( + arg_fields[index] + .as_ref() + .clone() + .with_data_type(value.data_type()), + ); + *arg = ColumnarValue::Scalar(value); + } + } + (args, arg_fields) +} diff --git a/datafusion/sqllogictest/test_files/expr.slt b/datafusion/sqllogictest/test_files/expr.slt index 32113890aadc0..0edfcef564722 100644 --- a/datafusion/sqllogictest/test_files/expr.slt +++ b/datafusion/sqllogictest/test_files/expr.slt @@ -1459,6 +1459,36 @@ SELECT encode(arrow_cast('tom', 'Binary'),'base64'); ---- dG9t +# a dictionary column reaches encode encoded, so repeated values are encoded +# once; the plan shows no cast away from the dictionary on the way in +query TT +EXPLAIN SELECT encode(column1, 'base64') FROM ( + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS column1 + FROM (VALUES ('tom'), ('tom'), ('jerry'), (NULL), ('tom'), ('jerry')) AS t(column1) +); +---- +logical_plan +01)Projection: encode(CAST(CAST(t.column1 AS Dictionary(Int32, Utf8)) AS Dictionary(Int32, BinaryView)), Utf8("base64")) AS encode(column1,Utf8("base64")) +02)--SubqueryAlias: t +03)----Projection: column1 AS column1 +04)------Values: (Utf8("tom")), (Utf8("tom")), (Utf8("jerry")), (Utf8(NULL)), (Utf8("tom"))... +physical_plan +01)ProjectionExec: expr=[encode(CAST(CAST(column1@0 AS Dictionary(Int32, Utf8)) AS Dictionary(Int32, BinaryView)), base64) as encode(column1,Utf8("base64"))] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query T +SELECT encode(column1, 'base64') FROM ( + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS column1 + FROM (VALUES ('tom'), ('tom'), ('jerry'), (NULL), ('tom'), ('jerry')) AS t(column1) +); +---- +dG9t +dG9t +amVycnk +NULL +dG9t +amVycnk + query T SELECT arrow_cast(decode(arrow_cast('dG9t', 'Binary'),'base64'), 'Utf8'); ---- diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 008be05852c85..6920963c51eb2 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -586,6 +586,30 @@ NULL NULL Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) statement ok DROP TABLE unicode_dictionary_test +# A batch referencing all 128 values an Int8 key can address, plus a null key: +# peeling has nowhere to put its NULL slot and must hand the dictionary to the +# function as it arrived, without changing the type the plan promised. +statement ok +CREATE TABLE narrow_key_dictionary_test AS +SELECT arrow_cast(CASE WHEN v = 128 THEN NULL ELSE 'value' || CAST(v AS VARCHAR) END, + 'Dictionary(Int8, Utf8)') AS dict_col +FROM generate_series(0, 128) AS t(v); + +query II +SELECT count(*), count(reverse(dict_col)) FROM narrow_key_dictionary_test +---- +129 128 + +query TT +SELECT reverse(dict_col), arrow_typeof(reverse(dict_col)) +FROM narrow_key_dictionary_test +WHERE dict_col = 'value42' +---- +24eulav Dictionary(Int8, Utf8) + +statement ok +DROP TABLE narrow_key_dictionary_test + query ? SELECT ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)')) ----