algeon-datafusion 0.1.0

use arrow::{array::Int32Array, record_batch::RecordBatch};
use cudf::testing;
use datafusion::datasource::memory::MemorySourceConfig;

use super::*;
use crate::native::ParquetReadStrategyConfig;

#[test]
fn benchmark_shape_lowers_to_one_native_query_plan() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering");
    let path = dir.join("input.parquet");
    write_float_measure_parquet_file(&path);
    let plan = aggregate_sort_limit_plan(parquet_scan(&path), Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.candidate_count(), 1);
    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    assert!(result.plan.is_some());
    assert_eq!(result.selected_plan().node_count(), 6);
    assert!(matches!(result.report.outcomes()[0], CandidateOutcome::Selected { .. }));

    let plan = result.selected_plan();
    let PlanNode::Limit(limit) = &plan.root_node().kind else {
        panic!("native root should be a limit node");
    };
    assert_eq!(limit.spec.fetch, Some(2));
    assert!(matches!(plan.sink(), ResultSink::HostArrow { input, .. } if *input == plan.root_id()));
    let sort_node = plan.node(limit.input).expect("limit input should exist");
    let PlanNode::Sort(sort) = &sort_node.kind else {
        panic!("limit input should be a sort node");
    };
    assert_eq!(sort.spec.expressions.len(), 2);
    assert_eq!(sort_node.facts.output_ordering.len(), 2);
    let aggregate_node = plan.node(sort.input).expect("sort input should exist");
    let PlanNode::Aggregate(aggregate) = &aggregate_node.kind else {
        panic!("sort input should be an aggregate node");
    };
    assert_eq!(aggregate.spec.aggr_exprs.len(), 1);
    let projection_node = plan.node(aggregate.input).expect("aggregate input should exist");
    let PlanNode::Projection(projection) = &projection_node.kind else {
        panic!("aggregate input should be a projection node");
    };
    assert_eq!(projection.outputs.len(), 2);
    let filter_node = plan.node(projection.input).expect("projection input should exist");
    let PlanNode::Filter(filter) = &filter_node.kind else {
        panic!("projection input should be a filter node");
    };
    assert!(matches!(
        &filter.predicate,
        Expr::Binary { right, .. } if matches!(right.as_ref(), Expr::Literal(LiteralValue::Int32(10)))
    ));
    let source_node = plan.node(filter.input).expect("filter input should exist");
    let PlanNode::Source(source_node) = &source_node.kind else {
        panic!("filter input should be a local parquet source");
    };
    let SourceNode::LocalParquet(source) = source_node.as_ref() else {
        panic!("filter input should be a local parquet source");
    };
    assert_eq!(source.file_count(), 1);
    assert_eq!(source.files[0].file_size_bytes, Some(fs::metadata(&path).unwrap().len()));
    assert_eq!(
        source.read_columns,
        vec!["key".to_string(), "m00".to_string(), "m01".to_string(), "payload".to_string()],
    );
    assert_eq!(source.read_column_file_indices, vec![0, 2, 3, 1]);
    assert_eq!(source.projected_output_columns, vec!["key".to_string(), "m00".to_string(), "m01".to_string()]);
    assert_eq!(
        source.projected_output_column_source_refs,
        vec![SourceColumnRef::File(0), SourceColumnRef::File(2), SourceColumnRef::File(3)]
    );
    assert_eq!(source.pure_filter_columns, vec!["payload".to_string()]);
    assert_eq!(source.pure_filter_column_source_refs, vec![SourceColumnRef::File(1)]);
    assert_eq!(source.diagnostics.read_column_count, 4);
    assert_eq!(source.diagnostics.projected_output_column_count, 3);
    assert_eq!(source.diagnostics.pure_filter_column_count, 1);
    assert_eq!(source.diagnostics.byte_range_disposition.as_str(), "full_files");
    assert_eq!(source.diagnostics.partition_column_disposition.as_str(), "absent");
    assert_eq!(source.diagnostics.zero_column_count_disposition.as_str(), "not_required");
    assert_eq!(source.diagnostics.schema_mapping_disposition.as_str(), "direct_by_name");
    assert!(result.selected_plan().to_string().contains("LocalParquetSource"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn multi_full_file_parquet_scan_lowers_to_one_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-multi-file");
    let first = dir.join("part-000.parquet");
    let second = dir.join("part-001.parquet");
    write_numeric_parquet_file(&first);
    write_numeric_parquet_file(&second);
    let plan = aggregate_sort_limit_plan(parquet_scan(&dir), Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();
    let source = plan
        .nodes()
        .iter()
        .find_map(|node| match &node.kind {
            PlanNode::Source(source_node) => match source_node.as_ref() {
                SourceNode::LocalParquet(source) => Some(source),
                SourceNode::GpuResident(_) => None,
                SourceNode::LiteralBatch(_) => None,
            },
            _ => None,
        })
        .expect("native plan should contain a local parquet source");
    assert_eq!(source.file_count(), 2);
    assert_eq!(
        source.diagnostics.known_bytes,
        Some(fs::metadata(&first).unwrap().len() + fs::metadata(&second).unwrap().len())
    );
    assert!(plan.to_string().contains("LocalParquetSource files=2"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn explicit_parquet_row_group_access_plan_lowers_to_native_row_groups() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-row-groups");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let access_plan = ParquetAccessPlan::new(vec![
        RowGroupAccess::Skip,
        RowGroupAccess::Scan,
        RowGroupAccess::Skip,
        RowGroupAccess::Scan,
    ]);
    let plan = aggregate_sort_limit_plan(parquet_scan_with_access_plan(&path, access_plan), Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1);
    assert_eq!(result.not_selected_count(), 0);
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.file_count(), 1);
    assert_eq!(source.files[0].row_groups.as_deref(), Some(&[1, 3][..]));
    assert_eq!(source.selected_row_group_count(), 2);
    assert_eq!(source.diagnostics.file_count, 1);
    assert_eq!(source.diagnostics.ranged_file_count, 0);
    assert_eq!(source.diagnostics.selected_row_groups, Some(2));
    assert_eq!(source.diagnostics.predicate_disposition.as_str(), "native_post_scan_only");
    assert_eq!(source.diagnostics.read_column_count, 4);
    assert_eq!(source.diagnostics.projected_output_column_count, 3);
    assert_eq!(source.diagnostics.pure_filter_column_count, 1);
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].file_count, 1);
    assert_eq!(result.source_diagnostics[0].ranged_file_count, 0);
    assert_eq!(result.source_diagnostics[0].selected_row_groups, Some(2));
    assert!(result.source_diagnostics[0].rejected_features.is_empty());
    assert!(result.selected_plan().to_string().contains("selected_row_groups=2"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn parquet_row_level_access_plan_is_rejected_with_source_reason() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-row-selection");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let selection = RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(1)]);
    let access_plan = ParquetAccessPlan::new(vec![RowGroupAccess::Selection(selection)]);
    let plan = aggregate_sort_limit_plan(parquet_scan_with_access_plan(&path, access_plan), Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::NativeSourceRowSelection, .. }
    ));
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].file_count, 1);
    assert_eq!(result.source_diagnostics[0].selected_row_groups, Some(1));
    assert!(result.source_diagnostics[0].rejected_features.iter().any(|feature| feature.as_str() == "row_selection"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn integer_modulo_projection_expression_lowers_to_native_planning() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-integer-modulo-shape");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = aggregate_sort_limit_plan(parquet_scan(&path), Operator::Modulo);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn modulo_with_floating_point_operands_lowers_to_native_planning() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-modulo-float");
    let path = dir.join("input.parquet");
    write_float_measure_parquet_file(&path);
    let plan = aggregate_sort_limit_plan(parquet_scan(&path), Operator::Modulo);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn modulo_with_decimal_operands_is_rejected_during_native_planning() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-modulo-decimal");
    let path = dir.join("input.parquet");
    write_decimal_measure_parquet_file(&path);
    let plan = aggregate_sort_limit_plan(parquet_scan(&path), Operator::Modulo);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::UnsupportedType, .. }
    ));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn tpch_like_source_types_do_not_reject_native_planning() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-tpch-types");
    let path = dir.join("input.parquet");
    write_parquet_file_with_tpch_like_columns(&path);
    let scan = parquet_scan(&path);
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(
        source.read_columns,
        vec!["key".to_string(), "m00".to_string(), "m01".to_string(), "payload".to_string()]
    );

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn order_insensitive_pushed_down_scan_limit_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-scan-property");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let scan = parquet_scan(&path).with_fetch(Some(1)).unwrap();
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.scan_limit, Some(1));
    assert_eq!(source.diagnostics.scan_limit_disposition.as_str(), "represented");
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].scan_limit_disposition.as_str(), "represented");

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn order_sensitive_single_file_group_scan_limit_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-order-sensitive-scan-limit");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let scan = parquet_scan_with_order_sensitive_limit(&path, 1);
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.scan_limit, Some(1));
    assert_eq!(source.diagnostics.scan_limit_disposition.as_str(), "represented");
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].scan_limit_disposition.as_str(), "represented");
    assert!(result.source_diagnostics[0].rejected_features.iter().all(|feature| feature.as_str() != "preserve_order"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn parquet_sorting_columns_scan_limit_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sorting-columns-scan-limit");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file_with_sorting_columns(&path);
    let scan = parquet_scan(&path).with_fetch(Some(1)).unwrap();
    assert!(scan.properties().output_ordering().is_some());

    let result = lower_physical_plan_to_native(scan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.scan_limit, Some(1));
    assert_eq!(source.diagnostics.scan_limit_disposition.as_str(), "represented");
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].scan_limit_disposition.as_str(), "represented");
    assert!(result.source_diagnostics[0].rejected_features.iter().all(|feature| feature.as_str() != "file_ordering"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn order_sensitive_multi_file_group_scan_limit_remains_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-order-sensitive-multi-group");
    write_numeric_parquet_file(&dir.join("part-000.parquet"));
    write_numeric_parquet_file(&dir.join("part-001.parquet"));
    let scan = parquet_scan_with_order_sensitive_limit_one_file_per_group(&dir, 1);
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::NativeSourcePreserveOrder, .. }
    ));
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].scan_limit_disposition.as_str(), "rejected");
    assert!(result.source_diagnostics[0].rejected_features.iter().any(|feature| feature.as_str() == "preserve_order"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn core_shapes_tc_declord_cases_1() {
    testing::run_named_cases(&[
        (
            "preserve_order_multi_file_group_without_scan_limit_lowers_when_no_parent_reads_the_order",
            preserve_order_multi_file_group_without_scan_limit_lowers_when_no_parent_reads_the_order,
        ),
        (
            "declared_output_ordering_single_file_group_multi_file_scan_lowers_to_native_source",
            declared_output_ordering_single_file_group_multi_file_scan_lowers_to_native_source,
        ),
        (
            "parquet_sorting_columns_inferred_output_ordering_lowers_to_native_source",
            parquet_sorting_columns_inferred_output_ordering_lowers_to_native_source,
        ),
        (
            "declared_output_ordering_sorted_runs_lower_when_no_parent_reads_the_scan_order",
            declared_output_ordering_sorted_runs_lower_when_no_parent_reads_the_scan_order,
        ),
        (
            "declared_output_ordering_single_file_group_scan_lowers_to_native_source",
            declared_output_ordering_single_file_group_scan_lowers_to_native_source,
        ),
    ]);
}

fn preserve_order_multi_file_group_without_scan_limit_lowers_when_no_parent_reads_the_order() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-preserve-order-multi-group");
    write_numeric_parquet_file(&dir.join("part-000.parquet"));
    write_numeric_parquet_file(&dir.join("part-001.parquet"));
    let scan = parquet_scan_with_preserve_order_one_file_per_group(&dir);
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(native_local_parquet_source(result.selected_plan()).file_count(), 2);
    assert_eq!(result.source_diagnostics.len(), 1);
    assert_eq!(result.source_diagnostics[0].scan_limit_disposition.as_str(), "absent");
    assert!(result.source_diagnostics[0].rejected_features.is_empty());

    fs::remove_dir_all(dir).unwrap();
}

fn declared_output_ordering_single_file_group_scan_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-declared-output-ordering");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file(&path);
    let scan = parquet_scan_with_declared_output_ordering(&path);
    assert!(scan.properties().output_ordering().is_some());
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.file_count(), 1);
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert!(result.source_diagnostics[0].rejected_features.iter().all(|feature| feature.as_str() != "output_ordering"));

    fs::remove_dir_all(dir).unwrap();
}

