sqlframe 3.10.1__py3-none-any.whl → 3.12.0__py3-none-any.whl

sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.10.1'
-__version_tuple__ = version_tuple = (3, 10, 1)
+__version__ = version = '3.12.0'
+__version_tuple__ = version_tuple = (3, 12, 0)

sqlframe/base/column.py

@@ -9,6 +9,7 @@ import typing as t
 import sqlglot
 from sqlglot import Dialect
 from sqlglot import expressions as exp
+from sqlglot.expressions import paren
 from sqlglot.helper import flatten, is_iterable
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
@@ -63,10 +64,10 @@ class Column:
         return self.binary_op(exp.LTE, other)
 
     def __and__(self, other: ColumnOrLiteral) -> Column:
-        return self.binary_op(exp.And, other)
+        return self.binary_op(exp.And, other, paren=True)
 
     def __or__(self, other: ColumnOrLiteral) -> Column:
-        return self.binary_op(exp.Or, other)
+        return self.binary_op(exp.Or, other, paren=True)
 
     def __mod__(self, other: ColumnOrLiteral) -> Column:
         return self.binary_op(exp.Mod, other, paren=True)
@@ -181,7 +182,7 @@ class Column:
     ) -> Column:
         columns = [] if column is None else [cls.ensure_col(column)]
         column_args = [cls.ensure_col(arg) for arg in args]
-        expressions = [x.expression for x in columns + column_args]
+        expressions = [x.column_expression for x in columns + column_args]
         new_expression = exp.Anonymous(this=func_name.upper(), expressions=expressions)
         return Column(new_expression)
 
@@ -192,9 +193,9 @@ class Column:
         ensured_column = None if column is None else cls.ensure_col(column)
         ensure_expression_values = {
             k: (
-                [cls.ensure_col(x).expression for x in v]
+                [cls.ensure_col(x).column_expression for x in v]
                 if is_iterable(v)
-                else cls.ensure_col(v).expression
+                else cls.ensure_col(v).column_expression
             )
             for k, v in kwargs.items()
             if v is not None
@@ -316,10 +317,12 @@ class Column:
         from sqlframe.base.functions import when
 
         column_with_if = when(condition, value)
-        if not isinstance(self.expression, exp.Case):
+        if not isinstance(self.column_expression, exp.Case):
             return column_with_if
         new_column = self.copy()
-        new_column.expression.args["ifs"].extend(column_with_if.expression.args["ifs"])
+        new_column.column_expression.args["ifs"].extend(
+            column_with_if.column_expression.args["ifs"]
+        )
         return new_column
 
     def otherwise(self, value: t.Any) -> Column:
@@ -327,7 +330,7 @@ class Column:
 
         true_value = value if isinstance(value, Column) else lit(value)
         new_column = self.copy()
-        new_column.expression.set("default", true_value.column_expression)
+        new_column.column_expression.set("default", true_value.column_expression)
         return new_column
 
     def isNull(self) -> Column:

sqlframe/base/dataframe.py

@@ -80,7 +80,7 @@ JOIN_HINTS = {
 }
 
 
-DF = t.TypeVar("DF", bound="_BaseDataFrame")
+DF = t.TypeVar("DF", bound="BaseDataFrame")
 
 
 class OpenAIMode(enum.Enum):
@@ -198,7 +198,7 @@ class _BaseDataFrameStatFunctions(t.Generic[DF]):
 STAT = t.TypeVar("STAT", bound=_BaseDataFrameStatFunctions)
 
 
-class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
+class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     _na: t.Type[NA]
     _stat: t.Type[STAT]
     _group_data: t.Type[GROUP_DATA]
@@ -1624,7 +1624,9 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             raise NotImplementedError("Vertical show is not yet supported")
         if truncate:
             logger.warning("Truncate is ignored so full results will be displayed")
-        result = self.limit(n).collect()
+        # Make sure that the limit we add doesn't affect the results
+        df = self._convert_leaf_to_cte()
+        result = df.limit(n).collect()
         table = PrettyTable()
         if row := seq_get(result, 0):
             table.field_names = row._unique_field_names

sqlframe/base/decorators.py

@@ -1,15 +1,54 @@
 from __future__ import annotations
 
+import re
 import typing as t
 
+from sqlglot import exp
 from sqlglot.helper import ensure_list
 
+from sqlframe.base.column import Column
+
 CALLING_CLASS = t.TypeVar("CALLING_CLASS")
 
 
 def func_metadata(unsupported_engines: t.Optional[t.Union[str, t.List[str]]] = None) -> t.Callable:
     def _metadata(func: t.Callable) -> t.Callable:
-        func.unsupported_engines = ensure_list(unsupported_engines) if unsupported_engines else []  # type: ignore
-        return func
+        def wrapper(*args, **kwargs):
+            funcs_to_not_auto_alias = [
+                "posexplode",
+                "explode_outer",
+                "json_tuple",
+                "posexplode_outer",
+                "stack",
+                "inline",
+                "inline_outer",
+                "window",
+                "session_window",
+                "window_time",
+            ]
+
+            result = func(*args, **kwargs)
+            if (
+                isinstance(result, Column)
+                and isinstance(result.column_expression, exp.Func)
+                and not isinstance(result.expression, exp.Alias)
+                and func.__name__ not in funcs_to_not_auto_alias
+            ):
+                col_name = result.column_expression.find(exp.Identifier)
+                if col_name:
+                    col_name = col_name.name
+                else:
+                    col_name = result.column_expression.find(exp.Literal)
+                    if col_name:
+                        col_name = col_name.this
+                alias_name = f"{func.__name__}__{col_name or ''}__"
+                # BigQuery has restrictions on alias names so we constrain it to alphanumeric characters and underscores
+                return result.alias(re.sub("\W", "_", alias_name))  # type: ignore
+            return result
+
+        wrapper.unsupported_engines = (  # type: ignore
+            ensure_list(unsupported_engines) if unsupported_engines else []
+        )
+        return wrapper
 
     return _metadata