Skip to main content

algeon_datafusion/native_exec/
fragment.rs

1//! The one DataFusion node for a native execution unit.
2//!
3//! A unit runs one admitted native plan. Its children are the GPU-resident
4//! producers that plan reads as inputs. Its output either streams host Arrow
5//! batches to DataFusion or stays on the GPU for a consumer that binds it by
6//! input name; such a producer cannot stream through
7//! [`ExecutionPlan::execute`].
8
9use std::{collections::BTreeSet, sync::Arc};
10
11use cudf::frame::dataframe::HostBatchCursorConfig;
12use datafusion::{
13    common::{DataFusionError, Result as DataFusionResult, tree_node::TreeNodeRecursion},
14    execution::context::TaskContext,
15    physical_expr::{EquivalenceProperties, PhysicalExpr},
16    physical_plan::{
17        DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
18        metrics::{ExecutionPlanMetricsSet, MetricsSet},
19    },
20};
21use futures::future::BoxFuture;
22#[cfg(test)]
23use query_engine::exec::admit_query_plan;
24use query_engine::exec::{QueryExecutionRequest, execute_query_plan, output::QueryOutput};
25#[cfg(feature = "cugraph")]
26use query_planning::plan::PlanNode;
27use query_planning::{
28    admission::{AdmittedPlan, AttemptCapabilityEvidence},
29    plan::{QueryPlan, ResultSink, row_bounds::query_plan_hard_row_upper_bound},
30};
31#[cfg(test)]
32use query_planning::{mode::StatisticalAggregateExecutionMode, policy::PlanningPolicy};
33use query_runtime::runtime::{ExecutionContext, resident::GpuResultHandle};
34
35#[cfg(test)]
36use super::metrics::NativeMetricsCapture;
37use super::{
38    admission::QueryAdmissionTicket,
39    common::{
40        NATIVE_ALGEON_EXEC_HOST_OUTPUT_MEMORY_CONSUMER, NativeHostOutputCursor, comma_separated,
41        conservative_properties, gpu_resident_source_names, mark_native_runtime_error, native_host_output_stream,
42        native_plan_one_line, native_query_runtime_error, properties_with_equivalence, run_gpu_execution_task,
43        write_native_plan_tree,
44    },
45    config::AlgeonNativeExecutionConfig,
46    graph::graph_execution_plan,
47    metrics::{
48        AlgeonNativeMetrics, GraphClaimCalibrationCapture, NativeExecutionMetricsCapture, SealedAttemptMetricsTarget,
49    },
50    producer::{GpuResidentProducer, gpu_resident_producer_input_name},
51    shared_resident::shared_resident_results,
52};
53use crate::error::{Error, ErrorCode, Result};
54
55/// Where a native unit's result goes.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub(crate) enum NativeOutput {
58    /// Host Arrow batches streamed to DataFusion.
59    Host,
60    /// A GPU-resident result a consumer binds under `input_name`.
61    GpuResident { input_name: String },
62}
63
64/// The four unit shapes, {leaf, composed} by {host, GPU-resident} output.
65/// The shape names the node and the operation of its errors.
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67enum Shape {
68    Host,
69    Composed,
70    Producer,
71    ComposedProducer,
72}
73
74impl Shape {
75    fn of(composed: bool, output: &NativeOutput) -> Self {
76        match (composed, output) {
77            (false, NativeOutput::Host) => Self::Host,
78            (true, NativeOutput::Host) => Self::Composed,
79            (false, NativeOutput::GpuResident { .. }) => Self::Producer,
80            (true, NativeOutput::GpuResident { .. }) => Self::ComposedProducer,
81        }
82    }
83
84    const fn name(self) -> &'static str {
85        match self {
86            Self::Host => "AlgeonNativeExec",
87            Self::Composed => "AlgeonNativeComposedExec",
88            Self::Producer => "AlgeonNativeGpuProducerExec",
89            Self::ComposedProducer => "AlgeonNativeComposedGpuProducerExec",
90        }
91    }
92
93    const fn try_new_operation(self) -> &'static str {
94        match self {
95            Self::Host => "algeon_datafusion.native_exec.try_new",
96            Self::Composed => "algeon_datafusion.native_composed.try_new",
97            Self::Producer => "algeon_datafusion.native_gpu_producer.try_new",
98            Self::ComposedProducer => "algeon_datafusion.native_composed_gpu_producer.try_new",
99        }
100    }
101
102    const fn with_new_children_operation(self) -> &'static str {
103        match self {
104            Self::Host => "algeon_datafusion.native_exec.with_new_children",
105            Self::Composed => "algeon_datafusion.native_composed.with_new_children",
106            Self::Producer => "algeon_datafusion.native_gpu_producer.with_new_children",
107            Self::ComposedProducer => "algeon_datafusion.native_composed_gpu_producer.with_new_children",
108        }
109    }
110}
111
112/// One native execution unit: an admitted native plan, the GPU-resident
113/// producers it reads, and its host or GPU-resident output.
114///
115/// Created by the Algeon physical optimizer and adapter composition; the
116/// public type identity serves [`ExecutionPlan`] inspection, downcasting,
117/// metrics collection, and plan diagnostics. Construction is crate-private
118/// and validates the sink against the output, the producer children against
119/// the plan's GPU-resident sources, and the equivalence properties. Each node
120/// exposes one output partition. A host unit executes its producers and its
121/// plan in one admitted native context and streams the result as Arrow.
122#[derive(Debug, Clone)]
123pub struct AlgeonNativeExec {
124    pub(super) admitted: Arc<AdmittedPlan>,
125    pub(super) gpu_inputs: Vec<Arc<dyn ExecutionPlan>>,
126    output: NativeOutput,
127    pub(super) execution_config: AlgeonNativeExecutionConfig,
128    pub(super) metrics_capture: Option<NativeExecutionMetricsCapture>,
129    pub(super) graph_claim_calibration_capture: GraphClaimCalibrationCapture,
130    /// Whether a composed GPU-resident unit resolved its plan and producers
131    /// before the composed root acquired the native context.
132    pub(super) prepared: bool,
133    /// Units of the composed island that read this GPU-resident output. Above
134    /// one, the island computes the output once and shares it.
135    pub(super) shared_consumer_count: usize,
136    pub(super) metrics: ExecutionPlanMetricsSet,
137    cache: Arc<PlanProperties>,
138}
139
140impl AlgeonNativeExec {
141    /// A native unit for `admitted`, reading `gpu_inputs` and writing
142    /// `output`. `eq_properties` are the replaced DataFusion subtree's
143    /// properties, which the native plan must already reproduce; `None`
144    /// declares no ordering.
145    pub(crate) fn try_new(
146        admitted: Arc<AdmittedPlan>,
147        gpu_inputs: Vec<Arc<dyn ExecutionPlan>>,
148        output: NativeOutput,
149        execution_config: AlgeonNativeExecutionConfig,
150        eq_properties: Option<EquivalenceProperties>,
151    ) -> Result<Self> {
152        let metrics_capture = execution_config.new_metrics_capture();
153        Self::try_new_with_metrics_capture(
154            admitted,
155            gpu_inputs,
156            output,
157            execution_config,
158            eq_properties,
159            metrics_capture,
160        )
161    }
162
163    fn try_new_with_metrics_capture(
164        admitted: Arc<AdmittedPlan>,
165        gpu_inputs: Vec<Arc<dyn ExecutionPlan>>,
166        output: NativeOutput,
167        execution_config: AlgeonNativeExecutionConfig,
168        eq_properties: Option<EquivalenceProperties>,
169        metrics_capture: Option<NativeExecutionMetricsCapture>,
170    ) -> Result<Self> {
171        let shape = Shape::of(!gpu_inputs.is_empty(), &output);
172        let operation = shape.try_new_operation();
173        let plan = admitted.plan();
174        let schema = match &output {
175            // The host unit executes the plan in the deferred form it was
176            // admitted in and materializes the host batches itself.
177            NativeOutput::Host => {
178                if !matches!(plan.sink(), ResultSink::DeferredHostArrow { .. }) {
179                    return Err(Error::new(
180                        ErrorCode::InvalidExecutionPlan,
181                        operation,
182                        "host native execution requires a deferred host-arrow result sink",
183                    )
184                    .with_fact("result_sink", result_sink_label(plan.sink())));
185                }
186                plan.caller_visible_output_schema()
187            },
188            NativeOutput::GpuResident { input_name } => {
189                if input_name.is_empty() {
190                    return Err(Error::new(
191                        ErrorCode::InvalidExecutionPlan,
192                        operation,
193                        "GPU-resident producer input name must not be empty",
194                    ));
195                }
196                if plan.sink().materializes_host_boundary() {
197                    return Err(Error::new(
198                        ErrorCode::InvalidExecutionPlan,
199                        operation,
200                        "GPU-resident producer requires a GPU-resident result sink",
201                    )
202                    .with_fact("result_sink", "host_arrow"));
203                }
204                plan.output_schema()
205            },
206        };
207        validate_gpu_inputs(plan, &gpu_inputs, operation)?;
208        let cache = match eq_properties {
209            Some(eq_properties) => properties_with_equivalence(Arc::clone(schema), eq_properties)?,
210            None => conservative_properties(Arc::clone(schema)),
211        };
212        // A host root records its producers' native execution under its own
213        // capture, so one attempt projection reads the whole unit.
214        let gpu_inputs = if output == NativeOutput::Host {
215            gpu_inputs
216                .into_iter()
217                .map(|input| match input.downcast_ref::<Self>() {
218                    Some(producer) if producer.gpu_inputs.is_empty() => {
219                        let mut producer = producer.clone();
220                        producer.metrics_capture.clone_from(&metrics_capture);
221                        Arc::new(producer) as Arc<dyn ExecutionPlan>
222                    },
223                    _ => input,
224                })
225                .collect()
226        } else {
227            gpu_inputs
228        };
229        Ok(Self {
230            admitted,
231            gpu_inputs,
232            output,
233            execution_config,
234            metrics_capture,
235            graph_claim_calibration_capture: GraphClaimCalibrationCapture::default(),
236            prepared: false,
237            shared_consumer_count: 1,
238            metrics: ExecutionPlanMetricsSet::new(),
239            cache,
240        })
241    }
242
243    /// A GPU-resident producer for `admitted` whose output `consumer_count`
244    /// units of one composed island read. The island computes it once, when
245    /// its first reading unit runs, and releases it after the last one.
246    pub(crate) fn try_new_shared_producer(
247        admitted: Arc<AdmittedPlan>,
248        input_name: String,
249        consumer_count: usize,
250        execution_config: AlgeonNativeExecutionConfig,
251    ) -> Result<Self> {
252        let mut producer =
253            Self::try_new(admitted, Vec::new(), NativeOutput::GpuResident { input_name }, execution_config, None)?;
254        producer.shared_consumer_count = consumer_count;
255        Ok(producer)
256    }
257
258    /// This unit running `admitted` over `gpu_inputs`, with its output,
259    /// configuration, metrics capture and declared properties.
260    pub(crate) fn with_rebound_plan(
261        &self,
262        admitted: Arc<AdmittedPlan>,
263        gpu_inputs: Vec<Arc<dyn ExecutionPlan>>,
264    ) -> Result<Self> {
265        let mut unit = Self::try_new_with_metrics_capture(
266            admitted,
267            gpu_inputs,
268            self.output.clone(),
269            self.execution_config.clone(),
270            Some(self.cache.equivalence_properties().clone()),
271            self.metrics_capture.clone(),
272        )?;
273        unit.shared_consumer_count = self.shared_consumer_count;
274        Ok(unit)
275    }
276
277    /// A host unit for `plan` admitted as planning would.
278    #[cfg(test)]
279    pub(crate) fn try_new_host_for_test(
280        plan: QueryPlan,
281        execution_config: AlgeonNativeExecutionConfig,
282    ) -> Result<Self> {
283        Self::try_new_composed_for_test(plan, Vec::new(), execution_config)
284    }
285
286    /// A host unit reading `gpu_inputs`, for `plan` admitted as planning would.
287    #[cfg(test)]
288    pub(crate) fn try_new_composed_for_test(
289        plan: QueryPlan,
290        gpu_inputs: Vec<Arc<dyn ExecutionPlan>>,
291        execution_config: AlgeonNativeExecutionConfig,
292    ) -> Result<Self> {
293        let admitted = super::admit_for_test(&plan.clone_with_deferred_host_materialization_sink(), &execution_config)?;
294        Self::try_new(admitted, gpu_inputs, NativeOutput::Host, execution_config, None)
295    }
296
297    /// A GPU-resident leaf producer for `plan`.
298    #[cfg(test)]
299    pub(crate) fn try_new_producer_for_test(input_name: impl Into<String>, plan: QueryPlan) -> Result<Self> {
300        Self::try_new_producer_with_capture_for_test(input_name, plan, None, None, false)
301    }
302
303    /// A GPU-resident leaf producer recording into `metrics_capture`.
304    #[cfg(test)]
305    pub(crate) fn try_new_producer_with_capture_for_test(
306        input_name: impl Into<String>,
307        plan: QueryPlan,
308        metrics_capture: Option<NativeMetricsCapture>,
309        plan_task_mapping_datafusion_shape: Option<&str>,
310        lakehouse_footer_pruning: bool,
311    ) -> Result<Self> {
312        let metrics_capture = Some(metrics_capture.unwrap_or_else(NativeMetricsCapture::standalone).new_execution());
313        let admitted = admit_query_plan(&plan, StatisticalAggregateExecutionMode::default(), PlanningPolicy::default())
314            .map_err(|error| {
315                Error::from_native("algeon_datafusion.native_exec.test_admission", "test plan is not admissible", error)
316            })?;
317        let mut execution_config =
318            AlgeonNativeExecutionConfig::default().with_lakehouse_footer_pruning(lakehouse_footer_pruning);
319        if let Some(shape) = plan_task_mapping_datafusion_shape {
320            execution_config = execution_config.with_plan_task_mapping_datafusion_shape(shape);
321        }
322        Self::try_new_with_metrics_capture(
323            admitted,
324            Vec::new(),
325            NativeOutput::GpuResident { input_name: input_name.into() },
326            execution_config,
327            None,
328            metrics_capture,
329        )
330    }
331
332    fn shape(&self) -> Shape {
333        Shape::of(!self.gpu_inputs.is_empty(), &self.output)
334    }
335
336    pub(crate) fn native_plan(&self) -> &QueryPlan {
337        self.admitted.plan()
338    }
339
340    #[cfg(feature = "bench-support")]
341    #[doc(hidden)]
342    pub fn admitted_plan_for_internal_consumer(&self) -> &Arc<AdmittedPlan> {
343        &self.admitted
344    }
345
346    /// Whether the unit streams host Arrow batches to DataFusion; otherwise
347    /// its output stays on the GPU for a consumer unit.
348    pub fn materializes_host(&self) -> bool {
349        self.output == NativeOutput::Host
350    }
351
352    /// The input name a consumer binds this unit's GPU-resident output under.
353    pub fn gpu_resident_input_name(&self) -> Option<&str> {
354        match &self.output {
355            NativeOutput::Host => None,
356            NativeOutput::GpuResident { input_name } => Some(input_name),
357        }
358    }
359
360    pub(crate) fn gpu_inputs(&self) -> &[Arc<dyn ExecutionPlan>] {
361        &self.gpu_inputs
362    }
363
364    #[cfg(feature = "cugraph")]
365    pub(crate) fn execution_config(&self) -> &AlgeonNativeExecutionConfig {
366        &self.execution_config
367    }
368
369    pub(crate) fn execution_metrics(&self) -> &ExecutionPlanMetricsSet {
370        &self.metrics
371    }
372
373    /// Returns the one post-seal attempt projection target for this host
374    /// unit. Its producers are not separate owners; their graph calibration
375    /// captures join this target. The attempt owner elects the target before
376    /// execution; the node never reads or retains terminal attempt state.
377    pub(crate) fn sealed_attempt_metrics_target(&self) -> SealedAttemptMetricsTarget {
378        let producer_captures = self.gpu_inputs.iter().filter_map(|input| {
379            super::sealed_attempt_metrics_target(input.as_ref())
380                .map(|target| target.graph_claim_calibration_captures().collect::<Vec<_>>())
381        });
382        SealedAttemptMetricsTarget::new(&self.metrics, self.graph_claim_calibration_capture.clone())
383            .with_graph_claim_calibration_captures(producer_captures.flatten())
384    }
385
386    /// Resolves footer-pruned metadata for the plan and its producers before
387    /// execution acquires the native context.
388    pub(crate) async fn prepare_lakehouse_metadata(&self, task_context: Arc<TaskContext>) -> DataFusionResult<Self> {
389        let mut prepared = self.clone();
390        prepared.admitted = super::lakehouse::resolve_admitted_plan(
391            self.execution_config.lakehouse_footer_pruning(),
392            &self.admitted,
393            task_context.as_ref(),
394        )
395        .await?;
396        let mut prepared_inputs = Vec::with_capacity(self.gpu_inputs.len());
397        for input in &self.gpu_inputs {
398            prepared_inputs.push(
399                super::prepare_gpu_resident_producer_for_execution(Arc::clone(input), Arc::clone(&task_context))
400                    .await?,
401            );
402        }
403        prepared.gpu_inputs = prepared_inputs;
404        prepared.execution_config = prepared.execution_config.with_lakehouse_footer_pruning(false);
405        Ok(prepared)
406    }
407
408    /// Streams the host result of a unit without producers.
409    fn execute_host_leaf(
410        &self,
411        partition: usize,
412        task_context: Arc<TaskContext>,
413    ) -> DataFusionResult<SendableRecordBatchStream> {
414        // Runtime, allocator, boundary, and memory-control state come from
415        // `AlgeonNativeExecutionConfig`; the DataFusion task context scopes
416        // backend admission across native fragments in one execution.
417        let schema = self.schema();
418        let admitted = Arc::clone(&self.admitted);
419        let execution_config = self.execution_config.clone();
420        let account_memory_pool = execution_config.host_output_memory_pool_accounting();
421        let metrics = AlgeonNativeMetrics::new_with_capture_and_shape(
422            &self.metrics,
423            partition,
424            self.metrics_capture.clone(),
425            self.execution_config.plan_task_mapping_datafusion_shape(),
426        )
427        .with_graph_claim_calibration_capture(self.graph_claim_calibration_capture.clone());
428        let output_schema = Arc::clone(&schema);
429        let output = native_host_output_stream(
430            NATIVE_ALGEON_EXEC_HOST_OUTPUT_MEMORY_CONSUMER,
431            schema,
432            Arc::clone(&task_context),
433            account_memory_pool,
434            async move {
435                let unresolved_capability_evidence = AttemptCapabilityEvidence::for_query_plan(admitted.plan());
436                let admitted = super::lakehouse::resolve_admitted_plan(
437                    execution_config.lakehouse_footer_pruning(),
438                    &admitted,
439                    task_context.as_ref(),
440                )
441                .await?;
442                let plan = admitted.plan();
443                // Footer pruning resolves the plan at execution, after the
444                // ticket's evidence was aggregated at planning. The resolved
445                // plan must refine a scope the ticket already attests to; a
446                // different plan would run under evidence that never
447                // described it.
448                if execution_config.lakehouse_footer_pruning()
449                    && let Some(ticket) =
450                        execution_config.query_admission_ticket_for_task_context(task_context.as_ref())
451                {
452                    let evidence = AttemptCapabilityEvidence::for_query_plan(plan);
453                    if let Some(existing) = ticket.attempt_capability_evidence_for_rebind()
454                        && unresolved_capability_evidence
455                            .execution_scopes()
456                            .iter()
457                            .any(|scope| !existing.execution_scopes().iter().any(|bound| bound == scope))
458                    {
459                        return Err(DataFusionError::from(
460                            Error::new(
461                                ErrorCode::InvalidExecutionPlan,
462                                "algeon_datafusion.native_exec.admission_evidence",
463                                "the active native host admission ticket does not attest to this unresolved fragment",
464                            )
465                            .with_fact("reason", "native_host_admission_evidence_scope_mismatch")
466                            .with_fact("ticket_evidence_scope", "different_native_plan"),
467                        ));
468                    }
469                    ticket.bind_fragment_evidence([evidence]).map_err(|error| {
470                        DataFusionError::from(Error::from_native(
471                            "algeon_datafusion.native_exec.admission_evidence",
472                            "failed to bind the native host fragment to its attempt admission",
473                            error,
474                        ))
475                    })?;
476                }
477                // Context construction is the configuration/validation stage. Failures here happen
478                // before the native query starts.
479                let context_lease = execution_config
480                    .execution_context_async(Arc::clone(&task_context))
481                    .await
482                    .map_err(|error| DataFusionError::from(mark_native_runtime_error(error, "runtime_context")))?;
483                if let Some(admission_facts) = context_lease.admission_facts() {
484                    metrics.record_admission_facts(admission_facts);
485                }
486                let gpu_executor = context_lease.context().gpu_executor().map_err(|error| {
487                    DataFusionError::from(mark_native_runtime_error(
488                        Error::from_native(
489                            "algeon_datafusion.native_exec.gpu_execution_pool",
490                            "query service closed before native execution acquired its GPU pool",
491                            error,
492                        ),
493                        "gpu_execution_pool",
494                    ))
495                })?;
496                let observation = context_lease.context().runtime_observation().clone();
497                if let Some(ticket) = super::admission::query_admission_ticket_from_task_context(task_context.as_ref())
498                {
499                    ticket.record_attempt_gpu_execution_started();
500                }
501                let cursor_executor = gpu_executor.clone();
502                let cursor_observation = observation.clone();
503                // The whole schedule runs in this one pool job; the cursor it
504                // returns performs the chunked device-to-host egress on later
505                // polls.
506                run_gpu_execution_task(
507                    gpu_executor,
508                    observation,
509                    "algeon_datafusion.native_exec.gpu_execution_pool",
510                    move |gpu_stream_scope| async move {
511                        let context = context_lease.context().with_runtime_gpu_stream_scope(gpu_stream_scope);
512                        let _elapsed_compute = metrics.baseline.elapsed_compute().timer();
513                        // Once native execution starts, the attempt is terminal. Errors are surfaced
514                        // with stable facts by `native_query_runtime_error` and are not replayed
515                        // through DataFusion.
516                        let outcome = execute_query_plan(&context, QueryExecutionRequest::admitted(&admitted));
517                        let (output, deferred_metrics) = metrics
518                            .finalize_query_outcome_for_deferred_host(outcome, context_lease)
519                            .map_err(|error| {
520                                DataFusionError::from(native_query_runtime_error(
521                                    "algeon_datafusion.native_exec.execute",
522                                    "Algeon native query execution failed",
523                                    error,
524                                ))
525                            })?;
526                        let QueryOutput::GpuResident(handle) = output else {
527                            return Err(DataFusionError::from(
528                                Error::new(
529                                    ErrorCode::InvalidExecutionPlan,
530                                    "algeon_datafusion.native_exec.execute",
531                                    "deferred native host execution did not produce a GPU-resident cursor input",
532                                )
533                                .with_fact("result_sink", "host_arrow"),
534                            ));
535                        };
536                        let cursor = handle
537                            .into_host_batch_cursor_with_host_schema(
538                                output_schema,
539                                context.cancellation().clone(),
540                                HostBatchCursorConfig::default(),
541                            )
542                            .map_err(|error| {
543                                DataFusionError::from(native_query_runtime_error(
544                                    "algeon_datafusion.native_exec.host_batch_cursor",
545                                    "failed to bind and create deferred native host-output cursor",
546                                    error,
547                                ))
548                            })?;
549                        Ok::<NativeHostOutputCursor, DataFusionError>(NativeHostOutputCursor::new(
550                            cursor,
551                            cursor_executor,
552                            cursor_observation,
553                            deferred_metrics,
554                        ))
555                    },
556                )
557                .await
558            },
559        );
560
561        Ok(output)
562    }
563
564    /// Executes a producer without inputs inside a composed root's shared
565    /// native context. The result stays on the GPU and is returned under its
566    /// input name, so the consumer binds it without host materialization.
567    fn execute_leaf_gpu_result(
568        &self,
569        native_context: &ExecutionContext,
570        partition: usize,
571        admitted: &AdmittedPlan,
572    ) -> DataFusionResult<(String, GpuResultHandle)> {
573        if partition != 0 {
574            return Err(Error::new(
575                ErrorCode::InvalidExecutionPlan,
576                "algeon_datafusion.native_gpu_producer.execute_gpu",
577                "native GPU-resident producer exposes exactly one output partition",
578            )
579            .with_fact("partition", partition)
580            .into());
581        }
582
583        let metrics = AlgeonNativeMetrics::new_with_capture_and_shape(
584            &self.metrics,
585            partition,
586            self.metrics_capture.clone(),
587            self.execution_config.plan_task_mapping_datafusion_shape(),
588        );
589        let _elapsed_compute = metrics.baseline.elapsed_compute().timer();
590        let outcome = execute_query_plan(native_context, QueryExecutionRequest::admitted(admitted));
591        let output = metrics.finalize_query_outcome(outcome).map_err(|error| {
592            DataFusionError::from(native_query_runtime_error(
593                "algeon_datafusion.native_gpu_producer.execute_gpu",
594                "algeon native GPU-resident fragment execution failed",
595                error,
596            ))
597        })?;
598        let QueryOutput::GpuResident(handle) = output else {
599            return Err(DataFusionError::from(
600                Error::new(
601                    ErrorCode::InvalidExecutionPlan,
602                    "algeon_datafusion.native_gpu_producer.execute_gpu",
603                    "native GPU-resident producer materialized a host result",
604                )
605                .with_fact("result_sink", "host_arrow"),
606            ));
607        };
608        Ok((self.gpu_result_input_name().to_string(), handle))
609    }
610}
611
612fn result_sink_label(sink: &ResultSink) -> &'static str {
613    match sink {
614        ResultSink::HostArrow { .. } => "host_arrow",
615        ResultSink::GpuResident { .. } => "gpu_resident",
616        ResultSink::DeferredHostArrow { .. } => "deferred_host_arrow",
617    }
618}
619
620/// Checks that the producer children supply the GPU-resident sources the plan
621/// declares and no others, each under a distinct input name.
622fn validate_gpu_inputs(plan: &QueryPlan, gpu_inputs: &[Arc<dyn ExecutionPlan>], operation: &'static str) -> Result<()> {
623    let required_inputs = gpu_resident_source_names(plan);
624    if gpu_inputs.is_empty() {
625        if required_inputs.is_empty() {
626            return Ok(());
627        }
628        return Err(Error::new(
629            ErrorCode::InvalidExecutionPlan,
630            operation,
631            "native execution is missing the GPU-resident producers its plan reads",
632        )
633        .with_fact("gpu_resident_source_count", required_inputs.len())
634        .with_fact("gpu_resident_sources", comma_separated(required_inputs.iter().map(String::as_str))));
635    }
636    let mut provided_inputs = BTreeSet::new();
637    for child in gpu_inputs {
638        let Some(input_name) = gpu_resident_producer_input_name(child.as_ref()) else {
639            if let Some(rejection) =
640                graph_execution_plan(child.as_ref()).and_then(|graph| graph.gpu_resident_output_rejection())
641            {
642                return Err(rejection);
643            }
644            return Err(Error::new(
645                ErrorCode::InvalidExecutionPlan,
646                operation,
647                "composed native execution children must be GPU-resident producers",
648            )
649            .with_fact("child_name", child.name()));
650        };
651        if !provided_inputs.insert(input_name.clone()) {
652            return Err(Error::new(
653                ErrorCode::InvalidExecutionPlan,
654                operation,
655                "composed native execution received duplicate GPU-resident input names",
656            )
657            .with_fact("input_name", input_name));
658        }
659    }
660    let missing_inputs = required_inputs.difference(&provided_inputs).map(String::as_str).collect::<Vec<_>>();
661    if !missing_inputs.is_empty() {
662        return Err(Error::new(
663            ErrorCode::InvalidExecutionPlan,
664            operation,
665            "composed native execution is missing GPU-resident producer inputs",
666        )
667        .with_fact("missing_inputs", missing_inputs.join(",")));
668    }
669    let unused_inputs = provided_inputs.difference(&required_inputs).map(String::as_str).collect::<Vec<_>>();
670    if !unused_inputs.is_empty() {
671        return Err(Error::new(
672            ErrorCode::InvalidExecutionPlan,
673            operation,
674            "composed native execution received GPU-resident producer inputs the plan does not declare",
675        )
676        .with_fact("unused_inputs", unused_inputs.join(",")));
677    }
678    Ok(())
679}
680
681impl GpuResidentProducer for AlgeonNativeExec {
682    #[cfg(feature = "cugraph")]
683    fn gpu_result_is_repeatable(&self) -> bool {
684        self.native_plan().nodes().iter().all(|node| match &node.kind {
685            PlanNode::Source(_)
686            | PlanNode::Filter(_)
687            | PlanNode::Projection(_)
688            | PlanNode::Aggregate(_)
689            | PlanNode::Window(_)
690            | PlanNode::Sort(_)
691            | PlanNode::Limit(_)
692            | PlanNode::Join(_)
693            | PlanNode::Union(_)
694            | PlanNode::EdgeNormalize(_) => true,
695            PlanNode::GraphAlgorithm(_) | PlanNode::VectorAlgorithm(_) | PlanNode::MlAlgorithm(_) => false,
696        }) && self.gpu_inputs.iter().all(|input| {
697            super::producer::gpu_resident_producer(input.as_ref())
698                .is_some_and(GpuResidentProducer::gpu_result_is_repeatable)
699        })
700    }
701
702    /// The name of a GPU-resident unit's output. Only GPU-resident units are
703    /// classified as producers, so a host unit never answers.
704    fn gpu_result_input_name(&self) -> &str {
705        self.gpu_resident_input_name().unwrap_or_default()
706    }
707
708    fn gpu_result_capability_evidence(&self) -> Result<AttemptCapabilityEvidence> {
709        if self.gpu_inputs.is_empty() {
710            return Ok(AttemptCapabilityEvidence::for_query_plan(self.native_plan()));
711        }
712        super::composed::resolve_composed_admission_capability_evidence(self.native_plan(), &self.gpu_inputs).map_err(
713            |error| {
714                Error::from_native(
715                    "algeon_datafusion.native_composed_gpu_producer.admission_evidence",
716                    "failed to derive capability evidence for a composed GPU-resident output",
717                    error,
718                )
719            },
720        )
721    }
722
723    fn hard_output_row_upper_bound(&self) -> Option<u64> {
724        self.gpu_inputs.is_empty().then(|| query_plan_hard_row_upper_bound(self.native_plan())).flatten()
725    }
726
727    fn execution_metrics(&self) -> Option<&ExecutionPlanMetricsSet> {
728        Some(&self.metrics)
729    }
730
731    fn prepared_admission_tickets(&self) -> Vec<Arc<QueryAdmissionTicket>> {
732        if self.gpu_inputs.is_empty() {
733            return Vec::new();
734        }
735        self.composed_prepared_admission_tickets()
736    }
737
738    fn prepare_execution_plan<'a>(
739        &'a self,
740        plan: Arc<dyn ExecutionPlan>,
741        task_context: Arc<TaskContext>,
742    ) -> BoxFuture<'a, DataFusionResult<Arc<dyn ExecutionPlan>>> {
743        Box::pin(async move {
744            if self.gpu_inputs.is_empty() && !self.execution_config.lakehouse_footer_pruning() {
745                return Ok(plan);
746            }
747            let mut prepared = self.prepare_lakehouse_metadata(task_context).await?;
748            prepared.prepared = true;
749            Ok(Arc::new(prepared) as Arc<dyn ExecutionPlan>)
750        })
751    }
752
753    fn execute_gpu_result_with_context<'a>(
754        &'a self,
755        native_context: &'a ExecutionContext,
756        partition: usize,
757        task_context: Arc<TaskContext>,
758    ) -> BoxFuture<'a, DataFusionResult<(String, GpuResultHandle)>> {
759        Box::pin(async move {
760            if !self.gpu_inputs.is_empty() {
761                return self.execute_composed_gpu_result(native_context, partition, task_context).await;
762            }
763            // A shared producer runs once per island execution. The first
764            // reading unit computes and retains the result; later units take
765            // a clone from the island's result scope.
766            let shared_results = if self.shared_consumer_count > 1 {
767                let results = shared_resident_results(&task_context).ok_or_else(|| {
768                    Error::new(
769                        ErrorCode::InvalidExecutionPlan,
770                        "algeon_datafusion.native_gpu_producer.execute_gpu",
771                        "a shared GPU-resident producer ran outside the composed island of its consumers",
772                    )
773                    .with_log_safe_fact("reason", "shared_resident_scope_missing")
774                    .with_fact("input_name", self.gpu_result_input_name())
775                })?;
776                if let Some(handle) = results.take(self.gpu_result_input_name()) {
777                    return Ok((self.gpu_result_input_name().to_string(), handle));
778                }
779                Some(results)
780            } else {
781                None
782            };
783            let admitted = super::lakehouse::resolve_admitted_plan(
784                self.execution_config.lakehouse_footer_pruning(),
785                &self.admitted,
786                task_context.as_ref(),
787            )
788            .await?;
789            let (input_name, handle) = self.execute_leaf_gpu_result(native_context, partition, &admitted)?;
790            if let Some(results) = shared_results {
791                results.retain(&input_name, &handle);
792            }
793            Ok((input_name, handle))
794        })
795    }
796}
797
798impl DisplayAs for AlgeonNativeExec {
799    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
800        let plan = self.native_plan();
801        let shape = self.shape();
802        let input_name = self.gpu_resident_input_name().unwrap_or_default();
803        let gpu_input_count = self.gpu_inputs.len();
804        let summary = match shape {
805            Shape::Host => format!(
806                "native_nodes={}, sink=host_arrow, native_pipeline_execution_mode=bounded, output_boundary=device_to_host, edge_normalize_inserted={}",
807                plan.node_count(),
808                plan.contains_edge_normalize(),
809            ),
810            Shape::Composed => format!(
811                "native_nodes={}, gpu_input_count={gpu_input_count}, sink=host_arrow, native_pipeline_execution_mode=bounded, input_boundary=gpu_resident, output_boundary=device_to_host, edge_normalize_inserted={}",
812                plan.node_count(),
813                plan.contains_edge_normalize(),
814            ),
815            Shape::Producer if self.shared_consumer_count > 1 => format!(
816                "input_name={input_name}, shared_consumers={}, native_nodes={}, sink=gpu_resident, output_boundary=device_resident, edge_normalize_inserted={}",
817                self.shared_consumer_count,
818                plan.node_count(),
819                plan.contains_edge_normalize(),
820            ),
821            Shape::Producer => format!(
822                "input_name={input_name}, native_nodes={}, sink=gpu_resident, output_boundary=device_resident, edge_normalize_inserted={}",
823                plan.node_count(),
824                plan.contains_edge_normalize(),
825            ),
826            Shape::ComposedProducer => format!(
827                "input_name={input_name}, native_nodes={}, gpu_input_count={gpu_input_count}, sink=gpu_resident, input_boundary=gpu_resident, output_boundary=device_resident",
828                plan.node_count(),
829            ),
830        };
831        match t {
832            DisplayFormatType::Default => write!(f, "{}: {summary}", shape.name()),
833            DisplayFormatType::Verbose => {
834                write!(f, "{}: {summary}, native_plan={}", shape.name(), native_plan_one_line(plan))
835            },
836            DisplayFormatType::TreeRender => {
837                for field in summary.split(", ") {
838                    writeln!(f, "{field}")?;
839                }
840                write_native_plan_tree(f, plan)
841            },
842        }
843    }
844}
845
846impl ExecutionPlan for AlgeonNativeExec {
847    fn name(&self) -> &str {
848        self.shape().name()
849    }
850
851    fn properties(&self) -> &Arc<PlanProperties> {
852        &self.cache
853    }
854
855    fn apply_expressions(
856        &self,
857        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DataFusionResult<TreeNodeRecursion>,
858    ) -> DataFusionResult<TreeNodeRecursion> {
859        Ok(TreeNodeRecursion::Continue)
860    }
861
862    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
863        self.gpu_inputs.iter().collect()
864    }
865
866    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
867        // Keep DataFusion distribution rules from inserting a row-stream
868        // wrapper between this unit and its device inputs.
869        vec![false; self.gpu_inputs.len()]
870    }
871
872    fn with_new_children(
873        self: Arc<Self>,
874        children: Vec<Arc<dyn ExecutionPlan>>,
875    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
876        if self.gpu_inputs.is_empty() {
877            if !children.is_empty() {
878                return Err(Error::new(
879                    ErrorCode::InvalidChildCount,
880                    self.shape().with_new_children_operation(),
881                    "native execution without GPU-resident inputs expected no children",
882                )
883                .with_fact("actual_child_count", children.len())
884                .into());
885            }
886            return Ok(self);
887        }
888        self.with_rebound_plan(Arc::clone(&self.admitted), children)
889            .map(|plan| Arc::new(plan) as Arc<dyn ExecutionPlan>)
890            .map_err(DataFusionError::from)
891    }
892
893    fn execute(&self, partition: usize, task_context: Arc<TaskContext>) -> DataFusionResult<SendableRecordBatchStream> {
894        let shape = self.shape();
895        let operation = match shape {
896            Shape::Host => "algeon_datafusion.native_exec.execute",
897            Shape::Composed => "algeon_datafusion.native_composed.execute",
898            Shape::Producer => "algeon_datafusion.native_gpu_producer.execute",
899            Shape::ComposedProducer => "algeon_datafusion.native_composed_gpu_producer.execute",
900        };
901        if !self.materializes_host() {
902            return Err(Error::new(
903                ErrorCode::InvalidExecutionPlan,
904                operation,
905                "native GPU-resident producer is an internal device boundary and cannot stream host Arrow directly",
906            )
907            .with_fact("result_sink", "gpu_resident")
908            .into());
909        }
910        if partition != 0 {
911            return Err(Error::new(
912                ErrorCode::InvalidExecutionPlan,
913                operation,
914                "native execution exposes exactly one output partition",
915            )
916            .with_fact("partition", partition)
917            .into());
918        }
919        if shape == Shape::Composed {
920            return self.execute_host_composed(partition, task_context);
921        }
922        self.execute_host_leaf(partition, task_context)
923    }
924
925    fn metrics(&self) -> Option<MetricsSet> {
926        Some(self.metrics.clone_inner())
927    }
928}