linkrs_query/optimizer/cost_based/aggregate_strategy/

selector.rs

//! Cost-based aggregation strategy selector.

use std::sync::Arc;

use crate::optimizer::cost::CostCalculator;
use crate::optimizer::cost::SelectivityEstimator;
use crate::optimizer::cost_based::ndv::estimate_group_count;
use crate::optimizer::cost_based::row_estimates::estimate_node_output_rows_logical;
use crate::optimizer::cost_based::traversal_logical::rewrite_children_logical;
use crate::optimizer::decision::OptimizationDecision;
use crate::optimizer::stats::StatsView;
use crate::planning::plan::logical::logical_node_traits::LogicalSingleInputNode;
use crate::planning::plan::logical::LogicalNodeEnum;
use linkrs_core::types::ContextualExpression;

use super::context::AggregateContext;
use super::strategy::{AggregateStrategy, AggregateStrategyDecision, SelectionReason};

/// Aggregation Policy Selector
///
/// Selecting the optimal aggregation execution strategy based on a cost model
#[derive(Debug)]
pub struct AggregateStrategySelector {
    cost_calculator: Arc<CostCalculator>,
}

impl AggregateStrategySelector {
    /// Create a new selector for aggregating policies.
    pub fn new(cost_calculator: Arc<CostCalculator>) -> Self {
        Self { cost_calculator }
    }

    /// Selecting the optimal aggregation strategy
    ///
    /// # Parameters
    /// Context: Information about the context in which the aggregation operation is taking place.
    ///
    /// # Return
    /// Aggregation policy decision-making, including the selected policy and the estimated cost.
    pub fn select_strategy(
        &self,
        space: &str,
        context: &AggregateContext,
    ) -> AggregateStrategyDecision {
        self.select_strategy_internal(space, context, false)
    }

    /// Selecting the optimal aggregation strategy with memory pressure awareness
    ///
    /// This method considers the global memory pressure threshold from the cost model configuration
    /// to make more robust decisions for large-scale aggregations.
    ///
    /// # Parameters
    /// - `context`: Information about the context in which the aggregation operation is taking place.
    ///
    /// # Return
    /// Aggregation policy decision-making, including the selected policy and the estimated cost.
    pub fn select_strategy_with_memory_pressure(
        &self,
        space: &str,
        context: &AggregateContext,
    ) -> AggregateStrategyDecision {
        self.select_strategy_internal(space, context, true)
    }

