query_planning/capability/
aggregate.rs

1//! Aggregate-strategy capability analysis.
2//!
3//! This module classifies grouped aggregate stages into a selected execution
4//! strategy plus locally rejected candidate strategies. It does not enable an
5//! optimized path by itself; executor admission must still have a concrete
6//! reservation and execution contract before selecting anything except the
7//! default libcudf groupby strategy.
8
9use std::collections::BTreeMap;
10
11use arrow::datatypes::{DataType, Schema};
12
13use crate::{
14    datatypes::is_arrow_fixed_width,
15    error::{QueryError, QueryErrorCode, Result as QueryResult},
16    expr::{AggregateExpr, AggregateFunction, AggregateSpec, Expr},
17};
18
19/// Aggregate execution stage being evaluated for strategy selection.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub enum AggregateExecutionStage {
22    /// One-shot grouped aggregate over a resident input frame.
23    Single,
24    /// Chunk-local partial grouped aggregate before final merge.
25    Partial,
26    /// Final grouped aggregate merge over partial aggregate rows.
27    FinalMerge,
28}
29
30fn aggregate_expr_is_exact_distinct_count(expr: &AggregateExpr) -> bool {
31    expr.function == AggregateFunction::NUnique || expr.distinct && expr.function == AggregateFunction::Count
32}
33
34/// Native aggregate state shape relevant to bounded admission.
35///
36/// Grouping sets, non-count distinct functions, and non-mergeable functions
37/// have no contract here (`None`): node capability already rejects those
38/// specs, so no edge may rely on their state.
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
40pub enum AggregateStateContract {
41    /// No distinct or group-only dedup state is required.
42    Plain,
43    /// SQL setop-style dedup over grouped keys that have no aggregate expressions.
44    DedupGroupOnly,
45    /// Exact count-distinct state represented as `COUNT(DISTINCT col)` or
46    /// native `NUnique(col)`.
47    ExactDistinctCount,
48}
49
50impl AggregateStateContract {
51    /// Stable diagnostic label for this aggregate state contract.
52    pub(crate) const fn as_str(self) -> &'static str {
53        match self {
54            Self::Plain => "plain",
55            Self::DedupGroupOnly => "dedup_group_only",
56            Self::ExactDistinctCount => "exact_distinct_count",
57        }
58    }
59}
60
61/// The state shape bounded admission may rely on for `spec`, or `None` when the
62/// spec has no mergeable state contract (see [`AggregateStateContract`]).
63pub(crate) fn aggregate_state_contract(spec: &AggregateSpec) -> Option<AggregateStateContract> {
64    if spec.grouping_sets.is_some() {
65        return None;
66    }
67    if spec.aggr_exprs.is_empty() && !spec.group_exprs.is_empty() {
68        return Some(AggregateStateContract::DedupGroupOnly);
69    }
70
71    let mut has_exact_distinct_count = false;
72    for expr in &spec.aggr_exprs {
73        if aggregate_expr_is_exact_distinct_count(expr) {
74            has_exact_distinct_count = true;
75        } else if expr.function == AggregateFunction::Median || expr.distinct {
76            // Non-mergeable functions and non-count distincts have no edge
77            // contract here; node capability already rejects those specs.
78            return None;
79        }
80    }
81
82    if has_exact_distinct_count {
83        Some(AggregateStateContract::ExactDistinctCount)
84    } else {
85        Some(AggregateStateContract::Plain)
86    }
87}
88
89impl AggregateExecutionStage {
90    /// Stable report label for this aggregate stage.
91    pub const fn as_str(self) -> &'static str {
92        match self {
93            Self::Single => "single",
94            Self::Partial => "partial",
95            Self::FinalMerge => "final_merge",
96        }
97    }
98}
99
100/// Native grouped aggregate execution strategy label.
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
102pub enum AggregateExecutionStrategy {
103    /// Existing default path using libcudf groupby directly on group keys.
104    LibcudfGroupby,
105    /// Candidate path that dictionary/factor encodes eligible group keys.
106    EncodedKeyPartial,
107    /// Candidate path using cuDF's prepared sort-backed groupby state.
108    PreparedSortGroupby,
109    /// Candidate path that partitions partial outputs before grouped merge.
110    PartitionedMerge,
111}
112
113impl AggregateExecutionStrategy {
114    /// Stable report label for this aggregate strategy.
115    pub const fn as_str(self) -> &'static str {
116        match self {
117            Self::LibcudfGroupby => "libcudf_groupby",
118            Self::EncodedKeyPartial => "encoded_key_partial",
119            Self::PreparedSortGroupby => "prepared_sort_groupby",
120            Self::PartitionedMerge => "partitioned_merge",
121        }
122    }
123}
124
125/// Runtime threshold policy for grouped aggregate strategy selection.
126///
127/// These thresholds choose between the default libcudf groupby path, encoded
128/// key partial aggregation, prepared sort-backed groupby, and exact
129/// cardinality probes. This is an expert Rust surface only; keep it out of
130/// server/session/user-facing configuration until the tuning model has a
131/// reviewed production contract.
132///
133/// Thresholds suffixed `_ppm` are ratios of distinct groups to input rows in
134/// parts per million; `_rows`, `_bytes`, `_row_groups`, `_fan_in` and
135/// `_key_columns` are absolute bounds. Zero disables the branch a threshold
136/// guards, which is why [`Self::validate`] accepts it.
137#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138#[non_exhaustive]
139pub struct AggregateStrategyExecutionOptions {
140    encoded_key_partial_low_cardinality_threshold_ppm: u64,
141    encoded_key_partial_min_input_rows: u64,
142    exact_grouped_key_cardinality_probe_max_key_columns: usize,
143    exact_single_string_avg_cardinality_probe_min_input_rows: u64,
144    exact_single_key_cardinality_probe_min_input_rows: u64,
145    exact_grouped_key_cardinality_probe_min_input_rows: u64,
146    exact_single_string_key_prepared_max_groups: u64,
147    exact_three_key_cardinality_probe_max_input_rows: u64,
148    exact_five_key_cardinality_probe_max_input_rows: u64,
149    prepared_sort_groupby_low_cardinality_threshold_ppm: u64,
150    source_adjacent_prepared_sort_groupby_max_row_groups: u64,
151    source_adjacent_prepared_sort_groupby_max_source_read_bytes: u64,
152    source_adjacent_prepared_sort_groupby_max_fan_in: u64,
153}
154
155impl AggregateStrategyExecutionOptions {
156    const MAX_RATIO_PPM: u64 = 1_000_000;
157
158    pub(crate) const DEFAULT_ENCODED_KEY_PARTIAL_LOW_CARDINALITY_THRESHOLD_PPM: u64 = 100_000;
159    pub(crate) const DEFAULT_ENCODED_KEY_PARTIAL_MIN_INPUT_ROWS: u64 = 100_000;
160    pub(crate) const DEFAULT_EXACT_GROUPED_KEY_CARDINALITY_PROBE_MAX_KEY_COLUMNS: usize = 5;
161    pub(crate) const DEFAULT_EXACT_SINGLE_STRING_AVG_CARDINALITY_PROBE_MIN_INPUT_ROWS: u64 = 100_000;
162    pub(crate) const DEFAULT_EXACT_SINGLE_KEY_CARDINALITY_PROBE_MIN_INPUT_ROWS: u64 = 500_000;
163    pub(crate) const DEFAULT_EXACT_GROUPED_KEY_CARDINALITY_PROBE_MIN_INPUT_ROWS: u64 = 500_000;
164    pub(crate) const DEFAULT_EXACT_SINGLE_STRING_KEY_PREPARED_MAX_GROUPS: u64 = 4_096;
165    pub(crate) const DEFAULT_EXACT_THREE_KEY_CARDINALITY_PROBE_MAX_INPUT_ROWS: u64 = 10_000_000;
166    pub(crate) const DEFAULT_EXACT_FIVE_KEY_CARDINALITY_PROBE_MAX_INPUT_ROWS: u64 = 7_000_000;
167    pub(crate) const DEFAULT_PREPARED_SORT_GROUPBY_LOW_CARDINALITY_THRESHOLD_PPM: u64 = 100_000;
168    /// Default maximum selected row groups for source-adjacent prepared sort grouping.
169    pub const DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_ROW_GROUPS: u64 = 64;
170    /// Default maximum source-read bytes for source-adjacent prepared sort grouping.
171    pub const DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_SOURCE_READ_BYTES: u64 = 384 * 1024 * 1024;
172    /// Default maximum chunk fan-in for source-adjacent prepared sort grouping.
173    pub const DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_FAN_IN: u64 = 64;
174
175    /// Maximum distinct-group/input-row ratio, in parts per million, for the
176    /// encoded-key partial aggregation strategy.
177    pub fn encoded_key_partial_low_cardinality_threshold_ppm(&self) -> u64 {
178        self.encoded_key_partial_low_cardinality_threshold_ppm
179    }
180
181    /// Minimum input row count for considering encoded-key partial aggregation.
182    pub fn encoded_key_partial_min_input_rows(&self) -> u64 {
183        self.encoded_key_partial_min_input_rows
184    }
185
186    /// Maximum number of grouping columns admitted to the exact cardinality probe.
187    pub fn exact_grouped_key_cardinality_probe_max_key_columns(&self) -> usize {
188        self.exact_grouped_key_cardinality_probe_max_key_columns
189    }
190
191    /// Minimum input row count for probing a single string key with AVG state.
192    pub fn exact_single_string_avg_cardinality_probe_min_input_rows(&self) -> u64 {
193        self.exact_single_string_avg_cardinality_probe_min_input_rows
194    }
195
196    /// Minimum input row count for the general single-key cardinality probe.
197    pub fn exact_single_key_cardinality_probe_min_input_rows(&self) -> u64 {
198        self.exact_single_key_cardinality_probe_min_input_rows
199    }
200
201    /// Minimum input row count for the multi-column grouping cardinality probe.
202    pub fn exact_grouped_key_cardinality_probe_min_input_rows(&self) -> u64 {
203        self.exact_grouped_key_cardinality_probe_min_input_rows
204    }
205
206    /// Maximum distinct string-key groups admitted to the prepared strategy.
207    pub fn exact_single_string_key_prepared_max_groups(&self) -> u64 {
208        self.exact_single_string_key_prepared_max_groups
209    }
210
211    /// Maximum input row count for an exact three-column grouping probe.
212    pub fn exact_three_key_cardinality_probe_max_input_rows(&self) -> u64 {
213        self.exact_three_key_cardinality_probe_max_input_rows
214    }
215
216    /// Maximum input row count for an exact five-column grouping probe.
217    pub fn exact_five_key_cardinality_probe_max_input_rows(&self) -> u64 {
218        self.exact_five_key_cardinality_probe_max_input_rows
219    }
220
221    /// Maximum distinct-group/input-row ratio, in parts per million, for
222    /// prepared sort grouping.
223    pub fn prepared_sort_groupby_low_cardinality_threshold_ppm(&self) -> u64 {
224        self.prepared_sort_groupby_low_cardinality_threshold_ppm
225    }
226
227    /// Maximum selected row groups for source-adjacent prepared sort grouping.
228    pub fn source_adjacent_prepared_sort_groupby_max_row_groups(&self) -> u64 {
229        self.source_adjacent_prepared_sort_groupby_max_row_groups
230    }
231
232    /// Maximum source-read bytes for source-adjacent prepared sort grouping.
233    pub fn source_adjacent_prepared_sort_groupby_max_source_read_bytes(&self) -> u64 {
234        self.source_adjacent_prepared_sort_groupby_max_source_read_bytes
235    }
236
237    /// Maximum chunk fan-in for source-adjacent prepared sort grouping.
238    pub fn source_adjacent_prepared_sort_groupby_max_fan_in(&self) -> u64 {
239        self.source_adjacent_prepared_sort_groupby_max_fan_in
240    }
241
242    /// Validates aggregate strategy policy values.
243    ///
244    /// Zero-valued thresholds remain accepted because this is an expert Rust
245    /// surface and callers may use zero to disable narrow strategy branches.
246    /// The checks here reject ratios above one hundred percent and contradictory
247    /// nonzero probe bounds so strategy facts remain interpretable.
248    pub fn validate(&self) -> QueryResult<()> {
249        validate_aggregate_strategy_execution_ppm(
250            "encoded_key_partial_low_cardinality_threshold_ppm",
251            self.encoded_key_partial_low_cardinality_threshold_ppm,
252        )?;
253        validate_aggregate_strategy_execution_ppm(
254            "prepared_sort_groupby_low_cardinality_threshold_ppm",
255            self.prepared_sort_groupby_low_cardinality_threshold_ppm,
256        )?;
257        validate_aggregate_strategy_execution_min_lte_max(
258            "exact_grouped_key_cardinality_probe_min_input_rows",
259            self.exact_grouped_key_cardinality_probe_min_input_rows,
260            "exact_three_key_cardinality_probe_max_input_rows",
261            self.exact_three_key_cardinality_probe_max_input_rows,
262            "three_key_probe_max_input_rows_below_grouped_probe_min_input_rows",
263        )?;
264        validate_aggregate_strategy_execution_min_lte_max(
265            "exact_grouped_key_cardinality_probe_min_input_rows",
266            self.exact_grouped_key_cardinality_probe_min_input_rows,
267            "exact_five_key_cardinality_probe_max_input_rows",
268            self.exact_five_key_cardinality_probe_max_input_rows,
269            "five_key_probe_max_input_rows_below_grouped_probe_min_input_rows",
270        )?;
271        validate_aggregate_strategy_execution_min_lte_max(
272            "exact_single_string_avg_cardinality_probe_min_input_rows",
273            self.exact_single_string_avg_cardinality_probe_min_input_rows,
274            "exact_single_key_cardinality_probe_min_input_rows",
275            self.exact_single_key_cardinality_probe_min_input_rows,
276            "single_string_avg_probe_min_input_rows_above_general_single_key_min_input_rows",
277        )?;
278        Ok(())
279    }
280
281    /// Sets the distinct-group/input-row ratio limit in parts per million.
282    /// [`Self::validate`] rejects values above 1,000,000.
283    pub fn with_encoded_key_partial_low_cardinality_threshold_ppm(mut self, threshold_ppm: u64) -> Self {
284        self.encoded_key_partial_low_cardinality_threshold_ppm = threshold_ppm;
285        self
286    }
287
288    /// Sets the minimum input rows for encoded-key partial aggregation.
289    /// This setter retains the supplied threshold without validation.
290    pub fn with_encoded_key_partial_min_input_rows(mut self, rows: u64) -> Self {
291        self.encoded_key_partial_min_input_rows = rows;
292        self
293    }
294
295    #[cfg(test)]
296    pub fn with_exact_grouped_key_cardinality_probe_min_input_rows(mut self, rows: u64) -> Self {
297        self.exact_grouped_key_cardinality_probe_min_input_rows = rows;
298        self
299    }
300
301    #[cfg(test)]
302    pub fn with_exact_three_key_cardinality_probe_max_input_rows(mut self, rows: u64) -> Self {
303        self.exact_three_key_cardinality_probe_max_input_rows = rows;
304        self
305    }
306
307    /// Sets the distinct-group/input-row ratio limit for prepared sort grouping
308    /// in parts per million. [`Self::validate`] rejects values above 1,000,000.
309    pub fn with_prepared_sort_groupby_low_cardinality_threshold_ppm(mut self, threshold_ppm: u64) -> Self {
310        self.prepared_sort_groupby_low_cardinality_threshold_ppm = threshold_ppm;
311        self
312    }
313
314    /// Sets the selected-row-group bound for source-adjacent prepared sort
315    /// grouping. Zero disables that strategy branch.
316    pub fn with_source_adjacent_prepared_sort_groupby_max_row_groups(mut self, row_groups: u64) -> Self {
317        self.source_adjacent_prepared_sort_groupby_max_row_groups = row_groups;
318        self
319    }
320
321    #[cfg(test)]
322    pub fn with_source_adjacent_prepared_sort_groupby_max_source_read_bytes(mut self, bytes: u64) -> Self {
323        self.source_adjacent_prepared_sort_groupby_max_source_read_bytes = bytes;
324        self
325    }
326
327    #[cfg(test)]
328    pub fn with_source_adjacent_prepared_sort_groupby_max_fan_in(mut self, fan_in: u64) -> Self {
329        self.source_adjacent_prepared_sort_groupby_max_fan_in = fan_in;
330        self
331    }
332}
333
334fn validate_aggregate_strategy_execution_ppm(field: &'static str, value: u64) -> QueryResult<()> {
335    if value > AggregateStrategyExecutionOptions::MAX_RATIO_PPM {
336        return Err(QueryError::new(
337            QueryErrorCode::InvalidConfiguration,
338            "query_engine.aggregate_strategy_execution_options.validate",
339            "aggregate strategy execution ratio threshold exceeds one hundred percent",
340        )
341        .with_fact("field", field)
342        .with_fact("value", value)
343        .with_fact("max_value", AggregateStrategyExecutionOptions::MAX_RATIO_PPM)
344        .with_fact("reason", "ratio_ppm_above_one_hundred_percent"));
345    }
346    Ok(())
347}
348
349fn validate_aggregate_strategy_execution_min_lte_max(
350    min_field: &'static str,
351    min_value: u64,
352    max_field: &'static str,
353    max_value: u64,
354    reason: &'static str,
355) -> QueryResult<()> {
356    if min_value != 0 && max_value != 0 && min_value > max_value {
357        return Err(QueryError::new(
358            QueryErrorCode::InvalidConfiguration,
359            "query_engine.aggregate_strategy_execution_options.validate",
360            "aggregate strategy execution threshold ordering is inconsistent",
361        )
362        .with_fact("field", min_field)
363        .with_fact("value", min_value)
364        .with_fact("bound_field", max_field)
365        .with_fact("bound_value", max_value)
366        .with_fact("reason", reason));
367    }
368    Ok(())
369}
370
371impl Default for AggregateStrategyExecutionOptions {
372    fn default() -> Self {
373        Self {
374            encoded_key_partial_low_cardinality_threshold_ppm:
375                Self::DEFAULT_ENCODED_KEY_PARTIAL_LOW_CARDINALITY_THRESHOLD_PPM,
376            encoded_key_partial_min_input_rows: Self::DEFAULT_ENCODED_KEY_PARTIAL_MIN_INPUT_ROWS,
377            exact_grouped_key_cardinality_probe_max_key_columns:
378                Self::DEFAULT_EXACT_GROUPED_KEY_CARDINALITY_PROBE_MAX_KEY_COLUMNS,
379            exact_single_string_avg_cardinality_probe_min_input_rows:
380                Self::DEFAULT_EXACT_SINGLE_STRING_AVG_CARDINALITY_PROBE_MIN_INPUT_ROWS,
381            exact_single_key_cardinality_probe_min_input_rows:
382                Self::DEFAULT_EXACT_SINGLE_KEY_CARDINALITY_PROBE_MIN_INPUT_ROWS,
383            exact_grouped_key_cardinality_probe_min_input_rows:
384                Self::DEFAULT_EXACT_GROUPED_KEY_CARDINALITY_PROBE_MIN_INPUT_ROWS,
385            exact_single_string_key_prepared_max_groups: Self::DEFAULT_EXACT_SINGLE_STRING_KEY_PREPARED_MAX_GROUPS,
386            exact_three_key_cardinality_probe_max_input_rows:
387                Self::DEFAULT_EXACT_THREE_KEY_CARDINALITY_PROBE_MAX_INPUT_ROWS,
388            exact_five_key_cardinality_probe_max_input_rows:
389                Self::DEFAULT_EXACT_FIVE_KEY_CARDINALITY_PROBE_MAX_INPUT_ROWS,
390            prepared_sort_groupby_low_cardinality_threshold_ppm:
391                Self::DEFAULT_PREPARED_SORT_GROUPBY_LOW_CARDINALITY_THRESHOLD_PPM,
392            source_adjacent_prepared_sort_groupby_max_row_groups:
393                Self::DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_ROW_GROUPS,
394            source_adjacent_prepared_sort_groupby_max_source_read_bytes:
395                Self::DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_SOURCE_READ_BYTES,
396            source_adjacent_prepared_sort_groupby_max_fan_in:
397                Self::DEFAULT_SOURCE_ADJACENT_PREPARED_SORT_GROUPBY_MAX_FAN_IN,
398        }
399    }
400}
401
402/// Local rejection for an aggregate strategy candidate.
403#[derive(Debug, Clone, PartialEq, Eq)]
404pub struct AggregateStrategyCandidateRejection {
405    /// Candidate strategy that was rejected.
406    pub strategy: AggregateExecutionStrategy,
407    /// Stable local rejection code.
408    pub reason: &'static str,
409}
410
411/// Capability decision for one aggregate strategy selection point.
412#[derive(Debug, Clone, PartialEq, Eq)]
413pub struct AggregateStrategyCapabilityDecision {
414    /// Strategy that is available for execution now.
415    pub selected_strategy: AggregateExecutionStrategy,
416    /// Candidate optimized strategies that were not admitted.
417    pub rejected_candidates: Vec<AggregateStrategyCandidateRejection>,
418    /// Stable shape and admission facts for reports.
419    pub facts: BTreeMap<&'static str, String>,
420}
421
422impl AggregateStrategyCapabilityDecision {
423    /// Formats rejected candidates as `strategy:reason` segments for TSV facts.
424    pub fn rejected_candidate_summary(&self) -> String {
425        if self.rejected_candidates.is_empty() {
426            return "none".to_string();
427        }
428        self.rejected_candidates
429            .iter()
430            .map(|rejection| format!("{}:{}", rejection.strategy.as_str(), rejection.reason))
431            .collect::<Vec<_>>()
432            .join(";")
433    }
434}
435
436/// Returns the aggregate-strategy capability decision for an aggregate stage.
437///
438/// Selection order is fixed: partitioned merge when the stage and spec allow
439/// it, otherwise encoded-key partial, otherwise the libcudf groupby default.
440/// Prepared sort groupby is never selected here because it needs grouped key
441/// cardinality the executor measures at runtime; it is always listed among the
442/// rejected candidates, with the shape reason when the shape rules it out and
443/// the runtime-evidence reason otherwise. The `facts` keys and every reason
444/// string are report contracts read by TSV consumers.
445pub fn aggregate_strategy_capability_decision(
446    spec: &AggregateSpec,
447    input_schema: &Schema,
448    stage: AggregateExecutionStage,
449) -> AggregateStrategyCapabilityDecision {
450    let encoded_key_partial_reason = encoded_key_partial_rejected_reason(spec, input_schema, stage);
451    let prepared_sort_groupby_reason = prepared_sort_groupby_rejected_reason(spec, input_schema, stage);
452    let partitioned_merge_rejected_reason = partitioned_merge_rejected_reason(spec, stage);
453    let selected_strategy = if partitioned_merge_rejected_reason.is_none() {
454        AggregateExecutionStrategy::PartitionedMerge
455    } else if encoded_key_partial_reason.is_none() {
456        AggregateExecutionStrategy::EncodedKeyPartial
457    } else {
458        AggregateExecutionStrategy::LibcudfGroupby
459    };
460    let mut rejected_candidates = Vec::new();
461    if let Some(reason) = encoded_key_partial_reason {
462        rejected_candidates.push(AggregateStrategyCandidateRejection {
463            strategy: AggregateExecutionStrategy::EncodedKeyPartial,
464            reason,
465        });
466    }
467    rejected_candidates.push(AggregateStrategyCandidateRejection {
468        strategy: AggregateExecutionStrategy::PreparedSortGroupby,
469        reason: prepared_sort_groupby_reason.unwrap_or("prepared_sort_groupby_requires_runtime_cardinality_evidence"),
470    });
471    if let Some(reason) = partitioned_merge_rejected_reason {
472        rejected_candidates.push(AggregateStrategyCandidateRejection {
473            strategy: AggregateExecutionStrategy::PartitionedMerge,
474            reason,
475        });
476    }
477    let mut facts = BTreeMap::from([
478        ("selected_strategy", selected_strategy.as_str().to_string()),
479        ("execution_stage", stage.as_str().to_string()),
480        ("group_key_count", spec.group_exprs.len().to_string()),
481        ("group_key_exprs", aggregate_group_key_exprs(spec)),
482        ("group_key_dtypes", aggregate_group_key_dtypes(spec, input_schema)),
483        ("aggregate_functions", aggregate_function_list(spec)),
484        ("aggregate_filter_count", aggregate_filter_count(spec).to_string()),
485        ("distinct_aggregate_count", distinct_aggregate_count(spec).to_string()),
486        ("grouping_mode", aggregate_grouping_mode(spec).to_string()),
487        ("grouping_sets_present", spec.grouping_sets.is_some().to_string()),
488        ("plain_column_group_keys", spec.group_exprs.iter().all(is_plain_column_expr).to_string()),
489        ("string_group_key_count", string_group_key_count(spec, input_schema).to_string()),
490        ("fixed_width_group_key_count", fixed_width_group_key_count(spec, input_schema).to_string()),
491        ("unsupported_group_key_dtypes", unsupported_group_key_dtypes(spec, input_schema)),
492        ("hidden_host_preparation", "false".to_string()),
493        ("encoded_key_partial_rejected_reason", encoded_key_partial_reason.unwrap_or("none").to_string()),
494        (
495            "encoded_key_partial_contract_status",
496            if encoded_key_partial_reason.is_none() { "available" } else { "blocked_by_shape_or_stage" }.to_string(),
497        ),
498        ("encoded_key_partial_required_contract", "encoded_key_partial_execution_contract".to_string()),
499        (
500            "prepared_sort_groupby_rejected_reason",
501            prepared_sort_groupby_reason
502                .unwrap_or("prepared_sort_groupby_requires_runtime_cardinality_evidence")
503                .to_string(),
504        ),
505        (
506            "prepared_sort_groupby_contract_status",
507            if prepared_sort_groupby_reason.is_none() { "available" } else { "blocked_by_shape_or_stage" }.to_string(),
508        ),
509        ("prepared_sort_groupby_required_contract", "prepared_sort_groupby_execution_contract".to_string()),
510        ("prepared_sort_groupby_cost_model", "exact_grouped_key_facts_required".to_string()),
511        (
512            "prepared_sort_groupby_cost_model_status",
513            if prepared_sort_groupby_reason.is_none() { "pending_runtime_evidence" } else { "not_applicable" }
514                .to_string(),
515        ),
516        (
517            "prepared_sort_groupby_selection_reason",
518            if prepared_sort_groupby_reason.is_none() {
519                "runtime_cardinality_evidence_required"
520            } else {
521                "capability_rejected"
522            }
523            .to_string(),
524        ),
525        ("prepared_sort_groupby_num_groups", "unknown".to_string()),
526        ("prepared_sort_groupby_grouped_key_output_bytes", "unknown".to_string()),
527        ("prepared_sort_groupby_max_group_size", "unknown".to_string()),
528        ("partitioned_merge_rejected_reason", partitioned_merge_rejected_reason.unwrap_or("none").to_string()),
529        (
530            "partitioned_merge_contract_status",
531            if partitioned_merge_rejected_reason.is_none() { "available" } else { "blocked_by_shape_or_stage" }
532                .to_string(),
533        ),
534        ("partitioned_merge_required_contract", "partitioned_merge_execution_contract".to_string()),
535        ("default_strategy_available", "true".to_string()),
536    ]);
537    facts.insert(
538        "rejected_candidate_strategies",
539        rejected_candidates
540            .iter()
541            .map(|rejection| format!("{}:{}", rejection.strategy.as_str(), rejection.reason))
542            .collect::<Vec<_>>()
543            .join(";"),
544    );
545    AggregateStrategyCapabilityDecision { selected_strategy, rejected_candidates, facts }
546}
547
548/// First shape or stage reason the encoded-key partial path cannot run `spec`,
549/// or `None`. The check order fixes which reason a spec with several defects
550/// reports.
551fn encoded_key_partial_rejected_reason(
552    spec: &AggregateSpec,
553    input_schema: &Schema,
554    stage: AggregateExecutionStage,
555) -> Option<&'static str> {
556    if stage != AggregateExecutionStage::Partial {
557        return Some("encoded_key_partial_requires_partial_stage");
558    }
559    if spec.grouping_sets.is_some() {
560        return Some("encoded_key_partial_grouping_sets_unsupported");
561    }
562    if spec.group_exprs.is_empty() {
563        return Some("encoded_key_partial_requires_grouped_aggregate");
564    }
565    if !spec.group_exprs.iter().all(is_plain_column_expr) {
566        return Some("encoded_key_partial_requires_plain_column_group_keys");
567    }
568    if aggregate_filter_count(spec) > 0 {
569        return Some("encoded_key_partial_aggregate_filter_unsupported");
570    }
571    if aggregate_spec_has_median(spec) {
572        return Some("encoded_key_partial_median_unsupported");
573    }
574    if distinct_aggregate_count(spec) > 0 {
575        return Some("encoded_key_partial_distinct_aggregate_unsupported");
576    }
577    let mut saw_string_key = false;
578    for expr in &spec.group_exprs {
579        let Some(data_type) = plain_column_group_key_dtype(expr, input_schema) else {
580            return Some("encoded_key_partial_group_key_column_missing");
581        };
582        if aggregate_encoded_key_string_dtype(data_type) {
583            saw_string_key = true;
584        } else if !aggregate_encoded_key_companion_dtype(data_type) {
585            return Some("encoded_key_partial_group_key_dtype_unsupported");
586        }
587    }
588    if !saw_string_key {
589        return Some("encoded_key_partial_requires_string_group_key");
590    }
591    None
592}
593
594/// First shape or stage reason prepared sort groupby cannot run `spec`, or
595/// `None`. `None` does not select it; the executor still needs runtime
596/// cardinality evidence.
597fn prepared_sort_groupby_rejected_reason(
598    spec: &AggregateSpec,
599    input_schema: &Schema,
600    stage: AggregateExecutionStage,
601) -> Option<&'static str> {
602    if !matches!(stage, AggregateExecutionStage::Partial | AggregateExecutionStage::Single) {
603        return Some("prepared_sort_groupby_requires_partial_or_single_stage");
604    }
605    if spec.grouping_sets.is_some() {
606        return Some("prepared_sort_groupby_grouping_sets_unsupported");
607    }
608    if spec.group_exprs.is_empty() {
609        return Some("prepared_sort_groupby_requires_grouped_aggregate");
610    }
611    if spec.aggr_exprs.is_empty() {
612        return Some("prepared_sort_groupby_requires_aggregate_expressions");
613    }
614    if !spec.group_exprs.iter().all(is_plain_column_expr) {
615        return Some("prepared_sort_groupby_requires_plain_column_group_keys");
616    }
617    if aggregate_filter_count(spec) > 0 {
618        return Some("prepared_sort_groupby_aggregate_filter_unsupported");
619    }
620    if stage != AggregateExecutionStage::Single && aggregate_spec_has_median(spec) {
621        return Some("prepared_sort_groupby_median_requires_single_stage");
622    }
623    if distinct_aggregate_count(spec) > 0 {
624        return Some("prepared_sort_groupby_distinct_aggregate_unsupported");
625    }
626    if aggregate_spec_has_floating_statistical_aggregate(spec) {
627        return Some("prepared_sort_groupby_floating_statistical_aggregate_unsupported");
628    }
629    if spec.group_exprs.iter().any(|expr| plain_column_group_key_dtype(expr, input_schema).is_none()) {
630        return Some("prepared_sort_groupby_group_key_column_missing");
631    }
632    None
633}
634
635/// First reason the partitioned merge cannot run `spec` at `stage`, or `None`.
636fn partitioned_merge_rejected_reason(spec: &AggregateSpec, stage: AggregateExecutionStage) -> Option<&'static str> {
637    if spec.grouping_sets.is_some() {
638        return Some("partitioned_merge_grouping_sets_unsupported");
639    }
640    if spec.group_exprs.is_empty() {
641        return Some("partitioned_merge_requires_grouped_aggregate");
642    }
643    if stage != AggregateExecutionStage::FinalMerge {
644        return Some("partitioned_merge_requires_final_merge_stage");
645    }
646    if aggregate_spec_has_floating_statistical_aggregate(spec) {
647        return Some("partitioned_merge_floating_statistical_aggregate_unsupported");
648    }
649    if aggregate_spec_has_median(spec) {
650        return Some("partitioned_merge_median_unsupported");
651    }
652    if distinct_aggregate_count(spec) > 0 {
653        return Some("partitioned_merge_distinct_aggregate_unsupported");
654    }
655    None
656}
657
658fn aggregate_spec_has_floating_statistical_aggregate(spec: &AggregateSpec) -> bool {
659    spec.aggr_exprs.iter().any(|expr| expr.function.is_floating_statistical())
660}
661
662fn aggregate_spec_has_median(spec: &AggregateSpec) -> bool {
663    spec.aggr_exprs.iter().any(|expr| expr.function == AggregateFunction::Median)
664}
665
666fn aggregate_group_key_exprs(spec: &AggregateSpec) -> String {
667    if spec.group_exprs.is_empty() {
668        return "none".to_string();
669    }
670    spec.group_exprs.iter().map(ToString::to_string).collect::<Vec<_>>().join(",")
671}
672
673fn aggregate_group_key_dtypes(spec: &AggregateSpec, input_schema: &Schema) -> String {
674    if spec.group_exprs.is_empty() {
675        return "none".to_string();
676    }
677    spec.group_exprs.iter().map(|expr| aggregate_group_key_dtype(expr, input_schema)).collect::<Vec<_>>().join(",")
678}
679
680fn aggregate_group_key_dtype(expr: &Expr, input_schema: &Schema) -> String {
681    match expr {
682        Expr::Column(column) => input_schema
683            .field_with_name(column.name())
684            .map(|field| format!("{:?}", field.data_type()))
685            .unwrap_or_else(|_| "missing_column".to_string()),
686        Expr::Alias { expr, .. } => aggregate_group_key_dtype(expr, input_schema),
687        _ => "non_column_expr".to_string(),
688    }
689}
690
691fn aggregate_grouping_mode(spec: &AggregateSpec) -> &'static str {
692    if spec.grouping_sets.is_some() {
693        "grouping_sets"
694    } else if spec.group_exprs.is_empty() {
695        "scalar"
696    } else {
697        "grouped"
698    }
699}
700
701fn string_group_key_count(spec: &AggregateSpec, input_schema: &Schema) -> usize {
702    spec.group_exprs
703        .iter()
704        .filter(|expr| plain_column_group_key_dtype(expr, input_schema).is_some_and(aggregate_encoded_key_string_dtype))
705        .count()
706}
707
708fn fixed_width_group_key_count(spec: &AggregateSpec, input_schema: &Schema) -> usize {
709    spec.group_exprs
710        .iter()
711        .filter(|expr| {
712            plain_column_group_key_dtype(expr, input_schema).is_some_and(aggregate_encoded_key_companion_dtype)
713        })
714        .count()
715}
716
717fn unsupported_group_key_dtypes(spec: &AggregateSpec, input_schema: &Schema) -> String {
718    let dtypes = spec
719        .group_exprs
720        .iter()
721        .filter_map(|expr| {
722            let dtype = plain_column_group_key_dtype(expr, input_schema)?;
723            (!aggregate_encoded_key_string_dtype(dtype) && !aggregate_encoded_key_companion_dtype(dtype))
724                .then(|| format!("{dtype:?}"))
725        })
726        .collect::<Vec<_>>();
727    if dtypes.is_empty() { "none".to_string() } else { dtypes.join(",") }
728}
729
730fn plain_column_group_key_dtype<'a>(expr: &Expr, input_schema: &'a Schema) -> Option<&'a DataType> {
731    match expr {
732        Expr::Column(column) => input_schema.field_with_name(column.name()).ok().map(|field| field.data_type()),
733        Expr::Alias { expr, .. } => plain_column_group_key_dtype(expr, input_schema),
734        _ => None,
735    }
736}
737
738fn is_plain_column_expr(expr: &Expr) -> bool {
739    match expr {
740        Expr::Column(_) => true,
741        Expr::Alias { expr, .. } => is_plain_column_expr(expr),
742        _ => false,
743    }
744}
745
746fn aggregate_function_list(spec: &AggregateSpec) -> String {
747    if spec.aggr_exprs.is_empty() {
748        return "none".to_string();
749    }
750    spec.aggr_exprs
751        .iter()
752        .map(|expr| {
753            let mut label = aggregate_function_label(expr.function).to_string();
754            if expr.distinct {
755                label.push_str("_distinct");
756            }
757            label
758        })
759        .collect::<Vec<_>>()
760        .join(",")
761}
762
763fn aggregate_function_label(function: AggregateFunction) -> &'static str {
764    function.as_str()
765}
766
767fn aggregate_filter_count(spec: &AggregateSpec) -> usize {
768    spec.aggr_exprs.iter().filter(|expr| expr.filter.is_some()).count()
769}
770
771fn distinct_aggregate_count(spec: &AggregateSpec) -> usize {
772    spec.aggr_exprs.iter().filter(|expr| aggregate_expr_is_exact_distinct_count(expr)).count()
773}
774
775fn aggregate_encoded_key_string_dtype(data_type: &DataType) -> bool {
776    matches!(data_type, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View)
777}
778
779fn aggregate_encoded_key_companion_dtype(data_type: &DataType) -> bool {
780    is_arrow_fixed_width(data_type) || matches!(data_type, DataType::Float16)
781}
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786
787    #[test]
788    fn aggregate_strategy_execution_options_reject_ratio_above_one_hundred_percent() {
789        let error = AggregateStrategyExecutionOptions::default()
790            .with_prepared_sort_groupby_low_cardinality_threshold_ppm(1_000_001)
791            .validate()
792            .unwrap_err();
793
794        assert_eq!(error.code(), QueryErrorCode::InvalidConfiguration);
795        assert_eq!(error.fact_value("field"), Some("prepared_sort_groupby_low_cardinality_threshold_ppm"));
796        assert_eq!(error.fact_value("max_value"), Some("1000000"));
797        assert_eq!(error.fact_value("reason"), Some("ratio_ppm_above_one_hundred_percent"));
798    }
799
800    #[test]
801    fn aggregate_strategy_execution_options_reject_contradictory_probe_bounds() {
802        let error = AggregateStrategyExecutionOptions::default()
803            .with_exact_grouped_key_cardinality_probe_min_input_rows(10)
804            .with_exact_three_key_cardinality_probe_max_input_rows(9)
805            .validate()
806            .unwrap_err();
807
808        assert_eq!(error.code(), QueryErrorCode::InvalidConfiguration);
809        assert_eq!(error.fact_value("field"), Some("exact_grouped_key_cardinality_probe_min_input_rows"));
810        assert_eq!(error.fact_value("bound_field"), Some("exact_three_key_cardinality_probe_max_input_rows"));
811        assert_eq!(
812            error.fact_value("reason"),
813            Some("three_key_probe_max_input_rows_below_grouped_probe_min_input_rows")
814        );
815    }
816
817    #[test]
818    fn aggregate_strategy_execution_options_accept_zero_threshold_overrides() {
819        AggregateStrategyExecutionOptions::default()
820            .with_encoded_key_partial_low_cardinality_threshold_ppm(0)
821            .with_prepared_sort_groupby_low_cardinality_threshold_ppm(0)
822            .with_exact_grouped_key_cardinality_probe_min_input_rows(0)
823            .with_exact_three_key_cardinality_probe_max_input_rows(0)
824            .with_source_adjacent_prepared_sort_groupby_max_row_groups(0)
825            .with_source_adjacent_prepared_sort_groupby_max_source_read_bytes(0)
826            .with_source_adjacent_prepared_sort_groupby_max_fan_in(0)
827            .validate()
828            .unwrap();
829    }
830}
query_planning/capability/aggregate.rs

query_planning/capability/
aggregate.rs
