Add a Polars backend

pixi.toml

@@ -9,6 +9,7 @@ numpy = ">=1.23.1"
 pandas = ">=1.4.3"
 SQLAlchemy = ">=1.4.27"
 pyarrow = ">=11.0.0"
+polars = ">=1.6.0"
 
 [feature.dev.dependencies]
 ruff = ">=0.5.6"
@@ -73,5 +74,4 @@ py312 = ["py312", "dev", "duckdb", "postgres", "mssql"]
 py39ibm = ["py39", "dev", "duckdb", "postgres", "mssql", "ibm-db"]
 py312ibm = ["py312", "dev", "duckdb", "postgres", "mssql", "ibm-db"]
 docs = ["docs"]
-release = { features=["release"], no-default-feature=true }
-
+release = { features = ["release"], no-default-feature = true }

src/pydiverse/transform/__init__.py

@@ -3,7 +3,7 @@
 from pydiverse.transform.core import functions
 from pydiverse.transform.core.alignment import aligned, eval_aligned
 from pydiverse.transform.core.dispatchers import verb
-from pydiverse.transform.core.expressions.lambda_getter import λ
+from pydiverse.transform.core.expressions.lambda_getter import C
 from pydiverse.transform.core.table import Table
 
 __all__ = [
@@ -12,5 +12,5 @@
     "eval_aligned",
     "functions",
     "verb",
-    "λ",
+    "C",
 ]

src/pydiverse/transform/core/dtypes.py

@@ -99,6 +99,14 @@ class DateTime(DType):
     name = "datetime"
 
 
+class Date(DType):
+    name = "date"
+
+
+class Duration(DType):
+    name = "duration"
+
+
 class Template(DType):
     name = None
 
@@ -167,8 +175,12 @@ def dtype_from_string(t: str) -> DType:
         return String(const=is_const, vararg=is_vararg)
     if base_type == "bool":
         return Bool(const=is_const, vararg=is_vararg)
+    if base_type == "date":
+        return Date(const=is_const, vararg=is_vararg)
     if base_type == "datetime":
         return DateTime(const=is_const, vararg=is_vararg)
+    if base_type == "duration":
+        return Duration(const=is_const, vararg=is_vararg)
     if base_type == "none":
         return NoneDType(const=is_const, vararg=is_vararg)
 

src/pydiverse/transform/core/expressions/expressions.py

@@ -106,7 +106,7 @@ class LambdaColumn(BaseExpression):
       The following fails because `table.a` gets referenced before it gets created.
         table >> mutate(a = table.x) >> mutate(b = table.a)
       Instead you can use a lambda column to achieve this:
-        table >> mutate(a = table.x) >> mutate(b = λ.a)
+        table >> mutate(a = table.x) >> mutate(b = C.a)
     """
 
     __slots__ = "name"
@@ -115,10 +115,10 @@ def __init__(self, name: str):
         self.name = name
 
     def __repr__(self):
-        return f"<λ.{self.name}>"
+        return f"<C.{self.name}>"
 
     def _expr_repr(self) -> str:
-        return f"λ.{self.name}"
+        return f"C.{self.name}"
 
     def __eq__(self, other):
         if isinstance(other, self.__class__):
@@ -129,7 +129,7 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
     def __hash__(self):
-        return hash(("λ", self.name))
+        return hash(("C", self.name))
 
 
 class LiteralColumn(BaseExpression, Generic[T]):

src/pydiverse/transform/core/expressions/lambda_getter.py

@@ -3,22 +3,16 @@
 from pydiverse.transform.core.expressions import LambdaColumn
 from pydiverse.transform.core.expressions.symbolic_expressions import SymbolicExpression
 
+__all__ = ["C"]
 
-class LambdaColumnGetter:
-    """
-    An instance of this object can be used to instantiate a LambdaColumn.
-    """
 
-    def __getattr__(self, item):
-        if item.startswith("__"):
-            raise AttributeError(
-                f"'{type(self).__name__}' object has no attribute '{item}'"
-            )
-        return SymbolicExpression(LambdaColumn(item))
+class MC(type):
+    def __getattr__(cls, name: str) -> SymbolicExpression:
+        return SymbolicExpression(LambdaColumn(name))
 
-    def __getitem__(self, item):
-        return SymbolicExpression(LambdaColumn(item))
+    def __getitem__(cls, name: str) -> SymbolicExpression:
+        return SymbolicExpression(LambdaColumn(name))
 
 
-# Global instance of LambdaColumnGetter.
-λ = LambdaColumnGetter()
+class C(metaclass=MC):
+    pass

src/pydiverse/transform/core/expressions/symbolic_expressions.py

@@ -34,12 +34,13 @@ def __getattr__(self, item) -> SymbolAttribute:
                 f"Invalid attribute {item}. Attributes can't begin and end with an"
                 " underscore."
             )
+
         return SymbolAttribute(item, self)
 
     def __getitem__(self, item):
         return SymbolicExpression(FunctionCall("__getitem__", self, item))
 
-    def case(self, *cases: tuple[Any, Any], default: Any = None):
+    def case(self, *cases: tuple[Any, Any], default: Any = None) -> SymbolicExpression:
         case_expression = CaseExpression(
             switching_on=self,
             cases=cases,

src/pydiverse/transform/core/expressions/translator.py

@@ -70,10 +70,7 @@ def _translate(self, expr, **kwargs):
 
         if isinstance(expr, FunctionCall):
             operator = self.operator_registry.get_operator(expr.name)
-
-            # Mutate function call arguments using operator
-            expr_args, expr_kwargs = operator.mutate_args(expr.args, expr.kwargs)
-            expr = FunctionCall(expr.name, *expr_args, **expr_kwargs)
+            expr = FunctionCall(expr.name, *expr.args, **expr.kwargs)
 
             op_args, op_kwargs, context_kwargs = self._translate_function_arguments(
                 expr, operator, **kwargs
@@ -88,7 +85,7 @@ def _translate(self, expr, **kwargs):
             )
 
             return self._translate_function(
-                expr, implementation, op_args, context_kwargs, **kwargs
+                implementation, op_args, context_kwargs, **kwargs
             )
 
         if isinstance(expr, CaseExpression):
@@ -118,15 +115,14 @@ def _translate(self, expr, **kwargs):
             f" {expr}."
         )
 
-    def _translate_col(self, expr: Column, **kwargs) -> T:
+    def _translate_col(self, col: Column, **kwargs) -> T:
         raise NotImplementedError
 
-    def _translate_literal_col(self, expr: LiteralColumn, **kwargs) -> T:
+    def _translate_literal_col(self, col: LiteralColumn, **kwargs) -> T:
         raise NotImplementedError
 
     def _translate_function(
         self,
-        expr: FunctionCall,
         implementation: registry.TypedOperatorImpl,
         op_args: list[T],
         context_kwargs: dict[str, Any],

src/pydiverse/transform/core/functions.py

@@ -18,15 +18,23 @@ def _sym_f_call(name, *args, **kwargs) -> SymbolicExpression[FunctionCall]:
     return SymbolicExpression(FunctionCall(name, *args, **kwargs))
 
 
-def count(expr: SymbolicExpression = None):
+def count(expr: SymbolicExpression | None = None):
     if expr is None:
         return _sym_f_call("count")
     else:
         return _sym_f_call("count", expr)
 
 
-def row_number(*, arrange: list):
-    return _sym_f_call("row_number", arrange=arrange)
+def row_number(*, arrange: list, partition_by: list | None = None):
+    return _sym_f_call("row_number", arrange=arrange, partition_by=partition_by)
+
+
+def rank(*, arrange: list, partition_by: list | None = None):
+    return _sym_f_call("rank", arrange=arrange, partition_by=partition_by)
+
+
+def dense_rank(*, arrange: list, partition_by: list | None = None):
+    return _sym_f_call("dense_rank", arrange=arrange, partition_by=partition_by)
 
 
 def case(*cases: tuple[Any, Any], default: Any = None):

src/pydiverse/transform/core/registry.py

@@ -195,9 +195,9 @@ class OperatorRegistry:
     def __init__(self, name, super_registry=None):
         self.name = name
         self.super_registry = super_registry
-        self.registered_ops = set()  # type: set[Operator]
-        self.implementations = dict()  # type: dict[str, OperatorImplementationStore]
-        self.check_super = dict()  # type: dict[str, bool]
+        self.registered_ops: set[Operator] = set()
+        self.implementations: dict[str, OperatorImplementationStore] = dict()
+        self.check_super: dict[str, bool] = dict()
 
     def register_op(self, operator: Operator, check_super=True):
         """

src/pydiverse/transform/core/table.py

@@ -10,6 +10,7 @@
     LambdaColumn,
     SymbolicExpression,
 )
+from pydiverse.transform.core.verbs import export
 
 
 class Table(Generic[ImplT]):
@@ -84,12 +85,10 @@ def __copy__(self):
         return self.__class__(impl_copy)
 
     def __str__(self):
-        from pydiverse.transform.core.verbs import collect
-
         try:
             return (
                 f"Table: {self._impl.name}, backend: {type(self._impl).__name__}\n"
-                f"{self >> collect()}"
+                f"{self >> export()}"
             )
         except Exception as e:
             return (
@@ -99,15 +98,13 @@ def __str__(self):
             )
 
     def _repr_html_(self) -> str | None:
-        from pydiverse.transform.core.verbs import collect
-
         html = (
             f"Table <code>{self._impl.name}</code> using"
             f" <code>{type(self._impl).__name__}</code> backend:</br>"
         )
         try:
             # TODO: For lazy backend only show preview (eg. take first 20 rows)
-            html += (self >> collect())._repr_html_()
+            html += (self >> export())._repr_html_()
         except Exception as e:
             html += (
                 "</br><pre>Failed to collect table due to an exception:\n"
@@ -117,3 +114,9 @@ def _repr_html_(self) -> str | None:
 
     def _repr_pretty_(self, p, cycle):
         p.text(str(self) if not cycle else "...")
+
+    def cols(self) -> list[Column]:
+        return [
+            self._impl.cols[uuid].as_column(name, self._impl)
+            for (name, uuid) in self._impl.selected_cols()
+        ]

src/pydiverse/transform/core/table_impl.py

@@ -77,13 +77,13 @@ def __init__(
         self.compiler = self.ExpressionCompiler(self)
         self.lambda_translator = self.LambdaTranslator(self)
 
-        self.selects = ordered_set()  # type: ordered_set[str]
-        self.named_cols = bidict()  # type: bidict[str: uuid.UUID]
-        self.available_cols = set()  # type: set[uuid.UUID]
-        self.cols = {}  # type: dict[uuid.UUID: ColumnMetaData]
+        self.selects: ordered_set[str] = ordered_set()  # subset of named_cols
+        self.named_cols: bidict[str, uuid.UUID] = bidict()
+        self.available_cols: set[uuid.UUID] = set()
+        self.cols: dict[uuid.UUID, ColumnMetaData] = dict()
 
-        self.grouped_by = ordered_set()  # type: ordered_set[Column]
-        self.intrinsic_grouped_by = ordered_set()  # type: ordered_set[Column]
+        self.grouped_by: ordered_set[Column] = ordered_set()
+        self.intrinsic_grouped_by: ordered_set[Column] = ordered_set()
 
         # Init Values
         for name, col in columns.items():
@@ -185,7 +185,7 @@ def select(self, *args): ...
 
     def mutate(self, **kwargs): ...
 
-    def join(self, right, on, how, *, validate=None): ...
+    def join(self, right, on, how, *, validate="m:m"): ...
 
     def filter(self, *args): ...
 
@@ -199,6 +199,8 @@ def summarise(self, **kwargs): ...
 
     def slice_head(self, n: int, offset: int): ...
 
+    def export(self): ...
+
     #### Symbolic Operators ####
 
     @classmethod
@@ -237,6 +239,10 @@ def _translate_literal(self, expr, **kwargs):
                 return TypedValue(literal, dtypes.String(const=True))
             if isinstance(expr, datetime.datetime):
                 return TypedValue(literal, dtypes.DateTime(const=True))
+            if isinstance(expr, datetime.date):
+                return TypedValue(literal, dtypes.Date(const=True))
+            if isinstance(expr, datetime.timedelta):
+                return TypedValue(literal, dtypes.Duration(const=True))
 
             if expr is None:
                 return TypedValue(literal, dtypes.NoneDType(const=True))
@@ -329,6 +335,10 @@ def _translate_literal(self, expr, **kwargs):
                 return TypedValue(expr, dtypes.String(const=True))
             if isinstance(expr, datetime.datetime):
                 return TypedValue(expr, dtypes.DateTime(const=True))
+            if isinstance(expr, datetime.date):
+                return TypedValue(expr, dtypes.Date(const=True))
+            if isinstance(expr, datetime.timedelta):
+                return TypedValue(expr, dtypes.Duration(const=True))
 
             if expr is None:
                 return TypedValue(expr, dtypes.NoneDType(const=True))
@@ -427,7 +437,7 @@ class ColumnMetaData:
 
     @classmethod
     def from_expr(cls, uuid, expr, table: AbstractTableImpl, **kwargs):
-        v = table.compiler.translate(expr, **kwargs)
+        v: TypedValue = table.compiler.translate(expr, **kwargs)
         return cls(
             uuid=uuid,
             expr=expr,
@@ -469,7 +479,7 @@ def _nulls_last(_):
     def _add(lhs, rhs):
         return lhs + rhs
 
-    @op.extension(ops.StringAdd)
+    @op.extension(ops.StrAdd)
     def _str_add(lhs, rhs):
         return lhs + rhs
 
@@ -480,7 +490,7 @@ def _str_add(lhs, rhs):
     def _radd(rhs, lhs):
         return lhs + rhs
 
-    @op.extension(ops.StringRAdd)
+    @op.extension(ops.StrRAdd)
     def _str_radd(lhs, rhs):
         return lhs + rhs