    /// Internal method for selecting aggregation strategy
    fn select_strategy_internal(
        &self,
        space: &str,
        context: &AggregateContext,
        consider_memory_pressure: bool,
    ) -> AggregateStrategyDecision {
        // If the input is already sorted and the sorting key matches the grouping key, stream aggregation should be used preferentially.
        if context.input_is_sorted && context.sort_keys_match_group_keys {
            return self.create_streaming_aggregate_decision(space, context);
        }

        // If the expression is non-deterministic, hash aggregation should be used preferentially (to avoid the uncertainties associated with sorting).
        if !context.is_deterministic {
            let group_by_cardinality = self.estimate_group_by_cardinality(space, context);
            let hash_cost = self.calculate_hash_aggregate_cost(context, group_by_cardinality);
            let hash_memory = self.estimate_hash_memory_usage(context, group_by_cardinality);
            return AggregateStrategyDecision {
                strategy: AggregateStrategy::HashAggregate,
                estimated_output_rows: group_by_cardinality.max(1),
                estimated_cost: hash_cost,
                estimated_memory_bytes: hash_memory,
                reason: SelectionReason::CostBased {
                    hash_cost,
                    sort_cost: hash_cost * 1.5, // Assume that the cost of sorting is higher.
                },
            };
        }

        // Estimating the cardinality of the group key
        let group_by_cardinality = self.estimate_group_by_cardinality(space, context);

        // Calculate the cost of each strategy.
        let hash_cost = self.calculate_hash_aggregate_cost(context, group_by_cardinality);
        let sort_cost = self.calculate_sort_aggregate_cost(context);

        // Check the memory limitations.
        let hash_memory = self.estimate_hash_memory_usage(context, group_by_cardinality);
        let memory_constrained = context.memory_limit > 0 && hash_memory > context.memory_limit;

        // Check global memory pressure if enabled
        let memory_pressure_threshold =
            self.cost_calculator.config().memory_pressure_threshold as u64;
        let under_memory_pressure =
            consider_memory_pressure && hash_memory > memory_pressure_threshold;

        // Decision-making logic
        let (strategy, reason) = if memory_constrained || under_memory_pressure {
            // Memory is limited or under global memory pressure
            // Priority should be given to sorting and aggregation operations
            if under_memory_pressure && sort_cost < hash_cost * 1.5 {
                // Allow up to 50% performance loss to avoid memory pressure
                (
                    AggregateStrategy::SortAggregate,
                    SelectionReason::MemoryConstrained,
                )
            } else if memory_constrained {
                (
                    AggregateStrategy::SortAggregate,
                    SelectionReason::MemoryConstrained,
                )
            } else {
                // Even under pressure, hash is significantly better
                (
                    AggregateStrategy::HashAggregate,
                    SelectionReason::CostBased {
                        hash_cost,
                        sort_cost,
                    },
                )
            }
        } else if context.input_rows
            < self
                .cost_calculator
                .config()
                .strategy_thresholds
                .small_dataset_threshold
        {
            // With small amounts of data, hash aggregation is simpler and more efficient.
            (
                AggregateStrategy::HashAggregate,
                SelectionReason::SmallDataSet,
            )
        } else if group_by_cardinality
            < self
                .cost_calculator
                .config()
                .strategy_thresholds
                .low_cardinality_threshold
        {
            // With a small base size, sorting and aggregation may be more advantageous (as the data becomes more localized after sorting).
            if sort_cost < hash_cost * 1.2 {
                (
                    AggregateStrategy::SortAggregate,
                    SelectionReason::LowCardinality,
                )
            } else {
                (
                    AggregateStrategy::HashAggregate,
                    SelectionReason::CostBased {
                        hash_cost,
                        sort_cost,
                    },
                )
            }
        } else if group_by_cardinality
            > (context.input_rows as f64
                * self
                    .cost_calculator
                    .config()
                    .strategy_thresholds
                    .high_cardinality_ratio) as u64
        {
            // For high cardinalities (close to unique values), hash aggregation is more advantageous.
            (
                AggregateStrategy::HashAggregate,
                SelectionReason::HighCardinality,
            )
        } else {
            // Based on cost comparison
            if hash_cost <= sort_cost {
                (
                    AggregateStrategy::HashAggregate,
                    SelectionReason::CostBased {
                        hash_cost,
                        sort_cost,
                    },
                )
            } else {
                (
                    AggregateStrategy::SortAggregate,
                    SelectionReason::CostBased {
                        hash_cost,
                        sort_cost,
                    },
                )
            }
        };

        let estimated_cost = match strategy {
            AggregateStrategy::HashAggregate => hash_cost,
            AggregateStrategy::SortAggregate => sort_cost,
            AggregateStrategy::StreamingAggregate => {
                self.calculate_streaming_aggregate_cost(context)
            }
        };

        AggregateStrategyDecision {
            strategy,
            estimated_output_rows: group_by_cardinality.max(1),
            estimated_cost,
            estimated_memory_bytes: match strategy {
                AggregateStrategy::HashAggregate => hash_memory,
                AggregateStrategy::SortAggregate => self.estimate_sort_memory_usage(context),
                AggregateStrategy::StreamingAggregate => {
                    self.estimate_streaming_memory_usage(context)
                }
            },
            reason,
        }
    }

    /// Rapid selection strategy (simplified version, for use in decision-making caching)
    pub fn select_strategy_quick(
        &self,
        input_rows: u64,
        group_key_count: usize,
        _agg_function_count: usize,
    ) -> AggregateStrategy {
        let thresholds = self.cost_calculator.config().strategy_thresholds;
        if input_rows < thresholds.small_dataset_threshold {
            return AggregateStrategy::HashAggregate;
        }

        // Estimating the cardinality of the grouping key
        let cardinality = self.estimate_cardinality_quick(input_rows, group_key_count);

        if cardinality < thresholds.low_cardinality_threshold {
            AggregateStrategy::SortAggregate
        } else {
            AggregateStrategy::HashAggregate
        }
    }