fn parquet_sorting_columns_inferred_output_ordering_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-parquet-sorting-columns");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file_with_sorting_columns(&path);
    let scan = parquet_scan(&path);
    assert!(scan.properties().output_ordering().is_some());

    let result = lower_physical_plan_to_native(scan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.file_count(), 1);
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert!(result.source_diagnostics[0].rejected_features.iter().all(|feature| feature.as_str() != "output_ordering"));

    fs::remove_dir_all(dir).unwrap();
}

fn declared_output_ordering_single_file_group_multi_file_scan_lowers_to_native_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-declared-output-ordering-multi-file");
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-000.parquet"), &[1, 1, 2]);
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-001.parquet"), &[3, 3, 4]);
    let scan = parquet_scan_with_declared_output_ordering(&dir);
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.file_count(), 2);
    assert!(source.diagnostics.rejected_features.is_empty());
    assert_eq!(result.source_diagnostics.len(), 1);
    assert!(result.source_diagnostics[0].rejected_features.iter().all(|feature| feature.as_str() != "output_ordering"));

    fs::remove_dir_all(dir).unwrap();
}

fn declared_output_ordering_sorted_runs_lower_when_no_parent_reads_the_scan_order() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sorted-runs-unordered-parent");
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-000.parquet"), &[1, 3, 5]);
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-001.parquet"), &[2, 3, 4]);
    let scan = parquet_scan_with_declared_output_ordering_one_file_per_group(&dir);
    assert!(scan.properties().output_ordering().is_some());
    let plan = aggregate_sort_limit_plan(scan, Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.file_count(), 2);
    assert!(source.diagnostics.rejected_features.is_empty(), "{:?}", source.diagnostics.rejected_features);

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn sort_preserving_merge_of_sorted_runs_lowers_to_one_native_sort() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-merge-sorted-runs");
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-000.parquet"), &[1, 3, 5]);
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-001.parquet"), &[2, 3, 4]);
    let scan = parquet_scan_with_declared_output_ordering_one_file_per_group(&dir);
    let ordering = scan.properties().output_ordering().cloned().expect("each run is sorted by key");
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(ordering, scan));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();
    assert!(matches!(plan.root_node().kind, PlanNode::Sort(_)), "the merge should become one native sort: {plan}");
    assert_eq!(plan.nodes().iter().filter(|node| matches!(node.kind, PlanNode::Sort(_))).count(), 1, "{plan}");
    assert_eq!(native_local_parquet_source(plan).file_count(), 2);

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn declared_output_ordering_of_sorted_runs_read_by_the_output_is_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sorted-runs-ordered-output");
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-000.parquet"), &[1, 3, 5]);
    write_sorted_numeric_parquet_file_with_keys(&dir.join("part-001.parquet"), &[2, 3, 4]);
    let scan = parquet_scan_with_declared_output_ordering_one_file_per_group(&dir);
    let scan_schema = scan.schema();
    let predicate =
        binary(col("key", &scan_schema).unwrap(), Operator::Gt, lit(ScalarValue::Int32(Some(1))), &scan_schema)
            .unwrap();
    // The filter keeps the scan's per-run order, which the plan declares as
    // its output ordering.
    let plan: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new(predicate, scan).unwrap());
    assert!(plan.properties().output_ordering().is_some());

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::NativeSourceOutputOrdering, .. }
    ));
    assert_eq!(result.source_diagnostics.len(), 1);
    assert!(result.source_diagnostics[0].rejected_features.iter().any(|feature| feature.as_str() == "output_ordering"));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn declared_output_ordering_claim_through_aggregate_is_reestablished_with_native_sort() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-ordered-aggregate-claim");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file(&path);
    let scan = parquet_scan_with_declared_output_ordering(&path);
    let plan = grouped_sum_aggregate_plan(scan);
    assert!(plan.properties().output_ordering().is_some());

    let (report, native_plan) = optimize_with_native_rule(plan);

    let [CandidateOutcome::Selected { candidate }] = report.outcomes() else {
        panic!("ordered aggregate should lower with an explicit native sort: {:?}", report.outcomes());
    };
    assert!(
        candidate.shape.to_string().starts_with("SortExec(AggregateExec("),
        "appended declared-ordering sort should be visible in candidate shape: {}",
        candidate.shape
    );
    let native_plan = native_plan.expect("the rule keeps the native plan");
    assert!(matches!(native_plan.root_node().kind, PlanNode::Sort(_)), "{native_plan}");

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn declared_output_ordering_limit_above_aggregate_is_rejected_instead_of_sorted_after_limit() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-ordered-limit-aggregate");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file(&path);
    let scan = parquet_scan_with_declared_output_ordering(&path);
    let aggregate = grouped_sum_aggregate_plan(scan);
    assert!(aggregate.properties().output_ordering().is_some());
    let plan: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(aggregate, 0, Some(5)));
    assert!(plan.properties().output_ordering().is_some());

    let (report, native_plan) = optimize_with_native_rule(plan);

    assert!(native_plan.is_none());
    assert_ordered_root_rejected_over_sorted_aggregate_fragment(&report);

    fs::remove_dir_all(dir).unwrap();
}

