feat: narwhals-compliant dataframe setup (#112)

Adds minimal setup for narwhals-complaint dataframe (lazyframe).

Author: tokoko

examples/dataframe_example.py

@@ -0,0 +1,17 @@
+from substrait.builders.plan import read_named_table
+from substrait.builders.type import i64, boolean, struct, named_struct
+from substrait.extension_registry import ExtensionRegistry
+import substrait.dataframe as sdf
+
+registry = ExtensionRegistry(load_default_extensions=True)
+
+ns = named_struct(
+    names=["id", "is_applicable"],
+    struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
+)
+
+table = read_named_table("example_table", ns)
+
+frame = sdf.DataFrame(read_named_table("example_table", ns))
+frame = frame.select(sdf.col("id"))
+print(frame.to_substrait(registry))

examples/narwhals_example.py

@@ -0,0 +1,36 @@
+# Install duckdb and pyarrow before running this example
+# /// script
+# dependencies = [
+#   "narwhals==2.9.0",
+#   "substrait[extensions] @ file:///${PROJECT_ROOT}/"
+# ]
+# ///
+
+from substrait.builders.plan import read_named_table
+from substrait.builders.type import i64, boolean, struct, named_struct
+from substrait.extension_registry import ExtensionRegistry
+
+from narwhals.typing import FrameT
+import narwhals as nw
+import substrait.dataframe as sdf
+
+
+registry = ExtensionRegistry(load_default_extensions=True)
+
+ns = named_struct(
+    names=["id", "is_applicable"],
+    struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
+)
+
+table = read_named_table("example_table", ns)
+
+
+lazy_frame: FrameT = nw.from_native(
+    sdf.DataFrame(read_named_table("example_table", ns))
+)
+
+lazy_frame = lazy_frame.select(nw.col("id").abs(), new_id=nw.col("id"))
+
+df: sdf.DataFrame = lazy_frame.to_native()
+
+print(df.to_substrait(registry))

src/substrait/dataframe/__init__.py

@@ -0,0 +1,16 @@
+import substrait.dataframe
+from substrait.builders.extended_expression import column
+
+from substrait.dataframe.dataframe import DataFrame
+from substrait.dataframe.expression import Expression
+
+__all__ = [DataFrame, Expression]
+
+
+def col(name: str) -> Expression:
+    """Column selection."""
+    return Expression(column(name))
+
+# TODO handle str_as_lit argument
+def parse_into_expr(expr, str_as_lit: bool):
+    return expr._to_compliant_expr(substrait.dataframe)

src/substrait/dataframe/dataframe.py

@@ -0,0 +1,36 @@
+from typing import Union, Iterable
+import substrait.dataframe
+from substrait.builders.plan import select
+from substrait.dataframe.expression import Expression
+
+
+class DataFrame:
+    def __init__(self, plan):
+        self.plan = plan
+        self._native_frame = self
+
+    def to_substrait(self, registry):
+        return self.plan(registry)
+
+    def __narwhals_lazyframe__(self) -> "DataFrame":
+        """Return object implementing CompliantDataFrame protocol."""
+        return self
+
+    def __narwhals_namespace__(self):
+        """
+        Return the namespace object that contains functions like col, lit, etc.
+        This is how Narwhals knows which backend's functions to use.
+        """
+        return substrait.dataframe
+
+    def select(
+        self, *exprs: Union[Expression, Iterable[Expression]], **named_exprs: Expression
+    ) -> "DataFrame":
+        expressions = [e.expr for e in exprs] + [
+            expr.alias(alias).expr for alias, expr in named_exprs.items()
+        ]
+        return DataFrame(select(self.plan, expressions=expressions))        
+
+    # TODO handle version
+    def _with_version(self, version):
+        return self

src/substrait/dataframe/expression.py

@@ -0,0 +1,36 @@
+from substrait.builders.extended_expression import (
+    UnboundExtendedExpression,
+    ExtendedExpressionOrUnbound,
+    resolve_expression,
+    scalar_function
+)
+import substrait.gen.proto.type_pb2 as stp
+import substrait.gen.proto.extended_expression_pb2 as stee
+from substrait.extension_registry import ExtensionRegistry
+
+
+def _alias(
+    expr: ExtendedExpressionOrUnbound,
+    alias: str = None,
+):
+    def resolve(
+        base_schema: stp.NamedStruct, registry: ExtensionRegistry
+    ) -> stee.ExtendedExpression:
+        bound_expression = resolve_expression(expr, base_schema, registry)
+        bound_expression.referred_expr[0].output_names[0] = alias
+        return bound_expression
+
+    return resolve
+
+
+class Expression:
+    def __init__(self, expr: UnboundExtendedExpression):
+        self.expr = expr
+
+    def alias(self, alias: str):
+        self.expr = _alias(self.expr, alias)
+        return self
+    
+    def abs(self):
+        self.expr = scalar_function("functions_arithmetic.yaml", "abs", expressions=[self.expr]) 
+        return self