diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 8785874c6f6..a1a74a38fa1 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -167,11 +167,14 @@ incorrect result. When any single window expression in a `WindowExec` falls back (`stddev`, `stddev_pop`, `stddev_samp`, `var_pop`, `var_samp`, `corr`, `covar_pop`, `covar_samp`). These run natively as plain aggregations but not as window functions ([#4766](https://github.com/apache/datafusion-comet/issues/4766)). -- `min` / `max` on timestamp-without-time-zone, interval, or nested (array / struct) input types, and `sum` / - `avg` on year-month or day-time interval input types. Windowed aggregates inherit the same input-type support as - the batch aggregates, so these fall back in both contexts. +- `min` / `max` on timestamp-without-time-zone, interval, or array input types, and `sum` / `avg` on year-month + or day-time interval input types. Windowed aggregates inherit the same input-type support as the batch + aggregates, so these fall back in both contexts. - `min` / `max` on string or binary input types. These run natively as plain aggregations but not as window functions. +- `min` / `max` on struct input types. A struct whose fields are all supported `min` / `max` types other than + `float` and `double` runs natively as a plain aggregation, but not as a window function. Other structs fall + back in both contexts. - `sum` or `avg` on `DECIMAL` with a sliding (non ever-expanding) frame, because the sliding path would wrap on overflow instead of returning Spark's `NULL`. - `avg` or `try_avg` on `DECIMAL(p, s)` with an expanding frame when `p >= 28`, because Spark can preserve diff --git a/native/core/src/execution/planner/aggregate.rs b/native/core/src/execution/planner/aggregate.rs index 4437141b302..c1dd517da85 100644 --- a/native/core/src/execution/planner/aggregate.rs +++ b/native/core/src/execution/planner/aggregate.rs @@ -51,10 +51,11 @@ use datafusion_comet_proto::{ spark_operator::{AggregateMode as ProtoAggregateMode, Operator}, }; use datafusion_comet_spark_expr::{ - ApproxPercentile, Avg, AvgDecimal, BloomFilterAgg, CometCollectList, CometCollectSet, - Correlation, Covariance, HllPlusPlus, HllSketchAgg, HllUnionAgg, MaxMinBy, Mode, Regr, - RegrType, SparkBloomFilterVersion, SparkListAgg, SparkMinMax, SparkPercentile, SparkPivotFirst, - Stddev, SumDecimal, SumInteger, Variance, ZeroOrder, + ApproxPercentile, Avg, AvgDecimal, BloomFilterAgg, Cast, CometCollectList, CometCollectSet, + Correlation, Covariance, EvalMode, HllPlusPlus, HllSketchAgg, HllUnionAgg, MaxMinBy, Mode, + Regr, RegrType, SparkBloomFilterVersion, SparkCastOptions, SparkListAgg, SparkMinMax, + SparkPercentile, SparkPivotFirst, SparkStructMinMax, Stddev, SumDecimal, SumInteger, Variance, + ZeroOrder, }; use jni::objects::{Global, JObject}; @@ -339,7 +340,7 @@ impl PhysicalPlanner { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap()); let func = min_max_udaf(&datatype, false); - let child = Arc::new(CastExpr::new(child, datatype.clone(), None)); + let child = Self::cast_min_max_child(child, &datatype, &schema)?; AggregateExprBuilder::new(func, vec![child]) .schema(schema) @@ -353,7 +354,7 @@ impl PhysicalPlanner { let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap()); let func = min_max_udaf(&datatype, true); - let child = Arc::new(CastExpr::new(child, datatype.clone(), None)); + let child = Self::cast_min_max_child(child, &datatype, &schema)?; AggregateExprBuilder::new(func, vec![child]) .schema(schema) @@ -807,6 +808,38 @@ impl PhysicalPlanner { Ok(Arc::new(CastExpr::new(child, nullable_type, None))) } + /// Cast a min/max input to the aggregate's declared type. Native expressions such as + /// `struct(...)` can report nullable fields where Spark declares them non-nullable, and + /// DataFusion's `CastExpr` rejects that narrowing. For a struct, only nullability may differ: + /// Comet's `Cast` then relabels the fields without touching the data (and errors if a field + /// declared non-nullable holds a null), so the input matches the state that shuffle and scans + /// stamp with Spark's type. Any other difference is a planning error. + fn cast_min_max_child( + child: Arc, + datatype: &DataType, + schema: &SchemaRef, + ) -> Result, ExecutionError> { + if !matches!(datatype, DataType::Struct(_)) { + return Ok(Arc::new(CastExpr::new(child, datatype.clone(), None))); + } + let child_type = child.data_type(schema.as_ref())?; + if &child_type == datatype { + return Ok(child); + } + if !equal_ignoring_nullability(&child_type, datatype) { + return Err(GeneralError(format!( + "min/max input type {child_type} does not match the declared type {datatype}" + ))); + } + Ok(Arc::new(Cast::new( + child, + datatype.clone(), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), + None, + None, + ))) + } + pub(super) fn create_aggr_func_expr( name: &str, schema: SchemaRef, @@ -823,6 +856,19 @@ impl PhysicalPlanner { } } +/// True if two types match apart from struct field nullability and metadata. +fn equal_ignoring_nullability(a: &DataType, b: &DataType) -> bool { + match (a, b) { + (DataType::Struct(fa), DataType::Struct(fb)) => { + fa.len() == fb.len() + && fa.iter().zip(fb.iter()).all(|(x, y)| { + x.name() == y.name() && equal_ignoring_nullability(x.data_type(), y.data_type()) + }) + } + _ => a == b, + } +} + /// Return a copy of `data_type` with every nested field marked nullable. Map key fields are left /// non-nullable to preserve Arrow's map invariant. Primitive types are returned unchanged. fn make_all_fields_nullable(data_type: &DataType) -> DataType { @@ -870,13 +916,15 @@ fn make_all_fields_nullable(data_type: &DataType) -> DataType { } } -/// `min` or `max` over `data_type`: Spark's version for floats, which DataFusion orders -/// differently, and DataFusion's for every other type. +/// `min` or `max` over `data_type`: Spark's version for floats and structs, which DataFusion +/// orders differently, and DataFusion's for every other type. pub(super) fn min_max_udaf(data_type: &DataType, is_max: bool) -> Arc { // Spark has only 32- and 64-bit floating-point types. Keep Float16 on DataFusion's path so // an internal plan using it does not reach SparkMinMax, whose accumulators reject that type. if matches!(data_type, DataType::Float32 | DataType::Float64) { Arc::new(AggregateUDF::new_from_impl(SparkMinMax::new(is_max))) + } else if matches!(data_type, DataType::Struct(_)) { + Arc::new(AggregateUDF::new_from_impl(SparkStructMinMax::new(is_max))) } else if is_max { max_udaf() } else { @@ -898,3 +946,53 @@ fn percentile_value(expr: &spark_expression::Expr) -> Result Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + is_max, + } + } +} + +impl AggregateUDFImpl for SparkStructMinMax { + fn name(&self) -> &str { + if self.is_max { + "max" + } else { + "min" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + fn accumulator(&self, args: AccumulatorArgs) -> Result> { + Ok(Box::new(StructMinMaxAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + Ok(vec![Arc::new(Field::new( + format_state_name(args.name, self.name()), + args.return_field.data_type().clone(), + true, + ))]) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result> { + Ok(Box::new(StructMinMaxGroupsAccumulator::try_new( + args.return_field.data_type(), + self.is_max, + )?)) + } +} + +/// Spark's struct order: ascending, with nulls (of the struct or of any field) first. +const SPARK_ORDER: SortOptions = SortOptions { + descending: false, + nulls_first: true, +}; + +/// Encodes struct values to Spark-ordered row bytes and decodes them back. +#[derive(Debug)] +struct StructRows { + converter: RowConverter, + /// The encoding of a null struct, emitted for a group that saw no value. + null_row: OwnedRow, + is_max: bool, +} + +impl StructRows { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + if !matches!(data_type, DataType::Struct(_)) { + return internal_err!("Spark struct min/max expects a struct, got {data_type}"); + } + let converter = RowConverter::new(vec![SortField::new_with_options( + data_type.clone(), + SPARK_ORDER, + )])?; + let null_row = converter + .convert_columns(&[new_null_array(data_type, 1)])? + .row(0) + .owned(); + Ok(Self { + converter, + null_row, + is_max, + }) + } + + fn encode(&self, values: &ArrayRef) -> Result { + Ok(self.converter.convert_columns(&[Arc::clone(values)])?) + } + + /// Compares rows of `values` with each other in Spark's struct order. + fn comparator(&self, values: &ArrayRef) -> Result { + Ok(make_comparator( + values.as_ref(), + values.as_ref(), + SPARK_ORDER, + )?) + } + + /// Whether a value that compares to the current one as `ordering` replaces it. Only a strictly + /// better value does, so of equal values the first one seen is kept. + #[inline] + fn beats(&self, ordering: Ordering) -> bool { + if self.is_max { + ordering == Ordering::Greater + } else { + ordering == Ordering::Less + } + } + + /// Keeps the better of `value` and `slot`, reusing the slot's buffer. Returns the bytes the + /// slot's buffer grew by, which is never negative since `clear` keeps the capacity. + #[inline] + fn update(&self, slot: &mut Option>, value: &[u8]) -> usize { + match slot { + None => { + let bytes = value.to_vec(); + let added = bytes.capacity(); + *slot = Some(bytes); + added + } + Some(current) if self.beats(value.cmp(current.as_slice())) => { + let before = current.capacity(); + current.clear(); + current.extend_from_slice(value); + current.capacity() - before + } + Some(_) => 0, + } + } + + fn decode<'a>(&self, values: impl Iterator>) -> Result { + let parser = self.converter.parser(); + let rows: Vec<_> = values + .map(|value| match value { + Some(bytes) => parser.parse(bytes), + None => self.null_row.row(), + }) + .collect(); + let mut columns = self.converter.convert_rows(rows)?; + Ok(columns.remove(0)) + } + + fn size(&self) -> usize { + self.converter.size() + self.null_row.as_ref().len() + } +} + +/// Whether row `i` takes part: a null struct or a row the filter drops does not. +#[inline] +fn is_skipped(values: &ArrayRef, filter: Option<&BooleanArray>, i: usize) -> bool { + values.is_null(i) || filter.is_some_and(|f| !f.is_valid(i) || !f.value(i)) +} + +#[derive(Debug)] +struct StructMinMaxAccumulator { + rows: StructRows, + value: Option>, +} + +impl StructMinMaxAccumulator { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + Ok(Self { + rows: StructRows::try_new(data_type, is_max)?, + value: None, + }) + } +} + +impl Accumulator for StructMinMaxAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let values = &values[0]; + let cmp = self.rows.comparator(values)?; + let mut best: Option = None; + for i in 0..values.len() { + if is_skipped(values, None, i) { + continue; + } + if best.is_none_or(|b| self.rows.beats(cmp(i, b))) { + best = Some(i); + } + } + if let Some(best) = best { + let encoded = self.rows.encode(&values.slice(best, 1))?; + self.rows.update(&mut self.value, encoded.row(0).as_ref()); + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.update_batch(states) + } + + fn state(&mut self) -> Result> { + Ok(vec![self.evaluate()?]) + } + + fn evaluate(&mut self) -> Result { + let array = self.rows.decode(std::iter::once(self.value.as_deref()))?; + ScalarValue::try_from_array(&array, 0) + } + + fn size(&self) -> usize { + size_of_val(self) + + self.rows.size() + + self.value.as_ref().map_or(0, |value| value.capacity()) + } +} + +#[derive(Debug)] +struct StructMinMaxGroupsAccumulator { + rows: StructRows, + /// The row bytes of each group's current winner, or `None` before it sees a value. + values: Vec>>, + /// The bytes allocated by the buffers in `values`. + value_bytes: usize, + /// Scratch space for `update_batch`: each group's best row in the current batch, or + /// `NO_ROW`, and the groups that have one. Reset after every batch. + batch_best: Vec, + touched: Vec, +} + +const NO_ROW: u32 = u32::MAX; + +impl StructMinMaxGroupsAccumulator { + fn try_new(data_type: &DataType, is_max: bool) -> Result { + Ok(Self { + rows: StructRows::try_new(data_type, is_max)?, + values: vec![], + value_bytes: 0, + batch_best: vec![], + touched: vec![], + }) + } +} + +impl GroupsAccumulator for StructMinMaxGroupsAccumulator { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + let values = &values[0]; + self.values.resize(total_num_groups, None); + self.batch_best.resize(total_num_groups, NO_ROW); + + // Find each group's best row in this batch without encoding anything. + let cmp = self.rows.comparator(values)?; + for (i, &group) in group_indices.iter().enumerate() { + if is_skipped(values, opt_filter, i) { + continue; + } + let best = &mut self.batch_best[group]; + if *best == NO_ROW { + *best = i as u32; + self.touched.push(group); + } else if self.rows.beats(cmp(i, *best as usize)) { + *best = i as u32; + } + } + if self.touched.is_empty() { + return Ok(()); + } + + // Encode only those rows, and keep each one that beats its group's stored winner. + let indices: UInt32Array = self + .touched + .iter() + .map(|&group| self.batch_best[group]) + .collect(); + let encoded = self.rows.encode(&take(values.as_ref(), &indices, None)?)?; + for (k, &group) in self.touched.iter().enumerate() { + self.value_bytes += self + .rows + .update(&mut self.values[group], encoded.row(k).as_ref()); + self.batch_best[group] = NO_ROW; + } + self.touched.clear(); + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let emitted = emit_to.take_needed(&mut self.values); + let freed: usize = emitted.iter().flatten().map(|value| value.capacity()).sum(); + self.value_bytes -= freed; + self.rows + .decode(emitted.iter().map(|value| value.as_deref())) + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + Ok(vec![self.evaluate(emit_to)?]) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + self.update_batch(values, group_indices, None, total_num_groups) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + let values = Arc::clone(&values[0]); + let Some(filter) = opt_filter else { + return Ok(vec![values]); + }; + // A row the filter drops, or whose filter value is null, contributes nothing. + let dropped = not(&prep_null_mask_filter(filter))?; + Ok(vec![nullif(&values, &dropped)?]) + } + + fn size(&self) -> usize { + size_of_val(self) + + self.rows.size() + + self.values.capacity() * size_of::>>() + + self.value_bytes + + self.batch_best.capacity() * size_of::() + + self.touched.capacity() * size_of::() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, Int64Array, StringArray, StructArray}; + use arrow::datatypes::{Fields, Int64Type}; + + fn fields(nullable: bool) -> Fields { + Fields::from(vec![ + Field::new("k", DataType::Int64, nullable), + Field::new("s", DataType::Utf8, nullable), + ]) + } + + fn structs(rows: &[Option<(Option, Option<&str>)>]) -> ArrayRef { + let k: Int64Array = rows.iter().map(|r| r.and_then(|(k, _)| k)).collect(); + let s: StringArray = rows.iter().map(|r| r.and_then(|(_, s)| s)).collect(); + let validity = rows.iter().map(Option::is_some).collect::>(); + Arc::new(StructArray::new( + fields(true), + vec![Arc::new(k), Arc::new(s)], + Some(validity.into()), + )) + } + + fn decoded(array: &ArrayRef) -> Vec, Option)>> { + let array = array.as_struct(); + let k = array.column(0).as_primitive::(); + let s = array.column(1).as_string::(); + (0..array.len()) + .map(|i| { + array.is_valid(i).then(|| { + ( + k.is_valid(i).then(|| k.value(i)), + s.is_valid(i).then(|| s.value(i).to_string()), + ) + }) + }) + .collect() + } + + fn groups(is_max: bool) -> StructMinMaxGroupsAccumulator { + StructMinMaxGroupsAccumulator::try_new(&DataType::Struct(fields(true)), is_max).unwrap() + } + + /// Spark orders a null field below every value, at any position, and skips a null struct. + #[test] + fn orders_null_fields_first_like_spark() -> Result<()> { + let input = structs(&[ + Some((Some(1), Some("z"))), + Some((None, Some("zz"))), + Some((Some(2), Some("a"))), + None, + Some((Some(2), None)), + ]); + for (is_max, expected) in [ + (true, Some((Some(2), Some("a".to_string())))), + (false, Some((None, Some("zz".to_string())))), + ] { + let mut acc = groups(is_max); + acc.update_batch(&[Arc::clone(&input)], &[0, 0, 0, 0, 0], None, 1)?; + assert_eq!(decoded(&acc.evaluate(EmitTo::All)?), vec![expected]); + } + Ok(()) + } + + /// A group with no value, or whose rows the filter drops, emits a null struct. + #[test] + fn empty_and_filtered_groups_are_null() -> Result<()> { + let input = structs(&[Some((Some(1), Some("a"))), Some((Some(5), Some("b")))]); + let filter = BooleanArray::from(vec![Some(true), None]); + let mut acc = groups(true); + acc.update_batch(&[input], &[0, 1], Some(&filter), 3)?; + assert_eq!( + decoded(&acc.evaluate(EmitTo::All)?), + vec![Some((Some(1), Some("a".to_string()))), None, None] + ); + Ok(()) + } + + /// Emitting a prefix of the groups frees exactly the bytes those groups held, so ordered + /// aggregation does not keep reserving memory for groups it has already emitted. + #[test] + fn emitting_groups_releases_their_bytes() -> Result<()> { + let payload = "x".repeat(64 * 1024); + let mut acc = groups(true); + let empty = acc.size(); + for round in 0..3 { + let rows: Vec<_> = (0..64) + .map(|k| Some((Some(round * 64 + k), Some(payload.as_str())))) + .collect(); + let group_indices: Vec = (0..64).collect(); + acc.update_batch(&[structs(&rows)], &group_indices, None, 64)?; + assert!(acc.size() > 64 * payload.len()); + acc.evaluate(EmitTo::First(32))?; + acc.evaluate(EmitTo::First(32))?; + assert_eq!(acc.value_bytes, 0); + assert!(acc.size() < empty + 64 * size_of::>>()); + } + Ok(()) + } + + /// Partial states merge into the same answer as one pass over every row. + #[test] + fn merge_matches_single_pass() -> Result<()> { + let first = structs(&[Some((Some(3), Some("b"))), Some((Some(1), None))]); + let second = structs(&[Some((Some(3), Some("c"))), None]); + let mut left = groups(true); + left.update_batch(&[first], &[0, 1], None, 2)?; + let mut right = groups(true); + right.update_batch(&[second], &[0, 1], None, 2)?; + let mut merged = groups(true); + merged.merge_batch(&left.state(EmitTo::All)?, &[0, 1], 2)?; + merged.merge_batch(&right.state(EmitTo::All)?, &[0, 1], 2)?; + assert_eq!( + decoded(&merged.evaluate(EmitTo::All)?), + vec![ + Some((Some(3), Some("c".to_string()))), + Some((Some(1), None)) + ] + ); + Ok(()) + } + + #[test] + fn ungrouped_accumulator_matches_grouped() -> Result<()> { + let input = structs(&[ + Some((Some(2), Some("a"))), + None, + Some((Some(2), Some("b"))), + Some((None, Some("zzz"))), + ]); + let mut acc = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + acc.update_batch(&[Arc::clone(&input)])?; + let state = acc.state()?; + let mut merged = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + merged.merge_batch(&[state[0].to_array()?])?; + let result = merged.evaluate()?.to_array()?; + assert_eq!( + decoded(&result), + vec![Some((Some(2), Some("b".to_string())))] + ); + + let mut empty = StructMinMaxAccumulator::try_new(&DataType::Struct(fields(true)), true)?; + assert_eq!(decoded(&empty.evaluate()?.to_array()?), vec![None]); + Ok(()) + } + + /// Non-nullable fields decode with the declared type. + #[test] + fn keeps_non_nullable_field_types() -> Result<()> { + let data_type = DataType::Struct(fields(false)); + let input: ArrayRef = Arc::new(StructArray::new( + fields(false), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["a", "b"])), + ], + None, + )); + let mut acc = StructMinMaxGroupsAccumulator::try_new(&data_type, false)?; + acc.update_batch(&[input], &[0, 0], None, 1)?; + let result = acc.evaluate(EmitTo::All)?; + assert_eq!(result.data_type(), &data_type); + Ok(()) + } + + /// Nested structs with nulls at every level, used to check the two orders agree. + fn nested_structs(n: usize) -> ArrayRef { + // Small value ranges so that ties on the leading fields are common. + let k: Int64Array = (0..n) + .map(|i| (i % 7 != 0).then_some((i * 7919 % 5) as i64 - 2)) + .collect(); + let s: StringArray = (0..n) + .map(|i| (i % 5 != 0).then(|| ["", "a", "ab", "日本", "\u{7f}"][i * 31 % 5])) + .collect(); + let inner_v: Int64Array = (0..n) + .map(|i| (i % 3 != 0).then_some((i * 104729 % 3) as i64)) + .collect(); + let inner_fields = Fields::from(vec![Field::new("v", DataType::Int64, true)]); + let inner_validity: Vec = (0..n).map(|i| i % 11 != 0).collect(); + let inner = StructArray::new( + inner_fields.clone(), + vec![Arc::new(inner_v)], + Some(inner_validity.into()), + ); + let outer_fields = Fields::from(vec![ + Field::new("k", DataType::Int64, true), + Field::new("s", DataType::Utf8, true), + Field::new("inner", DataType::Struct(inner_fields), true), + ]); + let outer_validity: Vec = (0..n).map(|i| i % 13 != 0).collect(); + Arc::new(StructArray::new( + outer_fields, + vec![Arc::new(k), Arc::new(s), Arc::new(inner)], + Some(outer_validity.into()), + )) + } + + /// Each batch picks candidates with Arrow's comparator, and winners are kept as row bytes. + /// The two must order every pair of values the same way. + #[test] + fn comparator_and_row_bytes_agree() -> Result<()> { + let values = nested_structs(200); + let rows = StructRows::try_new(values.data_type(), true)?; + let cmp = rows.comparator(&values)?; + let encoded = rows.encode(&values)?; + for i in 0..values.len() { + for j in 0..values.len() { + assert_eq!( + cmp(i, j), + encoded.row(i).cmp(&encoded.row(j)), + "rows {i} and {j}" + ); + } + } + Ok(()) + } + + /// Many rows per group in each batch, across several batches, against a brute-force answer. + #[test] + fn matches_brute_force_over_batches() -> Result<()> { + let values = nested_structs(600); + let data_type = values.data_type().clone(); + let groups: Vec = (0..values.len()).map(|i| i * 17 % 9).collect(); + let reference = StructRows::try_new(&data_type, true)?; + let encoded = reference.encode(&values)?; + for is_max in [true, false] { + let mut acc = StructMinMaxGroupsAccumulator::try_new(&data_type, is_max)?; + for start in (0..values.len()).step_by(100) { + acc.update_batch( + &[values.slice(start, 100)], + &groups[start..start + 100], + None, + 9, + )?; + } + let result = acc.evaluate(EmitTo::All)?; + let result_rows = reference.encode(&result)?; + for group in 0..9 { + let members = + (0..values.len()).filter(|&i| groups[i] == group && values.is_valid(i)); + let expected = if is_max { + members.max_by(|&a, &b| encoded.row(a).cmp(&encoded.row(b))) + } else { + members.min_by(|&a, &b| encoded.row(a).cmp(&encoded.row(b))) + }; + match expected { + Some(i) => assert_eq!(result_rows.row(group), encoded.row(i), "group {group}"), + None => assert!(result.is_null(group), "group {group}"), + } + } + } + Ok(()) + } + + /// Sorted aggregation emits finished groups after every batch. Struct max over 4,096 groups + /// of 4 KiB values (16 MiB in all), read in 64-row batches, must fit a 2 MiB pool, which it + /// only does if emitting a group releases the memory charged for it. DataFusion's own struct + /// max fails this with `ResourcesExhausted`, because it subtracts rows rather than bytes. + #[tokio::test] + async fn sorted_aggregation_fits_a_small_pool() -> Result<()> { + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::Schema; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::execution::TaskContext; + use datafusion::logical_expr::AggregateUDF; + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_expr::expressions::{col, Column}; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, + }; + use datafusion::physical_plan::{collect, ExecutionPlan, InputOrderMode}; + use datafusion::prelude::SessionConfig; + + const GROUPS: i64 = 4096; + const BATCH_ROWS: i64 = 64; + let payload = "x".repeat(4096); + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Struct(fields(true)), true), + ])); + let batches = (0..GROUPS / BATCH_ROWS) + .map(|b| { + let groups = (b * BATCH_ROWS..(b + 1) * BATCH_ROWS).collect::>(); + let rows: Vec<_> = groups + .iter() + .map(|&g| Some((Some(g), Some(payload.as_str())))) + .collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(groups)), structs(&rows)], + ) + }) + .collect::, _>>()?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("g", 0), + ))]) + .unwrap(); + let source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input: Arc = DataSourceExec::from_data_source(source); + + let max = Arc::new( + AggregateExprBuilder::new( + Arc::new(AggregateUDF::new_from_impl(SparkStructMinMax::new(true))), + vec![col("v", &schema)?], + ) + .schema(Arc::clone(&schema)) + .alias("max") + .build()?, + ); + // Comet splits a sort aggregate into a Partial and a Final stage. Both see sorted input, + // so DataFusion runs its ordered streams, which emit finished groups after every batch. + let group_by = PhysicalGroupBy::new_single(vec![(col("g", &schema)?, "g".to_string())]); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&max)], + vec![None], + input, + Arc::clone(&schema), + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![max], + vec![None], + Arc::clone(&partial) as Arc, + Arc::clone(&schema), + )?); + assert_eq!(partial.input_order_mode(), &InputOrderMode::Sorted); + // No disk, so an over-reserving accumulator fails instead of spilling. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(2 * 1024 * 1024, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(BATCH_ROWS as usize)) + .with_runtime(runtime), + ); + let output = collect(aggregate, context).await?; + assert_eq!( + output.iter().map(|b| b.num_rows()).sum::(), + GROUPS as usize + ); + Ok(()) + } +} diff --git a/native/spark-expr/src/agg_funcs/mod.rs b/native/spark-expr/src/agg_funcs/mod.rs index e7a72d60ec8..7133a9cff90 100644 --- a/native/spark-expr/src/agg_funcs/mod.rs +++ b/native/spark-expr/src/agg_funcs/mod.rs @@ -29,6 +29,7 @@ mod hll_union_agg; mod list_agg; mod max_min_by; mod min_max; +mod min_max_struct; mod mode; mod percentile; mod pivot_first; @@ -53,6 +54,7 @@ pub use hll_union_agg::HllUnionAgg; pub use list_agg::SparkListAgg; pub use max_min_by::MaxMinBy; pub use min_max::SparkMinMax; +pub use min_max_struct::SparkStructMinMax; pub use mode::Mode; pub use percentile::{SparkPercentile, ZeroOrder}; pub use pivot_first::SparkPivotFirst; diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 807c3ef3e35..dcff65311e6 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -552,11 +552,11 @@ fn cast_struct_to_struct( }) .collect::>>()?; - Ok(Arc::new(StructArray::new( + Ok(Arc::new(StructArray::try_new( to_fields.clone(), cast_fields, array.nulls().cloned(), - ))) + )?)) } _ => unreachable!(), } @@ -1248,6 +1248,70 @@ mod tests { } } + /// Struct min/max casts its input this way to relabel nullable fields as non-nullable, + /// so the values, including a UTC timestamp, must come through unchanged. + #[test] + fn test_cast_struct_relabels_nullable_fields_as_non_nullable() { + let ts_type = + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())); + let ts: ArrayRef = Arc::new( + PrimitiveArray::::from(vec![-1, 0, 1]).with_timezone("UTC"), + ); + let s: ArrayRef = Arc::new(StringArray::from(vec!["a", "中文", ""])); + let inner_from = Fields::from(vec![Field::new("s", DataType::Utf8, true)]); + let inner_to = Fields::from(vec![Field::new("s", DataType::Utf8, false)]); + let inner: ArrayRef = Arc::new(StructArray::new(inner_from.clone(), vec![s], None)); + let from_fields = Fields::from(vec![ + Field::new("ts", ts_type.clone(), true), + Field::new("inner", DataType::Struct(inner_from), true), + ]); + let to_fields = Fields::from(vec![ + Field::new("ts", ts_type, false), + Field::new("inner", DataType::Struct(inner_to), false), + ]); + // The outer struct itself may be null. + let outer: ArrayRef = Arc::new(StructArray::new( + from_fields, + vec![ts, inner], + Some(vec![true, false, true].into()), + )); + let to_type = DataType::Struct(to_fields); + + let result = cast_array( + Arc::clone(&outer), + &to_type, + &SparkCastOptions::new_without_timezone(EvalMode::Legacy), + ) + .unwrap(); + + assert_eq!(result.data_type(), &to_type); + assert_eq!(result.nulls(), outer.nulls()); + let result = result.as_struct(); + let original = outer.as_struct(); + assert_eq!(result.column(0).as_ref(), original.column(0).as_ref()); + assert_eq!( + result.column(1).as_struct().column(0).as_ref(), + original.column(1).as_struct().column(0).as_ref() + ); + } + + /// A null under a field relabeled as non-nullable must be an error, not a panic. + #[test] + fn test_cast_struct_to_non_nullable_field_with_null_returns_error() { + let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None])); + let from_fields = Fields::from(vec![Field::new("v", DataType::Int64, true)]); + let to_fields = Fields::from(vec![Field::new("v", DataType::Int64, false)]); + let array: ArrayRef = Arc::new(StructArray::new(from_fields, vec![values], None)); + + let result = cast_array( + array, + &DataType::Struct(to_fields), + &SparkCastOptions::new_without_timezone(EvalMode::Legacy), + ); + + assert!(result.is_err()); + } + #[test] fn test_cast_nested_struct_to_struct_ansi_overflow_returns_error() { let inner_values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), Some(128), None])); diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 17ab15d1810..a521d410ade 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, isSpark42Plus, withFallbackReason} import org.apache.comet.expressions.CometEvalMode -import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, isStringCollationType, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType} import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil, CometTypeShim} object CometMin extends CometAggregateExpressionSerde[Min] { @@ -1589,7 +1589,7 @@ object CometPivotFirst extends CometAggregateExpressionSerde[PivotFirst] with Co } } -object AggSerde { +object AggSerde extends CometTypeShim { import org.apache.spark.sql.types._ def minMaxDataTypeSupported(dt: DataType): Boolean = { @@ -1639,14 +1639,55 @@ object AggSerde { Unsupported(Some(ComparisonUtils.nonDefaultCollationReason("min/max"))) // DataFusion's grouped hash aggregate under-reports the memory these values hold (#6936). case _: StringType | BinaryType => Compatible() + case st: StructType => minMaxStructSupportLevel(st) case _ if minMaxDataTypeSupported(dt) => Compatible() case _ => Unsupported(Some(s"Unsupported data type: $dt")) } + /** + * The native accumulator compares structs by their Arrow row encoding, ascending with nulls + * first, which orders fields the way Spark does, a null field below every value included. The + * encoding orders floats by IEEE total order, which differs from Spark for -0.0 and NaN with + * the sign bit set, so float and double fields fall back. + */ + private def minMaxStructSupportLevel(st: StructType): SupportLevel = { + if (st.fields.isEmpty) { + Unsupported(Some("Struct min/max over an empty struct")) + } else if (hasNonDefaultStringCollation(st)) { + Unsupported(Some(ComparisonUtils.nonDefaultCollationReason("min/max"))) + } else if (SupportLevel.containsType(st, classOf[FloatType], classOf[DoubleType])) { + Unsupported( + Some( + "Struct min/max over a float or double field. Comet orders -0.0 below 0.0 and" + + " a NaN with the sign bit set below every value, while Spark treats -0.0 and 0.0" + + " as equal and every NaN as the largest value")) + } else { + unsupportedStructLeaf(st) match { + case Some(dt) => Unsupported(Some(s"Struct min/max over a field of type $dt")) + case None => Compatible() + } + } + } + + /** The first leaf type of a struct that native min/max cannot compare, if any. */ + private def unsupportedStructLeaf(dt: DataType): Option[DataType] = dt match { + case st: StructType if st.fields.isEmpty => Some(st) + case st: StructType => + st.fields.view.flatMap(f => unsupportedStructLeaf(f.dataType)).headOption + case _ => + minMaxSupportLevel(dt) match { + case _: Compatible => None + case _ => Some(dt) + } + } + /** The inputs that `minMaxSupportLevel` rejects, for the generated compatibility guide. */ val minMaxUnsupportedReasons: Seq[String] = Seq( "Input of a type other than boolean, integral, floating-point, decimal, date, timestamp," + - " string or binary, such as timestamp without time zone, intervals, arrays and structs.", + " string or binary, such as timestamp without time zone, intervals, arrays and maps." + + " A struct is supported only if all of its fields are of supported types other than" + + " float and double, since Comet orders -0.0 and NaN with the sign bit set" + + " differently from Spark inside a struct.", "A string with a non-default collation (Spark 4.0+), which Spark orders by its collation" + " while Comet compares raw bytes.") diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql index 02d96805fb1..97a43f236ea 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_collation.sql @@ -31,3 +31,8 @@ INSERT INTO test_min_max_collation VALUES ('x', 'a'), ('x', 'B'), ('x', NULL), ( query expect_fallback(min/max does not support non-UTF8_BINARY collated operands) SELECT grp, min(CAST(s AS STRING COLLATE UTF8_LCASE)), max(CAST(s AS STRING COLLATE UTF8_LCASE)) FROM test_min_max_collation GROUP BY grp ORDER BY grp + +-- A collated string inside a struct falls back for the same reason. +query expect_fallback(min/max does not support non-UTF8_BINARY collated operands) +SELECT grp, max(struct(coalesce(grp, '') AS g, CAST(s AS STRING COLLATE UTF8_LCASE) AS s)) +FROM test_min_max_collation GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql new file mode 100644 index 00000000000..d5ffc6334cf --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct.sql @@ -0,0 +1,175 @@ +-- 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. + +-- min/max over a struct. Spark orders structs field by field, with a null field (or a null +-- nested struct) below every value, and skips a null struct entirely. Float and double fields +-- fall back. + +statement +CREATE TABLE test_mm_struct(g int, ts bigint, s string, d decimal(38,0), dt date, t timestamp, n string) +USING parquet + +statement +INSERT INTO test_mm_struct VALUES + (1, 5, 'b', 99999999999999999999999999999999999999, DATE'1980-01-01', TIMESTAMP'1980-01-01 00:00:00', 'x'), + (1, 9, '日本語', -99999999999999999999999999999999999999, DATE'9999-12-31', TIMESTAMP'2100-01-01 00:00:00', 'y'), + (1, 9, '中文', 0, DATE'2000-02-29', TIMESTAMP'2000-02-29 12:00:00', NULL), + (2, 9223372036854775807, 'a', 1, DATE'1980-01-01', TIMESTAMP'1999-12-31 23:59:59', 'z'), + (2, -9223372036854775808, 'c', 2, DATE'1980-01-02', TIMESTAMP'1999-12-31 23:59:58', 'w'), + (3, 1, 'only', 3, DATE'2020-01-01', TIMESTAMP'2020-01-01 00:00:00', NULL) + +-- Non-nullable fields. +query +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct GROUP BY g ORDER BY g + +query +SELECT g, min(struct(coalesce(d, 0BD) AS d, coalesce(dt, DATE'1970-01-01') AS dt, coalesce(t, TIMESTAMP'1970-01-01 00:00:00') AS t)), + max(struct(coalesce(d, 0BD) AS d, coalesce(dt, DATE'1970-01-01') AS dt, coalesce(t, TIMESTAMP'1970-01-01 00:00:00') AS t)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Nested struct, and a whole-struct null from a guarding CASE WHEN. +query +SELECT g, max(CASE WHEN ts > 0 THEN struct(coalesce(ts, 0L) AS ts, struct(coalesce(s, '') AS s, true AS b) AS inner) END) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- No grouping. +query +SELECT min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct + +-- Float and double fields fall back: Comet orders -0.0 below 0.0, and a NaN with the sign bit +-- set below every value. +query expect_fallback(float or double field) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, 1.5D AS d)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Array field: unsupported. +query expect_fallback(Struct min/max over a field of type ArrayType) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, array(1, 2) AS a)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Spark cannot order maps, so it rejects a struct holding one during analysis. +query expect_error(INVALID_ORDERING_TYPE) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, map('k', 1) AS m)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- timestamp_ntz is not supported by native min/max, so a struct holding one falls back too. +query expect_fallback(Struct min/max over a field of type TimestampNTZType) +SELECT g, max(struct(coalesce(ts, 0L) AS ts, TIMESTAMP_NTZ'2020-01-01 00:00:00' AS t)) FROM test_mm_struct GROUP BY g ORDER BY g + +-- Boolean, integral and binary fields. +query +SELECT g, + min(struct(coalesce(ts, 0L) > 5 AS b, coalesce(CAST(ts % 100 AS TINYINT), 0Y) AS i8, + coalesce(CAST(ts % 30000 AS SMALLINT), 0S) AS i16, coalesce(CAST(ts % 100000 AS INT), 0) AS i32, + coalesce(CAST(s AS BINARY), X'') AS bin)), + max(struct(coalesce(ts, 0L) > 5 AS b, coalesce(CAST(ts % 100 AS TINYINT), 0Y) AS i8, + coalesce(CAST(ts % 30000 AS SMALLINT), 0S) AS i16, coalesce(CAST(ts % 100000 AS INT), 0) AS i32, + coalesce(CAST(s AS BINARY), X'') AS bin)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Ties on the first field are decided by the next one. Group 1 has two rows with ts = 9, and +-- '中文' sorts before '日本語' as raw UTF-8 bytes. +query +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct WHERE ts = 9 GROUP BY g ORDER BY g + +-- count(DISTINCT) makes Spark plan a PartialMerge stage, which merges the struct buffers again. +query +SELECT g, count(DISTINCT s), max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- A VALUES input keeps the partial aggregate on Spark, while the final one merges Spark's struct +-- buffers natively. The plan is not fully native, so only the answer is checked. +query spark_answer_only +SELECT k, min(struct(o, v)), max(struct(o, v)) +FROM VALUES (1, 5L, 'b'), (1, 9L, 'ア'), (2, -1L, '😀'), (2, 3L, 'a') AS t(k, o, v) +GROUP BY k ORDER BY k + +-- Empty input: one NULL row without grouping, no rows with grouping. +query +SELECT max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct WHERE g < 0 + +query +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct WHERE g < 0 GROUP BY g + +-- FILTER turns filtered-out rows into nulls before they reach the struct accumulator. +query +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FILTER (WHERE ts < 9) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- A nested struct that can be null orders below every non-null nested struct. +query +SELECT g, max(struct(coalesce(ts, 0L) AS ts, CASE WHEN ts > 5 THEN struct(coalesce(s, '') AS s) END AS inner)) +FROM test_mm_struct GROUP BY g ORDER BY g + +-- Null fields at every level. +statement +CREATE TABLE test_mm_struct_null(g int, ts bigint, s string, n string, inner_v int) USING parquet + +statement +INSERT INTO test_mm_struct_null VALUES + (1, 5, 'b', 'x', 1), + (1, 9, NULL, NULL, NULL), + (1, 7, '日本語', NULL, 3), + (2, 1, 'a', 'z', NULL), + (2, 2, NULL, 'w', 2), + (3, 4, NULL, NULL, NULL), + (4, NULL, NULL, NULL, NULL), + (5, NULL, 'z', 'p', 1), + (5, 1, 'a', 'q', 2), + (5, 1, NULL, 'r', 3), + (6, 3, 'a', NULL, NULL), + (6, 3, 'a', 'b', NULL) + +-- Nulls in every position, including the deciding field: group 5 orders (NULL, ...) first and +-- breaks the tie on ts = 1 with the null s first. +query +SELECT g, min(struct(ts, s, n)), max(struct(ts, s, n)) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- Nested struct with nullable inner fields, and a nested struct that is itself null. +query +SELECT g, + min(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)), + max(struct(ts, CASE WHEN inner_v IS NOT NULL THEN struct(s, inner_v) END AS inner)) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- Whole-struct nulls are skipped. Group 4 has only a NULL struct, so the result is NULL. +query +SELECT g, max(CASE WHEN ts IS NOT NULL THEN struct(ts, s) END) +FROM test_mm_struct_null GROUP BY g ORDER BY g + +-- A null deciding field sorts first: Spark returns (7, 'a') for max and (NULL, '~') for min. +query +SELECT min(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)), + max(struct(CASE WHEN ts = 9 THEN NULL ELSE ts END AS ts, + CASE WHEN ts = 9 THEN '~' ELSE 'a' END AS s)) +FROM test_mm_struct_null WHERE g = 1 + +-- A struct column read from a table, whose fields Spark reads as nullable. +statement +CREATE TABLE test_mm_struct_null_col(g int, st struct) USING parquet + +statement +INSERT INTO test_mm_struct_null_col VALUES + (1, named_struct('ts', 5L, 's', 'b')), (1, named_struct('ts', 9L, 's', NULL)), + (1, named_struct('ts', 7L, 's', '日本語')), (2, named_struct('ts', 1L, 's', 'a')), (2, NULL), + (3, named_struct('ts', NULL, 's', 'x')), (3, named_struct('ts', 2L, 's', NULL)) + +query +SELECT g, min(st), max(st) FROM test_mm_struct_null_col GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql new file mode 100644 index 00000000000..87df5c13d09 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_many_batches.sql @@ -0,0 +1,33 @@ +-- 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. + +-- Struct min/max over many small batches and many groups, so the accumulator carries each group's +-- winner from batch to batch instead of seeing all rows at once. + +-- Config: spark.comet.batchSize=64 + +query +SELECT g, min(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)), + max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) +FROM (SELECT id, id % 1000 AS g FROM range(0, 200000, 1, 8)) GROUP BY g ORDER BY g + +query +SELECT max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) FROM range(0, 200000, 1, 8) + +query +SELECT g, max(struct(coalesce(id % 7, 0L) AS k, CAST(id AS STRING) AS s)) FILTER (WHERE id % 3 = 0) +FROM (SELECT id, id % 1000 AS g FROM range(0, 200000, 1, 8)) GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql new file mode 100644 index 00000000000..57503e070f8 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_mixed_stages.sql @@ -0,0 +1,55 @@ +-- 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. + +-- Struct min/max with its partial and final stages on different engines. The test-only configs +-- below keep one stage on Spark, so the struct buffer crosses between Spark and Comet in both +-- directions: native partial -> Spark final, and Spark partial -> native final. Min and Max allow +-- both. The plan is mixed by design, so only the answer is checked. + +-- ConfigMatrix: spark.comet.testing.aggregate.partialMode.enabled=true,false +-- ConfigMatrix: spark.comet.testing.aggregate.finalMode.enabled=true,false + +statement +CREATE TABLE test_mm_struct_mixed(g int, ts bigint, s string, x int) USING parquet + +statement +INSERT INTO test_mm_struct_mixed VALUES + (1, 5, 'b', 1), (1, 9, '日本語', 2), (1, 9, '中文', 3), + (2, 9223372036854775807, 'a', 4), (2, -9223372036854775808, 'c', 5), + (3, NULL, NULL, NULL) + +query spark_answer_only +SELECT g, min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), + max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +query spark_answer_only +SELECT g, max(CASE WHEN ts > 0 THEN struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s) END) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +query spark_answer_only +SELECT max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) FROM test_mm_struct_mixed + +-- count(DISTINCT) adds PartialMerge stages between the partial and final ones. +query spark_answer_only +SELECT g, count(DISTINCT s), min(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g + +-- collect_list makes Spark plan an ObjectHashAggregate instead of a sort aggregate. +query spark_answer_only +SELECT g, max(struct(coalesce(ts, 0L) AS ts, coalesce(s, '') AS s)), sort_array(collect_list(x)) +FROM test_mm_struct_mixed GROUP BY g ORDER BY g diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql new file mode 100644 index 00000000000..1677acb61ae --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max_struct_timezone.sql @@ -0,0 +1,56 @@ +-- 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. + +-- min/max over a struct with timestamp and date fields under several session time zones, +-- including values near the ends of the range and DST transitions. 1901 stays clear of the +-- INT96 rebase cutoff in every zone listed. + +-- ConfigMatrix: spark.sql.session.timeZone=UTC,Etc/UTC,America/Los_Angeles,Asia/Kolkata,Australia/Lord_Howe + +statement +CREATE TABLE test_mm_struct_tz(g int, t timestamp, dt date, s string) USING parquet + +statement +INSERT INTO test_mm_struct_tz VALUES + (1, TIMESTAMP'1901-01-01 00:00:00', DATE'1901-01-01', 'a'), + (1, TIMESTAMP'9999-12-30 23:59:59.999999', DATE'9999-12-31', 'b'), + (1, TIMESTAMP'1970-01-01 00:00:00', DATE'1970-01-01', 'c'), + (2, TIMESTAMP'2024-03-10 01:59:59', DATE'2024-03-10', 'before'), + (2, TIMESTAMP'2024-03-10 03:00:00', DATE'2024-03-10', 'after'), + (2, TIMESTAMP'2024-11-03 01:30:00', DATE'2024-11-03', 'overlap'), + (3, TIMESTAMP'1969-12-31 23:59:59.999999', DATE'1969-12-31', '한국어'), + (3, NULL, NULL, NULL) + +query +SELECT g, + min(struct(coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t, coalesce(dt, DATE'2000-01-01') AS dt, coalesce(s, '') AS s)), + max(struct(coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t, coalesce(dt, DATE'2000-01-01') AS dt, coalesce(s, '') AS s)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g + +-- The date leads, so ties on the date are decided by the timestamp. +query +SELECT g, + min(struct(coalesce(dt, DATE'2000-01-01') AS dt, coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t)), + max(struct(coalesce(dt, DATE'2000-01-01') AS dt, coalesce(t, TIMESTAMP'2000-01-01 00:00:00') AS t)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g + +-- The timestamp field is computed natively from a date, which depends on the session time zone. +query +SELECT g, + min(struct(CAST(coalesce(dt, DATE'2000-01-01') AS TIMESTAMP) AS t, coalesce(s, '') AS s)), + max(struct(CAST(coalesce(dt, DATE'2000-01-01') AS TIMESTAMP) AS t, coalesce(s, '') AS s)) +FROM test_mm_struct_tz GROUP BY g ORDER BY g diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 5c73433a552..deaeb7c2bd3 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -1416,6 +1416,43 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("struct min/max buffers cross engines in both directions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + // JVM columnar shuffle can carry Spark's partial rows to a native final. + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withTable("struct_min_max") { + sql("CREATE TABLE struct_min_max(g int, ts bigint, s string) USING parquet") + // Group 1 has a null deciding field, which sorts first on both engines. + sql( + "INSERT INTO struct_min_max VALUES (1, 5, 'b'), (1, 9, NULL), (1, NULL, 'z'), " + + "(2, 1, '日本語'), (2, 1, 'a'), (3, NULL, NULL)") + for { + (disabled, nativeMode) <- Seq( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE -> Partial, + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE -> Final) + } { + withSQLConf(disabled.key -> "false") { + val df = sql( + "SELECT g, min(struct(ts, s)), max(struct(ts, s)) FROM struct_min_max GROUP BY g") + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan + assert( + collect(plan) { case a: CometBaseAggregateExec => a.modes }.flatten == + Seq(nativeMode), + plan) + assert( + collect(plan) { case a: SortAggregateExec => + a.aggregateExpressions.map(_.mode).distinct + }.flatten == Seq(if (nativeMode == Partial) Final else Partial), + plan) + } + } + } + } + } + test("count, avg with null") { Seq(false, true).foreach { dictionary => withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala index 950ebaab7aa..be98e2dca9b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometAggregateExpressionBenchmark.scala @@ -199,6 +199,24 @@ object CometAggregateExpressionBenchmark extends CometBenchmarkBase { "min_by_high_card", "SELECT min_by(c_int, c_long) FROM parquetV1Table GROUP BY high_card_grp")) + // min / max over a struct of an ordering key and a payload. The key decides every comparison, + // so the narrow (64 byte) and wide (4 KiB) cases differ only in what has to be carried along. + private val structMinMaxAggregates = List( + AggExprConfig( + "max_struct_narrow", + "SELECT MAX(struct(k, narrow)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "max_struct_wide", + "SELECT MAX(struct(k, wide)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "min_struct_wide", + "SELECT MIN(struct(k, wide)) FROM parquetV1Table GROUP BY grp"), + AggExprConfig( + "max_struct_narrow_global", + "SELECT MAX(struct(k, narrow)) FROM parquetV1Table"), + AggExprConfig("max_struct_wide_global", "SELECT MAX(struct(k, wide)) FROM parquetV1Table"), + AggExprConfig("min_struct_wide_global", "SELECT MIN(struct(k, wide)) FROM parquetV1Table")) + override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024 @@ -231,5 +249,28 @@ object CometAggregateExpressionBenchmark extends CometBenchmarkBase { } } } + + // Fewer rows than above, since each wide payload is 4 KiB. + runBenchmarkWithTable("Struct min/max benchmarks", 200 * 1000) { v => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s""" + SELECT + CAST(value % 1000 AS INT) AS grp, + CAST(value AS LONG) AS k, + CASE WHEN value % 10 = 0 THEN NULL ELSE sha2(CAST(value AS STRING), 256) END AS narrow, + CASE WHEN value % 10 = 0 THEN NULL + ELSE repeat(sha2(CAST(value AS STRING), 256), 64) END AS wide + FROM $tbl + """)) + + structMinMaxAggregates.foreach { config => + runExpressionBenchmark(config.name, v.toLong, config.query, config.extraCometConfigs) + } + } + } + } } }