/// The root candidate misses its ordering contract, and the aggregate below
/// the DataFusion root runs as a native fragment that sorts its output.
fn assert_ordered_root_rejected_over_sorted_aggregate_fragment(report: &PlanningReport) {
    let [
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::NativeOrderedOutputUnsupported, .. },
        CandidateOutcome::Selected { candidate },
    ] = report.outcomes()
    else {
        panic!("outcomes={:?}", report.outcomes());
    };
    assert!(candidate.shape.to_string().starts_with("SortExec(AggregateExec("), "{}", candidate.shape);
}

#[test]
fn ordered_window_over_order_breaking_aggregate_reads_a_native_sort_of_its_keys() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-window-over-aggregate-order");
    let path = dir.join("input.parquet");
    write_sorted_numeric_parquet_file(&path);
    let scan = parquet_scan_with_declared_output_ordering(&path);
    let aggregate = grouped_sum_aggregate_plan(scan);
    assert!(aggregate.properties().output_ordering().is_some());
    let aggregate_schema = aggregate.schema();
    let window_expr: Arc<dyn WindowExpr> = Arc::new(StandardWindowExpr::new(
        create_udwf_window_expr(
            &row_number_udwf(),
            &[],
            aggregate_schema.as_ref(),
            "row_number() ORDER BY [key]".to_string(),
            false,
        )
        .unwrap(),
        &[],
        &[PhysicalSortExpr::new(col("key", &aggregate_schema).unwrap(), ArrowSortOptions::new(false, false))],
        Arc::new(WindowFrame::new(None)),
    ));
    let plan: Arc<dyn ExecutionPlan> =
        Arc::new(BoundedWindowAggExec::try_new(vec![window_expr], aggregate, InputOrderMode::Sorted, true).unwrap());

    let (report, native_plan) = optimize_with_native_rule(plan);

    assert!(matches!(report.outcomes(), [CandidateOutcome::Selected { .. }]), "outcomes={:?}", report.outcomes());
    let native_plan = native_plan.expect("the rule keeps the native plan");
    let PlanNode::Window(window) = &native_plan.root_node().kind else {
        panic!("the window is the plan root:\n{native_plan}");
    };
    let Some(PlanNode::Sort(sort)) = native_plan.node(window.input).map(|node| &node.kind) else {
        panic!("the window reads a sort of its keys:\n{native_plan}");
    };
    assert!(matches!(native_plan.node(sort.input).map(|node| &node.kind), Some(PlanNode::Aggregate(_))));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn declared_output_ordering_multi_partition_memory_source_is_rejected() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let first_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap();
    let second_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![3, 4]))]).unwrap();
    let ordering =
        LexOrdering::new([PhysicalSortExpr::new(col("key", &schema).unwrap(), ArrowSortOptions::new(false, false))])
            .unwrap();
    let source = MemorySourceConfig::try_new(&[vec![first_batch], vec![second_batch]], Arc::clone(&schema), None)
        .unwrap()
        .try_with_sort_information(vec![ordering])
        .unwrap();
    let plan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    assert!(plan.properties().output_ordering().is_some());

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported { capability_reason: CapabilityReason::NativeOrderedOutputUnsupported, .. }
    ));
}

