datafusion_python/
physical_plan.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable};
21use datafusion_proto::physical_plan::AsExecutionPlan;
22use prost::Message;
23use pyo3::exceptions::PyRuntimeError;
24use pyo3::prelude::*;
25use pyo3::types::PyBytes;
26
27use crate::codec::PythonPhysicalCodec;
28use crate::context::PySessionContext;
29use crate::errors::PyDataFusionResult;
30use crate::metrics::PyMetricsSet;
31
32#[pyclass(
33    from_py_object,
34    frozen,
35    name = "ExecutionPlan",
36    module = "datafusion",
37    subclass
38)]
39#[derive(Debug, Clone)]
40pub struct PyExecutionPlan {
41    pub plan: Arc<dyn ExecutionPlan>,
42}
43
44impl PyExecutionPlan {
45    /// creates a new PyPhysicalPlan
46    pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
47        Self { plan }
48    }
49}
50
51#[pymethods]
52impl PyExecutionPlan {
53    /// Get the inputs to this plan
54    pub fn children(&self) -> Vec<PyExecutionPlan> {
55        self.plan
56            .children()
57            .iter()
58            .map(|&p| p.to_owned().into())
59            .collect()
60    }
61
62    pub fn display(&self) -> String {
63        let d = displayable(self.plan.as_ref());
64        format!("{}", d.one_line())
65    }
66
67    pub fn display_indent(&self) -> String {
68        let d = displayable(self.plan.as_ref());
69        format!("{}", d.indent(false))
70    }
71
72    #[pyo3(signature = (ctx=None))]
73    pub fn to_bytes<'py>(
74        &'py self,
75        py: Python<'py>,
76        ctx: Option<PySessionContext>,
77    ) -> PyDataFusionResult<Bound<'py, PyBytes>> {
78        // Route through the session's physical codec when supplied so
79        // user FFI codecs registered via
80        // `with_physical_extension_codec` see the encode path.
81        let default_codec;
82        let codec: &dyn datafusion_proto::physical_plan::PhysicalExtensionCodec = match ctx {
83            Some(ref ctx) => ctx.physical_codec().as_ref(),
84            None => {
85                default_codec = PythonPhysicalCodec::default();
86                &default_codec
87            }
88        };
89        let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(
90            self.plan.clone(),
91            codec,
92        )?;
93
94        let bytes = proto.encode_to_vec();
95        Ok(PyBytes::new(py, &bytes))
96    }
97
98    #[staticmethod]
99    pub fn from_bytes(
100        ctx: PySessionContext,
101        proto_msg: Bound<'_, PyBytes>,
102    ) -> PyDataFusionResult<Self> {
103        let bytes: &[u8] = proto_msg.extract().map_err(Into::<PyErr>::into)?;
104        let proto_plan =
105            datafusion_proto::protobuf::PhysicalPlanNode::decode(bytes).map_err(|e| {
106                PyRuntimeError::new_err(format!(
107                    "Unable to decode physical node from serialized bytes: {e}"
108                ))
109            })?;
110
111        let codec = ctx.physical_codec();
112        let plan =
113            proto_plan.try_into_physical_plan(ctx.ctx.task_ctx().as_ref(), codec.as_ref())?;
114        Ok(Self::new(plan))
115    }
116
117    pub fn metrics(&self) -> Option<PyMetricsSet> {
118        self.plan.metrics().map(PyMetricsSet::new)
119    }
120
121    fn __repr__(&self) -> String {
122        self.display_indent()
123    }
124
125    #[getter]
126    pub fn partition_count(&self) -> usize {
127        self.plan.output_partitioning().partition_count()
128    }
129}
130
131impl From<PyExecutionPlan> for Arc<dyn ExecutionPlan> {
132    fn from(plan: PyExecutionPlan) -> Arc<dyn ExecutionPlan> {
133        plan.plan.clone()
134    }
135}
136
137impl From<Arc<dyn ExecutionPlan>> for PyExecutionPlan {
138    fn from(plan: Arc<dyn ExecutionPlan>) -> PyExecutionPlan {
139        PyExecutionPlan { plan: plan.clone() }
140    }
141}
datafusion_python/physical_plan.rs

datafusion_python/
physical_plan.rs
