databricks-labs-lakebridge 0.10.0__py3-none-any.whl

databricks/labs/lakebridge/transpiler/sqlglot/lca_utils.py

@@ -0,0 +1,138 @@
+import logging
+from collections.abc import Iterable
+from pathlib import Path
+
+from sqlglot import expressions as exp, Dialect
+from sqlglot import parse
+from sqlglot.errors import ErrorLevel, ParseError, TokenError, UnsupportedError
+from sqlglot.expressions import Expression, Select
+from sqlglot.optimizer.scope import Scope, build_scope
+
+from databricks.labs.lakebridge.transpiler.transpile_status import TranspileError, ErrorKind, ErrorSeverity
+from databricks.labs.lakebridge.transpiler.sqlglot.local_expression import AliasInfo
+
+logger = logging.getLogger(__name__)
+
+
+def check_for_unsupported_lca(
+    from_dialect: Dialect,
+    source_sql: str,
+    file_path: Path,
+) -> TranspileError | None:
+    """
+    Check for presence of unsupported lateral column aliases in window expressions and where clauses
+    :return: An error if found
+    """
+    try:
+        all_parsed_expressions: Iterable[Expression | None] = parse(
+            source_sql, read=from_dialect, error_level=ErrorLevel.RAISE
+        )
+        root_expressions: Iterable[Expression] = [pe for pe in all_parsed_expressions if pe is not None]
+    except (ParseError, TokenError, UnsupportedError) as e:
+        logger.warning(f"Error while preprocessing {file_path}: {e}")
+        return None
+
+    aliases_in_where = set()
+    aliases_in_window = set()
+
+    for expr in root_expressions:
+        for select in expr.find_all(exp.Select, bfs=False):
+            alias_info = _find_aliases_in_select(select)
+            aliases_in_where.update(_find_invalid_lca_in_where(select, alias_info))
+            aliases_in_window.update(_find_invalid_lca_in_window(select, alias_info))
+
+    if not (aliases_in_where or aliases_in_window):
+        return None
+
+    err_messages = [f"Unsupported operation found in file {file_path}. Needs manual review of transpiled query."]
+    if aliases_in_where:
+        err_messages.append(f"Lateral column aliases `{', '.join(aliases_in_where)}` found in where clause.")
+
+    if aliases_in_window:
+        err_messages.append(f"Lateral column aliases `{', '.join(aliases_in_window)}` found in window expressions.")
+
+    return TranspileError("UNSUPPORTED_LCA", ErrorKind.ANALYSIS, ErrorSeverity.ERROR, file_path, " ".join(err_messages))
+
+
+def unalias_lca_in_select(expr: exp.Expression) -> exp.Expression:
+    if not isinstance(expr, exp.Select):
+        return expr
+
+    root_select: Scope | None = build_scope(expr)
+    if not root_select:
+        return expr
+
+    # We won't search inside nested selects, they will be visited separately
+    nested_selects = {*root_select.derived_tables, *root_select.subqueries}
+    alias_info = _find_aliases_in_select(expr)
+    where_ast: Expression | None = expr.args.get("where")
+    if where_ast:
+        for column in where_ast.walk(prune=lambda n: n in nested_selects):
+            _replace_aliases(column, alias_info)
+    for window in _find_windows_in_select(expr):
+        for column in window.walk():
+            _replace_aliases(column, alias_info)
+    return expr
+
+
+def _replace_aliases(column: Expression, alias_info: dict[str, AliasInfo]):
+    if (
+        isinstance(column, exp.Column)
+        and column.name in alias_info
+        and not alias_info[column.name].is_same_name_as_column
+    ):
+        unaliased_expr = alias_info[column.name].expression
+        column.replace(unaliased_expr)
+        for col in unaliased_expr.walk():
+            _replace_aliases(col, alias_info)
+
+
+def _find_windows_in_select(select: Select) -> list[exp.Window]:
+    window_expressions = []
+    for expr in select.expressions:
+        window_expr = expr.find(exp.Window)
+        if window_expr:
+            window_expressions.append(window_expr)
+    return window_expressions
+
+
+def _find_aliases_in_select(select_expr: Select) -> dict[str, AliasInfo]:
+    aliases = {}
+    for expr in select_expr.expressions:
+        if isinstance(expr, exp.Alias):
+            alias_name = expr.output_name
+            is_same_name_as_column = False
+            for column in expr.find_all(exp.Column):
+                if column.name == alias_name:
+                    is_same_name_as_column = True
+                    break
+            aliases[alias_name] = AliasInfo(alias_name, expr.unalias().copy(), is_same_name_as_column)
+    return aliases
+
+
+def _find_invalid_lca_in_where(
+    select_expr: Select,
+    aliases: dict[str, AliasInfo],
+) -> set[str]:
+    aliases_in_where = set()
+    where_ast: Expression | None = select_expr.args.get("where")
+    if where_ast:
+        for column in where_ast.find_all(exp.Column):
+            if column.name in aliases and not aliases[column.name].is_same_name_as_column:
+                aliases_in_where.add(column.name)
+
+    return aliases_in_where
+
+
+def _find_invalid_lca_in_window(
+    select_expr: Select,
+    aliases: dict[str, AliasInfo],
+) -> set[str]:
+    aliases_in_window = set()
+    windows = _find_windows_in_select(select_expr)
+    for window in windows:
+        for column in window.find_all(exp.Column):
+            if column.name in aliases and not aliases[column.name].is_same_name_as_column:
+                aliases_in_window.add(column.name)
+
+    return aliases_in_window

