algeon-datafusion 0.1.2

use query_planning::plan::{ShiftDirection, WindowFunction};

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ExpectedGpuCoverageOutcome {
    Selected,
    /// `blocking` is the kind and one-line text of what blocked the candidate.
    NotSupported {
        reason: CapabilityReason,
        unsupported_category: &'static str,
        blocking: (&'static str, &'static str),
    },
}

#[derive(Debug, Clone, Copy)]
struct GpuGapCase {
    id: &'static str,

    sql: &'static str,
    expected: ExpectedGpuCoverageOutcome,
}

const AGGREGATE_FILTER_CASES: &[GpuGapCase] = &[
    GpuGapCase {
        id: "aggregate_filter_sum",

        sql: "SELECT key, SUM(payload) FILTER (WHERE flag) AS s FROM bench GROUP BY key",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    // TRY_CAST(string AS int) is admitted natively (Utf8/LargeUtf8/Utf8View →
    // any signed/unsigned integer); the resulting nullable predicate is a
    // valid aggregate FILTER expression.
    GpuGapCase {
        id: "aggregate_filter_try_cast",

        sql: "SELECT key, SUM(payload) FILTER (WHERE TRY_CAST(name AS INT) > 0) AS s FROM bench GROUP BY key",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
];

const AGGREGATE_FUNCTION_CASES: &[GpuGapCase] = &[
    GpuGapCase {
        id: "aggregate_approx_distinct",

        sql: "SELECT APPROX_DISTINCT(name) AS distinct_names FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "aggregate_approx_distinct",
            blocking: ("aggregate_function", "aggregate function approx_distinct"),
        },
    },
    GpuGapCase {
        id: "aggregate_approx_percentile_cont",

        sql: "SELECT APPROX_PERCENTILE_CONT(score, 0.5) AS median_score FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "aggregate_approx_percentile_cont",
            blocking: ("aggregate_function", "aggregate function approx_percentile_cont"),
        },
    },
];

const WINDOW_CASES: &[GpuGapCase] = &[
    GpuGapCase {
        id: "window_lag",

        sql: "SELECT id, LAG(payload) OVER (PARTITION BY key ORDER BY id) AS prev_payload FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "window_lag_offset_default",

        sql: "SELECT id, LAG(payload, 2, 0) OVER (PARTITION BY key ORDER BY id) AS prev_payload \
              FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "window_lead",

        sql: "SELECT id, LEAD(payload) OVER (PARTITION BY key ORDER BY id) AS next_payload FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "window_row_number",

        sql: "SELECT id, ROW_NUMBER() OVER (PARTITION BY key ORDER BY payload) AS rn FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "window_bounded_sum",

        sql: "SELECT id, SUM(payload) OVER (PARTITION BY key ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS s FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "window_frame_rows_1_preceding_to_current_row",
            blocking: ("window_function", "window function sum over frame ROWS BETWEEN 1 PRECEDING AND CURRENT ROW"),
        },
    },
    GpuGapCase {
        id: "window_whole_input_aggregate",

        sql: "SELECT id, COUNT(*) OVER () AS total FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "window_cumulative_count",

        sql: "SELECT id, COUNT(payload) OVER (PARTITION BY key ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS c FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "window_frame_rows_unbounded_preceding_to_current_row",
            blocking: (
                "window_function",
                "window function count over frame ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW",
            ),
        },
    },
    GpuGapCase {
        id: "window_lag_without_order_by",

        sql: "SELECT id, LAG(payload) OVER (PARTITION BY key) AS prev_payload FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "window_function_lag_without_order",
            blocking: ("window_function", "window function lag without ORDER BY"),
        },
    },
    GpuGapCase {
        id: "window_dense_rank",

        sql: "SELECT id, DENSE_RANK() OVER (PARTITION BY key ORDER BY payload) AS r FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "window_function_dense_rank",
            blocking: ("window_function", "window function dense_rank"),
        },
    },
];

