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
118

use std::sync::Arc;

use async_trait::async_trait;
use datafusion::{
    catalog::Session,
    common::Result,
    logical_expr::{LogicalPlan, UserDefinedLogicalNode, physical_planning_context::PhysicalPlanningContext},
    physical_plan::ExecutionPlan,
    physical_planner::{ExtensionPlanner, PhysicalPlanner},
};
#[cfg(feature = "cugraph")]
use datafusion::{common::DataFusionError, execution::SessionState};

#[cfg(not(feature = "cugraph"))]
use super::planning::CugraphFeatureDisabledExec;
#[cfg(feature = "cugraph")]
use super::{
    CugraphAlgoExec,
    planning::{literal_bfs_source_plan, validate_edge_source_schema},
};
use super::{CugraphLogicalNode, GpuIslandHostBoundaryConfig};
#[cfg(feature = "cugraph")]
use crate::{
    planner::native::InputPlanningSink,
    report::PlanningReport,
    session::{planning_native_execution_config, session_state_with_native_execution_config},
};

#[derive(Debug)]
pub struct CugraphExtensionPlanner {
    host_boundary_config: GpuIslandHostBoundaryConfig,
}

impl CugraphExtensionPlanner {
    pub(crate) fn new(host_boundary_config: GpuIslandHostBoundaryConfig) -> Self {
        Self { host_boundary_config }
    }
}

#[async_trait]
impl ExtensionPlanner for CugraphExtensionPlanner {
    async fn plan_extension(
        &self,
        planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        logical_inputs: &[&LogicalPlan],
        _physical_inputs: &[Arc<dyn ExecutionPlan>],
        session: &dyn Session,
        _planning_ctx: &PhysicalPlanningContext,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        let Some(node) = node.as_any().downcast_ref::<CugraphLogicalNode>() else {
            return Ok(None);
        };

        #[cfg(not(feature = "cugraph"))]
        {
            let _ = (planner, logical_inputs, session, self.host_boundary_config);
            let exec = CugraphFeatureDisabledExec::new(node.spec().algorithm, node.output_schema(), None, None)?;
            return Ok(Some(Arc::new(exec) as Arc<dyn ExecutionPlan>));
        }

        #[cfg(feature = "cugraph")]
        {
            let session_state = session.as_any().downcast_ref::<SessionState>().ok_or_else(|| {
                DataFusionError::Plan("cuGraph extension planning requires a DataFusion SessionState".to_string())
            })?;
            let native_execution_config = planning_native_execution_config(session);
            // Each input is planned on its own state because a native input
            // root becomes a GPU-resident producer named for its role. The
            // outcomes of that planning are labelled with the function and the
            // relation, such as `cugraph_pagerank.edges`, and travel with the
            // graph to the report of the query root.
            let function = node.spec().algorithm.function_name();
            let planning_state = |role: &str| {
                let sink = InputPlanningSink::default();
                let state = session_state_with_native_execution_config(
                    session_state,
                    native_execution_config.clone(),
                    format!("{}.{role}", node.gpu_result_input_name()),
                    sink.clone(),
                )
                .unwrap_or_else(|| session_state.clone());
                (state, sink)
            };
            let mut children = Vec::with_capacity(logical_inputs.len());
            let mut input_planning = PlanningReport::new();
            for (index, (input, relation)) in logical_inputs.iter().zip(node.roles()).enumerate() {
                let role = if index == 0 { "edge".to_string() } else { format!("predicate_{index}") };
                let (state, sink) = planning_state(&role);
                children.push(planner.create_physical_plan(input, &state).await?);
                let mut report = sink.take();
                report.label_relation_input(&format!("{function}.{}", relation.as_str()));
                input_planning.append(report);
            }
            let edge_source = children.remove(0);
            let vertex_domain = validate_edge_source_schema(node.spec(), &edge_source.schema())?;
            let mut predicate_sources = Vec::with_capacity(children.len() + 1);
            if let Some(source) =
                literal_bfs_source_plan(&planning_state("bfs_source").0, node.spec(), &vertex_domain).await?
            {
                predicate_sources.push(source);
            }
            predicate_sources.extend(children);
            let exec = CugraphAlgoExec::try_new_with_gpu_result_input_name(
                node.spec().clone(),
                edge_source,
                predicate_sources,
                None,
                None,
                native_execution_config,
                self.host_boundary_config,
                node.gpu_result_input_name().to_string(),
            )?
            .with_input_planning_report(input_planning);
            Ok(Some(Arc::new(exec) as Arc<dyn ExecutionPlan>))
        }
    }
}