#[test]
fn sort_preserving_merge_over_partition_preserving_memory_sort_lowers() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let first_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 3]))]).unwrap();
    let second_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![2, 4]))]).unwrap();
    let ordering =
        LexOrdering::new([PhysicalSortExpr::new(col("key", &schema).unwrap(), ArrowSortOptions::new(false, false))])
            .unwrap();
    let source = MemorySourceConfig::try_new(&[vec![first_batch], vec![second_batch]], Arc::clone(&schema), None)
        .unwrap()
        .try_with_sort_information(vec![ordering.clone()])
        .unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    let sorted: Arc<dyn ExecutionPlan> =
        Arc::new(SortExec::new(ordering.clone(), scan).with_preserve_partitioning(true));
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(ordering, sorted));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("partition-preserving SPM should lower to native");
    };
    assert_eq!(candidate.shape.to_string(), "SortPreservingMergeExec(SortExec(MemorySource))");
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Sort(_)));
}

#[test]
fn sort_preserving_merge_over_ordered_memory_source_lowers_without_child_sort() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let first_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 3]))]).unwrap();
    let second_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![2, 4]))]).unwrap();
    let ordering =
        LexOrdering::new([PhysicalSortExpr::new(col("key", &schema).unwrap(), ArrowSortOptions::new(false, false))])
            .unwrap();
    let source = MemorySourceConfig::try_new(&[vec![first_batch], vec![second_batch]], Arc::clone(&schema), None)
        .unwrap()
        .try_with_sort_information(vec![ordering.clone()])
        .unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(ordering, scan));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("SPM over an ordered input should lower to native");
    };
    assert_eq!(candidate.shape.to_string(), "SortPreservingMergeExec(MemorySource)");
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Sort(_)));
}