databricks/labs/lakebridge/transpiler/sqlglot/local_expression.py

@@ -0,0 +1,197 @@
+from dataclasses import dataclass
+
+from sqlglot import expressions as exp
+from sqlglot.expressions import AggFunc, Condition, Expression, Func
+
+
+class NthValue(AggFunc):
+    arg_types = {"this": True, "offset": False}
+
+
+class Parameter(Expression):
+    arg_types = {"this": True, "wrapped": False, "suffix": False}
+
+
+class Collate(Func):
+    arg_types = {"this": True, "expressions": True}
+
+
+class Bracket(Condition):
+    arg_types = {"this": True, "expressions": True}
+
+
+class Split(Func):
+    """
+    Redefined Split(sqlglot/expression) class with expression: False to handle default delimiter
+    Please refer the test case `test_strtok_to_array` -> `select STRTOK_TO_ARRAY('my text is divided')`
+    """
+
+    arg_types = {"this": True, "expression": False, "limit": False}
+
+
+class MakeDate(Func):
+    arg_types = {"this": True, "expression": False, "zone": False}
+
+
+class ConvertTimeZone(Func):
+    arg_types = {"srcTZ": True, "tgtTZ": True, "this": False}
+
+
+class TryToDate(Func):
+    arg_types = {"this": True, "format": False}
+
+
+class TryToTimestamp(Func):
+    arg_types = {"this": True, "format": False}
+
+
+class SplitPart(Func):
+    arg_types = {"this": True, "expression": False, "partNum": False}
+
+
+class StrTok(Func):
+    arg_types = {"this": True, "expression": False, "partNum": False}
+
+
+class TryToNumber(Func):
+    arg_types = {"this": True, "expression": False, "precision": False, "scale": False}
+
+    _sql_names = ["TRY_TO_DECIMAL", "TRY_TO_NUMBER", "TRY_TO_NUMERIC"]
+
+
+class DateFormat(Func):
+    arg_types = {"this": True, "expression": False}
+
+
+class IsInteger(Func):
+    pass
+
+
+class JsonExtractPathText(Func):
+    arg_types = {"this": True, "path_name": True}
+
+
+class BitOr(AggFunc):
+    pass
+
+
+class ArrayConstructCompact(Func):
+    arg_types = {"expressions": False}
+
+    is_var_len_args = True
+
+
+class ArrayIntersection(Func):
+    arg_types = {"this": True, "expression": True}
+
+
+class ArraySlice(Func):
+    arg_types = {"this": True, "from": True, "to": True}
+
+
+class ObjectKeys(Func):
+    arg_types = {"this": True}
+
+
+class ToBoolean(Func):
+    arg_types = {"this": True, "raise_error": False}
+
+
+class ToDouble(Func):
+    pass
+
+
+class ToObject(Func):
+    pass
+
+
+class ToNumber(Func):
+    arg_types = {"this": True, "expression": False, "precision": False, "scale": False}
+
+    _sql_names = ["TO_DECIMAL", "TO_NUMBER", "TO_NUMERIC"]
+
+
+class TimestampFromParts(Func):
+    arg_types = {
+        "this": True,
+        "expression": True,
+        "day": True,
+        "hour": True,
+        "min": True,
+        "sec": True,
+        "nanosec": False,
+        "Zone": False,
+    }
+
+
+class ToVariant(Func):
+    pass
+
+
+class UUID(Func):
+    arg_types = {"this": False, "name": False}
+
+
+class DateTrunc(Func):
+    arg_types = {"unit": False, "this": True, "zone": False}
+
+
+class Median(Func):
+    arg_types = {"this": True}
+
+
+class CumeDist(Func):
+    arg_types = {"this": False}
+
+
+class DenseRank(Func):
+    arg_types = {"this": False}
+
+
+class Rank(Func):
+    arg_types = {"this": False}
+
+
+class PercentRank(Func):
+    arg_types = {"this": False}
+
+
+class Ntile(Func):
+    arg_types = {"this": True, "is_string": False}
+
+
+class ToArray(Func):
+    arg_types = {"this": True, "expression": False}
+
+
+@dataclass
+class WithinGroupParams:
+    agg_col: exp.Column
+    order_cols: list[tuple[exp.Column, bool]]  # List of (column, is ascending)
+
+
+@dataclass
+class AliasInfo:
+    name: str
+    expression: exp.Expression
+    is_same_name_as_column: bool
+
+
+class MapKeys(Func):
+    arg_types = {"this": True}
+
+
+class ArrayExists(Func):
+    arg_types = {"this": True, "expression": True}
+
+
+class Locate(Func):
+    arg_types = {"substring": True, "this": True, "position": False}
+
+
+class NamedStruct(Func):
+    arg_types = {"expressions": True}
+
+
+class GetJsonObject(Func):
+    arg_types = {"this": True, "path": True}

