algeon-datafusion 0.1.2

DataFusion physical-plan adapter for native GPU execution.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use query_engine::{
    dimension_pruning::local_dimension_pruning_resolver, rapids_backend::rapids_analysis_context,
    shadow::estimate_compiled_pipeline,
};
use query_planning::{
    admission::admit,
    mode::StatisticalAggregateExecutionMode,
    optimizer::{NativeOptimizerOptions, optimize_query_plan},
    pipeline::ShadowBytes,
    plan::AdvisoryStatistic,
    policy::PlanningPolicy,
};

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

#[tokio::test]
async fn shadow_advisory_binds_projection_and_final_aggregate_outputs() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-shadow-advisory");
    let path = dir.join("input.parquet");
    write_numeric_parquet_file(&path);
    let plan = sql_plan(
        &path,
        direct_friendly_sql_session_config(),
        "SELECT key, SUM(payload) AS total FROM bench WHERE payload > 1 GROUP BY key",
    )
    .await;
    let lowered = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap();
    assert_eq!(lowered.selected_count(), 1, "{:?}", lowered.report.outcomes());
    let native = lowered.selected_plan();
    let aggregate = native.nodes().iter().find(|node| matches!(node.kind, PlanNode::Aggregate(_))).unwrap();
    let binding = native
        .advisory_outputs()
        .iter()
        .find(|binding| binding.node == aggregate.id)
        .expect("final aggregate output has frontend statistics even when unknown");
    assert_eq!(&binding.schema, aggregate.output_schema());
    assert_eq!(binding.frontend_version, datafusion::DATAFUSION_VERSION);

    let projected =
        sql_plan(&path, direct_friendly_sql_session_config(), "SELECT payload + 1 AS value FROM bench").await;
    let lowered = lower_physical_plan_to_native(projected, &ParquetReadStrategyConfig::default()).unwrap();
    assert_eq!(lowered.selected_count(), 1, "{:?}", lowered.report.outcomes());
    let native = lowered.selected_plan();
    let binding = native
        .advisory_outputs()
        .iter()
        .find(|binding| binding.node == native.root_id())
        .expect("root projection carries its own frontend estimate");
    assert!(!matches!(binding.statistics.rows, AdvisoryStatistic::Absent));
    assert_eq!(&binding.schema, native.root_node().output_schema());
    fs::remove_dir_all(dir).unwrap();
}

#[tokio::test]
async fn shadow_estimate_composes_frontend_advisory_with_the_compiled_schedule() {
    let dir = unique_temp_dir("algeon-datafusion-native-lowering-shadow-estimate");
    write_tiny_tpch_parquet_dataset(&dir).unwrap();
    let items = tpch_fallback_pressure_workload_items(&dir, tpch_report_session_config()).await.unwrap();
    let q3 = items.iter().find(|item| item.query_id == "q03").expect("TPC-H workload should include q03");
    let NativeFallbackPressureWorkloadInput::PhysicalPlan(plan) = &q3.input else {
        panic!("q03 should produce a physical plan");
    };
    let lowered = lower_physical_plan_to_native(Arc::clone(plan), &ParquetReadStrategyConfig::default()).unwrap();
    assert_eq!(lowered.selected_count(), 1, "{:?}", lowered.report.outcomes());
    assert!(!lowered.selected_plan().advisory_outputs().is_empty(), "q03 lowering attaches frontend statistics");

    // The estimate is composed over the schedule that executes, so the
    // advisory volumes must survive the optimizer rewrites that produce it.
    let native = &optimize_query_plan(
        lowered.selected_plan().clone(),
        &NativeOptimizerOptions::default(),
        rapids_analysis_context(PlanningPolicy::default()),
        local_dimension_pruning_resolver(),
    )
    .unwrap()
    .plan;
    assert!(!native.advisory_outputs().is_empty(), "optimizer rewrites carry frontend statistics forward");
    let admission = admit(
        native,
        StatisticalAggregateExecutionMode::default(),
        PlanningPolicy::default(),
        rapids_analysis_context(PlanningPolicy::default()),
    )
    .unwrap()
    .admission()
    .cloned()
    .expect("the lowered q03 join plan compiles to a bounded pipeline");
    let pipeline = admission.pipeline();

    let estimate = estimate_compiled_pipeline(native, pipeline).unwrap();
    assert_eq!(estimate.pipeline.envelopes.len(), pipeline.tasks().len());
    assert_eq!(estimate.pipeline.task_workspace.len(), pipeline.tasks().len());
    assert_eq!(estimate.pipeline.resources.len(), pipeline.resources().len());
    assert_eq!(estimate.frontend_outputs.len(), native.advisory_outputs().len());
    for (binding, captured) in native.advisory_outputs().iter().zip(&estimate.frontend_outputs) {
        assert_eq!(captured["node"].as_u64(), Some(binding.node.index() as u64));
        assert_eq!(captured["frontend"].as_str(), Some(binding.frontend.as_str()));
        assert_eq!(captured["frontend_version"].as_str(), Some(binding.frontend_version.as_str()));
    }

    // Local Parquet decode and every operator workspace still lack a component
    // allocation model, so the composed query-owned peak stays unknown with its
    // reason, and no zero is reported for it.
    assert!(
        estimate.pipeline.resources.iter().any(|resource| matches!(
            &resource.model.bytes,
            ShadowBytes::Unknown { reason } if reason.contains("decode")
        )),
        "{:?}",
        estimate.pipeline.resources
    );
    assert!(matches!(estimate.pipeline.query_owned_peak, ShadowBytes::Unknown { .. }));
    assert!(matches!(estimate.pipeline.borrowed_resident_peak, ShadowBytes::Estimated { bytes: 0, .. }));

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