#[test]
fn projection_over_sort_preserving_merge_and_ordered_union_lowers() {
    let input_schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Int32, false)]));
    let inputs = [1, 2]
        .into_iter()
        .map(|key| {
            let batch =
                RecordBatch::try_new(Arc::clone(&input_schema), vec![Arc::new(Int32Array::from(vec![key, key]))])
                    .unwrap();
            let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&input_schema), None).unwrap();
            let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
            Arc::new(ProjectionExec::try_new([(lit(ScalarValue::Int32(Some(key))), "key".to_string())], scan).unwrap())
                as Arc<dyn ExecutionPlan>
        })
        .collect();
    let union: Arc<dyn ExecutionPlan> = UnionExec::try_new(inputs).unwrap();
    assert!(union.properties().output_ordering().is_none());
    let ordering = LexOrdering::new([PhysicalSortExpr::new(
        col("key", &union.schema()).unwrap(),
        ArrowSortOptions::new(false, false),
    )])
    .unwrap();
    let merge: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(ordering, union));
    let plan: Arc<dyn ExecutionPlan> =
        Arc::new(ProjectionExec::try_new([(col("key", &merge.schema()).unwrap(), "key".to_string())], merge).unwrap());

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("SPM over an ordered union should lower to native");
    };
    assert_eq!(
        candidate.shape.to_string(),
        "ProjectionExec(SortPreservingMergeExec(UnionExec(ProjectionExec(MemorySource),ProjectionExec(MemorySource))))"
    );
    assert!(result.selected_plan().nodes().iter().any(|node| matches!(node.kind, PlanNode::Sort(_))));
    assert!(result.selected_plan().nodes().iter().any(|node| matches!(node.kind, PlanNode::Projection(_))));
}

#[test]
fn partition_preserving_sort_with_single_input_partition_lowers() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![2, 1]))]).unwrap();
    let ordering =
        LexOrdering::new([PhysicalSortExpr::new(col("key", &schema).unwrap(), ArrowSortOptions::new(false, false))])
            .unwrap();
    let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None).unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    assert_eq!(scan.properties().output_partitioning().partition_count(), 1);
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(ordering, scan).with_preserve_partitioning(true));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("single-partition partition-preserving sort should lower to native");
    };
    assert_eq!(candidate.shape.to_string(), "SortExec(MemorySource)");
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Sort(_)));
}

#[test]
fn partition_preserving_sort_with_multiple_input_partitions_is_rejected() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let first_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 3]))]).unwrap();
    let second_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![2, 4]))]).unwrap();
    let ordering =
        LexOrdering::new([PhysicalSortExpr::new(col("key", &schema).unwrap(), ArrowSortOptions::new(false, false))])
            .unwrap();
    let source =
        MemorySourceConfig::try_new(&[vec![first_batch], vec![second_batch]], Arc::clone(&schema), None).unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    assert_eq!(scan.properties().output_partitioning().partition_count(), 2);
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(ordering, scan).with_preserve_partitioning(true));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativePartitionPreservingSortUnsupported,
            ..
        }
    ));
}

#[test]
fn filter_over_computed_scan_projection_reads_the_projected_column() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-filter-computed-scan-projection");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let scan = parquet_scan(&path);
    let scan_schema = scan.schema();
    let projection = ProjectionExec::try_new(
        vec![
            (Arc::new(CastExpr::new(col("key", &scan_schema).unwrap(), DataType::Int64, None)) as _, "key".to_string()),
            (col("m00", &scan_schema).unwrap(), "m00".to_string()),
        ],
        Arc::clone(&scan),
    )
    .unwrap();
    let scan = scan
        .try_swapping_with_projection(&projection)
        .unwrap()
        .expect("the parquet scan takes the computed projection");
    let projected_schema = scan.schema();
    assert_eq!(projected_schema.field_with_name("key").unwrap().data_type(), &DataType::Int64);
    let predicate = binary(
        col("key", &projected_schema).unwrap(),
        Operator::Gt,
        lit(ScalarValue::Int64(Some(1))),
        &projected_schema,
    )
    .unwrap();
    let plan: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new(predicate, scan).unwrap());

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();
    let PlanNode::Filter(filter) = &plan.root_node().kind else {
        panic!("root should be the native filter: {plan}");
    };
    let PlanNode::Projection(projection) = &plan.node(filter.input).expect("filter input").kind else {
        panic!("the computed scan projection should run below the filter: {plan}");
    };
    assert!(
        projection.outputs.iter().any(|output| output.field.name() == "key"
            && output.field.data_type() == &DataType::Int64
            && matches!(&output.expr, Expr::Alias { expr, .. } if matches!(expr.as_ref(), Expr::Cast { .. }))),
        "projection should cast key before the filter reads it: {plan}"
    );

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn local_limit_over_multiple_partitions_is_rejected() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let first_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 3]))]).unwrap();
    let second_batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![2, 4]))]).unwrap();
    let source =
        MemorySourceConfig::try_new(&[vec![first_batch], vec![second_batch]], Arc::clone(&schema), None).unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    assert_eq!(scan.properties().output_partitioning().partition_count(), 2);
    // Each partition keeps one row, so DataFusion returns two rows where one
    // native global limit would return one.
    let plan: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(scan, 1));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert!(matches!(
        result.report.outcomes(),
        [CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativeLocalLimitPartitionUnsupported,
            candidate,
            ..
        }] if candidate.shape.to_string() == "LocalLimitExec"
    ));
}

#[test]
fn local_limit_over_one_partition_lowers_to_native_limit() {
    let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
    let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap();
    let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None).unwrap();
    let scan = DataSourceExec::from_data_source(source) as Arc<dyn ExecutionPlan>;
    let plan: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(scan, 2));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let PlanNode::Limit(limit) = &result.selected_plan().root_node().kind else {
        panic!("root should be the native limit: {}", result.selected_plan());
    };
    assert_eq!(limit.spec, LimitSpec::fetch(2));
}

