sqlframe 1.1.3__py3-none-any.whl

sqlframe/base/normalize.py

@@ -0,0 +1,84 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import expressions as exp
+from sqlglot.helper import ensure_list
+
+from sqlframe.base.column import Column
+from sqlframe.base.util import get_tables_from_expression_with_join
+
+if t.TYPE_CHECKING:
+    from sqlframe.base.dataframe import SESSION
+
+    NORMALIZE_INPUT = t.TypeVar("NORMALIZE_INPUT", bound=t.Union[str, exp.Expression, Column])
+
+
+def normalize(session: SESSION, expression_context: exp.Select, expr: t.List[NORMALIZE_INPUT]):
+    expr = ensure_list(expr)
+    expressions = _ensure_expressions(expr)
+    for expression in expressions:
+        identifiers = expression.find_all(exp.Identifier)
+        for identifier in identifiers:
+            identifier.transform(session.input_dialect.normalize_identifier)
+            replace_alias_name_with_cte_name(session, expression_context, identifier)
+            replace_branch_and_sequence_ids_with_cte_name(session, expression_context, identifier)
+
+
+def replace_alias_name_with_cte_name(
+    session: SESSION, expression_context: exp.Select, id: exp.Identifier
+):
+    normalized_id = session._normalize_string(id.alias_or_name)
+    if normalized_id in session.name_to_sequence_id_mapping:
+        for cte in reversed(expression_context.ctes):
+            if cte.args["sequence_id"] in session.name_to_sequence_id_mapping[normalized_id]:
+                _set_alias_name(id, cte.alias_or_name)
+                break
+
+
+def replace_branch_and_sequence_ids_with_cte_name(
+    session: SESSION, expression_context: exp.Select, id: exp.Identifier
+):
+    normalized_id = session._normalize_string(id.alias_or_name)
+    if normalized_id in session.known_ids:
+        # Check if we have a join and if both the tables in that join share a common branch id
+        # If so we need to have this reference the left table by default unless the id is a sequence
+        # id then it keeps that reference. This handles the weird edge case in spark that shouldn't
+        # be common in practice
+        if expression_context.args.get("joins") and normalized_id in session.known_branch_ids:
+            join_table_aliases = [
+                x.alias_or_name for x in get_tables_from_expression_with_join(expression_context)
+            ]
+            ctes_in_join = [
+                cte for cte in expression_context.ctes if cte.alias_or_name in join_table_aliases
+            ]
+            if ctes_in_join[0].args["branch_id"] == ctes_in_join[1].args["branch_id"]:
+                assert len(ctes_in_join) == 2
+                _set_alias_name(id, ctes_in_join[0].alias_or_name)
+                return
+
+        for cte in reversed(expression_context.ctes):
+            if normalized_id in (cte.args["branch_id"], cte.args["sequence_id"]):
+                _set_alias_name(id, cte.alias_or_name)
+                return
+
+
+def _set_alias_name(id: exp.Identifier, name: str):
+    id.set("this", name)
+    id.set("quoted", False)
+
+
+def _ensure_expressions(values: t.List[NORMALIZE_INPUT]) -> t.List[exp.Expression]:
+    results = []
+    for value in values:
+        if isinstance(value, str):
+            results.append(Column.ensure_col(value).expression)
+        elif isinstance(value, Column):
+            results.append(value.expression)
+        elif isinstance(value, exp.Expression):
+            results.append(value)
+        else:
+            raise ValueError(f"Got an invalid type to normalize: {type(value)}")
+    return results

sqlframe/base/operations.py

@@ -0,0 +1,87 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import functools
+import typing as t
+from enum import IntEnum
+
+if t.TYPE_CHECKING:
+    from sqlframe.base.dataframe import _BaseDataFrame
+    from sqlframe.base.group import _BaseGroupedData
+
+
+class Operation(IntEnum):
+    INIT = -1
+    NO_OP = 0
+    FROM = 1
+    WHERE = 2
+    GROUP_BY = 3
+    HAVING = 4
+    SELECT = 5
+    ORDER_BY = 6
+    LIMIT = 7
+
+
+def operation(op: Operation) -> t.Callable[[t.Callable], t.Callable]:
+    """
+    Decorator used around DataFrame methods to indicate what type of operation is being performed from the
+    ordered Operation enums. This is used to determine which operations should be performed on a CTE vs.
+    included with the previous operation.
+
+    Ex: After a user does a join we want to allow them to select which columns for the different
+    tables that they want to carry through to the following operation. If we put that join in
+    a CTE preemptively then the user would not have a chance to select which column they want
+    in cases where there is overlap in names.
+    """
+
+    def decorator(func: t.Callable) -> t.Callable:
+        @functools.wraps(func)
+        def wrapper(self: _BaseDataFrame, *args, **kwargs) -> _BaseDataFrame:
+            if self.last_op == Operation.INIT:
+                self = self._convert_leaf_to_cte()
+                self.last_op = Operation.NO_OP
+            last_op = self.last_op
+            new_op = op if op != Operation.NO_OP else last_op
+            if new_op < last_op or (last_op == new_op == Operation.SELECT):
+                self = self._convert_leaf_to_cte()
+            df: t.Union[_BaseDataFrame, _BaseGroupedData] = func(self, *args, **kwargs)
+            df.last_op = new_op  # type: ignore
+            return df  # type: ignore
+
+        wrapper.__wrapped__ = func  # type: ignore
+        return wrapper
+
+    return decorator
+
+
+def group_operation(op: Operation) -> t.Callable[[t.Callable], t.Callable]:
+    """
+    Decorator used around DataFrame methods to indicate what type of operation is being performed from the
+    ordered Operation enums. This is used to determine which operations should be performed on a CTE vs.
+    included with the previous operation.
+
+    Ex: After a user does a join we want to allow them to select which columns for the different
+    tables that they want to carry through to the following operation. If we put that join in
+    a CTE preemptively then the user would not have a chance to select which column they want
+    in cases where there is overlap in names.
+    """
+
+    def decorator(func: t.Callable) -> t.Callable:
+        @functools.wraps(func)
+        def wrapper(self: _BaseGroupedData, *args, **kwargs) -> _BaseDataFrame:
+            if self._df.last_op == Operation.INIT:
+                self._df = self._df._convert_leaf_to_cte()
+                self._df.last_op = Operation.NO_OP
+            last_op = self._df.last_op
+            new_op = op if op != Operation.NO_OP else last_op
+            if new_op < last_op or (last_op == new_op == Operation.SELECT):
+                self._df = self._df._convert_leaf_to_cte()
+            df: _BaseDataFrame = func(self, *args, **kwargs)
+            df.last_op = new_op  # type: ignore
+            return df
+
+        wrapper.__wrapped__ = func  # type: ignore
+        return wrapper
+
+    return decorator