1use 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#[derive(Debug, Clone, PartialEq, Eq)]
57pub(crate) enum NativeOutput {
58 Host,
60 GpuResident { input_name: String },
62}
63
64#[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#[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 pub(super) prepared: bool,
133 pub(super) shared_consumer_count: usize,
136 pub(super) metrics: ExecutionPlanMetricsSet,
137 cache: Arc<PlanProperties>,
138}
139
140impl AlgeonNativeExec {
141 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 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 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 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 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 #[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 #[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 #[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 #[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 pub fn materializes_host(&self) -> bool {
349 self.output == NativeOutput::Host
350 }
351
352 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 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 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 fn execute_host_leaf(
410 &self,
411 partition: usize,
412 task_context: Arc<TaskContext>,
413 ) -> DataFusionResult<SendableRecordBatchStream> {
414 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 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 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 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 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 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
620fn 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 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 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 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}