    /// Estimating the cardinality of the grouping key
    ///
    /// First tries to use property combination statistics if available,
    /// then falls back to heuristic estimation.
    fn estimate_group_by_cardinality(&self, space: &str, context: &AggregateContext) -> u64 {
        // Try to use property combination statistics for better estimation
        if let Some(table_name) = self.get_table_name_from_context(context) {
            if let Some(cardinality) = self
                .cost_calculator
                .statistics_manager()
                .get_combined_cardinality(space, Some(&table_name), &context.group_keys)
            {
                return cardinality.min(context.input_rows).max(1);
            }
        }

        // Fall back to heuristic estimation
        self.estimate_cardinality_quick(context.input_rows, context.group_keys.len())
    }

    /// Get table name from aggregate context
    /// This is a helper method to extract table name for statistics lookup
    fn get_table_name_from_context(&self, context: &AggregateContext) -> Option<String> {
        context.table_name.clone()
    }

    /// Quick estimation of the base number
    fn estimate_cardinality_quick(&self, input_rows: u64, key_count: usize) -> u64 {
        estimate_group_count(input_rows, key_count)
    }

    /// Calculating the cost of hash aggregation
    fn calculate_hash_aggregate_cost(
        &self,
        context: &AggregateContext,
        _group_by_cardinality: u64,
    ) -> f64 {
        self.cost_calculator.calculate_aggregate_cost(
            context.input_rows,
            context.agg_function_count,
            context.group_keys.len(),
        )
    }

    /// Calculating the cost of sorting and aggregation operations
    fn calculate_sort_aggregate_cost(&self, context: &AggregateContext) -> f64 {
        // Sorting cost + Aggregation cost
        let sort_cost = self.cost_calculator.calculate_sort_cost(
            context.input_rows,
            context.group_keys.len(),
            None,
        );

        // The aggregated cost after sorting is lower (the data has been grouped).
        let agg_cost = context.input_rows as f64
            * context.agg_function_count as f64
            * self.cost_calculator.config().cpu_operator_cost
            * 0.5; // The cost of aggregation is reduced by half after sorting.

        sort_cost + agg_cost
    }

    /// Calculating the cost of stream aggregation
    fn calculate_streaming_aggregate_cost(&self, context: &AggregateContext) -> f64 {
        // Stream aggregation only requires the cost of aggregation; sorting is not necessary.
        context.input_rows as f64
            * context.agg_function_count as f64
            * self.cost_calculator.config().cpu_operator_cost
    }

    /// Estimating the memory usage of hash aggregation
    fn estimate_hash_memory_usage(
        &self,
        _context: &AggregateContext,
        group_by_cardinality: u64,
    ) -> u64 {
        // Single row-width source shared with the cost calculator and the
        // memory budget allocator; hash overhead is a fixed factor.
        group_by_cardinality * self.cost_calculator.config().row_width_bytes as u64 * 2
    }

    /// Estimating the memory usage for sorting and aggregation operations
    fn estimate_sort_memory_usage(&self, context: &AggregateContext) -> u64 {
        // Sorting may require caching all the data; width from shared config.
        let row_size = self.cost_calculator.config().row_width_bytes as u64;
        context.input_rows * row_size
    }

    /// Estimating the memory usage of stream aggregation
    fn estimate_streaming_memory_usage(&self, _context: &AggregateContext) -> u64 {
        // Streaming keeps a small window; width still from the shared source.
        self.cost_calculator.config().row_width_bytes as u64 * 2
    }

    /// Creating streaming aggregation decisions
    fn create_streaming_aggregate_decision(
        &self,
        space: &str,
        context: &AggregateContext,
    ) -> AggregateStrategyDecision {
        let estimated_cost = self.calculate_streaming_aggregate_cost(context);
        let group_by_cardinality = self.estimate_group_by_cardinality(space, context);

        AggregateStrategyDecision {
            strategy: AggregateStrategy::StreamingAggregate,
            estimated_output_rows: group_by_cardinality.max(1),
            estimated_cost,
            estimated_memory_bytes: self.estimate_streaming_memory_usage(context),
            reason: SelectionReason::InputAlreadySorted,
        }
    }

