algeon_datafusion/
native_exec.rs

1//! DataFusion `ExecutionPlan` wrappers around native cuDF query plans.
2//!
3//! Host-sink wrappers materialize native results back into Arrow streams for
4//! DataFusion. Composed wrappers can pass GPU-resident fragment results between
5//! native subplans, but the final DataFusion-facing boundary still streams
6//! host Arrow.
7//!
8//! `ExecutionPlan::execute` on these wrappers performs no GPU work and claims
9//! no grant; it returns a stream whose first poll admits the attempt, runs the
10//! fragment's whole schedule on the device's GPU execution pool, and keeps the
11//! result on the device. Each later poll moves one bounded chunk to host
12//! memory under a DataFusion memory-pool reservation. Every failure crosses
13//! into DataFusion once, as `DataFusionError::External` around an adapter
14//! `Error` whose facts keep the native classification; a failure after native
15//! execution started ends the attempt without a CPU replay. A GPU-resident
16//! producer's output belongs to the context of the composed root that runs it
17//! and is released with that root's stream. The lifetime contract is in
18//! `website/docs/design/admission.mdx#resident-handoff-and-release`.
19
20use std::{collections::HashSet, sync::Arc};
21
22use datafusion::{
23    common::{DataFusionError, Result as DataFusionResult},
24    execution::context::TaskContext,
25    physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet},
26};
27#[cfg(test)]
28use query_engine::exec::admit_query_plan;
29#[cfg(test)]
30use query_planning::{admission::AdmittedPlan, plan::QueryPlan};
31pub use query_planning::{
32    capability::{AggregateStrategyExecutionOptions, KeyedHashJoinStrategyOptions},
33    chunk::SourceChunkPlanningOptions,
34    correlation::{QueryId, TraceId},
35    dynamic_filter::DynamicFilterExecutionOptions,
36    execution_contracts::SourceChunkPlanningMode,
37    mode::StatisticalAggregateExecutionMode,
38    policy::PlanningPolicy,
39};
40pub use query_runtime::{
41    observability::{RuntimeObservationConfig, RuntimeObservationLimits},
42    runtime::{LakehouseCacheFillConfig, S3SourceStorageProfile, SourceSchedulerPolicy},
43};
44
45#[cfg(test)]
46use crate::error::Result;
47use crate::error::{Error, ErrorCode};
48
49/// Admits a plan a test builds for a native node, as planning would.
50#[cfg(test)]
51pub(crate) fn admit_for_test(
52    plan: &QueryPlan,
53    execution_config: &AlgeonNativeExecutionConfig,
54) -> Result<Arc<AdmittedPlan>> {
55    admit_query_plan(plan, execution_config.statistical_aggregate_execution_mode(), execution_config.planning_policy())
56        .map_err(|error| {
57            Error::from_native("algeon_datafusion.native_exec.test_admission", "test plan is not admissible", error)
58        })
59}
60
61#[cfg(test)]
62mod adapter_tests;
63pub(crate) mod admission;
64pub(crate) mod child_inputs;
65mod common;
66mod composed;
67mod config;
68mod fragment;
69pub(crate) mod graph;
70pub(crate) mod lakehouse;
71mod metrics;
72mod native_node;
73mod occurrence;
74mod producer;
75pub(crate) mod relation_function;
76mod shared_resident;
77
78pub(crate) use common::{NativeHostOutputCursor, native_host_output_stream, run_gpu_execution_task};
79
80#[cfg(feature = "cugraph")]
81#[cfg(feature = "cugraph")]
82pub(crate) use self::composed::composed_admission_capability_evidence;
83pub(crate) use self::{
84    admission::{QueryAdmissionSource, plan_is_query_admission_wrapped},
85    composed::producer_bound_capability_evidence,
86    fragment::NativeOutput,
87    metrics::{
88        AlgeonNativeMetrics, GraphClaimCalibrationCapture, NativeExecutionMetricsCapture, SealedAttemptMetricsTarget,
89    },
90    occurrence::{GpuFunctionHostOccurrenceExec, gpu_function_occurrence_from_task_context},
91    producer::{
92        GpuResidentProducer, execute_gpu_resident_producer_with_context, gpu_resident_producer,
93        gpu_resident_producer_hard_output_row_upper_bound, gpu_resident_producer_input_name,
94        gpu_resident_producer_unique_output_key, prepare_gpu_resident_producer_for_execution,
95        prepared_gpu_resident_producer_admission_tickets, rewrite_gpu_resident_producer_for_subsumed_composed_consumer,
96    },
97};
98pub use self::{
99    config::{AlgeonNativeExecutionConfig, LakehouseByteCache, LakehouseByteCacheFill, SourceChunking},
100    fragment::AlgeonNativeExec,
101    metrics::{
102        GraphClaimCalibrationExecution, GraphClaimCalibrationNonPhysicalReason, NATIVE_ALGEON_DATAFUSION_METRIC_NAMES,
103        SERVER_CAPTURE_LOSS_METRIC_NAME, SealedAttemptGraphClaimCalibrationSnapshot,
104    },
105};
106
107/// The plan below the admission node the native optimizer places at a
108/// query's root, or `plan` itself when it has none.
109pub fn below_native_admission(plan: &Arc<dyn ExecutionPlan>) -> &Arc<dyn ExecutionPlan> {
110    match native_node::native_node(plan.as_ref()) {
111        Some(native_node::NativeNode::Admission(node)) => node.input(),
112        _ => plan,
113    }
114}
115
116/// Return the input of an adapter-owned root node (native admission or a GPU
117/// function host occurrence), which adds no rows or GPU work of its own.
118pub(crate) fn transparent_root_wrapper_input(plan: &dyn ExecutionPlan) -> Option<&Arc<dyn ExecutionPlan>> {
119    native_node::native_node(plan).and_then(|node| node.transparent_root_wrapper_input())
120}
121
122/// Validates every host-materializing GPU descendant against the outer
123/// final-plan occurrence manifest.
124pub(crate) fn validate_host_wrapped_gpu_execution(
125    plan: &dyn ExecutionPlan,
126    task_context: &TaskContext,
127) -> DataFusionResult<()> {
128    let ticket = admission::invocation_admission_ticket(task_context)
129        .filter(|ticket| ticket.has_execution_manifest())
130        .ok_or_else(|| {
131            DataFusionError::from(
132                Error::new(
133                    ErrorCode::InvalidExecutionPlan,
134                    "algeon_datafusion.native_exec.host_input_staging",
135                    "host staging cannot execute a GPU descendant without a final-plan occurrence manifest",
136                )
137                .with_log_safe_fact("reason", "host_wrapped_gpu_occurrence_manifest_missing"),
138            )
139        })?;
140
141    fn validate(plan: &dyn ExecutionPlan, ticket: &Arc<admission::QueryAdmissionTicket>) -> DataFusionResult<bool> {
142        if let Some(producer) = gpu_resident_producer(plan) {
143            let host_plan = producer.host_result_query_plan()?.ok_or_else(|| {
144                DataFusionError::from(
145                    Error::new(
146                        ErrorCode::InvalidExecutionPlan,
147                        "algeon_datafusion.native_exec.host_input_staging",
148                        "a host-wrapped GPU producer does not expose its host execution identity",
149                    )
150                    .with_log_safe_fact("reason", "host_wrapped_gpu_occurrence_identity_missing")
151                    .with_log_safe_fact("occurrence_id", producer.gpu_result_input_name()),
152                )
153            })?;
154            ticket.validate_host_occurrence(producer.gpu_result_input_name(), &host_plan)?.ok_or_else(|| {
155                DataFusionError::from(
156                    Error::new(
157                        ErrorCode::InvalidExecutionPlan,
158                        "algeon_datafusion.native_exec.host_input_staging",
159                        "a host-wrapped GPU producer is not covered by the final-plan occurrence manifest",
160                    )
161                    .with_log_safe_fact("reason", "host_wrapped_gpu_occurrence_not_admitted")
162                    .with_log_safe_fact("occurrence_id", producer.gpu_result_input_name()),
163                )
164            })?;
165            return Ok(true);
166        }
167        let mut found = false;
168        for child in plan.children() {
169            found |= validate(child.as_ref(), ticket)?;
170        }
171        Ok(found)
172    }
173
174    if !validate(plan, &ticket)? {
175        return Err(Error::new(
176            ErrorCode::InvalidExecutionPlan,
177            "algeon_datafusion.native_exec.host_input_staging",
178            "a host-wrapped native subtree has no admitted GPU function occurrence",
179        )
180        .with_log_safe_fact("reason", "host_wrapped_gpu_occurrence_identity_missing")
181        .into());
182    }
183    Ok(())
184}
185
186/// Finds the single native execution wrapper closest to an invocation root
187/// and returns its post-seal attempt-metric target. Transparent root wrappers
188/// are skipped, and mixed DataFusion plans select their first native subtree in
189/// stable child order. The target is chosen before execution and receives only
190/// the already-sealed report.
191pub(crate) fn sealed_attempt_metrics_target(plan: &dyn ExecutionPlan) -> Option<SealedAttemptMetricsTarget> {
192    if let Some(exec) = plan.downcast_ref::<AlgeonNativeExec>()
193        && exec.materializes_host()
194    {
195        return Some(exec.sealed_attempt_metrics_target());
196    }
197    if let Some(exec) = graph::graph_execution_plan(plan) {
198        return Some(exec.sealed_attempt_metrics_target());
199    }
200    if let Some(input) = transparent_root_wrapper_input(plan) {
201        return sealed_attempt_metrics_target(input.as_ref());
202    }
203    let descendant = plan.children().into_iter().find_map(|child| sealed_attempt_metrics_target(child.as_ref()));
204    // A producer reached here is either a relation function materializing its
205    // own host result, which makes it the native root, or a composed root's
206    // GPU input whose captures alone are read. It owns no graph calibration
207    // capture, so captures elected below it still reach the sealed snapshot.
208    if let Some(metrics) = gpu_resident_producer(plan).and_then(GpuResidentProducer::execution_metrics) {
209        let captures = descendant.iter().flat_map(SealedAttemptMetricsTarget::graph_claim_calibration_captures);
210        return Some(
211            SealedAttemptMetricsTarget::new(metrics, GraphClaimCalibrationCapture::default())
212                .with_graph_claim_calibration_captures(captures),
213        );
214    }
215    descendant
216}
217
218/// Returns the live metric set of every native execution wrapper in `plan`,
219/// including GPU-resident producers, relation functions, and graph nodes below
220/// a composed root. Execution clones share these sets and wrappers register
221/// their metrics only while executing; callers inspect them after completion.
222/// Readers do not use an `ExecutionPlan::metrics` snapshot.
223pub(crate) fn native_execution_metrics(plan: &dyn ExecutionPlan) -> Vec<ExecutionPlanMetricsSet> {
224    fn visit(plan: &dyn ExecutionPlan, seen: &mut HashSet<*const ()>, sets: &mut Vec<ExecutionPlanMetricsSet>) {
225        // Physical plans can be DAGs; visit each shared node once.
226        if !seen.insert(std::ptr::from_ref(plan).cast::<()>()) {
227            return;
228        }
229        sets.extend(native_execution_wrapper_metrics(plan).cloned());
230        for child in plan.children() {
231            visit(child.as_ref(), seen, sets);
232        }
233    }
234
235    let mut sets = Vec::new();
236    visit(plan, &mut HashSet::new(), &mut sets);
237    sets
238}
239
240/// Whether `plan` contains a native execution wrapper.
241#[cfg(feature = "server")]
242pub(crate) fn contains_native_execution_wrapper(plan: &dyn ExecutionPlan) -> bool {
243    native_execution_wrapper_metrics(plan).is_some()
244        || plan.children().into_iter().any(|child| contains_native_execution_wrapper(child.as_ref()))
245}
246
247fn native_execution_wrapper_metrics(plan: &dyn ExecutionPlan) -> Option<&ExecutionPlanMetricsSet> {
248    if let Some(exec) = plan.downcast_ref::<AlgeonNativeExec>() {
249        return Some(exec.execution_metrics());
250    }
251    gpu_resident_producer(plan).and_then(GpuResidentProducer::execution_metrics)
252}
algeon_datafusion/native_exec.rs

algeon_datafusion/
native_exec.rs
