diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index e7f2d8b7eac47..3bc5c255235a9 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -4,6 +4,7 @@ //! Defines the FIRST_VALUE/LAST_VALUE aggregations. +use std::borrow::Cow; use std::fmt::Debug; use std::hash::Hash; use std::mem::size_of_val; @@ -431,6 +432,16 @@ struct FirstLastGroupsAccumulator { // and `update_batch` takes the comparison-free fast path. is_input_pre_ordered: bool, default_orderings: Vec, + // The time zone of each ordering column that is a time zone aware + // timestamp, `None` for every other column. + // + // Such columns are stored in `orderings` without their time zone. + // Comparison does not depend on it (the values are UTC instants), and a + // `ScalarValue::Timestamp*(_, Some(tz))` holds an `Arc` shared by + // every partition, so cloning and dropping one per updated group makes + // the threads contend on its reference count. The time zone is put back + // on the arrays returned by `state`. + ordering_timezones: Vec>>, } impl FirstLastGroupsAccumulator { @@ -442,9 +453,21 @@ impl FirstLastGroupsAccumulator { pick_first_in_group: bool, is_input_pre_ordered: bool, ) -> Result { + let ordering_timezones: Vec<_> = ordering_dtypes + .iter() + .map(|dt| match dt { + DataType::Timestamp(_, tz) => tz.clone(), + _ => None, + }) + .collect(); let default_orderings = ordering_dtypes .iter() - .map(ScalarValue::try_from) + .map(|dt| match dt { + DataType::Timestamp(unit, _) => { + ScalarValue::try_from(DataType::Timestamp(*unit, None)) + } + _ => ScalarValue::try_from(dt), + }) .collect::>()?; let sort_options = get_sort_options(&ordering_req); @@ -463,9 +486,29 @@ impl FirstLastGroupsAccumulator { touched_groups_buf: Vec::new(), pick_first_in_group, is_input_pre_ordered, + ordering_timezones, }) } + /// Returns `columns` with the time zone removed from the ordering columns + /// (`columns[1..=ordering_req.len()]`) that have one. See + /// `ordering_timezones`. + fn strip_ordering_timezones<'a>( + &self, + columns: &'a [ArrayRef], + ) -> Cow<'a, [ArrayRef]> { + if self.ordering_timezones.iter().all(Option::is_none) { + return Cow::Borrowed(columns); + } + let mut columns = columns.to_vec(); + for (column, tz) in columns[1..].iter_mut().zip(&self.ordering_timezones) { + if tz.is_some() { + *column = with_timezone_opt(column, None); + } + } + Cow::Owned(columns) + } + fn should_update_state( &self, group_idx: usize, @@ -767,6 +810,8 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator Result<()> { self.resize_states(total_num_groups); + let values_and_order_cols = self.strip_ordering_timezones(values_and_order_cols); + let values_and_order_cols = values_and_order_cols.as_ref(); if self.is_input_pre_ordered { return self.update_batch_pre_ordered( @@ -829,8 +874,13 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator with_timezone_opt(&array, Some(Arc::clone(tz))), + None => array, + }); } result.push(Arc::new(BooleanArray::new(is_sets, None))); @@ -845,6 +895,8 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator Result<()> { self.resize_states(total_num_groups); + let values = self.strip_ordering_timezones(values); + let values = values.as_ref(); let mut ordering_buf = Vec::with_capacity(self.ordering_req.len()); @@ -881,6 +933,7 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator() + self.extreme_of_each_group_buf.1.capacity() / 8 + + self.ordering_timezones.capacity() * size_of::>>() } fn shrink_to_fit(&mut self) { @@ -1579,6 +1632,38 @@ fn filter_states_according_to_is_set( .collect() } +/// Returns a timestamp `array` with its time zone replaced by `tz`, without +/// changing its values (unlike a cast); other arrays are returned unchanged. +fn with_timezone_opt(array: &ArrayRef, tz: Option>) -> ArrayRef { + match array.data_type() { + DataType::Timestamp(TimeUnit::Second, _) => Arc::new( + array + .as_primitive::() + .clone() + .with_timezone_opt(tz), + ), + DataType::Timestamp(TimeUnit::Millisecond, _) => Arc::new( + array + .as_primitive::() + .clone() + .with_timezone_opt(tz), + ), + DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new( + array + .as_primitive::() + .clone() + .with_timezone_opt(tz), + ), + DataType::Timestamp(TimeUnit::Nanosecond, _) => Arc::new( + array + .as_primitive::() + .clone() + .with_timezone_opt(tz), + ), + _ => Arc::clone(array), + } +} + /// Combines array refs and their corresponding orderings to construct `SortColumn`s. fn convert_to_sort_cols(arrs: &[ArrayRef], sort_exprs: &LexOrdering) -> Vec { arrs.iter() @@ -1595,7 +1680,10 @@ mod tests { use std::iter::repeat_with; use arrow::{ - array::{BooleanArray, Int64Array, ListArray, PrimitiveArray, StringArray}, + array::{ + BooleanArray, Int64Array, ListArray, PrimitiveArray, StringArray, + TimestampMicrosecondArray, + }, buffer::NullBuffer, compute::SortOptions, datatypes::Schema, @@ -1815,6 +1903,80 @@ mod tests { Ok(()) } + #[test] + fn test_group_acc_timezone_ordering() -> Result<()> { + let tz: Arc = Arc::from("+02:00"); + let ts_type = DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::clone(&tz))); + let schema = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Int64, true), + Field::new("ts", ts_type.clone(), true), + ])); + let sort_keys = [PhysicalSortExpr { + expr: col("ts", &schema)?, + options: SortOptions::default(), + }]; + let new_acc = |pick_first_in_group| { + FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::::new(DataType::Int64), + sort_keys.clone().into(), + false, + std::slice::from_ref(&ts_type), + pick_first_in_group, + false, + ) + }; + let ts = |values: Vec>| -> ArrayRef { + Arc::new( + TimestampMicrosecondArray::from(values) + .with_timezone_opt(Some(Arc::clone(&tz))), + ) + }; + + for (pick_first_in_group, expected) in [(true, [1, 5]), (false, [3, 4])] { + let mut partial = new_acc(pick_first_in_group)?; + partial.update_batch( + &[ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + ts(vec![Some(10), Some(20), Some(30), Some(15), Some(5)]), + ], + &[0, 0, 0, 1, 1], + None, + 2, + )?; + // The ordering values are kept without the time zone ... + assert!( + partial.orderings.iter().flatten().all(|v| matches!( + v, + ScalarValue::TimestampMicrosecond(Some(_), None) + )) + ); + assert_eq!( + partial.size_of_orderings, + partial.compute_size_of_orderings() + ); + + // ... and the state carries it again, with unchanged values, also + // when it is emitted in parts. + let first = partial.state(EmitTo::First(1))?; + let rest = partial.state(EmitTo::All)?; + assert_eq!(first[1].data_type(), &ts_type); + assert_eq!(rest[1].data_type(), &ts_type); + + let mut last = new_acc(pick_first_in_group)?; + last.merge_batch(&first, &[0], 1)?; + last.merge_batch(&rest, &[1], 2)?; + let other = vec![ + Arc::new(Int64Array::from(vec![6, 7])) as ArrayRef, + ts(vec![Some(25), Some(10)]), + Arc::new(BooleanArray::from(vec![true, true])), + ]; + last.merge_batch(&other, &[0, 1], 2)?; + let result = last.evaluate(EmitTo::All)?; + assert_eq!(result.as_primitive::().values(), &expected); + } + Ok(()) + } + #[test] fn test_group_acc_size_of_ordering() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 605e54d1543cb..a0fd1dadf4d23 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -10458,6 +10458,152 @@ GROUP BY id ORDER BY id; statement ok DROP TABLE first_last_value_str_tests; +################# +# first_value / last_value ordered by a time zone aware timestamp +################# + +# The grouped accumulator keeps the ordering values without the time zone +# and puts it back on its intermediate state. The rows are split over two +# inputs so that partial states from different partitions are merged. +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +CREATE TABLE first_last_tz_a AS +SELECT column1 AS g, column2 AS v, arrow_cast(column3, 'Timestamp(Microsecond, Some("+05:00"))') AS ts +FROM (VALUES + (1, 1, '2024-01-01T00:00:00Z'), + (1, 2, '2024-01-01T02:00:00Z'), + (2, 3, NULL), + (2, 4, '2024-01-01T03:00:00Z'), + (3, 5, '2024-01-01T01:00:00Z')); + +statement ok +CREATE TABLE first_last_tz_b AS +SELECT column1 AS g, column2 AS v, arrow_cast(column3, 'Timestamp(Microsecond, Some("+05:00"))') AS ts +FROM (VALUES + (1, 6, '2023-12-31T23:00:00Z'), + (2, 7, '2024-01-01T04:00:00Z'), + (3, 8, '2024-01-01T01:00:00Z'), + (4, 9, NULL)); + +statement ok +CREATE VIEW first_last_tz AS +SELECT * FROM first_last_tz_a UNION ALL SELECT * FROM first_last_tz_b; + +query IIIPP +SELECT g, + first_value(v ORDER BY ts, v), last_value(v ORDER BY ts, v), + first_value(ts ORDER BY ts), last_value(ts ORDER BY ts) +FROM first_last_tz GROUP BY g ORDER BY g; +---- +1 6 2 2024-01-01T04:00:00+05:00 2024-01-01T07:00:00+05:00 +2 4 3 2024-01-01T08:00:00+05:00 NULL +3 5 8 2024-01-01T06:00:00+05:00 2024-01-01T06:00:00+05:00 +4 9 9 NULL NULL + +# The time zone aware column is not the first ordering column +query III +SELECT g, first_value(v ORDER BY v % 2, ts), last_value(v ORDER BY v % 2, ts) +FROM first_last_tz GROUP BY g ORDER BY g; +---- +1 6 1 +2 4 3 +3 8 5 +4 9 9 + +query IIP +SELECT g, first_value(v ORDER BY ts DESC NULLS FIRST, v DESC), last_value(ts ORDER BY ts DESC NULLS FIRST, v DESC) +FROM first_last_tz GROUP BY g ORDER BY g; +---- +1 2 2024-01-01T04:00:00+05:00 +2 3 2024-01-01T08:00:00+05:00 +3 8 2024-01-01T06:00:00+05:00 +4 9 NULL + +query TT +SELECT arrow_typeof(first_value(ts ORDER BY ts)), arrow_typeof(last_value(v ORDER BY ts)) +FROM first_last_tz GROUP BY g ORDER BY g LIMIT 1; +---- +Timestamp(µs, "+05:00") Int64 + +# DISTINCT ON is planned as first_value per column +query IIP +SELECT DISTINCT ON (g) g, v, ts FROM first_last_tz ORDER BY g, ts DESC, v; +---- +1 2 2024-01-01T07:00:00+05:00 +2 3 NULL +3 5 2024-01-01T06:00:00+05:00 +4 9 NULL + +statement ok +DROP VIEW first_last_tz; + +statement ok +DROP TABLE first_last_tz_a; + +statement ok +DROP TABLE first_last_tz_b; + +# Other time units, several ordering columns with different time zones, and +# IGNORE NULLS +statement ok +CREATE TABLE first_last_tz_u_a AS +SELECT column1 AS g, column2 AS v, + arrow_cast(column3, 'Timestamp(Second, Some("UTC"))') AS ts_s, + arrow_cast(column4, 'Timestamp(Millisecond, Some("Europe/Berlin"))') AS ts_ms, + arrow_cast(column4, 'Timestamp(Nanosecond, Some("-03:00"))') AS ts_ns +FROM (VALUES + (1, 1, '2024-01-01T00:00:00Z', '2024-01-01T00:00:00.300Z'), + (1, NULL, '2024-01-01T00:00:00Z', '2024-01-01T00:00:00.100Z'), + (2, 3, '2024-01-02T00:00:00Z', NULL), + (2, NULL, '2024-01-01T00:00:00Z', '2024-01-01T00:00:00.500Z')); + +statement ok +CREATE TABLE first_last_tz_u_b AS +SELECT column1 AS g, column2 AS v, + arrow_cast(column3, 'Timestamp(Second, Some("UTC"))') AS ts_s, + arrow_cast(column4, 'Timestamp(Millisecond, Some("Europe/Berlin"))') AS ts_ms, + arrow_cast(column4, 'Timestamp(Nanosecond, Some("-03:00"))') AS ts_ns +FROM (VALUES + (1, 5, '2024-01-01T00:00:00Z', '2024-01-01T00:00:00.200Z'), + (2, 6, '2024-01-03T00:00:00Z', '2024-01-01T00:00:00.400Z')); + +statement ok +CREATE VIEW first_last_tz_u AS +SELECT * FROM first_last_tz_u_a UNION ALL SELECT * FROM first_last_tz_u_b; + +query IIIIP +SELECT g, first_value(v ORDER BY ts_s, ts_ms DESC), last_value(v ORDER BY ts_s, ts_ms DESC), + first_value(v ORDER BY ts_ns), last_value(ts_ns ORDER BY ts_s, ts_ns) +FROM first_last_tz_u GROUP BY g ORDER BY g; +---- +1 1 NULL NULL 2023-12-31T21:00:00.300-03:00 +2 NULL 6 6 2023-12-31T21:00:00.400-03:00 + +query IIIII +SELECT g, first_value(v ORDER BY ts_ms) IGNORE NULLS, last_value(v ORDER BY ts_ms) IGNORE NULLS, + first_value(v ORDER BY ts_ms), last_value(v ORDER BY ts_ms) +FROM first_last_tz_u GROUP BY g ORDER BY g; +---- +1 5 1 NULL 1 +2 6 3 6 3 + +query TT +SELECT arrow_typeof(last_value(ts_ns ORDER BY ts_s, ts_ns)), arrow_typeof(first_value(ts_ms ORDER BY ts_ns)) +FROM first_last_tz_u GROUP BY g ORDER BY g LIMIT 1; +---- +Timestamp(ns, "-03:00") Timestamp(ms, "Europe/Berlin") + +statement ok +DROP VIEW first_last_tz_u; + +statement ok +DROP TABLE first_last_tz_u_a; + +statement ok +DROP TABLE first_last_tz_u_b; + # MIN/MAX over a CAST projection folds from parquet statistics via the # aggregate-statistics optimizer. Non-CAST projections (multi-column or # self-referenced expressions) must keep returning correct results — guard