    /// Update and optimize the information on the aggregation strategies used in decision-making processes.
    pub fn update_decision(
        &self,
        space: &str,
        decision: &mut OptimizationDecision,
        context: &AggregateContext,
    ) {
        let strategy_decision = self.select_strategy(space, context);

        // Encode the aggregation policy information into the `rewrite_rules` used for decision-making.
        // Use a specific rule ID to indicate the selection of the aggregation policy.
        let rule_id = match strategy_decision.strategy {
            AggregateStrategy::HashAggregate => {
                crate::optimizer::decision::RewriteRuleId::AggregateOptimization
            }
            AggregateStrategy::SortAggregate => {
                crate::optimizer::decision::RewriteRuleId::AggregateOptimization
            }
            AggregateStrategy::StreamingAggregate => {
                crate::optimizer::decision::RewriteRuleId::AggregateOptimization
            }
        };

        if !decision.rewrite_rules.contains(&rule_id) {
            decision.rewrite_rules.push(rule_id);
        }
    }

    /// Analyze aggregation expression complexity using expression cost
    ///
    /// Returns the total expression cost for all aggregation functions.
    /// Higher cost indicates more complex expressions that benefit from
    /// HashAggregate (to avoid redundant computations).
    pub fn analyze_aggregation_complexity(&self, agg_functions: &[ContextualExpression]) -> f64 {
        agg_functions
            .iter()
            .filter_map(|f| f.get_expression())
            .map(|expr| self.cost_calculator.calculate_expression_cost(&expr))
            .sum()
    }

    /// Select strategy considering both memory pressure and expression complexity
    ///
    /// This is the most comprehensive strategy selection method that considers:
    /// - Input sorting
    /// - Memory constraints and pressure
    /// - Expression complexity
    /// - Cost comparison
    pub fn select_strategy_comprehensive(
        &self,
        space: &str,
        context: &AggregateContext,
        agg_functions: &[ContextualExpression],
    ) -> AggregateStrategyDecision {
        // First check if we can use streaming aggregation
        if context.input_is_sorted && context.sort_keys_match_group_keys {
            return self.create_streaming_aggregate_decision(space, context);
        }

        // Analyze expression complexity
        let total_expression_cost = self.analyze_aggregation_complexity(agg_functions);
        let avg_expression_cost = if agg_functions.is_empty() {
            0.0
        } else {
            total_expression_cost / agg_functions.len() as f64
        };

        // Get base decision with memory pressure awareness
        let mut decision = self.select_strategy_with_memory_pressure(space, context);

        // If expressions are complex and we're using SortAggregate,
        // consider switching to HashAggregate to avoid redundant computations
        if decision.strategy == AggregateStrategy::SortAggregate
            && avg_expression_cost > self.cost_calculator.config().function_call_base_cost * 2.0
        {
            let group_by_cardinality = self.estimate_group_by_cardinality(space, context);
            let hash_cost = self.calculate_hash_aggregate_cost(context, group_by_cardinality);
            let sort_cost = decision.estimated_cost;

            // If hash aggregation is not too much more expensive, use it for complex expressions
            if hash_cost < sort_cost * 1.3 {
                let hash_memory = self.estimate_hash_memory_usage(context, group_by_cardinality);
                decision = AggregateStrategyDecision {
                    strategy: AggregateStrategy::HashAggregate,
                    estimated_output_rows: group_by_cardinality.max(1),
                    estimated_cost: hash_cost,
                    estimated_memory_bytes: hash_memory,
                    reason: SelectionReason::CostBased {
                        hash_cost,
                        sort_cost,
                    },
                };
            }
        }

        decision
    }
}

// =====================================================================
// Logical-plan variant (PlanNodeEnum logic/physical separation).
//
// The aggregation strategy decision is taken on the pure logical tree
// (`LogicalNodeEnum`) and emits the same `aggregate:` note as the physical
// walker in the engine. The tree is returned unchanged — the strategy is a
// decision note consumed by the physical planner.
// =====================================================================

