diff --git a/datafusion/__init__.py b/datafusion/__init__.py index 367cda382..9de141636 100644 --- a/datafusion/__init__.py +++ b/datafusion/__init__.py @@ -77,6 +77,7 @@ TryCast, Between, Explain, + CreateMemoryTable, SubqueryAlias, Extension, ) @@ -133,6 +134,7 @@ "Explain", "SubqueryAlias", "Extension", + "CreateMemoryTable", ] diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py index 0e3a122bd..76fbc481c 100644 --- a/datafusion/tests/test_imports.py +++ b/datafusion/tests/test_imports.py @@ -80,6 +80,7 @@ Between, Explain, Extension, + CreateMemoryTable, ) @@ -149,6 +150,7 @@ def test_class_module_is_datafusion(): SubqueryAlias, Explain, Extension, + CreateMemoryTable, ]: assert klass.__module__ == "datafusion.expr" diff --git a/src/expr.rs b/src/expr.rs index 9dea30da8..5cf712d09 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -47,6 +47,7 @@ pub mod bool_expr; pub mod case; pub mod cast; pub mod column; +pub mod create_memory_table; pub mod cross_join; pub mod empty_relation; pub mod exists; @@ -278,6 +279,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/src/expr/create_memory_table.rs b/src/expr/create_memory_table.rs new file mode 100644 index 000000000..509bf2168 --- /dev/null +++ b/src/expr/create_memory_table.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; + +use datafusion_expr::CreateMemoryTable; +use pyo3::prelude::*; + +use crate::sql::logical::PyLogicalPlan; + +use super::logical_node::LogicalNode; + +#[pyclass(name = "CreateMemoryTable", module = "datafusion.expr", subclass)] +#[derive(Clone)] +pub struct PyCreateMemoryTable { + create: CreateMemoryTable, +} + +impl From for CreateMemoryTable { + fn from(create: PyCreateMemoryTable) -> Self { + create.create + } +} + +impl From for PyCreateMemoryTable { + fn from(create: CreateMemoryTable) -> PyCreateMemoryTable { + PyCreateMemoryTable { create } + } +} + +impl Display for PyCreateMemoryTable { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "CreateMemoryTable + Name: {:?} + Input: {:?} + if_not_exists: {:?} + or_replace: {:?}", + &self.create.name, + &self.create.input, + &self.create.if_not_exists, + &self.create.or_replace, + ) + } +} + +#[pymethods] +impl PyCreateMemoryTable { + fn name(&self) -> PyResult { + Ok(self.create.name.to_string()) + } + + fn input(&self) -> PyResult> { + Ok(Self::inputs(self)) + } + + fn if_not_exists(&self) -> bool { + self.create.if_not_exists + } + + fn or_replace(&self) -> bool { + self.create.or_replace + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateMemoryTable({})", self)) + } + + fn __name__(&self) -> PyResult { + Ok("CreateMemoryTable".to_string()) + } +} + +impl LogicalNode for PyCreateMemoryTable { + fn inputs(&self) -> Vec { + vec![PyLogicalPlan::from((*self.create.input).clone())] + } + + fn to_variant(&self, py: Python) -> PyResult { + Ok(self.clone().into_py(py)) + } +}