#[test]
fn topk_sort_fetch_lowers_to_native_sort_plus_limit() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-topk");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = aggregate_topk_plan(parquet_scan(&path), Operator::Plus, 2);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.candidate_count(), 1);
    assert_eq!(result.selected_count(), 1);
    assert_eq!(result.not_selected_count(), 0);
    let PlanNode::Limit(limit) = &result.selected_plan().root_node().kind else {
        panic!("native root should be a limit node");
    };
    assert_eq!(limit.spec.fetch, Some(2));
    let sort_node = result.selected_plan().node(limit.input).expect("limit input should exist");
    assert!(matches!(sort_node.kind, PlanNode::Sort(_)));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn input_tree_sort_fetch_lowers_to_native_sort_plus_limit() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-input-topk");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = count_over_topk_sort_fetch_input_plan(parquet_scan(&path), 2);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.candidate_count(), 1);
    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("input-tree top-k aggregate should lower to native");
    };
    assert_eq!(candidate.shape.to_string(), "AggregateExec(SortExec(fetch)(LocalParquetScan))");
    let PlanNode::Aggregate(_aggregate) = &result.selected_plan().root_node().kind else {
        panic!("native root should be an aggregate node");
    };
    let has_sort_fetch_limit = result.selected_plan().nodes().iter().any(|node| {
        let PlanNode::Limit(limit) = &node.kind else {
            return false;
        };
        if limit.spec.fetch != Some(2) {
            return false;
        }
        let Some(input) = result.selected_plan().node(limit.input) else {
            return false;
        };
        match &input.kind {
            PlanNode::Sort(_) => true,
            PlanNode::Projection(projection) => result
                .selected_plan()
                .node(projection.input)
                .is_some_and(|input| matches!(input.kind, PlanNode::Sort(_))),
            _ => false,
        }
    });
    assert!(
        has_sort_fetch_limit,
        "native plan should keep SortExec.fetch as a limit above sort:\n{}",
        result.selected_plan()
    );

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn sort_without_fetch_over_aggregate_lowers_without_native_limit() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-no-fetch");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.candidate_count(), 1);
    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    assert!(matches!(result.report.outcomes()[0], CandidateOutcome::Selected { .. }));
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Sort(_)));
    assert!(!result.selected_plan().nodes().iter().any(|node| matches!(&node.kind, PlanNode::Limit(_))));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn aggregate_root_global_limit_does_not_push_scan_limit_to_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-aggregate-root-limit-no-source-pushdown");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sorted = aggregate_sort_plan_without_filter(parquet_scan(&path), Operator::Plus);
    let sort = sorted.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let aggregate = Arc::clone(sort.input());
    let plan: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(aggregate, 0, Some(2)));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.scan_limit, None, "LIMIT after aggregate must not be applied before aggregate input");
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Limit(_)));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn nested_limit_offsets_push_enough_rows_to_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-nested-limit-source-pushdown");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let scan = parquet_scan(&path);
    let inner: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(scan, 50, Some(100)));
    let outer: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(inner, 10, Some(10)));

    let result = lower_physical_plan_to_native(outer, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(
        source.scan_limit,
        Some(70),
        "source must keep the inner offset plus the rows needed by the outer limit"
    );

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn zero_fetch_global_limit_pushes_zero_scan_limit_to_source() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-zero-limit-source-pushdown");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(parquet_scan(&path), 100, Some(0)));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    let source = native_local_parquet_source(result.selected_plan());
    assert_eq!(source.scan_limit, Some(0), "LIMIT 0 can be answered without reading source rows");

    fs::remove_dir_all(dir).unwrap();
}

fn sort_preserving_merge_without_fetch_lowers_to_native_sort() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort_plan = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort_plan.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let merge: Arc<dyn ExecutionPlan> =
        Arc::new(SortPreservingMergeExec::new(sort.expr().clone(), Arc::clone(&sort_plan)));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let CandidateOutcome::Selected { candidate } = &result.report.outcomes()[0] else {
        panic!("non-fetch SPM should lower to native");
    };
    assert!(
        candidate.shape.to_string().starts_with("SortPreservingMergeExec(SortExec(AggregateExec("),
        "unexpected candidate shape: {}",
        candidate.shape
    );
    assert!(matches!(result.selected_plan().root_node().kind, PlanNode::Sort(_)));
    assert!(!result.selected_plan().nodes().iter().any(|node| matches!(node.kind, PlanNode::Limit(_))));

    fs::remove_dir_all(dir).unwrap();
}

fn sort_preserving_merge_fetch_with_matching_child_sort_lowers_to_native_sort_plus_limit() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge-fetch-matching-child-sort");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort_plan = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort_plan.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let merge: Arc<dyn ExecutionPlan> =
        Arc::new(SortPreservingMergeExec::new(sort.expr().clone(), Arc::clone(&sort_plan)).with_fetch(Some(2)));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let PlanNode::Limit(limit) = &result.selected_plan().root_node().kind else {
        panic!("SPM fetch should lower to a native limit above sort:\n{}", result.selected_plan());
    };
    assert_eq!(limit.spec.fetch, Some(2));
    let sort_node = result.selected_plan().node(limit.input).expect("limit input should exist");
    assert!(matches!(sort_node.kind, PlanNode::Sort(_)));

    fs::remove_dir_all(dir).unwrap();
}

fn sort_preserving_merge_fetch_without_child_ordering_is_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge-fetch-no-child-ordering");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let merge: Arc<dyn ExecutionPlan> =
        Arc::new(SortPreservingMergeExec::new(sort.expr().clone(), Arc::clone(sort.input())).with_fetch(Some(2)));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativeSortPreservingMergeUnsupported,
            ..
        }
    ));

    fs::remove_dir_all(dir).unwrap();
}

fn sort_preserving_merge_fetch_with_mismatched_child_sort_is_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge-fetch-mismatched-child-sort");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let first_sort_expr = &sort.expr()[0];
    let child_ordering = LexOrdering::new([PhysicalSortExpr::new(
        Arc::clone(&first_sort_expr.expr),
        ArrowSortOptions::new(!first_sort_expr.options.descending, first_sort_expr.options.nulls_first),
    )])
    .expect("mismatched child ordering should be non-empty");
    let child_sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(child_ordering, Arc::clone(sort.input())));
    let merge: Arc<dyn ExecutionPlan> =
        Arc::new(SortPreservingMergeExec::new(sort.expr().clone(), child_sort).with_fetch(Some(2)));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativeSortPreservingMergeUnsupported,
            ..
        }
    ));

    fs::remove_dir_all(dir).unwrap();
}