/// Walk a logical plan tree and record the aggregation strategy decision
/// for every aggregate node as a CBO note (cost-driven, observable in
/// EXPLAIN). The tree itself is returned unchanged.
pub fn walk_aggregate_strategies_logical(
    node: &LogicalNodeEnum,
    stats: &StatsView,
    selector: &AggregateStrategySelector,
    selectivity: &SelectivityEstimator,
    notes: &mut Vec<String>,
) -> LogicalNodeEnum {
    use LogicalNodeEnum::*;

    if let Aggregate(aggregate) = node {
        let input_rows = estimate_node_output_rows_logical(aggregate.input(), stats, selectivity);
        let context = AggregateContext {
            input_rows,
            group_keys: aggregate
                .group_key_exprs
                .iter()
                .map(|e| e.to_expression_string())
                .collect(),
            agg_function_count: aggregate.aggregation_functions.len(),
            memory_limit: 0,
            input_is_sorted: false,
            sort_keys_match_group_keys: false,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };
        let decision = selector.select_strategy(stats.space().unwrap_or(""), &context);
        notes.push(format!(
            "aggregate: strategy={:?} (reason={:?}, est_rows={})",
            decision.strategy, decision.reason, decision.estimated_output_rows
        ));
    }

    let mut closure = |child: &LogicalNodeEnum| {
        walk_aggregate_strategies_logical(child, stats, selector, selectivity, notes)
    };
    rewrite_children_logical(node, &mut closure)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::optimizer::stats::StatisticsManager;

    fn create_test_selector() -> AggregateStrategySelector {
        let stats_manager = Arc::new(StatisticsManager::new());
        let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
        AggregateStrategySelector::new(cost_calculator)
    }

    #[test]
    fn test_streaming_aggregate_when_sorted() {
        let selector = create_test_selector();
        let context = AggregateContext {
            input_rows: 10000,
            group_keys: vec!["category".to_string()],
            agg_function_count: 2,
            memory_limit: 0,
            input_is_sorted: true,
            sort_keys_match_group_keys: true,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };

        let decision = selector.select_strategy("test", &context);
        assert_eq!(decision.strategy, AggregateStrategy::StreamingAggregate);
        matches!(decision.reason, SelectionReason::InputAlreadySorted);
    }

    #[test]
    fn test_hash_aggregate_for_small_data() {
        let selector = create_test_selector();
        let context = AggregateContext {
            input_rows: 500,
            group_keys: vec!["category".to_string()],
            agg_function_count: 1,
            memory_limit: 0,
            input_is_sorted: false,
            sort_keys_match_group_keys: false,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };

        let decision = selector.select_strategy("test", &context);
        assert_eq!(decision.strategy, AggregateStrategy::HashAggregate);
        matches!(decision.reason, SelectionReason::SmallDataSet);
    }

    #[test]
    fn test_memory_constrained_fallback() {
        let selector = create_test_selector();
        let context = AggregateContext {
            input_rows: 100000,
            group_keys: vec!["category".to_string(), "subcategory".to_string()],
            agg_function_count: 3,
            memory_limit: 1024, // Memory limit of 1 KB (very small)
            input_is_sorted: false,
            sort_keys_match_group_keys: false,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };

        let decision = selector.select_strategy("test", &context);
        assert_eq!(decision.strategy, AggregateStrategy::SortAggregate);
        matches!(decision.reason, SelectionReason::MemoryConstrained);
    }

    #[test]
    fn test_quick_selection() {
        let selector = create_test_selector();

        // For small amounts of data, hash aggregation should be chosen.
        let strategy = selector.select_strategy_quick(500, 1, 1);
        assert_eq!(strategy, AggregateStrategy::HashAggregate);

        // For large amounts of data with multiple keys (and a low cardinality), sorting and aggregation should be chosen as the appropriate methods for processing the data.
        let strategy = selector.select_strategy_quick(100000, 10, 1);
        assert_eq!(strategy, AggregateStrategy::SortAggregate);

        // For large volumes of data with a high number of unique keys (high cardinality), hash aggregation should be chosen.
        let strategy = selector.select_strategy_quick(100000, 1, 1);
        assert_eq!(strategy, AggregateStrategy::HashAggregate);
    }

    #[test]
    fn test_cardinality_estimation() {
        let selector = create_test_selector();

        // If there is no grouping key, the result should be 1.
        let cardinality = selector.estimate_cardinality_quick(1000, 0);
        assert_eq!(cardinality, 1);

        // The cardinality of a single key should be high.
        let cardinality = selector.estimate_cardinality_quick(10000, 1);
        assert!(cardinality > 100);

        // The cardinality of multiple keys should be relatively low.
        let cardinality = selector.estimate_cardinality_quick(10000, 5);
        assert!(cardinality < 1000);
    }

    #[test]
    fn test_memory_pressure_aware_selection() {
        let selector = create_test_selector();

        // Create a context that would normally choose HashAggregate
        // but with high memory usage that exceeds the threshold
        let context = AggregateContext {
            input_rows: 100000,
            group_keys: vec!["category".to_string()],
            agg_function_count: 2,
            memory_limit: 0, // No explicit limit
            input_is_sorted: false,
            sort_keys_match_group_keys: false,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };

        // Without memory pressure, should likely choose HashAggregate for high cardinality
        let _normal_decision = selector.select_strategy("test", &context);

        // With memory pressure awareness, check if it considers memory threshold
        let pressure_decision = selector.select_strategy_with_memory_pressure("test", &context);

        // Both should return valid decisions
        assert!(pressure_decision.estimated_cost > 0.0);
        assert!(pressure_decision.estimated_memory_bytes > 0);
    }

    #[test]
    fn test_analyze_aggregation_complexity() {
        let selector = create_test_selector();

        // Test with empty functions
        let empty_complexity = selector.analyze_aggregation_complexity(&[]);
        assert_eq!(empty_complexity, 0.0);

        // Test with simple expressions
        use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
        use linkrs_core::types::expr::Expression;
        use linkrs_core::value::Value;

        let ctx = ExpressionAnalysisContext::new();
        let expr = Expression::Literal(Value::Int(42));
        let id = ctx.register_expression(linkrs_core::types::expr::ExpressionMeta::new(expr));
        let simple_expr = ContextualExpression::new(id, std::sync::Arc::new(ctx));
        let simple_complexity = selector.analyze_aggregation_complexity(&[simple_expr]);
        assert!(simple_complexity >= 0.0);
    }

    #[test]
    fn test_comprehensive_strategy_selection() {
        let selector = create_test_selector();

        let context = AggregateContext {
            input_rows: 10000,
            group_keys: vec!["category".to_string()],
            agg_function_count: 2,
            memory_limit: 0,
            input_is_sorted: false,
            sort_keys_match_group_keys: false,
            is_deterministic: true,
            complexity_score: 0,
            table_name: None,
        };

        // Test comprehensive selection with empty functions
        let decision = selector.select_strategy_comprehensive("test", &context, &[]);
        assert!(decision.estimated_cost > 0.0);
        assert!(decision.estimated_memory_bytes > 0);
    }

    #[test]
    fn logical_walk_emits_aggregate_strategy_note() {
        use crate::optimizer::cost::SelectivityEstimator;
        use crate::optimizer::stats::StatsView;
        use crate::planning::plan::logical::logical_nodes::access::LogicalScanVerticesNode;
        use crate::planning::plan::logical::logical_nodes::operation::LogicalAggregateNode;
        use crate::planning::plan::logical::LogicalNodeEnum;
        use linkrs_core::types::operators::AggregateFunction;

        let stats_manager = Arc::new(StatisticsManager::new());
        let cost_calculator = Arc::new(CostCalculator::new(stats_manager.clone()));
        let selector = AggregateStrategySelector::new(cost_calculator);
        let selectivity = SelectivityEstimator::new(stats_manager.clone());
        let stats = StatsView::new(&stats_manager, Some("test"));

        let scan = LogicalNodeEnum::ScanVertices(LogicalScanVerticesNode {
            id: 1,
            space_id: 1,
            space_name: "test".to_string(),
            tag: Some("person".to_string()),
            expression: None,
            limit: None,
            projected_properties: vec![],
            index_hint: None,
            estimated_cardinality: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        });
        let ctx = std::sync::Arc::new(
            linkrs_core::types::expr::expression_context::ExpressionAnalysisContext::new(),
        );
        let expr = linkrs_core::Expression::Variable("n.age".to_string());
        let meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
        let id = ctx.register_expression(meta);
        let ctx_expr = linkrs_core::types::expr::contextual::ContextualExpression::new(id, ctx);
        let aggregate = LogicalNodeEnum::Aggregate(LogicalAggregateNode {
            id: 2,
            input: Some(Box::new(scan)),
            group_key_exprs: vec![ctx_expr],
            aggregation_functions: vec![AggregateFunction::Count],
            aggregation_args: vec![],
            aggregation_distinct: vec![],
            aggregation_filters: vec![],
            grouping_sets: vec![],
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        });

        let mut notes = Vec::new();
        let rewritten = walk_aggregate_strategies_logical(
            &aggregate,
            &stats,
            &selector,
            &selectivity,
            &mut notes,
        );

        // The logical tree is returned unchanged; the decision is a note.
        assert!(matches!(rewritten, LogicalNodeEnum::Aggregate(_)));
        assert_eq!(notes.len(), 1);
        assert!(notes[0].starts_with("aggregate: strategy="));
        assert!(notes[0].contains("est_rows="));
    }
}