databricks/labs/lakebridge/transpiler/sqlglot/parsers/oracle.py

@@ -0,0 +1,23 @@
+from sqlglot.dialects.oracle import Oracle as Orc
+from sqlglot.tokens import TokenType
+
+
+class Oracle(Orc):
+    # Instantiate Oracle Dialect
+    oracle = Orc()
+
+    class Tokenizer(Orc.Tokenizer):
+        KEYWORDS = {
+            **Orc.Tokenizer.KEYWORDS,
+            'LONG': TokenType.TEXT,
+            'NCLOB': TokenType.TEXT,
+            'ROWID': TokenType.TEXT,
+            'UROWID': TokenType.TEXT,
+            'ANYTYPE': TokenType.TEXT,
+            'ANYDATA': TokenType.TEXT,
+            'ANYDATASET': TokenType.TEXT,
+            'XMLTYPE': TokenType.TEXT,
+            'SDO_GEOMETRY': TokenType.TEXT,
+            'SDO_TOPO_GEOMETRY': TokenType.TEXT,
+            'SDO_GEORASTER': TokenType.TEXT,
+        }

databricks/labs/lakebridge/transpiler/sqlglot/parsers/presto.py

@@ -0,0 +1,202 @@
+import logging
+from sqlglot.dialects.presto import Presto as presto
+from sqlglot import exp
+from sqlglot.helper import seq_get
+from sqlglot.errors import ParseError
+from sqlglot.tokens import TokenType
+
+from databricks.labs.lakebridge.transpiler.sqlglot import local_expression
+
+
+logger = logging.getLogger(__name__)
+
+
+def _build_approx_percentile(args: list) -> exp.Expression:
+    if len(args) == 4:
+        arg3 = seq_get(args, 3)
+        try:
+            number = float(arg3.this) if arg3 is not None else 0
+            return exp.ApproxQuantile(
+                this=seq_get(args, 0),
+                weight=seq_get(args, 1),
+                quantile=seq_get(args, 2),
+                accuracy=exp.Literal(this=f'{int((1/number) * 100)} ', is_string=False),
+            )
+        except ValueError as exc:
+            raise ParseError(f"Expected a string representation of a number for argument 2, but got {arg3}") from exc
+    if len(args) == 3:
+        arg2 = seq_get(args, 2)
+        try:
+            number = float(arg2.this) if arg2 is not None else 0
+            return exp.ApproxQuantile(
+                this=seq_get(args, 0),
+                quantile=seq_get(args, 1),
+                accuracy=exp.Literal(this=f'{int((1/number) * 100)}', is_string=False),
+            )
+        except ValueError as exc:
+            raise ParseError(f"Expected a string representation of a number for argument 2, but got {arg2}") from exc
+    return exp.ApproxQuantile.from_arg_list(args)
+
+
+def _build_any_keys_match(args: list) -> local_expression.ArrayExists:
+    return local_expression.ArrayExists(
+        this=local_expression.MapKeys(this=seq_get(args, 0)), expression=seq_get(args, 1)
+    )
+
+
+def _build_str_position(args: list) -> local_expression.Locate:
+    # TODO the 3rd param in presto strpos and databricks locate has different implementation.
+    # For now we haven't implemented the logic same as presto for 3rd param.
+    # Users should be vigilant when using 3 param function in presto strpos.
+    if len(args) == 3:
+        msg = (
+            "*Warning:: The third parameter in Presto's `strpos` function and Databricks' `locate` function "
+            "have different implementations. Please exercise caution when using the three-parameter version "
+            "of the `strpos` function in Presto."
+        )
+        logger.warning(msg)
+        return local_expression.Locate(substring=seq_get(args, 1), this=seq_get(args, 0), position=seq_get(args, 2))
+    return local_expression.Locate(substring=seq_get(args, 1), this=seq_get(args, 0))
+
+
+def _build_array_average(args: list) -> exp.Reduce:
+    return exp.Reduce(
+        this=exp.ArrayFilter(
+            this=seq_get(args, 0),
+            expression=exp.Lambda(
+                this=exp.Not(this=exp.Is(this=exp.Identifier(this="x", quoted=False), expression=exp.Null())),
+                expressions=[exp.Identifier(this="x", quoted=False)],
+            ),
+        ),
+        initial=local_expression.NamedStruct(
+            expressions=[
+                exp.Literal(this="sum", is_string=True),
+                exp.Cast(this=exp.Literal(this="0", is_string=False), to=exp.DataType(this="DOUBLE")),
+                exp.Literal(this="cnt", is_string=True),
+                exp.Literal(this="0", is_string=False),
+            ],
+        ),
+        merge=exp.Lambda(
+            this=local_expression.NamedStruct(
+                expressions=[
+                    exp.Literal(this="sum", is_string=True),
+                    exp.Add(
+                        this=exp.Dot(
+                            this=exp.Identifier(this="acc", quoted=False),
+                            expression=exp.Identifier(this="sum", quoted=False),
+                        ),
+                        expression=exp.Identifier(this="x", quoted=False),
+                    ),
+                    exp.Literal(this="cnt", is_string=True),
+                    exp.Add(
+                        this=exp.Dot(
+                            this=exp.Identifier(this="acc", quoted=False),
+                            expression=exp.Identifier(this="cnt", quoted=False),
+                        ),
+                        expression=exp.Literal(this="1", is_string=False),
+                    ),
+                ],
+            ),
+            expressions=[exp.Identifier(this="acc", quoted=False), exp.Identifier(this="x", quoted=False)],
+        ),
+        finish=exp.Lambda(
+            this=exp.Anonymous(
+                this="try_divide",
+                expressions=[
+                    exp.Dot(
+                        this=exp.Identifier(this="acc", quoted=False),
+                        expression=exp.Identifier(this="sum", quoted=False),
+                    ),
+                    exp.Dot(
+                        this=exp.Identifier(this="acc", quoted=False),
+                        expression=exp.Identifier(this="cnt", quoted=False),
+                    ),
+                ],
+            ),
+            expressions=[exp.Identifier(this="acc", quoted=False)],
+        ),
+    )
+
+
+def _build_json_size(args: list):
+    return exp.Case(
+        ifs=[
+            exp.If(
+                this=exp.Like(
+                    this=local_expression.GetJsonObject(
+                        this=exp.Column(this=seq_get(args, 0)),
+                        path=exp.Column(this=seq_get(args, 1)),
+                    ),
+                    expression=exp.Literal(this="{%", is_string=True),
+                ),
+                true=exp.ArraySize(
+                    this=exp.Anonymous(
+                        this="from_json",
+                        expressions=[
+                            local_expression.GetJsonObject(
+                                this=exp.Column(this=seq_get(args, 0)),
+                                path=exp.Column(this=seq_get(args, 1)),
+                            ),
+                            exp.Literal(this="map<string,string>", is_string=True),
+                        ],
+                    )
+                ),
+            ),
+            exp.If(
+                this=exp.Like(
+                    this=local_expression.GetJsonObject(
+                        this=exp.Column(this=seq_get(args, 0)),
+                        path=exp.Column(this=seq_get(args, 1)),
+                    ),
+                    expression=exp.Literal(this="[%", is_string=True),
+                ),
+                true=exp.ArraySize(
+                    this=exp.Anonymous(
+                        this="from_json",
+                        expressions=[
+                            local_expression.GetJsonObject(
+                                this=exp.Column(this=seq_get(args, 0)),
+                                path=exp.Column(this=seq_get(args, 1)),
+                            ),
+                            exp.Literal(this="array<string>", is_string=True),
+                        ],
+                    )
+                ),
+            ),
+            exp.If(
+                this=exp.Not(
+                    this=exp.Is(
+                        this=local_expression.GetJsonObject(
+                            this=exp.Column(this=seq_get(args, 0)),
+                            path=exp.Column(this=seq_get(args, 1)),
+                        ),
+                        expression=exp.Null(),
+                    )
+                ),
+                true=exp.Literal(this="0", is_string=False),
+            ),
+        ],
+        default=exp.Null(),
+    )
+
+
+class Presto(presto):
+
+    class Parser(presto.Parser):
+        VALUES_FOLLOWED_BY_PAREN = False
+
+        FUNCTIONS = {
+            **presto.Parser.FUNCTIONS,
+            "APPROX_PERCENTILE": _build_approx_percentile,
+            "STRPOS": _build_str_position,
+            "ANY_KEYS_MATCH": _build_any_keys_match,
+            "ARRAY_AVERAGE": _build_array_average,
+            "JSON_SIZE": _build_json_size,
+            "FORMAT_DATETIME": local_expression.DateFormat.from_arg_list,
+        }
+
+    class Tokenizer(presto.Tokenizer):
+        KEYWORDS = {
+            **presto.Tokenizer.KEYWORDS,
+            "JSON": TokenType.TEXT,
+        }