[arrow-datafusion-python] branch main updated: Make from_substrait_plan return DataFrame instead of LogicalPlan (#164)

[agrove]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git


The following commit(s) were added to refs/heads/main by this push:
     new 8c618f9  Make from_substrait_plan return DataFrame instead of 
LogicalPlan (#164)
8c618f9 is described below

commit 8c618f905a39d2651fc6585e1583a72a89018642
Author: Andy Grove <andygrov...@gmail.com>
AuthorDate: Sun Feb 5 11:59:43 2023 -0700

    Make from_substrait_plan return DataFrame instead of LogicalPlan (#164)
    
    * revert changes, add context method 'create_dataframe_from_logical_plan'
    
    * lint
---
 datafusion/tests/test_substrait.py | 8 ++++++--
 src/context.rs                     | 6 ++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/datafusion/tests/test_substrait.py 
b/datafusion/tests/test_substrait.py
index 4c6c8a3..9a08b76 100644
--- a/datafusion/tests/test_substrait.py
+++ b/datafusion/tests/test_substrait.py
@@ -45,7 +45,11 @@ def test_substrait_serialization(ctx):
         "SELECT * FROM t", ctx
     )
     substrait_plan = ss.substrait.serde.deserialize_bytes(substrait_bytes)
-    df_logical_plan = ss.substrait.consumer.from_substrait_plan(
+    logical_plan = ss.substrait.consumer.from_substrait_plan(
         ctx, substrait_plan
     )
-    substrait_plan = ss.substrait.producer.to_substrait_plan(df_logical_plan)
+
+    # demonstrate how to create a DataFrame from a deserialized logical plan
+    df = ctx.create_dataframe_from_logical_plan(logical_plan)
+
+    substrait_plan = ss.substrait.producer.to_substrait_plan(df.logical_plan())
diff --git a/src/context.rs b/src/context.rs
index 915ac45..cf19b8e 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -29,6 +29,7 @@ use crate::catalog::{PyCatalog, PyTable};
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
 use crate::errors::DataFusionError;
+use crate::logical::PyLogicalPlan;
 use crate::store::StorageContexts;
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
@@ -175,6 +176,11 @@ impl PySessionContext {
         Ok(df)
     }
 
+    /// Create a DataFrame from an existing logical plan
+    fn create_dataframe_from_logical_plan(&mut self, plan: PyLogicalPlan) -> 
PyDataFrame {
+        PyDataFrame::new(DataFrame::new(self.ctx.state(), 
plan.plan.as_ref().clone()))
+    }
+
     fn register_table(&mut self, name: &str, table: &PyTable) -> PyResult<()> {
         self.ctx
             .register_table(name, table.table())