fn sort_preserving_merge_with_same_name_different_index_ordering_is_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge-same-name-different-index");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort_plan = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort_plan.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let merge_ordering =
        LexOrdering::new([PhysicalSortExpr::new(Arc::new(Column::new("sum_x00", 0)), sort.expr()[0].options)])
            .expect("SPM ordering should be non-empty");
    let merge: Arc<dyn ExecutionPlan> = Arc::new(SortPreservingMergeExec::new(merge_ordering, sort_plan));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativeSortPreservingMergeUnsupported,
            ..
        }
    ));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn core_shapes_tc_spm_cases_1() {
    testing::run_named_cases(&[
        (
            "sort_preserving_merge_with_child_fetch_smaller_than_root_fetch_is_rejected",
            sort_preserving_merge_with_child_fetch_smaller_than_root_fetch_is_rejected,
        ),
        (
            "sort_preserving_merge_with_same_name_different_index_ordering_is_rejected",
            sort_preserving_merge_with_same_name_different_index_ordering_is_rejected,
        ),
        (
            "sort_preserving_merge_without_fetch_lowers_to_native_sort",
            sort_preserving_merge_without_fetch_lowers_to_native_sort,
        ),
        (
            "sort_preserving_merge_fetch_with_mismatched_child_sort_is_rejected",
            sort_preserving_merge_fetch_with_mismatched_child_sort_is_rejected,
        ),
        (
            "sort_preserving_merge_fetch_with_matching_child_sort_lowers_to_native_sort_plus_limit",
            sort_preserving_merge_fetch_with_matching_child_sort_lowers_to_native_sort_plus_limit,
        ),
        (
            "sort_preserving_merge_fetch_without_child_ordering_is_rejected",
            sort_preserving_merge_fetch_without_child_ordering_is_rejected,
        ),
    ]);
}

fn sort_preserving_merge_with_child_fetch_smaller_than_root_fetch_is_rejected() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-sort-preserving-merge-fetch");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let sort = aggregate_sort_plan(parquet_scan(&path), Operator::Plus);
    let sort = sort.downcast_ref::<SortExec>().expect("helper returns SortExec");
    let child_sort: Arc<dyn ExecutionPlan> =
        Arc::new(SortExec::new(sort.expr().clone(), Arc::clone(sort.input())).with_fetch(Some(1)));
    let merge: Arc<dyn ExecutionPlan> =
        Arc::new(SortPreservingMergeExec::new(sort.expr().clone(), child_sort).with_fetch(Some(2)));

    let result = lower_physical_plan_to_native(merge, &ParquetReadStrategyConfig::default()).unwrap();

    assert!(result.plan.is_none());
    assert_eq!(result.not_selected_count(), 1);
    assert!(matches!(
        result.report.outcomes()[0],
        CandidateOutcome::NotSupported {
            capability_reason: CapabilityReason::NativeSortPreservingMergeUnsupported,
            ..
        }
    ));

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn grouped_aggregate_late_original_index_columns_bind_after_source_compaction() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-aggregate-compacted-bindings");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = grouped_late_column_multi_sum_sort_limit_plan(parquet_scan(&path));

    let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();
    let aggregate = native_aggregate_nodes(plan).into_iter().next().expect("native plan should contain an aggregate");
    let column_name = |expr: &Expr| match expr {
        Expr::Alias { expr, .. } => match expr.as_ref() {
            Expr::Column(column) => Some(column.name().to_string()),
            _ => None,
        },
        Expr::Column(column) => Some(column.name().to_string()),
        _ => None,
    };
    assert_eq!(
        column_name(&aggregate.spec.group_exprs[0]).as_deref(),
        Some("payload"),
        "group expression should bind original payload, not compacted ordinal m00: {:?}",
        aggregate.spec.group_exprs
    );
    let aggregate_exprs = native_aggregate_exprs(plan);
    assert_eq!(aggregate_exprs.len(), 2);
    assert_eq!(
        column_name(&aggregate_exprs[0].expr).as_deref(),
        Some("m00"),
        "first aggregate arg should bind original m00, not compacted ordinal m01: {aggregate_exprs:?}"
    );
    assert_eq!(
        column_name(&aggregate_exprs[1].expr).as_deref(),
        Some("m01"),
        "second aggregate arg should bind original m01: {aggregate_exprs:?}"
    );

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn sort_projection_sort_nested_shape_preserves_both_sort_orderings_and_fetches() {
    let dir = unique_temp_dir("algeon-datafusion-sort-proj-sort");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);

    let scan = parquet_scan(&path);
    let scan_schema = scan.schema();

    let inner_sort_exprs = LexOrdering::new([PhysicalSortExpr::new(
        col("m00", &scan_schema).unwrap(),
        ArrowSortOptions::new(false, false),
    )])
    .unwrap();
    let inner_sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(inner_sort_exprs, scan).with_fetch(Some(5)));

    let inner_schema = inner_sort.schema();
    let projection: Arc<dyn ExecutionPlan> = Arc::new(
        ProjectionExec::try_new(
            vec![
                (col("key", &inner_schema).unwrap(), "key".to_string()),
                (col("m00", &inner_schema).unwrap(), "m00".to_string()),
                (col("m01", &inner_schema).unwrap(), "m01".to_string()),
            ],
            inner_sort,
        )
        .unwrap(),
    );

    let proj_schema = projection.schema();
    let outer_sort_exprs = LexOrdering::new([PhysicalSortExpr::new(
        col("m01", &proj_schema).unwrap(),
        ArrowSortOptions::new(true, false),
    )])
    .unwrap();
    let outer_sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(outer_sort_exprs, projection).with_fetch(Some(2)));

    let result = lower_physical_plan_to_native(outer_sort, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();

    let PlanNode::Limit(outer_limit) = &plan.root_node().kind else {
        panic!("root should be the outer limit node: {plan}");
    };
    assert_eq!(outer_limit.spec.fetch, Some(2));

    let outer_sort_node = plan.node(outer_limit.input).expect("outer sort node should exist");
    let PlanNode::Sort(outer_sort) = &outer_sort_node.kind else {
        panic!("outer limit input should be outer sort: {plan}");
    };
    assert_eq!(outer_sort.spec.expressions.len(), 1);
    assert_eq!(outer_sort.spec.expressions[0].expr.to_string(), "m01");
    assert!(outer_sort.spec.expressions[0].descending);

    let proj_node = plan.node(outer_sort.input).expect("projection node should exist");
    let PlanNode::Projection(proj) = &proj_node.kind else {
        panic!("outer sort input should be projection: {plan}");
    };
    assert_eq!(proj.outputs.len(), 3);

    let inner_limit_node = plan.node(proj.input).expect("inner limit node should exist");
    let PlanNode::Limit(inner_limit) = &inner_limit_node.kind else {
        panic!("projection input should be inner limit: {plan}");
    };
    assert_eq!(inner_limit.spec.fetch, Some(5));

    let inner_sort_node = plan.node(inner_limit.input).expect("inner sort node should exist");
    let PlanNode::Sort(inner_sort) = &inner_sort_node.kind else {
        panic!("inner limit input should be inner sort: {plan}");
    };
    assert_eq!(inner_sort.spec.expressions.len(), 1);
    assert_eq!(inner_sort.spec.expressions[0].expr.to_string(), "m00");
    assert!(!inner_sort.spec.expressions[0].descending);

    fs::remove_dir_all(dir).unwrap();
}

