datafusion_python/
table.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 arrow::datatypes::SchemaRef;
21use arrow::pyarrow::ToPyArrow;
22use async_trait::async_trait;
23use datafusion::catalog::{Session, TableProviderFactory};
24use datafusion::common::Column;
25use datafusion::datasource::{TableProvider, TableType};
26use datafusion::logical_expr::{
27    CreateExternalTable, Expr, LogicalPlanBuilder, TableProviderFilterPushDown,
28};
29use datafusion::physical_plan::ExecutionPlan;
30use datafusion::prelude::DataFrame;
31use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
32use datafusion_python_util::{create_logical_extension_capsule, table_provider_from_pycapsule};
33use pyo3::IntoPyObjectExt;
34use pyo3::prelude::*;
35
36use crate::context::PySessionContext;
37use crate::dataframe::PyDataFrame;
38use crate::dataset::Dataset;
39use crate::errors;
40use crate::expr::create_external_table::PyCreateExternalTable;
41
42/// This struct is used as a common method for all TableProviders,
43/// whether they refer to an FFI provider, an internally known
44/// implementation, a dataset, or a dataframe view.
45#[pyclass(
46    from_py_object,
47    frozen,
48    name = "RawTable",
49    module = "datafusion.catalog",
50    subclass
51)]
52#[derive(Clone)]
53pub struct PyTable {
54    pub table: Arc<dyn TableProvider>,
55}
56
57impl PyTable {
58    pub fn table(&self) -> Arc<dyn TableProvider> {
59        self.table.clone()
60    }
61}
62
63#[pymethods]
64impl PyTable {
65    /// Instantiate from any Python object that supports any of the table
66    /// types. We do not know a priori when using this method if the object
67    /// will be passed a wrapped or raw class. Here we handle all of the
68    /// following object types:
69    ///
70    /// - PyTable (essentially a clone operation), but either raw or wrapped
71    /// - DataFrame, either raw or wrapped
72    /// - FFI Table Providers via PyCapsule
73    /// - PyArrow Dataset objects
74    #[new]
75    pub fn new(obj: Bound<'_, PyAny>, session: Option<Bound<PyAny>>) -> PyResult<Self> {
76        let py = obj.py();
77        if let Ok(py_table) = obj.extract::<PyTable>() {
78            Ok(py_table)
79        } else if let Ok(py_table) = obj
80            .getattr("_inner")
81            .and_then(|inner| inner.extract::<PyTable>().map_err(Into::<PyErr>::into))
82        {
83            Ok(py_table)
84        } else if let Ok(py_df) = obj.extract::<PyDataFrame>() {
85            let provider = py_df.inner_df().as_ref().clone().into_view();
86            Ok(PyTable::from(provider))
87        } else if let Ok(py_df) = obj
88            .getattr("df")
89            .and_then(|inner| inner.extract::<PyDataFrame>().map_err(Into::<PyErr>::into))
90        {
91            let provider = py_df.inner_df().as_ref().clone().into_view();
92            Ok(PyTable::from(provider))
93        } else if let Some(provider) = {
94            let session = match session {
95                Some(session) => session,
96                None => PySessionContext::global_ctx()?.into_bound_py_any(obj.py())?,
97            };
98            table_provider_from_pycapsule(obj.clone(), session)?
99        } {
100            Ok(PyTable::from(provider))
101        } else {
102            let provider = Arc::new(Dataset::new(&obj, py)?) as Arc<dyn TableProvider>;
103            Ok(PyTable::from(provider))
104        }
105    }
106
107    /// Get a reference to the schema for this table
108    #[getter]
109    fn schema<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
110        self.table.schema().to_pyarrow(py)
111    }
112
113    /// Get the type of this table for metadata/catalog purposes.
114    #[getter]
115    fn kind(&self) -> &str {
116        match self.table.table_type() {
117            TableType::Base => "physical",
118            TableType::View => "view",
119            TableType::Temporary => "temporary",
120        }
121    }
122
123    fn __repr__(&self) -> PyResult<String> {
124        let kind = self.kind();
125        Ok(format!("Table(kind={kind})"))
126    }
127}
128
129impl From<Arc<dyn TableProvider>> for PyTable {
130    fn from(table: Arc<dyn TableProvider>) -> Self {
131        Self { table }
132    }
133}
134
135#[derive(Clone, Debug)]
136pub(crate) struct TempViewTable {
137    df: Arc<DataFrame>,
138}
139
140/// This is nearly identical to `DataFrameTableProvider`
141/// except that it is for temporary tables.
142/// Remove when https://github.com/apache/datafusion/issues/18026
143/// closes.
144impl TempViewTable {
145    pub(crate) fn new(df: Arc<DataFrame>) -> Self {
146        Self { df }
147    }
148}
149
150#[async_trait]
151impl TableProvider for TempViewTable {
152    fn schema(&self) -> SchemaRef {
153        Arc::new(self.df.schema().as_arrow().clone())
154    }
155
156    fn table_type(&self) -> TableType {
157        TableType::Temporary
158    }
159
160    async fn scan(
161        &self,
162        state: &dyn Session,
163        projection: Option<&Vec<usize>>,
164        filters: &[Expr],
165        limit: Option<usize>,
166    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
167        let filter = filters.iter().cloned().reduce(|acc, new| acc.and(new));
168        let plan = self.df.logical_plan().clone();
169        let mut plan = LogicalPlanBuilder::from(plan);
170
171        if let Some(filter) = filter {
172            plan = plan.filter(filter)?;
173        }
174
175        let mut plan = if let Some(projection) = projection {
176            // avoiding adding a redundant projection (e.g. SELECT * FROM view)
177            let current_projection = (0..plan.schema().fields().len()).collect::<Vec<usize>>();
178            if projection == &current_projection {
179                plan
180            } else {
181                let fields: Vec<Expr> = projection
182                    .iter()
183                    .map(|i| {
184                        Expr::Column(Column::from(
185                            self.df.logical_plan().schema().qualified_field(*i),
186                        ))
187                    })
188                    .collect();
189                plan.project(fields)?
190            }
191        } else {
192            plan
193        };
194
195        if let Some(limit) = limit {
196            plan = plan.limit(0, Some(limit))?;
197        }
198
199        state.create_physical_plan(&plan.build()?).await
200    }
201
202    fn supports_filters_pushdown(
203        &self,
204        filters: &[&Expr],
205    ) -> datafusion::common::Result<Vec<TableProviderFilterPushDown>> {
206        Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
207    }
208}
209
210#[derive(Debug)]
211pub(crate) struct RustWrappedPyTableProviderFactory {
212    pub(crate) table_provider_factory: Py<PyAny>,
213    pub(crate) codec: Arc<FFI_LogicalExtensionCodec>,
214}
215
216impl RustWrappedPyTableProviderFactory {
217    pub fn new(table_provider_factory: Py<PyAny>, codec: Arc<FFI_LogicalExtensionCodec>) -> Self {
218        Self {
219            table_provider_factory,
220            codec,
221        }
222    }
223
224    fn create_inner(
225        &self,
226        cmd: CreateExternalTable,
227        codec: Bound<PyAny>,
228    ) -> PyResult<Arc<dyn TableProvider>> {
229        Python::attach(|py| {
230            let provider = self.table_provider_factory.bind(py);
231            let cmd = PyCreateExternalTable::from(cmd);
232
233            provider
234                .call_method1("create", (cmd,))
235                .and_then(|t| PyTable::new(t, Some(codec)))
236                .map(|t| t.table())
237        })
238    }
239}
240
241#[async_trait]
242impl TableProviderFactory for RustWrappedPyTableProviderFactory {
243    async fn create(
244        &self,
245        _: &dyn Session,
246        cmd: &CreateExternalTable,
247    ) -> datafusion::common::Result<Arc<dyn TableProvider>> {
248        Python::attach(|py| {
249            let codec = create_logical_extension_capsule(py, self.codec.as_ref())
250                .map_err(errors::to_datafusion_err)?;
251
252            self.create_inner(cmd.clone(), codec.into_any())
253                .map_err(errors::to_datafusion_err)
254        })
255    }
256}
datafusion_python/table.rs

datafusion_python/
table.rs