const SCALAR_EXPRESSION_CASES: &[GpuGapCase] = &[
    // TRY_CAST(string AS integer) is admitted natively. DataFusion projection
    // pushdown collapses the cast into the DataSourceExec projection; the
    // adapter still recognises it as a supported try-cast.
    GpuGapCase {
        id: "try_cast_string_to_int",

        sql: "SELECT id, TRY_CAST(name AS INT) AS maybe_int FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    // TRY_CAST to a non-integer target (Float64) remains unsupported until
    // libcudf provides a string→float null-on-failure primitive; the structured
    // reject category pins the unsupported type pair.
    GpuGapCase {
        id: "try_cast_string_to_float",

        sql: "SELECT id, TRY_CAST(name AS DOUBLE) AS maybe_float FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedExpression,
            unsupported_category: "expression_try_cast_utf8view_to_float64",
            blocking: ("cast", "try cast from Utf8View to Float64 in `TRY_CAST(name AS Float64)`"),
        },
    },
    // TRY_CAST whose source is not a string column (here a numeric column)
    // also stays unsupported; the lower-level primitive accepts
    // string→integer only.
    GpuGapCase {
        id: "try_cast_int_to_int",

        sql: "SELECT id, TRY_CAST(payload AS SMALLINT) AS narrowed FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedExpression,
            unsupported_category: "expression_try_cast_int64_to_int16",
            blocking: ("cast", "try cast from Int64 to Int16 in `TRY_CAST(payload AS Int16)`"),
        },
    },
    // Casting a float to text has no native conversion.
    GpuGapCase {
        id: "float_to_string_cast",

        sql: "SELECT id, CAST(score AS VARCHAR) AS score_text FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedType,
            unsupported_category: "type_cast_float64_to_utf8view",
            blocking: ("cast", "cast from Float64 to Utf8View in `CAST(score AS Utf8View)`"),
        },
    },
    GpuGapCase {
        id: "integer_to_timestamp_cast",

        sql: "SELECT id, CAST(payload AS TIMESTAMP) AS ts FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "sqrt_scalar_function",

        sql: "SELECT id, sqrt(score) AS root_score FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "ln_scalar_function",

        sql: "SELECT id, CAST(LN(1.0 + payload) AS REAL) AS log_payload FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "exp_scalar_function",

        sql: "SELECT id, EXP(score) AS exp_score FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "floor_ceil_scalar_functions",

        sql: "SELECT id, FLOOR(score) AS floor_score, CEIL(score) AS ceil_score FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    // The CASE yields Int64; narrowing it to Int32 is a range-checked native cast.
    GpuGapCase {
        id: "narrowing_cast_of_case",

        sql: "SELECT id, CAST(CASE WHEN score >= 3.0 THEN 1 ELSE 0 END AS INT) AS label FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "integer_to_string_cast",

        sql: "SELECT id, CAST(payload AS VARCHAR) AS payload_text FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "integer_modulo_and_division",

        sql: "SELECT id, payload / 20 AS bucket FROM bench WHERE payload % 2 = 1",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    // DataFusion divides only on the rows that reach the THEN branch; lowering restricts the divisor
    // to those rows, so the guarded division stays native.
    GpuGapCase {
        id: "variable_divisor_inside_case_branch",

        sql: "SELECT id, CASE WHEN key <> 0 THEN payload / key ELSE 0 END AS ratio FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    // Only the decimal modulo is rejected; the supported SQRT beside it must
    // not be reported as the unsupported expression.
    GpuGapCase {
        id: "supported_function_beside_decimal_modulo",

        sql: "SELECT id, sqrt(score) AS root_score FROM bench \
              WHERE CAST(payload AS DECIMAL(18, 2)) % CAST(2 AS DECIMAL(18, 2)) = CAST(1 AS DECIMAL(18, 2))",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedType,
            unsupported_category: "type",
            blocking: (
                "binary_operator",
                "binary operator `%` over Decimal128(18, 2) and Decimal128(18, 2) in `CAST(payload AS Decimal128(18, 2)) % 2.00`",
            ),
        },
    },
    GpuGapCase {
        id: "floating_point_modulo",

        sql: "SELECT id, score % 2.0 AS fraction FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "string_concat_operator",

        sql: "SELECT id, name || '-x' AS label FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "is_distinct_from",

        sql: "SELECT id FROM bench WHERE name IS DISTINCT FROM 'alpha'",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
];

const ROOT_SOURCE_CASES: &[GpuGapCase] = &[
    GpuGapCase {
        id: "offset_without_limit",

        sql: "SELECT id, payload FROM bench OFFSET 2",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "explain_query",

        sql: "EXPLAIN SELECT id FROM bench WHERE key = 10",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "operator_explain",
            blocking: ("operator", "operator ExplainExec"),
        },
    },
    GpuGapCase {
        id: "scalar_subquery_in_projection",

        sql: "SELECT id, (SELECT COUNT(*) FROM bench) AS total FROM bench",
        expected: ExpectedGpuCoverageOutcome::NotSupported {
            reason: CapabilityReason::UnsupportedOperator,
            unsupported_category: "operator_scalarsubquery",
            blocking: ("operator", "operator ScalarSubqueryExec"),
        },
    },
    GpuGapCase {
        id: "literal_only_projection",

        sql: "SELECT 1 AS one",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "values_table",

        sql: "SELECT a, b FROM (VALUES (1, 10), (2, 20)) AS t(a, b)",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
];

const EXISTING_COVERAGE_CASES: &[GpuGapCase] = &[
    GpuGapCase {
        id: "filter_projection",

        sql: "SELECT id, payload + 1 AS payload_plus FROM bench WHERE key = 10",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "grouped_aggregate_order_limit",

        sql: "SELECT key, SUM(payload) AS s FROM bench GROUP BY key ORDER BY s DESC, key ASC LIMIT 2",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "having_after_aggregate",

        sql: "SELECT key, SUM(payload) AS s FROM bench GROUP BY key HAVING SUM(payload) > 10",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "self_hash_join_aggregate",

        sql: "SELECT b.key, COUNT(*) AS rows FROM bench b JOIN bench s ON b.key = s.key GROUP BY b.key",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "left_join",

        sql: "SELECT b.id, s.payload AS matched_payload FROM bench b LEFT JOIN bench s ON b.key = s.key",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "non_equi_inner_join",

        sql: "SELECT b.id, s.id AS sid FROM bench b JOIN bench s ON b.payload < s.payload",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "union_all",

        sql: "SELECT key FROM bench WHERE flag UNION ALL SELECT key FROM bench WHERE payload > 10",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "count_distinct",

        sql: "SELECT COUNT(DISTINCT name) AS distinct_names FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "avg_distinct",

        sql: "SELECT AVG(DISTINCT payload) AS avg_payload FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "rank_window",

        sql: "SELECT id, RANK() OVER (PARTITION BY key ORDER BY payload) AS r FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "starts_with_like_rewrite",

        sql: "SELECT id FROM bench WHERE starts_with(name, 'a')",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "regexp_replace",

        sql: "SELECT id, regexp_replace(name, 'a', 'x') AS rewritten_name FROM bench",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "order_limit_offset",

        sql: "SELECT id, payload FROM bench ORDER BY payload ASC NULLS LAST LIMIT 2 OFFSET 1",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "intersect_rewrite",

        sql: "SELECT key FROM bench WHERE flag INTERSECT SELECT key FROM bench WHERE payload > 10",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "except_rewrite",

        sql: "SELECT key FROM bench EXCEPT SELECT key FROM bench WHERE flag",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "exists_rewrite",

        sql: "SELECT id FROM bench b WHERE EXISTS (SELECT 1 FROM bench s WHERE s.key = b.key AND s.flag)",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
    GpuGapCase {
        id: "rollup",

        sql: "SELECT key, SUM(payload) AS s FROM bench GROUP BY ROLLUP(key)",
        expected: ExpectedGpuCoverageOutcome::Selected,
    },
];

#[tokio::test]
async fn full_gpu_gap_analysis_sql_planning_coverage_is_stable() {
    let dir = unique_temp_dir("algeon-datafusion-full-gpu-gap-analysis");
    let path = dir.join("bench.parquet");
    write_full_gpu_gap_analysis_parquet_file(&path);

    for case in gpu_gap_cases() {
        let plan = sql_plan(&path, direct_friendly_sql_session_config(), case.sql).await;
        let lowering = lower_physical_plan_to_native(Arc::clone(&plan), &ParquetReadStrategyConfig::default()).unwrap();
        assert_planning_report(&case, &lowering);
        assert_native_ir_details(&case, &lowering);

        let report = NativeFallbackPressureReport::collect([NativeFallbackPressureWorkloadItem::new(case.id, plan)]);
        assert_eq!(report.rows.len(), 1, "{}", case.id);
        let row = &report.rows[0];
        match case.expected {
            ExpectedGpuCoverageOutcome::Selected => {
                assert_eq!(row.outcome, Some(CandidateOutcomeKind::Selected), "{}", case.id);
                assert_eq!(row.reason_code, None, "{}", case.id);
                assert_eq!(row.unsupported_category, None, "{}", case.id);
                assert!(!row.rejected_before_execution, "{}", case.id);
            },
            ExpectedGpuCoverageOutcome::NotSupported { reason, unsupported_category, .. } => {
                assert_eq!(row.outcome, Some(CandidateOutcomeKind::NotSupported), "{}", case.id);
                assert_eq!(row.reason_code.as_deref(), Some(reason.code()), "{}", case.id);
                assert_eq!(row.unsupported_category.as_deref(), Some(unsupported_category), "{}", case.id);
                assert!(!row.rejected_before_execution, "{}", case.id);
            },
        }
    }

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

fn gpu_gap_cases() -> Vec<GpuGapCase> {
    [
        AGGREGATE_FILTER_CASES,
        AGGREGATE_FUNCTION_CASES,
        WINDOW_CASES,
        SCALAR_EXPRESSION_CASES,
        ROOT_SOURCE_CASES,
        EXISTING_COVERAGE_CASES,
    ]
    .into_iter()
    .flatten()
    .copied()
    .collect()
}

fn assert_planning_report(case: &GpuGapCase, lowering: &NativeLoweringResult) {
    assert_eq!(lowering.candidate_count(), 1, "{} outcomes={:?}", case.id, lowering.report.outcomes());
    match case.expected {
        ExpectedGpuCoverageOutcome::Selected => {
            assert_eq!(lowering.selected_count(), 1, "{} outcomes={:?}", case.id, lowering.report.outcomes());
            assert_eq!(lowering.not_selected_count(), 0, "{} outcomes={:?}", case.id, lowering.report.outcomes());
            assert!(lowering.plan.is_some(), "{}", case.id);
            assert!(matches!(lowering.report.outcomes()[0], CandidateOutcome::Selected { .. }), "{}", case.id);
        },
        ExpectedGpuCoverageOutcome::NotSupported { reason, blocking, .. } => {
            assert!(lowering.plan.is_none(), "{}", case.id);
            assert_eq!(lowering.selected_count(), 0, "{} outcomes={:?}", case.id, lowering.report.outcomes());
            assert_eq!(lowering.not_selected_count(), 1, "{} outcomes={:?}", case.id, lowering.report.outcomes());
            let CandidateOutcome::NotSupported { capability_reason, detail, .. } = &lowering.report.outcomes()[0]
            else {
                panic!("{} outcomes={:?}", case.id, lowering.report.outcomes());
            };
            assert_eq!(*capability_reason, reason, "{} outcomes={:?}", case.id, lowering.report.outcomes());
            let actual = detail.blocking().unwrap_or_else(|| panic!("{} names no blocking subject", case.id));
            assert_eq!((actual.kind(), actual.text()), blocking, "{}", case.id);
        },
    }
}

fn assert_native_ir_details(case: &GpuGapCase, lowering: &NativeLoweringResult) {
    match case.id {
        "window_row_number" => {
            let plan = lowering.plan.as_ref().expect("row_number case should lower to a native plan");
            assert!(
                plan.nodes().iter().any(|node| {
                    matches!(
                        &node.kind,
                        PlanNode::Window(window)
                            if window
                                .spec
                                .functions
                                .iter()
                                .any(|function| matches!(function, WindowFunction::RowNumber { .. }))
                    )
                }),
                "window_row_number should lower to native RowNumber IR"
            );
        },
        "window_lag" => {
            let plan = lowering.plan.as_ref().expect("lag case should lower to a native plan");
            let observed = collect_lag_lead_functions(plan);
            assert!(
                observed.iter().any(|function| matches!(
                    function,
                    WindowFunction::LagLead { offset: 1, default: None, direction: ShiftDirection::Lag, .. }
                )),
                "window_lag should lower to LagLead with offset=1, default=NULL, direction=Lag; got {observed:#?}"
            );
        },
        "window_lag_offset_default" => {
            let plan = lowering.plan.as_ref().expect("LAG(payload, 2, 0) case should lower to a native plan");
            let observed = collect_lag_lead_functions(plan);
            assert!(
                observed.iter().any(|function| matches!(
                    function,
                    WindowFunction::LagLead {
                        offset: 2,
                        default: Some(LiteralValue::Int64(0)),
                        direction: ShiftDirection::Lag,
                        ..
                    }
                )),
                "window_lag_offset_default should lower to LagLead with offset=2, default=Int64(0), direction=Lag; got {observed:#?}"
            );
        },
        "window_lead" => {
            let plan = lowering.plan.as_ref().expect("lead case should lower to a native plan");
            let observed = collect_lag_lead_functions(plan);
            assert!(
                observed.iter().any(|function| matches!(
                    function,
                    WindowFunction::LagLead { offset: 1, default: None, direction: ShiftDirection::Lead, .. }
                )),
                "window_lead should lower to LagLead with offset=1, default=NULL, direction=Lead; got {observed:#?}"
            );
        },
        "offset_without_limit" => {
            let plan = lowering.plan.as_ref().expect("offset case should lower to a native plan");
            assert!(
                matches!(
                    &plan.root_node().kind,
                    PlanNode::Limit(limit) if limit.spec.skip == 2 && limit.spec.fetch.is_none()
                ),
                "offset_without_limit should lower to an offset-only native Limit"
            );
        },
        "literal_only_projection" | "values_table" => {
            let plan = lowering.plan.as_ref().expect("literal source case should lower to a native plan");
            assert!(
                plan.nodes().iter().any(|node| {
                    matches!(
                        &node.kind,
                        PlanNode::Source(source) if matches!(source.as_ref(), SourceNode::LiteralBatch(_))
                    )
                }),
                "{} should lower to a native literal batch source",
                case.id
            );
        },
        _ => {},
    }
}

fn collect_lag_lead_functions(plan: &QueryPlan) -> Vec<WindowFunction> {
    plan.nodes()
        .iter()
        .filter_map(|node| match &node.kind {
            PlanNode::Window(window) => Some(window.spec.functions.clone()),
            _ => None,
        })
        .flatten()
        .filter(|function| matches!(function, WindowFunction::LagLead { .. }))
        .collect()
}