diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..58c4d02d9f394 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -98,6 +98,11 @@ env_logger = { workspace = true } rand = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync"] } +[[bench]] +harness = false +name = "round_dense" +required-features = ["math_expressions"] + [[bench]] harness = false name = "ascii" diff --git a/datafusion/functions/benches/date_trunc.rs b/datafusion/functions/benches/date_trunc.rs index 0668a1cc5085c..e2372fff2a02e 100644 --- a/datafusion/functions/benches/date_trunc.rs +++ b/datafusion/functions/benches/date_trunc.rs @@ -18,52 +18,64 @@ use std::hint::black_box; use std::sync::Arc; -use arrow::array::{Array, ArrayRef, TimestampSecondArray}; +use arrow::array::{Array, ArrayRef, TimestampNanosecondArray, TimestampSecondArray}; use arrow::datatypes::Field; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs}; use datafusion_functions::datetime::date_trunc; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; -fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { - let mut seconds = vec![]; - for _ in 0..1000 { - seconds.push(rng.random_range(0..1_000_000)); - } +const NUM_ROWS: usize = 1000; +const NANOS_PER_SECOND: i64 = 1_000_000_000; +/// Roughly 30 years, so that values span many months, quarters and years. +const RANGE_SECONDS: i64 = 30 * 365 * 24 * 60 * 60; - TimestampSecondArray::from(seconds) +fn seedable_rng() -> StdRng { + StdRng::seed_from_u64(42) } -fn criterion_benchmark(c: &mut Criterion) { - c.bench_function("date_trunc_minute_1000", |b| { - let mut rng = rand::rng(); - let timestamps_array = Arc::new(timestamps(&mut rng)) as ArrayRef; - let batch_len = timestamps_array.len(); - let precision = - ColumnarValue::Scalar(ScalarValue::Utf8(Some("minute".to_string()))); - let timestamps = ColumnarValue::Array(timestamps_array); - let udf = date_trunc(); - let args = vec![precision, timestamps]; - let arg_fields = args - .iter() - .enumerate() - .map(|(idx, arg)| { - Field::new(format!("arg_{idx}"), arg.data_type(), true).into() - }) - .collect::>(); +fn second_timestamps() -> TimestampSecondArray { + let mut rng = seedable_rng(); + (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..1_000_000i64))) + .collect() +} - let scalar_arguments = vec![None; arg_fields.len()]; - let return_field = udf - .return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &scalar_arguments, - }) - .unwrap(); - let config_options = Arc::new(ConfigOptions::default()); +fn nanosecond_timestamps() -> TimestampNanosecondArray { + let mut rng = seedable_rng(); + (0..NUM_ROWS) + .map(|_| { + let seconds = rng.random_range(-RANGE_SECONDS..RANGE_SECONDS); + Some(seconds * NANOS_PER_SECOND + rng.random_range(0..NANOS_PER_SECOND)) + }) + .collect() +} + +fn run_benchmark(c: &mut Criterion, name: &str, granularity: &str, array: ArrayRef) { + let batch_len = array.len(); + let precision = + ColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string()))); + let udf = date_trunc(); + let args = vec![precision, ColumnarValue::Array(array)]; + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) + .collect::>(); + let scalar_arguments = vec![None; arg_fields.len()]; + let return_field = udf + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + }) + .unwrap(); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { b.iter(|| { black_box( udf.invoke_with_args(ScalarFunctionArgs { @@ -79,5 +91,23 @@ fn criterion_benchmark(c: &mut Criterion) { }); } +fn criterion_benchmark(c: &mut Criterion) { + let seconds: ArrayRef = Arc::new(second_timestamps()); + run_benchmark(c, "date_trunc_minute_1000", "minute", Arc::clone(&seconds)); + run_benchmark(c, "date_trunc_month_second_1000", "month", seconds); + + // Coarse granularities on an untimezoned array: these need calendar + // arithmetic rather than a plain division. + let nanos: ArrayRef = Arc::new(nanosecond_timestamps()); + for granularity in ["week", "month", "quarter", "year"] { + run_benchmark( + c, + &format!("date_trunc_{granularity}_nanos_1000"), + granularity, + Arc::clone(&nanos), + ); + } +} + criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/datafusion/functions/benches/round_dense.rs b/datafusion/functions/benches/round_dense.rs new file mode 100644 index 0000000000000..2c37849bde489 --- /dev/null +++ b/datafusion/functions/benches/round_dense.rs @@ -0,0 +1,94 @@ +// 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. + +//! Microbenchmark for `round(float_array, scalar_decimal_places)` over a +//! Float column with no NULLs — the dense elementwise-rounding path. + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field, Float32Type, Float64Type}; +use arrow::util::bench_util::create_primitive_array; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::math::round::RoundFunc; +use std::hint::black_box; +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let round_fn = RoundFunc::new(); + let config_options = Arc::new(ConfigOptions::default()); + + for size in [1024usize, 4096, 8192] { + // Float64, no nulls. + let f64_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f64_args = vec![ + ColumnarValue::Array(Arc::clone(&f64_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f64", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float64, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float64, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + + // Float32, no nulls. + let f32_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f32_args = vec![ + ColumnarValue::Array(Arc::clone(&f32_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f32", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float32, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float32, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/datetime/date_part.rs b/datafusion/functions/src/datetime/date_part.rs index 80f6bc0f66a70..e3f67db905615 100644 --- a/datafusion/functions/src/datetime/date_part.rs +++ b/datafusion/functions/src/datetime/date_part.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::iter::repeat_n; use std::str::FromStr; use std::sync::Arc; @@ -390,7 +391,7 @@ fn part_normalization(part: &str) -> &str { /// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the /// result to a total number of seconds, milliseconds, microseconds or -/// nanoseconds +/// nanoseconds as an `Int32Array` fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result { // Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing // with overflow and precision issue we don't support nanosecond @@ -398,6 +399,19 @@ fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result { return not_impl_err!("Date part {unit:?} not supported"); } + // Fast path with seconds - no need to compute nanoseconds + if unit == Second { + return Ok(date_part(array, DatePart::Second)?); + } + + // Fast path for Date32 and Date64 - no seconds + if array.data_type() == &Date32 || array.data_type() == &Date64 { + return Ok(Arc::new(Int32Array::from_iter_values_with_nulls( + repeat_n(0, array.len()), + array.nulls().cloned(), + ))); + } + let conversion_factor = match unit { Second => 1_000_000_000, Millisecond => 1_000_000, @@ -539,6 +553,14 @@ fn epoch(array: &dyn Array) -> Result { /// `nanosecond`s in each second, so representing up to 60 seconds as /// nanoseconds can be values up to 60 billion, which does not fit in Int32. fn seconds_ns(array: &dyn Array) -> Result { + // Fast path for Date32 and Date64 - no nanoseconds + if array.data_type() == &Date32 || array.data_type() == &Date64 { + return Ok(Arc::new(Int64Array::from_iter_values_with_nulls( + repeat_n(0, array.len()), + array.nulls().cloned(), + ))); + } + let secs = date_part(array, DatePart::Second)?; // This assumes array is primitive and not a dictionary let secs = as_int32_array(secs.as_ref())?; diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index a4b244405cc22..6dcd7a666d0a6 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -23,7 +23,6 @@ use std::sync::Arc; use arrow::array::temporal_conversions::{ MICROSECONDS, MILLISECONDS, NANOSECONDS, as_datetime_with_timezone, - timestamp_ns_to_datetime, }; use arrow::array::timezone::Tz; use arrow::array::types::{ @@ -462,6 +461,7 @@ const NANOS_PER_MILLISECOND: i64 = NANOSECONDS / MILLISECONDS; const NANOS_PER_SECOND: i64 = NANOSECONDS; const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; const NANOS_PER_HOUR: i64 = 60 * NANOS_PER_MINUTE; +const NANOS_PER_DAY: i64 = 24 * NANOS_PER_HOUR; const MICROS_PER_MILLISECOND: i64 = MICROSECONDS / MILLISECONDS; const MICROS_PER_SECOND: i64 = MICROSECONDS; @@ -591,52 +591,143 @@ where fn _date_trunc_coarse_with_tz( granularity: DateTruncGranularity, - value: Option>, + value: DateTime, ) -> Result> { - if let Some(value) = value { - let local = value.naive_local(); - let truncated = _date_trunc_coarse::(granularity, Some(local))?; - let truncated = truncated.and_then(|truncated| { - match truncated.and_local_timezone(value.timezone()) { - LocalResult::None => { - // This can happen if the date_trunc operation moves the time into - // an hour that doesn't exist due to daylight savings. On known example where - // this can happen is with historic dates in the America/Sao_Paulo time zone. - // To account for this adjust the time by a few hours, convert to local time, - // and then adjust the time back. - truncated - .sub(TimeDelta::try_hours(3).unwrap()) - .and_local_timezone(value.timezone()) - .single() - .map(|v| v.add(TimeDelta::try_hours(3).unwrap())) - } - LocalResult::Single(datetime) => Some(datetime), - LocalResult::Ambiguous(datetime1, datetime2) => { - // Because we are truncating from an equally or more specific time - // the original time must have been within the ambiguous local time - // period. Therefore the offset of one of these times should match the - // offset of the original time. - if datetime1.offset().fix() == value.offset().fix() { - Some(datetime1) - } else { - Some(datetime2) - } + let local = value.naive_local(); + let truncated = _date_trunc_coarse::(granularity, Some(local))?; + let truncated = truncated.and_then(|truncated| { + match truncated.and_local_timezone(value.timezone()) { + LocalResult::None => { + // This can happen if the date_trunc operation moves the time into + // an hour that doesn't exist due to daylight savings. On known example where + // this can happen is with historic dates in the America/Sao_Paulo time zone. + // To account for this adjust the time by a few hours, convert to local time, + // and then adjust the time back. + truncated + .sub(TimeDelta::try_hours(3).unwrap()) + .and_local_timezone(value.timezone()) + .single() + .map(|v| v.add(TimeDelta::try_hours(3).unwrap())) + } + LocalResult::Single(datetime) => Some(datetime), + LocalResult::Ambiguous(datetime1, datetime2) => { + // Because we are truncating from an equally or more specific time + // the original time must have been within the ambiguous local time + // period. Therefore the offset of one of these times should match the + // offset of the original time. + if datetime1.offset().fix() == value.offset().fix() { + Some(datetime1) + } else { + Some(datetime2) } } - }); - Ok(truncated.and_then(|value| value.timestamp_nanos_opt())) + } + }); + Ok(truncated.and_then(|value| value.timestamp_nanos_opt())) +} + +// The two helpers below duplicate `chrono::NaiveDate::{from_epoch_days, +// to_epoch_days}`. They are kept separate because chrono's versions round trip +// through a validated `NaiveDate`: `from_epoch_days` computes year flags and +// returns an `Option`, and reading the year/month/day back out decodes them from +// its packed representation. These helpers stay in plain integers, which is all +// the truncation below needs. + +/// Days from the Unix epoch to 0000-03-01, the epoch used by the civil calendar +/// conversions below. +const DAYS_EPOCH_SHIFT: i64 = 719_468; + +/// Days in a 400 year era of the proleptic Gregorian calendar. +const DAYS_PER_ERA: i64 = 146_097; + +/// Splits a day count relative to the Unix epoch into a proleptic Gregorian +/// year, month (1-12) and day of month (1-31). +/// +/// This is a port of Howard Hinnant's `civil_from_days`, which documents the +/// derivation of the constants and the March-based year used below: +/// +fn civil_from_days(days: i64) -> (i64, i64, i64) { + let z = days + DAYS_EPOCH_SHIFT; + let era = z.div_euclid(DAYS_PER_ERA); + let day_of_era = z.rem_euclid(DAYS_PER_ERA); + let year_of_era = (day_of_era - day_of_era / 1460 + day_of_era / 36524 + - day_of_era / 146_096) + / 365; + let day_of_year = + day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + // Month index with March as 0, so that the leap day falls at the end of the year. + let month_index = (5 * day_of_year + 2) / 153; + let day = day_of_year - (153 * month_index + 2) / 5 + 1; + let month = if month_index < 10 { + month_index + 3 } else { - _date_trunc_coarse::(granularity, None)?; - Ok(None) - } + month_index - 9 + }; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +/// Inverse of [`civil_from_days`]: the day count relative to the Unix epoch for +/// the given proleptic Gregorian date. +/// +/// This is a port of Howard Hinnant's `days_from_civil`, which documents the +/// derivation of the constants and the March-based year used below: +/// +fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { + let year = year - i64::from(month <= 2); + let era = year.div_euclid(400); + let year_of_era = year.rem_euclid(400); + let month_index = if month > 2 { month - 3 } else { month + 9 }; + let day_of_year = (153 * month_index + 2) / 5 + day - 1; + let day_of_era = + year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + era * DAYS_PER_ERA + day_of_era - DAYS_EPOCH_SHIFT } +/// Truncates a UTC nanosecond timestamp with integer arithmetic. Truncating on +/// the calendar directly avoids converting every value to a `NaiveDateTime` and +/// rebuilding it field by field. +/// +/// Returns `None` when the truncated timestamp is no longer representable as +/// nanoseconds since the epoch, which the caller reports as an out of range +/// error. fn _date_trunc_coarse_without_tz( granularity: DateTruncGranularity, - value: Option, -) -> Result> { - let value = _date_trunc_coarse::(granularity, value)?; - Ok(value.and_then(|value| value.and_utc().timestamp_nanos_opt())) + value: i64, +) -> Option { + let truncate_to = |unit: i64| value.checked_sub(value.rem_euclid(unit)); + let days = || value.div_euclid(NANOS_PER_DAY); + let nanos_from_days = |days: i64| days.checked_mul(NANOS_PER_DAY); + + match granularity { + // Sub-second granularities are applied by the caller, which rescales + // the nanoseconds to the time unit of the array. + DateTruncGranularity::Millisecond | DateTruncGranularity::Microsecond => { + Some(value) + } + DateTruncGranularity::Second => truncate_to(NANOS_PER_SECOND), + DateTruncGranularity::Minute => truncate_to(NANOS_PER_MINUTE), + DateTruncGranularity::Hour => truncate_to(NANOS_PER_HOUR), + DateTruncGranularity::Day => nanos_from_days(days()), + DateTruncGranularity::Week => { + let days = days(); + // `Weekday::num_days_from_monday` for the epoch (a Thursday) is 3. + nanos_from_days(days - (days + 3).rem_euclid(7)) + } + DateTruncGranularity::Month => { + let days = days(); + let (_, _, day_of_month) = civil_from_days(days); + nanos_from_days(days - (day_of_month - 1)) + } + DateTruncGranularity::Quarter => { + let (year, month, _) = civil_from_days(days()); + nanos_from_days(days_from_civil(year, 1 + 3 * ((month - 1) / 3), 1)) + } + DateTruncGranularity::Year => { + let (year, _, _) = civil_from_days(days()); + nanos_from_days(days_from_civil(year, 1, 1)) + } + } } /// Truncates the single `value`, expressed in nanoseconds since the @@ -655,15 +746,10 @@ fn date_trunc_coarse( // and NaiveDateTime (ISO 8601) has no concept of timezones let value = as_datetime_with_timezone::(value, tz) .ok_or(exec_datafusion_err!("Timestamp {value} out of range"))?; - _date_trunc_coarse_with_tz(granularity, Some(value)) - } - None => { - // Use chrono NaiveDateTime to clear the various fields, if we don't have a timezone. - let value = timestamp_ns_to_datetime(value) - .ok_or_else(|| exec_datafusion_err!("Timestamp {value} out of range"))?; - _date_trunc_coarse_without_tz(granularity, Some(value)) + _date_trunc_coarse_with_tz(granularity, value)? } - }?; + None => _date_trunc_coarse_without_tz(granularity, value), + }; value.ok_or_else(|| { exec_datafusion_err!( diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs index 62f1c3540b9ce..49385d087af0d 100644 --- a/datafusion/functions/src/math/round.rs +++ b/datafusion/functions/src/math/round.rs @@ -23,9 +23,9 @@ use arrow::datatypes::DataType::{ Int64, UInt8, UInt16, UInt32, UInt64, }; use arrow::datatypes::{ - ArrowNativeTypeOp, DataType, Decimal32Type, Decimal64Type, Decimal128Type, - Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, - Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + ArrowNativeTypeOp, ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use arrow::datatypes::{Field, FieldRef}; use arrow::error::ArrowError; @@ -495,22 +495,8 @@ fn round_columnar( )?, } } - (Float64, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ - } - (Float32, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ - } + (Float64, _) => round_float_column::(&value_array, decimal_places)?, + (Float32, _) => round_float_column::(&value_array, decimal_places)?, (Decimal32(input_precision, scale), Decimal32(precision, new_scale)) => { // reduce scale to reclaim integer precision let result = calculate_binary_decimal_math_cast::< @@ -859,15 +845,59 @@ fn round_integer_array( } } -fn round_float(value: T, decimal_places: i32) -> Result +/// Rounds a float array to `decimal_places`. +/// +/// The shared `calculate_binary_math` kernel routes through `try_unary` and +/// re-evaluates `round_float` (including `10f64.powi(decimal_places)` and a +/// `Result` check) for every element. When `decimal_places` is a non-null +/// scalar, the scaling factor can instead be hoisted out of the loop and the +/// infallible `unary` kernel used, which the compiler can autovectorize. +/// `unary` also computes over null slots, but it carries the input null buffer +/// through to the output, so those values stay masked. +fn round_float_column( + value_array: &ArrayRef, + decimal_places: &ColumnarValue, +) -> Result where - T: num_traits::Float, + PT: ArrowPrimitiveType, + PT::Native: num_traits::Float, { - let factor = T::from(10_f64.powi(decimal_places)).ok_or_else(|| { + // Bring `Float` into scope so `.round()` resolves on the `PT::Native` + // projection below. + use num_traits::Float; + + if let ColumnarValue::Scalar(ScalarValue::Int32(Some(decimal_places))) = + decimal_places + { + let factor = round_factor::(*decimal_places)?; + let result = value_array + .as_primitive::() + .unary::<_, PT>(|value| (value * factor).round() / factor); + return Ok(Arc::new(result) as ArrayRef); + } + + let result = calculate_binary_math::( + value_array.as_ref(), + decimal_places, + round_float::, + )?; + Ok(result as _) +} + +/// Computes the power-of-ten scaling factor used to round to `decimal_places`. +fn round_factor(decimal_places: i32) -> Result { + T::from(10_f64.powi(decimal_places)).ok_or_else(|| { ArrowError::ComputeError(format!( "Invalid value for decimal places: {decimal_places}" )) - })?; + }) +} + +fn round_float(value: T, decimal_places: i32) -> Result +where + T: num_traits::Float, +{ + let factor = round_factor::(decimal_places)?; Ok((value * factor).round() / factor) } @@ -957,6 +987,7 @@ mod test { use std::sync::Arc; use arrow::array::{ArrayRef, Float32Array, Float64Array, Int64Array}; + use arrow::datatypes::DataType; use datafusion_common::DataFusionError; use datafusion_common::ScalarValue; use datafusion_common::cast::{as_float32_array, as_float64_array}; @@ -1022,6 +1053,35 @@ mod test { assert_eq!(floats, &expected); } + /// A scalar `decimal_places` takes the hoisted-factor `unary` path, which + /// computes over null slots as well. The nulls must survive into the output. + #[test] + fn test_round_f64_scalar_decimal_places_preserves_nulls() { + let value: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(125.2345), + None, + Some(-1.555), + None, + ])); + + let result = super::round_columnar( + &ColumnarValue::Array(value), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + 4, + &DataType::Float64, + ) + .expect("failed to initialize function round"); + let ColumnarValue::Array(result) = result else { + panic!("expected an array result"); + }; + let floats = + as_float64_array(&result).expect("failed to initialize function round"); + + let expected = Float64Array::from(vec![Some(125.23), None, Some(-1.56), None]); + + assert_eq!(floats, &expected); + } + #[test] fn test_round_f32_one_input() { let args: Vec = vec![ diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 7a5d2573ec257..b46353b609f1a 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -27,7 +27,7 @@ use arrow::array::{ Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; -use arrow::buffer::{Buffer, ScalarBuffer}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; @@ -262,6 +262,65 @@ fn trim_and_append_view( } } +/// Builds the trimmed output array by writing the trimmed slices straight into +/// the value buffer, rather than collecting through a string builder. +/// +/// Every trimmed value is a substring of its input, so the byte range the input +/// spans bounds the output's. Reserving that much up front means one allocation +/// and no growth during the copy, and it also guarantees the running offset stays +/// within `T` (the input array's own offsets already fit). +/// +/// `nulls` becomes the output null buffer; null rows contribute no bytes. +/// `trim_row` is called only for non-null rows, with the row index and its value, +/// and must return a subslice of the value it is given. +fn build_trimmed( + string_array: &GenericStringArray, + nulls: Option, + mut trim_row: F, +) -> ArrayRef +where + F: for<'a> FnMut(usize, &'a str) -> &'a str, +{ + let len = string_array.len(); + let input_offsets = string_array.value_offsets(); + let start = input_offsets.first().unwrap().as_usize(); + let end = input_offsets.last().unwrap().as_usize(); + + let mut values: Vec = Vec::with_capacity(end - start); + let mut offsets: Vec = Vec::with_capacity(len + 1); + offsets.push(T::usize_as(0)); + + match &nulls { + // Keeping the null check out of the all-valid path leaves it branch-free. + None => { + for i in 0..len { + // SAFETY: `i` is in bounds. + let s = unsafe { string_array.value_unchecked(i) }; + values.extend_from_slice(trim_row(i, s).as_bytes()); + offsets.push(T::usize_as(values.len())); + } + } + Some(validity) => { + for i in 0..len { + if validity.is_valid(i) { + // SAFETY: `i` is in bounds. + let s = unsafe { string_array.value_unchecked(i) }; + values.extend_from_slice(trim_row(i, s).as_bytes()); + } + offsets.push(T::usize_as(values.len())); + } + } + } + + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + // SAFETY: trimming splits `s` on char boundaries, so the value buffer is a + // concatenation of valid UTF-8; the offsets are monotonic and end at its length. + let array = unsafe { + GenericStringArray::::new_unchecked(offsets, Buffer::from_vec(values), nulls) + }; + Arc::new(array) +} + /// Applies the trim function to the given string array(s) /// and returns a new string array with the trimmed values. /// @@ -273,12 +332,11 @@ fn string_trim(args: &[ArrayRef]) -> Result { // Trim spaces by default - let result = string_array - .iter() - .map(|string| string.map(|s| Tr::trim_ascii_char(s, b' ').0)) - .collect::>(); - - Ok(Arc::new(result) as ArrayRef) + Ok(build_trimmed( + string_array, + string_array.nulls().cloned(), + |_, s| Tr::trim_ascii_char(s, b' ').0, + )) } 2 => { let characters_array = as_generic_string_array::(&args[1])?; @@ -293,29 +351,31 @@ fn string_trim(args: &[ArrayRef]) -> Result = characters_array.value(0).chars().collect(); - let result = string_array - .iter() - .map(|item| item.map(|s| Tr::trim(s, &pattern).0)) - .collect::>(); - return Ok(Arc::new(result) as ArrayRef); + return Ok(build_trimmed( + string_array, + string_array.nulls().cloned(), + |_, s| Tr::trim(s, &pattern).0, + )); + } + + // Indexing `characters_array` per row below requires the two arguments + // to line up. + if characters_array.len() != string_array.len() { + return exec_err!( + "Function TRIM was called with mismatched argument lengths" + ); } + // A row is null if either argument is null. + let nulls = NullBuffer::union(string_array.nulls(), characters_array.nulls()); + // Per-row pattern - must compute pattern chars for each row let mut pattern: Vec = Vec::new(); - let result = string_array - .iter() - .zip(characters_array.iter()) - .map(|(string, characters)| match (string, characters) { - (Some(s), Some(c)) => { - pattern.clear(); - pattern.extend(c.chars()); - Some(Tr::trim(s, &pattern).0) - } - _ => None, - }) - .collect::>(); - - Ok(Arc::new(result) as ArrayRef) + Ok(build_trimmed(string_array, nulls, |i, s| { + pattern.clear(); + pattern.extend(characters_array.value(i).chars()); + Tr::trim(s, &pattern).0 + })) } other => { exec_err!( diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs index dbea192d4947d..f97aa2cd73974 100644 --- a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs +++ b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs @@ -35,6 +35,10 @@ use datafusion_physical_expr_common::physical_expr::DynHash; mod tracker; pub use tracker::{DynamicFilterTracker, DynamicFilterTracking}; +/// Per-generation cache of the remapped current expression for +/// [`DynamicFilterPhysicalExpr::current`]. See the field docs there. +type CurrentExprCache = Arc)>>>; + /// State of a dynamic filter, tracking both updates and completion. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum FilterState { @@ -62,7 +66,6 @@ impl FilterState { /// For more background, please also see the [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog] /// /// [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog]: https://datafusion.apache.org/blog/2025/09/10/dynamic-filters -#[derive(Debug)] pub struct DynamicFilterPhysicalExpr { /// The original children of this PhysicalExpr, if any. /// This is necessary because the dynamic filter may be initialized with a placeholder (e.g. `lit(true)`) @@ -72,6 +75,16 @@ pub struct DynamicFilterPhysicalExpr { /// If any of the children were remapped / modified (e.g. to adjust for projections) we need to keep track of the new children /// so that when we update `current()` in subsequent iterations we can re-apply the replacements. remapped_children: Option>>, + /// Cache of the last (generation, remapped-expression) pair returned by + /// [`Self::current`]. `current()` is hot on the per-batch RowFilter path; + /// when the inner generation hasn't changed (common — updates fire once + /// per HashJoin build or once per TopK threshold refresh, but `evaluate` + /// is called per batch), the cache serves the remapped expression + /// without re-running the `transform_up` tree walk in + /// [`Self::remap_children`]. Reset on `update()` (by generation bump) + /// and populated with `None` on `with_new_children` (each derived + /// filter owns its own cache). + current_cache: CurrentExprCache, /// The source of dynamic filters. inner: Arc>, /// Broadcasts filter state (updates and completion) to all waiters. @@ -83,6 +96,23 @@ pub struct DynamicFilterPhysicalExpr { nullable: Arc>>, } +impl std::fmt::Debug for DynamicFilterPhysicalExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Manual impl deliberately omits `current_cache`: it is a pure + // optimization artifact whose contents depend on whether + // `current()` has been called, and roundtrip tests (e.g. in + // `datafusion-proto`) compare `format!("{:?}", ..)` output. + f.debug_struct("DynamicFilterPhysicalExpr") + .field("children", &self.children) + .field("remapped_children", &self.remapped_children) + .field("inner", &self.inner) + .field("state_watch", &self.state_watch) + .field("data_type", &self.data_type) + .field("nullable", &self.nullable) + .finish() + } +} + /// Atomic internal state of a [`DynamicFilterPhysicalExpr`]. /// /// `expression_id` lives here because it identifies the actual filter expression `expr`. @@ -189,6 +219,7 @@ impl DynamicFilterPhysicalExpr { children, remapped_children: None, // Initially no remapped children inner: Arc::new(RwLock::new(Inner::new(inner))), + current_cache: Arc::new(RwLock::new(None)), state_watch, data_type: Arc::new(RwLock::new(None)), nullable: Arc::new(RwLock::new(None)), @@ -232,9 +263,48 @@ impl DynamicFilterPhysicalExpr { /// Get the current expression. /// This will return the current expression with any children /// remapped to match calls to [`PhysicalExpr::with_new_children`]. + /// + /// Called per batch on the RowFilter path (via + /// [`PhysicalExpr::evaluate`]). The remap walk is O(tree size) and, for + /// dynamic filters that carry a large `InListExpr` (join key IN list), + /// dominated by `InListExpr::with_new_children` cloning the whole list. + /// The inner generation only changes when [`Self::update`] fires, so we + /// cache the remapped expression per generation and return it directly + /// on subsequent per-batch calls. pub fn current(&self) -> Result> { - let expr = Arc::clone(self.inner.read().expr()); - Self::remap_children(&self.children, self.remapped_children.as_ref(), expr) + // Fast path: cache hit for the current generation. + let (expr, generation) = { + let inner = self.inner.read(); + (Arc::clone(inner.expr()), inner.generation) + }; + if let Some((cached_gen, cached_expr)) = self.current_cache.read().as_ref() + && *cached_gen == generation + { + return Ok(Arc::clone(cached_expr)); + } + // Slow path: (re)compute the remap and store it under a write lock. + let remapped = + Self::remap_children(&self.children, self.remapped_children.as_ref(), expr)?; + // Only publish our result if it is strictly newer than whatever is + // currently cached. Without this guard a slow computation that + // observed an older `inner` could clobber a newer entry that a + // concurrent caller has already published (see #23532 review), which + // would force subsequent readers to redo the remap for the newer + // generation. Same-generation writes are also skipped: the cached + // and about-to-write remaps are semantically identical (same input + // expression, same remapped_children), so overwriting is redundant + // and only wastes a write-lock take. + { + let mut cache = self.current_cache.write(); + let should_write = match cache.as_ref() { + Some((cached_gen, _)) => generation > *cached_gen, + None => true, + }; + if should_write { + *cache = Some((generation, Arc::clone(&remapped))); + } + } + Ok(remapped) } /// Update the current expression and notify all waiters. @@ -411,6 +481,7 @@ impl DynamicFilterPhysicalExpr { children, remapped_children, inner: Arc::new(RwLock::new(inner)), + current_cache: Arc::new(RwLock::new(None)), state_watch, data_type: Arc::new(RwLock::new(None)), nullable: Arc::new(RwLock::new(None)), @@ -436,6 +507,9 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr { remapped_children: Some(children), // Note: expression_id is preserved inner: Arc::clone(&self.inner), + // Fresh cache per derived filter — remap depends on this + // instance's `remapped_children`, which just changed. + current_cache: Arc::new(RwLock::new(None)), state_watch: self.state_watch.clone(), data_type: Arc::clone(&self.data_type), nullable: Arc::clone(&self.nullable), @@ -721,6 +795,25 @@ impl ExpressionIdAtomicCounter { /// file and be made public for other expressions to use. static EXPR_ID_SOURCE: ExpressionIdAtomicCounter = ExpressionIdAtomicCounter::new(); +#[cfg(test)] +impl DynamicFilterPhysicalExpr { + /// Test-only clone that produces a fresh outer instance sharing the + /// same `inner`. Used by the concurrent stress test to obtain a + /// standalone `Arc` without going through `with_new_children` + /// (which would clear `remapped_children`). + fn clone_with_remapped_children_for_test(&self) -> Self { + Self { + children: self.children.clone(), + remapped_children: self.remapped_children.clone(), + inner: Arc::clone(&self.inner), + current_cache: Arc::new(RwLock::new(None)), + state_watch: self.state_watch.clone(), + data_type: Arc::clone(&self.data_type), + nullable: Arc::clone(&self.nullable), + } + } +} + #[cfg(test)] mod test { use crate::{ @@ -1259,4 +1352,238 @@ mod test { "mark_complete() must not change expression_id", ); } + + /// Repeated `current()` at the same generation must return the exact same + /// `Arc` — the cache serves without re-running `remap_children`. + #[test] + fn test_current_cache_hits_within_generation() { + let table_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &table_schema).unwrap(); + let expr = Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(10) as Arc, + )); + // Force the remap path to actually run: give the filter a distinct + // `remapped_children`. Without this, `remap_children` returns the + // input Arc unchanged and every call is trivially pointer-equal. + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + expr as Arc, + )); + let remapped_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let derived = reassign_expr_columns( + Arc::clone(&filter) as Arc, + &remapped_schema, + ) + .unwrap(); + let derived = derived + .downcast_ref::() + .expect("derived filter must be a DynamicFilterPhysicalExpr"); + + // First call populates the cache. Second and third must return the + // *same* Arc — proving `remap_children` did not run again. + let first = derived.current().unwrap(); + let second = derived.current().unwrap(); + let third = derived.current().unwrap(); + assert!( + Arc::ptr_eq(&first, &second), + "current() should return the cached Arc within a generation", + ); + assert!(Arc::ptr_eq(&second, &third)); + } + + /// `update()` bumps the generation; the next `current()` must return a + /// fresh remapped expression, not the stale cached one. + #[test] + fn test_current_cache_invalidates_on_update() { + let table_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &table_schema).unwrap(); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(10) as Arc, + )); + // Remap to force the cache path. + let derived = reassign_expr_columns( + Arc::clone(&filter) as Arc, + &table_schema, + ) + .unwrap(); + let derived = derived + .downcast_ref::() + .expect("derived filter must be a DynamicFilterPhysicalExpr"); + + let before = derived.current().unwrap(); + // Bump the generation with a distinct expression. + filter + .update(Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(42) as Arc, + )) as Arc) + .unwrap(); + let after = derived.current().unwrap(); + assert!( + !Arc::ptr_eq(&before, &after), + "current() must return a fresh Arc after update() bumps the generation", + ); + assert_ne!(format!("{before:?}"), format!("{after:?}")); + } + + /// `with_new_children` produces a derived filter with its own cache slot; + /// populating one filter's cache must not leak into the other. + #[test] + fn test_current_cache_is_per_derived_filter() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + // Original expression references `a`. Each derived filter remaps `a` + // to a *different* column so remap_children returns distinct exprs + // per filter (and thus distinct cached Arcs). + let col_a = col("a", &schema).unwrap(); + let expr = Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(10) as Arc, + )); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + expr as Arc, + )); + + let d1 = Arc::clone(&filter) + .with_new_children(vec![col("b", &schema).unwrap()]) + .unwrap(); + let d2 = Arc::clone(&filter) + .with_new_children(vec![col("c", &schema).unwrap()]) + .unwrap(); + let d1 = d1.downcast_ref::().unwrap(); + let d2 = d2.downcast_ref::().unwrap(); + + let d1_first = d1.current().unwrap(); + let d2_first = d2.current().unwrap(); + // Distinct remap_children paths produce distinct cached Arcs. + assert!(!Arc::ptr_eq(&d1_first, &d2_first)); + assert_ne!(format!("{d1_first:?}"), format!("{d2_first:?}")); + // Subsequent calls each hit their own cache. + assert!(Arc::ptr_eq(&d1_first, &d1.current().unwrap())); + assert!(Arc::ptr_eq(&d2_first, &d2.current().unwrap())); + } + + /// Stress-test the cache under concurrent readers and periodic writes. + /// + /// Motivation: prod scans run with tens/hundreds of partitions, each + /// calling `current()` on the same `Arc` per + /// batch, while the producer (HashJoin build / TopK) fires `update()` + /// on a separate task. A caching bug that only shows up under + /// contention (torn read, ABA-style Arc lifetime issue, cache monotonicity + /// violation) would be invisible in single-threaded tests. This test + /// hot-loops many readers against a writer and asserts the invariants + /// that matter: + /// 1. `current()` never panics and always returns a valid `Arc`. + /// 2. Cache generation never regresses (monotonic non-decreasing). + /// 3. After the writer stops, the cache eventually converges to the + /// final `inner.generation`. + #[test] + fn test_current_cache_concurrent_readers_and_writer() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::thread; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &schema).unwrap(); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(true) as Arc, + )); + // Force the remap path: give the derived filter distinct + // remapped_children so `current()` actually runs `remap_children` + // instead of the short-circuit `Arc::clone(&expr)`. + let derived = + reassign_expr_columns(Arc::clone(&filter) as Arc, &schema) + .unwrap(); + // Re-wrap in Arc for cross-thread sharing. + let derived: Arc = Arc::new( + derived + .downcast_ref::() + .expect("derived is DynamicFilterPhysicalExpr") + .clone_with_remapped_children_for_test(), + ); + + let stop = Arc::new(AtomicBool::new(false)); + const READERS: usize = 8; + const READER_ITERS: usize = 5_000; + const WRITER_ITERS: i32 = 200; + + let mut readers = Vec::with_capacity(READERS); + for _ in 0..READERS { + let d = Arc::clone(&derived); + let stop = Arc::clone(&stop); + readers.push(thread::spawn(move || { + let mut last_seen_gen: u64 = 0; + for _ in 0..READER_ITERS { + if stop.load(Ordering::Relaxed) { + break; + } + let expr = d.current().expect("current must not fail"); + // Cheap sanity: the returned Arc's Debug must be + // formattable — proves it's a valid PhysicalExpr. + let _ = format!("{expr:?}"); + // Cache generation observed by this reader must never + // decrease across successive calls on the same filter. + let cached = d + .current_cache + .read() + .as_ref() + .map(|(g, _)| *g) + .unwrap_or(0); + assert!( + cached >= last_seen_gen, + "cache generation regressed: {cached} < {last_seen_gen}", + ); + last_seen_gen = cached; + } + })); + } + + let f_writer = Arc::clone(&filter); + let writer = thread::spawn(move || { + for i in 0..WRITER_ITERS { + f_writer + .update(lit(i) as Arc) + .expect("update must succeed"); + // Yield to give readers a chance to see intermediate states. + thread::yield_now(); + } + }); + + writer.join().expect("writer thread panicked"); + stop.store(true, Ordering::Relaxed); + for h in readers { + h.join().expect("reader thread panicked"); + } + + // After the writer is done, one final `current()` should sync the + // cache to the latest generation. + let _ = derived.current().unwrap(); + let (cache_gen, _) = derived + .current_cache + .read() + .as_ref() + .expect("cache populated after final current()") + .clone(); + let inner_gen = derived.inner.read().generation; + assert_eq!( + cache_gen, inner_gen, + "final cache generation must match inner.generation", + ); + // Writer bumps generation once per update, so final generation is + // starting-generation + WRITER_ITERS. Starting is 1, so final is + // WRITER_ITERS + 1. + assert_eq!(inner_gen, WRITER_ITERS as u64 + 1); + } } diff --git a/datafusion/physical-plan/src/common.rs b/datafusion/physical-plan/src/common.rs index 0dafcf6bd3390..734ec96debc85 100644 --- a/datafusion/physical-plan/src/common.rs +++ b/datafusion/physical-plan/src/common.rs @@ -181,7 +181,8 @@ pub fn project_plan_to_schema( } /// If running in a tokio context spawns the execution of `stream` to a separate task -/// allowing it to execute in parallel with an intermediate buffer of size `buffer` +/// allowing it to execute in parallel with an intermediate buffer of size `buffer`. +/// At most `buffer` record batches will be produced ahead of the consumer. pub fn spawn_buffered( mut input: SendableRecordBatchStream, buffer: usize, @@ -196,11 +197,22 @@ pub fn spawn_buffered( let sender = builder.tx(); builder.spawn(async move { - while let Some(item) = input.next().await { - if sender.send(item).await.is_err() { - // Receiver dropped when query is shutdown early (e.g., limit) or error, - // no need to return propagate the send error. - return Ok(()); + // We call `reserve` (which waits until there's room for at least 1 message in the + // channel buffer) **before** polling from input to ensure we hold a maximum of + // `buffer` record batches in memory. + // Polling from input and then calling send() would block when the channel is full + // so it would essentially hold `buffer` + 1 record batches: + // * `buffer`: this many elements would live inside the channel, since this is the + // channel's capacity + // * 1 extra RecordBatch which was produced, but there was no room for it in the + // channel, so it's being owned by the send() future, which keeps the batch in + // memory while it waits for a slot to free up + while let Ok(permit) = sender.reserve().await { + // Receiver dropped when query is shutdown early (e.g., limit) or error, + // no need to return propagate the send error. + match input.next().await { + Some(item) => permit.send(item), + None => break, } } @@ -298,10 +310,13 @@ mod tests { use crate::empty::EmptyExec; use crate::projection::ProjectionExec; + use crate::stream::RecordBatchStreamAdapter; + use futures::stream; use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::{ - array::{Float32Array, Float64Array, UInt64Array}, + array::{Float32Array, Float64Array, Int32Array, UInt64Array}, datatypes::{DataType, Field, Schema}, }; @@ -554,4 +569,55 @@ mod tests { let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); assert!(err.to_string().contains("schema metadata differ")); } + + /// Verifies that `spawn_buffered` holds exactly `buffer` record batches in memory + /// when no receiver is polling + async fn spawn_buffered_max_in_flight_batches(buffer_size: usize) { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let num_batches = 10; + + let produced_count = Arc::new(AtomicUsize::new(0)); + let produced_clone = Arc::clone(&produced_count); + let schema_clone = Arc::clone(&schema); + + // Stream increments the counter each time a batch is pulled by the producer. + let input_stream = stream::unfold(0usize, move |i| { + let schema = Arc::clone(&schema_clone); + let counter = Arc::clone(&produced_clone); + async move { + if i >= num_batches { + return None; + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![i as i32]))], + ) + .unwrap(); + counter.fetch_add(1, Ordering::SeqCst); + Some((Ok(batch), i + 1)) + } + }); + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + input_stream, + )); + // Drop the returned stream immediately so no receiver is ever polled. + let _buffered = spawn_buffered(input, buffer_size); + + // Give the producer task time to fill the channel and stall on send(). + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + + assert_eq!( + produced_count.load(Ordering::SeqCst), + buffer_size, + "expected exactly {buffer_size} batch(es) in memory with no receiver polling" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_spawn_buffered_max_in_flight_batches() { + spawn_buffered_max_in_flight_batches(1).await; + spawn_buffered_max_in_flight_batches(2).await; + } } diff --git a/datafusion/sqllogictest/test_files/datetime/date_part.slt b/datafusion/sqllogictest/test_files/datetime/date_part.slt index 891319f9e2cd2..0a992b2d78a22 100644 --- a/datafusion/sqllogictest/test_files/datetime/date_part.slt +++ b/datafusion/sqllogictest/test_files/datetime/date_part.slt @@ -838,6 +838,40 @@ SELECT extract(millisecond from arrow_cast('23:32:50.123456789'::time, 'Time64(N ---- 50123 +# date32 and date64 + +statement ok +CREATE TABLE source_dt AS +with t as (values + ('1970-01-01'), + ('2020-06-02'), + ('2026-02-28'), + (NULL) +) +SELECT + arrow_cast(column1, 'Date32') as date32, + arrow_cast(column1, 'Date64') as date64, +FROM t; + +query IIIIIIIIII +SELECT date_part('year', date32), date_part('month', date32), date_part('week', date32), date_part('day', date32), date_part('hour', date32), date_part('minute', date32), date_part('second', date32), date_part('millisecond', date32), date_part('microsecond', date32), date_part('nanosecond', date32) FROM source_dt; +---- +1970 1 1 1 0 0 0 0 0 0 +2020 6 23 2 0 0 0 0 0 0 +2026 2 9 28 0 0 0 0 0 0 +NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL + +query IIIIIIIIII +SELECT date_part('year', date64), date_part('month', date64), date_part('week', date64), date_part('day', date64), date_part('hour', date64), date_part('minute', date64), date_part('second', date64), date_part('millisecond', date64), date_part('microsecond', date64), date_part('nanosecond', date64) FROM source_dt; +---- +1970 1 1 1 0 0 0 0 0 0 +2020 6 23 2 0 0 0 0 0 0 +2026 2 9 28 0 0 0 0 0 0 +NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL + +statement ok +drop table source_dt; + # just some floating point stuff happening in the result here query I SELECT date_part('microsecond', arrow_cast('23:32:50.123456789'::time, 'Time64(Nanosecond)'))