Add unit tests for to_proto and from_proto for logical and physical plans

Author: timsaucer

python/datafusion/plan.py

@@ -85,11 +85,17 @@ def display_graphviz(self) -> str:
 
     @staticmethod
     def from_proto(ctx: SessionContext, data: bytes) -> LogicalPlan:
-        """Create a LogicalPlan from protobuf bytes."""
+        """Create a LogicalPlan from protobuf bytes.
+
+        Tables created in memory from record batches are currently not supported.
+        """
         return LogicalPlan(df_internal.LogicalPlan.from_proto(ctx.ctx, data))
 
     def to_proto(self) -> bytes:
-        """Convert a LogicalPlan to protobuf bytes."""
+        """Convert a LogicalPlan to protobuf bytes.
+
+        Tables created in memory from record batches are currently not supported.
+        """
         return self._raw_plan.to_proto()
 
 
@@ -127,9 +133,15 @@ def partition_count(self) -> int:
 
     @staticmethod
     def from_proto(ctx: SessionContext, data: bytes) -> ExecutionPlan:
-        """Create an ExecutionPlan from protobuf bytes."""
+        """Create an ExecutionPlan from protobuf bytes.
+
+        Tables created in memory from record batches are currently not supported.
+        """
         return ExecutionPlan(df_internal.ExecutionPlan.from_proto(ctx.ctx, data))
 
     def to_proto(self) -> bytes:
-        """Convert an ExecutionPlan into protobuf bytes."""
+        """Convert an ExecutionPlan into protobuf bytes.
+
+        Tables created in memory from record batches are currently not supported.
+        """
         return self._raw_plan.to_proto()

python/datafusion/tests/test_plans.py

@@ -0,0 +1,42 @@
+# 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.
+
+from datafusion import SessionContext, LogicalPlan, ExecutionPlan
+import pytest
+
+
+# Note: We must use CSV because memory tables are currently not supported for
+# conversion to/from protobuf.
+@pytest.fixture
+def df():
+    ctx = SessionContext()
+    return ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select("c1")
+
+
+def test_logical_plan_to_proto(ctx, df) -> None:
+    logical_plan_bytes = df.logical_plan().to_proto()
+    logical_plan = LogicalPlan.from_proto(ctx, logical_plan_bytes)
+
+    df_round_trip = ctx.create_dataframe_from_logical_plan(logical_plan)
+
+    assert df.collect() == df_round_trip.collect()
+
+    original_execution_plan = df.execution_plan()
+    execution_plan_bytes = original_execution_plan.to_proto()
+    execution_plan = ExecutionPlan.from_proto(ctx, execution_plan_bytes)
+
+    assert str(original_execution_plan) == str(execution_plan)