@altimateai/altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/group.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot.dataframe.sql import functions as F
+from sqlglot.dataframe.sql.column import Column
+from sqlglot.dataframe.sql.operations import Operation, operation
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql.dataframe import DataFrame
+
+
+class GroupedData:
+    def __init__(self, df: DataFrame, group_by_cols: t.List[Column], last_op: Operation):
+        self._df = df.copy()
+        self.spark = df.spark
+        self.last_op = last_op
+        self.group_by_cols = group_by_cols
+
+    def _get_function_applied_columns(
+        self, func_name: str, cols: t.Tuple[str, ...]
+    ) -> t.List[Column]:
+        func_name = func_name.lower()
+        return [getattr(F, func_name)(name).alias(f"{func_name}({name})") for name in cols]
+
+    @operation(Operation.SELECT)
+    def agg(self, *exprs: t.Union[Column, t.Dict[str, str]]) -> DataFrame:
+        columns = (
+            [Column(f"{agg_func}({column_name})") for column_name, agg_func in exprs[0].items()]
+            if isinstance(exprs[0], dict)
+            else exprs
+        )
+        cols = self._df._ensure_and_normalize_cols(columns)
+
+        expression = self._df.expression.group_by(
+            *[x.expression for x in self.group_by_cols]
+        ).select(*[x.expression for x in self.group_by_cols + cols], append=False)
+        return self._df.copy(expression=expression)
+
+    def count(self) -> DataFrame:
+        return self.agg(F.count("*").alias("count"))
+
+    def mean(self, *cols: str) -> DataFrame:
+        return self.avg(*cols)
+
+    def avg(self, *cols: str) -> DataFrame:
+        return self.agg(*self._get_function_applied_columns("avg", cols))
+
+    def max(self, *cols: str) -> DataFrame:
+        return self.agg(*self._get_function_applied_columns("max", cols))
+
+    def min(self, *cols: str) -> DataFrame:
+        return self.agg(*self._get_function_applied_columns("min", cols))
+
+    def sum(self, *cols: str) -> DataFrame:
+        return self.agg(*self._get_function_applied_columns("sum", cols))
+
+    def pivot(self, *cols: str) -> DataFrame:
+        raise NotImplementedError("Sum distinct is not currently implemented")

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/normalize.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import expressions as exp
+from sqlglot.dataframe.sql.column import Column
+from sqlglot.dataframe.sql.util import get_tables_from_expression_with_join
+from sqlglot.helper import ensure_list
+
+NORMALIZE_INPUT = t.TypeVar("NORMALIZE_INPUT", bound=t.Union[str, exp.Expression, Column])
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql.session import SparkSession
+
+
+def normalize(spark: SparkSession, 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(spark.dialect.normalize_identifier)
+            replace_alias_name_with_cte_name(spark, expression_context, identifier)
+            replace_branch_and_sequence_ids_with_cte_name(spark, expression_context, identifier)
+
+
+def replace_alias_name_with_cte_name(
+    spark: SparkSession, expression_context: exp.Select, id: exp.Identifier
+):
+    if id.alias_or_name in spark.name_to_sequence_id_mapping:
+        for cte in reversed(expression_context.ctes):
+            if cte.args["sequence_id"] in spark.name_to_sequence_id_mapping[id.alias_or_name]:
+                _set_alias_name(id, cte.alias_or_name)
+                break
+
+
+def replace_branch_and_sequence_ids_with_cte_name(
+    spark: SparkSession, expression_context: exp.Select, id: exp.Identifier
+):
+    if id.alias_or_name in spark.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 id.alias_or_name in spark.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 id.alias_or_name 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)
+
+
+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

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/operations.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import functools
+import typing as t
+from enum import IntEnum
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql.dataframe import DataFrame
+    from sqlglot.dataframe.sql.group import GroupedData
+
+
+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):
+    """
+    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):
+        @functools.wraps(func)
+        def wrapper(self: DataFrame, *args, **kwargs):
+            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[DataFrame, GroupedData] = func(self, *args, **kwargs)
+            df.last_op = new_op  # type: ignore
+            return df
+
+        wrapper.__wrapped__ = func  # type: ignore
+        return wrapper
+
+    return decorator

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/readwriter.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import typing as t
+
+import sqlglot as sqlglot
+from sqlglot import expressions as exp
+from sqlglot.helper import object_to_dict
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql.dataframe import DataFrame
+    from sqlglot.dataframe.sql.session import SparkSession
+
+
+class DataFrameReader:
+    def __init__(self, spark: SparkSession):
+        self.spark = spark
+
+    def table(self, tableName: str) -> DataFrame:
+        from sqlglot.dataframe.sql.dataframe import DataFrame
+        from sqlglot.dataframe.sql.session import SparkSession
+
+        sqlglot.schema.add_table(tableName, dialect=SparkSession().dialect)
+
+        return DataFrame(
+            self.spark,
+            exp.Select()
+            .from_(
+                exp.to_table(tableName, dialect=SparkSession().dialect).transform(
+                    SparkSession().dialect.normalize_identifier
+                )
+            )
+            .select(
+                *(
+                    column
+                    for column in sqlglot.schema.column_names(
+                        tableName, dialect=SparkSession().dialect
+                    )
+                )
+            ),
+        )
+
+
+class DataFrameWriter:
+    def __init__(
+        self,
+        df: DataFrame,
+        spark: t.Optional[SparkSession] = None,
+        mode: t.Optional[str] = None,
+        by_name: bool = False,
+    ):
+        self._df = df
+        self._spark = spark or df.spark
+        self._mode = mode
+        self._by_name = by_name
+
+    def copy(self, **kwargs) -> DataFrameWriter:
+        return DataFrameWriter(
+            **{
+                k[1:] if k.startswith("_") else k: v
+                for k, v in object_to_dict(self, **kwargs).items()
+            }
+        )
+
+    def sql(self, **kwargs) -> t.List[str]:
+        return self._df.sql(**kwargs)
+
+    def mode(self, saveMode: t.Optional[str]) -> DataFrameWriter:
+        return self.copy(_mode=saveMode)
+
+    @property
+    def byName(self):
+        return self.copy(by_name=True)
+
+    def insertInto(self, tableName: str, overwrite: t.Optional[bool] = None) -> DataFrameWriter:
+        from sqlglot.dataframe.sql.session import SparkSession
+
+        output_expression_container = exp.Insert(
+            **{
+                "this": exp.to_table(tableName),
+                "overwrite": overwrite,
+            }
+        )
+        df = self._df.copy(output_expression_container=output_expression_container)
+        if self._by_name:
+            columns = sqlglot.schema.column_names(
+                tableName, only_visible=True, dialect=SparkSession().dialect
+            )
+            df = df._convert_leaf_to_cte().select(*columns)
+
+        return self.copy(_df=df)
+
+    def saveAsTable(self, name: str, format: t.Optional[str] = None, mode: t.Optional[str] = None):
+        if format is not None:
+            raise NotImplementedError("Providing Format in the save as table is not supported")
+        exists, replace, mode = None, None, mode or str(self._mode)
+        if mode == "append":
+            return self.insertInto(name)
+        if mode == "ignore":
+            exists = True
+        if mode == "overwrite":
+            replace = True
+        output_expression_container = exp.Create(
+            this=exp.to_table(name),
+            kind="TABLE",
+            exists=exists,
+            replace=replace,
+        )
+        return self.copy(_df=self._df.copy(output_expression_container=output_expression_container))

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/session.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import typing as t
+import uuid
+from collections import defaultdict
+
+import sqlglot as sqlglot
+from sqlglot import Dialect, expressions as exp
+from sqlglot.dataframe.sql import functions as F
+from sqlglot.dataframe.sql.dataframe import DataFrame
+from sqlglot.dataframe.sql.readwriter import DataFrameReader
+from sqlglot.dataframe.sql.types import StructType
+from sqlglot.dataframe.sql.util import get_column_mapping_from_schema_input
+from sqlglot.helper import classproperty
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql._typing import ColumnLiterals, SchemaInput
+
+
+class SparkSession:
+    DEFAULT_DIALECT = "spark"
+    _instance = None
+
+    def __init__(self):
+        if not hasattr(self, "known_ids"):
+            self.known_ids = set()
+            self.known_branch_ids = set()
+            self.known_sequence_ids = set()
+            self.name_to_sequence_id_mapping = defaultdict(list)
+            self.incrementing_id = 1
+            self.dialect = Dialect.get_or_raise(self.DEFAULT_DIALECT)()
+
+    def __new__(cls, *args, **kwargs) -> SparkSession:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @property
+    def read(self) -> DataFrameReader:
+        return DataFrameReader(self)
+
+    def table(self, tableName: str) -> DataFrame:
+        return self.read.table(tableName)
+
+    def createDataFrame(
+        self,
+        data: t.Sequence[t.Union[t.Dict[str, ColumnLiterals], t.List[ColumnLiterals], t.Tuple]],
+        schema: t.Optional[SchemaInput] = None,
+        samplingRatio: t.Optional[float] = None,
+        verifySchema: bool = False,
+    ) -> DataFrame:
+        from sqlglot.dataframe.sql.dataframe import DataFrame
+
+        if samplingRatio is not None or verifySchema:
+            raise NotImplementedError("Sampling Ratio and Verify Schema are not supported")
+        if schema is not None and (
+            not isinstance(schema, (StructType, str, list))
+            or (isinstance(schema, list) and not isinstance(schema[0], str))
+        ):
+            raise NotImplementedError("Only schema of either list or string of list supported")
+        if not data:
+            raise ValueError("Must provide data to create into a DataFrame")
+
+        column_mapping: t.Dict[str, t.Optional[str]]
+        if schema is not None:
+            column_mapping = get_column_mapping_from_schema_input(schema)
+        elif isinstance(data[0], dict):
+            column_mapping = {col_name.strip(): None for col_name in data[0]}
+        else:
+            column_mapping = {f"_{i}": None for i in range(1, len(data[0]) + 1)}
+
+        data_expressions = [
+            exp.Tuple(
+                expressions=list(
+                    map(
+                        lambda x: F.lit(x).expression,
+                        row if not isinstance(row, dict) else row.values(),
+                    )
+                )
+            )
+            for row in data
+        ]
+
+        sel_columns = [
+            F.col(name).cast(data_type).alias(name).expression
+            if data_type is not None
+            else F.col(name).expression
+            for name, data_type in column_mapping.items()
+        ]
+
+        select_kwargs = {
+            "expressions": sel_columns,
+            "from": exp.From(
+                this=exp.Values(
+                    expressions=data_expressions,
+                    alias=exp.TableAlias(
+                        this=exp.to_identifier(self._auto_incrementing_name),
+                        columns=[exp.to_identifier(col_name) for col_name in column_mapping],
+                    ),
+                ),
+            ),
+        }
+
+        sel_expression = exp.Select(**select_kwargs)
+        return DataFrame(self, sel_expression)
+
+    def sql(self, sqlQuery: str) -> DataFrame:
+        expression = sqlglot.parse_one(sqlQuery, read=self.dialect)
+        if isinstance(expression, exp.Select):
+            df = DataFrame(self, expression)
+            df = df._convert_leaf_to_cte()
+        elif isinstance(expression, (exp.Create, exp.Insert)):
+            select_expression = expression.expression.copy()
+            if isinstance(expression, exp.Insert):
+                select_expression.set("with", expression.args.get("with"))
+                expression.set("with", None)
+            del expression.args["expression"]
+            df = DataFrame(self, select_expression, output_expression_container=expression)  # type: ignore
+            df = df._convert_leaf_to_cte()
+        else:
+            raise ValueError(
+                "Unknown expression type provided in the SQL. Please create an issue with the SQL."
+            )
+        return df
+
+    @property
+    def _auto_incrementing_name(self) -> str:
+        name = f"a{self.incrementing_id}"
+        self.incrementing_id += 1
+        return name
+
+    @property
+    def _random_branch_id(self) -> str:
+        id = self._random_id
+        self.known_branch_ids.add(id)
+        return id
+
+    @property
+    def _random_sequence_id(self):
+        id = self._random_id
+        self.known_sequence_ids.add(id)
+        return id
+
+    @property
+    def _random_id(self) -> str:
+        id = "r" + uuid.uuid4().hex
+        self.known_ids.add(id)
+        return id
+
+    @property
+    def _join_hint_names(self) -> t.Set[str]:
+        return {"BROADCAST", "MERGE", "SHUFFLE_HASH", "SHUFFLE_REPLICATE_NL"}
+
+    def _add_alias_to_mapping(self, name: str, sequence_id: str):
+        self.name_to_sequence_id_mapping[name].append(sequence_id)
+
+    class Builder:
+        SQLFRAME_DIALECT_KEY = "sqlframe.dialect"
+
+        def __init__(self):
+            self.dialect = "spark"
+
+        def __getattr__(self, item) -> SparkSession.Builder:
+            return self
+
+        def __call__(self, *args, **kwargs):
+            return self
+
+        def config(
+            self,
+            key: t.Optional[str] = None,
+            value: t.Optional[t.Any] = None,
+            *,
+            map: t.Optional[t.Dict[str, t.Any]] = None,
+            **kwargs: t.Any,
+        ) -> SparkSession.Builder:
+            if key == self.SQLFRAME_DIALECT_KEY:
+                self.dialect = value
+            elif map and self.SQLFRAME_DIALECT_KEY in map:
+                self.dialect = map[self.SQLFRAME_DIALECT_KEY]
+            return self
+
+        def getOrCreate(self) -> SparkSession:
+            spark = SparkSession()
+            spark.dialect = Dialect.get_or_raise(self.dialect)()
+            return spark
+
+    @classproperty
+    def builder(cls) -> Builder:
+        return cls.Builder()

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/transforms.py

@@ -0,0 +1,9 @@
+import typing as t
+
+from sqlglot import expressions as exp
+
+
+def replace_id_value(node, replacement_mapping: t.Dict[exp.Identifier, exp.Identifier]):
+    if isinstance(node, exp.Identifier) and node in replacement_mapping:
+        node = node.replace(replacement_mapping[node].copy())
+    return node