Create memory table (#271)

jdye64 · web-flow · commit 325b7cedb9c8 · 2023-03-13T13:27:25.000-06:00
* Introduce to_variant trait function to LogicalNode and create Explain LogicalNode bindings

* Cargo fmt

* bindings for Extension LogicalNode

* Add missing classes to list of exports so test_imports will pass

* Update to point to proper repo

* Update pytest to adhere to aggregate calls being wrapped in projections

* Address linter change which causes a pytest to fail

* Add bindings for LogicalPlan::CreateMemoryTable
diff --git 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
@@ -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
@@ -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::<filter::PyFilter>()?;
     m.add_class::<projection::PyProjection>()?;
     m.add_class::<table_scan::PyTableScan>()?;
+    m.add_class::<create_memory_table::PyCreateMemoryTable>()?;
     m.add_class::<subquery_alias::PySubqueryAlias>()?;
     Ok(())
 }
diff --git a/src/expr/create_memory_table.rs 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<PyCreateMemoryTable> for CreateMemoryTable {
+    fn from(create: PyCreateMemoryTable) -> Self {
+        create.create
+    }
+}
+
+impl From<CreateMemoryTable> 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<String> {
+        Ok(self.create.name.to_string())
+    }
+
+    fn input(&self) -> PyResult<Vec<PyLogicalPlan>> {
+        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<String> {
+        Ok(format!("CreateMemoryTable({})", self))
+    }
+
+    fn __name__(&self) -> PyResult<String> {
+        Ok("CreateMemoryTable".to_string())
+    }
+}
+
+impl LogicalNode for PyCreateMemoryTable {
+    fn inputs(&self) -> Vec<PyLogicalPlan> {
+        vec![PyLogicalPlan::from((*self.create.input).clone())]
+    }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -77,6 +77,7 @@`
`77`	`77`	`TryCast,`
`78`	`78`	`Between,`
`79`	`79`	`Explain,`
	`80`	`+ CreateMemoryTable,`
`80`	`81`	`SubqueryAlias,`
`81`	`82`	`Extension,`
`82`	`83`	`)`
`@@ -133,6 +134,7 @@`
`133`	`134`	`"Explain",`
`134`	`135`	`"SubqueryAlias",`
`135`	`136`	`"Extension",`
	`137`	`+ "CreateMemoryTable",`
`136`	`138`	`]`
`137`	`139`
`138`	`140`