[DO NOT MERGE] feat: remove generated protobuf code from substrait-python

examples/builder_example.py

@@ -1,4 +1,5 @@
-import substrait.gen.proto.algebra_pb2 as stalg
+import substrait.algebra_pb2 as stalg
+
 from substrait.builders.extended_expression import (
     column,
     literal,

gen_proto.sh

@@ -2,23 +2,8 @@
 
 set -eou pipefail
 
-namespace=substrait.gen.proto
 submodule_dir=./third_party/substrait
-src_dir="$submodule_dir"/proto
-tmp_dir=./buf_work_dir
-dest_dir=./src
-extension_dir=./src/substrait/extensions
-
-# Prefix the protobuf files with a unique configuration to prevent namespace conflicts
-# with other substrait packages. Save output to the work dir.
-python "$submodule_dir"/tools/proto_prefix.py "$tmp_dir" "$namespace" "$src_dir"
-
-# Remove the old python protobuf files
-rm -rf "$dest_dir/substrait/gen/proto"
-
-# Generate the new python protobuf files
-buf generate
-find "$dest_dir/substrait/gen" -type d -exec touch {}/__init__.py \;
+extension_dir=./src/substrait/extension_files
 
 # Remove the old extension files
 rm -rf "$extension_dir"
@@ -29,6 +14,3 @@ find "$extension_dir" -type f -exec chmod u+rw {} +
 
 # Ensure there's an __init__.py file in the extension directory
 touch $extension_dir/__init__.py
-
-# Remove the temporary work dir
-rm -rf "$tmp_dir"

pyproject.toml

@@ -5,7 +5,10 @@ authors = [{name = "Substrait contributors", email = "[email protected]
 license = {text = "Apache-2.0"}
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
-dependencies = ["protobuf >=5,<7"]
+dependencies = [
+    "protobuf >=5,<7",
+    "vbarua-substrait-protobuf>=0.1.0.dev0",
+]
 dynamic = ["version"]
 
 [tool.setuptools_scm]

src/substrait/builders/extended_expression.py

@@ -2,10 +2,11 @@
 from datetime import date
 from typing import Any, Callable, Iterable, Union
 
-import substrait.gen.proto.algebra_pb2 as stalg
-import substrait.gen.proto.extended_expression_pb2 as stee
-import substrait.gen.proto.extensions.extensions_pb2 as ste
-import substrait.gen.proto.type_pb2 as stp
+import substrait.algebra_pb2 as stalg
+import substrait.extended_expression_pb2 as stee
+import substrait.extensions.extensions_pb2 as ste
+import substrait.type_pb2 as stp
+
 from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_extended_expression_schema
 from substrait.utils import (

src/substrait/builders/plan.py

@@ -8,23 +8,24 @@
 import re
 from typing import Callable, Iterable, Optional, Union
 
-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
-from substrait import __substrait_version__
+import substrait.algebra_pb2 as stalg
+import substrait.extended_expression_pb2 as stee
+import substrait.plan_pb2 as stp
+import substrait.type_pb2 as stt
+from substrait.extensions.extensions_pb2 import AdvancedExtension
+
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     resolve_expression,
 )
 from substrait.extension_registry import ExtensionRegistry
-from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
 from substrait.type_inference import infer_plan_schema
 from substrait.utils import (
     merge_extension_declarations,
     merge_extension_uris,
     merge_extension_urns,
 )
+from substrait.version import substrait_version
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 
@@ -33,7 +34,7 @@
 
 def _create_default_version():
     p = re.compile(r"(\d+)\.(\d+)\.(\d+)")
-    m = p.match(__substrait_version__)
+    m = p.match(substrait_version)
     global default_version
     default_version = stp.Version(
         major_number=int(m.group(1)),

src/substrait/builders/type.py

@@ -1,6 +1,6 @@
 from typing import Iterable
 
-import substrait.gen.proto.type_pb2 as stt
+import substrait.type_pb2 as stt
 
 
 def boolean(nullable=True) -> stt.Type:

src/substrait/dataframe/expression.py

@@ -1,5 +1,6 @@
-import substrait.gen.proto.extended_expression_pb2 as stee
-import substrait.gen.proto.type_pb2 as stp
+import substrait.extended_expression_pb2 as stee
+import substrait.type_pb2 as stp
+
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     UnboundExtendedExpression,

src/substrait/derivation_expression.py

@@ -1,10 +1,10 @@
 from typing import Optional
 
 from antlr4 import CommonTokenStream, InputStream
+from substrait.type_pb2 import NamedStruct, Type
 
 from substrait.gen.antlr.SubstraitTypeLexer import SubstraitTypeLexer
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
-from substrait.gen.proto.type_pb2 import NamedStruct, Type
 
 
 def _evaluate(x, values: dict):

src/substrait/extension_registry/function_entry.py

@@ -3,9 +3,10 @@
 from enum import Enum
 from typing import Optional, Union
 
+from substrait.type_pb2 import Type
+
 from substrait.derivation_expression import _parse, evaluate
 from substrait.gen.json import simple_extensions as se
-from substrait.gen.proto.type_pb2 import Type
 
 from .signature_checker_helpers import covers, normalize_substrait_type_names
 

src/substrait/extension_registry/registry.py

@@ -8,9 +8,9 @@
 from typing import Optional, Union
 
 import yaml
+from substrait.type_pb2 import Type
 
 from substrait.bimap import UriUrnBiDiMap
-from substrait.gen.proto.type_pb2 import Type
 from substrait.simple_extension_utils import build_simple_extensions
 
 from .function_entry import FunctionEntry, FunctionType
@@ -32,7 +32,7 @@ def __init__(self, load_default_extensions=True) -> None:
         # Bidirectional URI <-> URN mapping (temporary during migration)
         self._uri_urn_bimap = UriUrnBiDiMap()
         if load_default_extensions:
-            for fpath in importlib_files("substrait.extensions").glob(  # type: ignore
+            for fpath in importlib_files("substrait.extension_files").glob(  # type: ignore
                 "functions*.yaml"
             ):
                 # Derive URI from DEFAULT_URN_PREFIX and filename

src/substrait/extension_registry/signature_checker_helpers.py

@@ -2,9 +2,10 @@
 
 from typing import Dict
 
+from substrait.type_pb2 import Type
+
 from substrait.derivation_expression import _evaluate
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
-from substrait.gen.proto.type_pb2 import Type
 
 from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError