feat(write_table): added the write table builder

Author: giospada

src/substrait/builders/plan.py

@@ -7,16 +7,16 @@
 
 from typing import Iterable, Optional, Union, Callable
 
-import substrait.gen.proto.algebra_pb2 as stalg
 from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
+import substrait.gen.proto.algebra_pb2 as stalg
+import substrait.gen.proto.extended_expression_pb2 as stee
 import substrait.gen.proto.plan_pb2 as stp
 import substrait.gen.proto.type_pb2 as stt
-import substrait.gen.proto.extended_expression_pb2 as stee
-from substrait.extension_registry import ExtensionRegistry
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     resolve_expression,
 )
+from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_plan_schema
 from substrait.utils import (
     merge_extension_declarations,
@@ -379,3 +379,33 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
     return resolve
+
+
+def write_table(
+    table_names: Union[str, Iterable[str]],
+    input: PlanOrUnbound,
+    create_mode: Union[stalg.WriteRel.CreateMode.ValueType, None] = None,
+) -> UnboundPlan:
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        bound_input = input if isinstance(input, stp.Plan) else input(registry)
+        ns = infer_plan_schema(bound_input)
+        _table_names = [table_names] if isinstance(table_names, str) else table_names
+        _create_mode = create_mode or stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS
+
+        write_rel = stalg.Rel(
+            write=stalg.WriteRel(
+                input=bound_input.relations[-1].root.input,
+                table_schema=ns,
+                op=stalg.WriteRel.WRITE_OP_CTAS,
+                create_mode=_create_mode,
+                named_table=stalg.NamedObjectWrite(names=_table_names),
+            )
+        )
+        return stp.Plan(
+            relations=[
+                stp.PlanRel(root=stalg.RelRoot(input=write_rel, names=ns.names))
+            ],
+            **_merge_extensions(bound_input),
+        )
+
+    return resolve

tests/builders/plan/test_write.py

@@ -1,57 +1,41 @@
-import substrait.gen.proto.type_pb2 as stt
-import substrait.gen.proto.plan_pb2 as stp
 import substrait.gen.proto.algebra_pb2 as stalg
+import substrait.gen.proto.plan_pb2 as stp
+import substrait.gen.proto.type_pb2 as stt
+from substrait.builders.plan import read_named_table, write_table
 from substrait.builders.type import boolean, i64
-from substrait.builders.plan import read_named_table
 
 struct = stt.Type.Struct(types=[i64(nullable=False), boolean()])
 
 named_struct = stt.NamedStruct(names=["id", "is_applicable"], struct=struct)
 
 
 def test_write_rel():
-    actual = read_named_table("example_table", named_struct)(None)
+    actual = write_table(
+        "example_table_write_test",
+        read_named_table("example_table", named_struct),
+    )(None)
 
-    # write example table test
-    stp.Plan(
+    expected = stp.Plan(
         relations=[
             stp.PlanRel(
                 root=stalg.RelRoot(
                     input=stalg.Rel(
                         write=stalg.WriteRel(
                             input=stalg.Rel(
                                 read=stalg.ReadRel(
-                                    common=stalg.RelCommon(direct=stalg.RelCommon.Direct()),
+                                    common=stalg.RelCommon(
+                                        direct=stalg.RelCommon.Direct()
+                                    ),
                                     base_schema=named_struct,
                                     named_table=stalg.ReadRel.NamedTable(
                                         names=["example_table"]
                                     ),
                                 )
                             ),
-                            common=stalg.RelCommon(direct=stalg.RelCommon.Direct()),
+                            op=stalg.WriteRel.WRITE_OP_CTAS,
                             table_schema=named_struct,
-                            create_mode=stalg.WriteRel.CreateMode.CREATE_MODE_REPLACE_IF_EXISTS,
-                            named_table=stalg.NamedTable(
-                                names=["example_table_write_test"]
-                            ),
-                        )
-                    ),
-                    names=["id", "is_applicable"],
-                )
-            )
-        ]
-    )
-
-    # read back the table
-    expected = stp.Plan(
-        relations=[
-            stp.PlanRel(
-                root=stalg.RelRoot(
-                    input=stalg.Rel(
-                        read=stalg.ReadRel(
-                            common=stalg.RelCommon(direct=stalg.RelCommon.Direct()),
-                            base_schema=named_struct,
-                            named_table=stalg.ReadRel.NamedTable(
+                            create_mode=stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS,
+                            named_table=stalg.NamedObjectWrite(
                                 names=["example_table_write_test"]
                             ),
                         )
@@ -61,30 +45,4 @@ def test_write_rel():
             )
         ]
     )
-
-    assert actual == expected
-
-
-def test_write_rel_db():
-    actual = read_named_table(["example_db", "example_table"], named_struct)(None)
-
-    expected = stp.Plan(
-        relations=[
-            stp.PlanRel(
-                root=stalg.RelRoot(
-                    input=stalg.Rel(
-                        read=stalg.ReadRel(
-                            common=stalg.RelCommon(direct=stalg.RelCommon.Direct()),
-                            base_schema=named_struct,
-                            named_table=stalg.ReadRel.NamedTable(
-                                names=["example_db", "example_table"]
-                            ),
-                        )
-                    ),
-                    names=["id", "is_applicable"],
-                )
-            )
-        ]
-    )
-
     assert actual == expected