#[test]
fn projection_sort_projection_aggregate_nested_shape_preserves_both_projections() {
    let dir = unique_temp_dir("algeon-datafusion-proj-sort-proj-agg");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);

    let scan = parquet_scan(&path);
    let agg = grouped_sum_aggregate_plan(scan);
    let agg_schema = agg.schema();

    let derived_expr =
        binary(col("sum_m00", &agg_schema).unwrap(), Operator::Plus, lit(ScalarValue::Int64(Some(1))), &agg_schema)
            .unwrap();
    let inner_projection: Arc<dyn ExecutionPlan> = Arc::new(
        ProjectionExec::try_new(
            vec![
                (col("key", &agg_schema).unwrap(), "key".to_string()),
                (col("sum_m00", &agg_schema).unwrap(), "sum_m00".to_string()),
                (derived_expr, "derived".to_string()),
            ],
            agg,
        )
        .unwrap(),
    );

    let inner_proj_schema = inner_projection.schema();
    let sort_exprs = LexOrdering::new([PhysicalSortExpr::new(
        col("derived", &inner_proj_schema).unwrap(),
        ArrowSortOptions::new(true, false),
    )])
    .unwrap();
    let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_exprs, inner_projection));

    let sort_schema = sort.schema();
    let final_out_expr =
        binary(col("derived", &sort_schema).unwrap(), Operator::Plus, lit(ScalarValue::Int64(Some(2))), &sort_schema)
            .unwrap();
    let outer_projection: Arc<dyn ExecutionPlan> = Arc::new(
        ProjectionExec::try_new(
            vec![
                (col("key", &sort_schema).unwrap(), "key".to_string()),
                (col("derived", &sort_schema).unwrap(), "derived".to_string()),
                (final_out_expr, "final_out".to_string()),
            ],
            sort,
        )
        .unwrap(),
    );

    let result = lower_physical_plan_to_native(outer_projection, &ParquetReadStrategyConfig::default()).unwrap();

    assert_eq!(result.selected_count(), 1, "outcomes={:?}", result.report.outcomes());
    assert_eq!(result.not_selected_count(), 0, "outcomes={:?}", result.report.outcomes());
    let plan = result.selected_plan();

    let PlanNode::Projection(outer_proj_node) = &plan.root_node().kind else {
        panic!("root should be outer projection: {plan}");
    };
    assert_eq!(outer_proj_node.outputs.len(), 3);
    assert_eq!(outer_proj_node.outputs[2].field.name(), "final_out");

    let sort_node = plan.node(outer_proj_node.input).expect("sort node should exist");
    let PlanNode::Sort(sort_kind) = &sort_node.kind else {
        panic!("outer projection input should be sort: {plan}");
    };
    assert_eq!(sort_kind.spec.expressions.len(), 1);
    assert_eq!(sort_kind.spec.expressions[0].expr.to_string(), "derived");

    let inner_proj_node = plan.node(sort_kind.input).expect("inner projection node should exist");
    let PlanNode::Projection(inner_proj_kind) = &inner_proj_node.kind else {
        panic!("sort input should be inner projection: {plan}");
    };
    // `sum_m00` has no reader above the inner projection.
    assert_eq!(
        inner_proj_kind.outputs.iter().map(|output| output.field.name().as_str()).collect::<Vec<_>>(),
        vec!["key", "derived"]
    );

    let agg_node = plan.node(inner_proj_kind.input).expect("aggregate node should exist");
    let PlanNode::Aggregate(agg_kind) = &agg_node.kind else {
        panic!("inner projection input should be aggregate: {plan}");
    };
    assert_eq!(agg_kind.spec.group_exprs.len(), 1);
    assert_eq!(agg_kind.spec.aggr_exprs.len(), 1);

    fs::remove_dir_all(dir).unwrap();
}