ydb-sqlglot-plugin 0.2.1__tar.gz → 0.2.3__tar.gz

{ydb_sqlglot_plugin-0.2.1 → ydb_sqlglot_plugin-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ydb-sqlglot-plugin
-Version: 0.2.1
+Version: 0.2.3
 Summary: YDB dialect plugin for sqlglot
 Author: YDB Team
 License: Apache-2.0
@@ -102,6 +102,31 @@ LEFT JOIN (
 
 The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subqueries.
 
+#### GROUP BY aliases
+
+YDB accepts aliases directly inside `GROUP BY` items. The generator uses this
+form for grouped columns so later clauses and decorrelated subqueries can refer
+to a stable grouping name:
+
+```sql
+-- input
+SELECT user_id, COUNT(*) FROM events GROUP BY user_id
+
+-- output
+SELECT user_id, COUNT(*) FROM `events` GROUP BY user_id AS user_id
+```
+
+If a grouped column is selected under a generated alias, the `GROUP BY` item uses
+that alias as well:
+
+```sql
+SELECT user_id AS _u_1, COUNT(*) FROM `events` GROUP BY user_id AS _u_1
+```
+
+Positional `GROUP BY` references are expanded before generation. When a
+positional reference points to a constant expression, the grouping item is
+removed because YDB rejects grouping by constants.
+
 ---
 
 ### YDB → any SQL
@@ -119,7 +144,10 @@ The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-
 | `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
 | Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
 | `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| `GROUP BY expr AS alias` | `SELECT v, COUNT(*) FROM t GROUP BY v AS v` |
 | Named expressions | `$t = (SELECT 1 AS x)` |
+| Lambda expressions | `($x, $y?) -> ($x + COALESCE($y, 0))`, `($y) -> { $p = "x"; RETURN $p \|\| $y }` |
+| `IN COMPACT` | `WHERE key IN COMPACT $values` |
 | `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
 
 Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
@@ -204,11 +232,30 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 |---|---|
 | `ARRAY(v1, v2, ...)` | `AsList(v1, v2, ...)` |
 | `ARRAY_LENGTH(x)` / `ARRAY_SIZE(x)` | `ListLength(x)` |
-| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> {RETURN cond})` |
-| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> {RETURN cond}))` |
+| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> (cond))` |
+| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> (cond)))` |
 | `ARRAY_AGG(x)` | `AGGREGATE_LIST(x)` |
 | `UNNEST(x)` | `FLATTEN BY x` |
 
+Lambda expressions are represented with sqlglot's standard `exp.Lambda` AST node.
+When a source dialect parses lambdas, the YDB generator emits YQL lambda syntax:
+
+```sql
+-- DuckDB input
+SELECT list_filter(arr, x -> x > 0) FROM t
+
+-- YDB output
+SELECT ListFilter(arr, ($x) -> ($x > 0)) FROM `t`
+```
+
+YDB input also supports documented YQL lambda forms, including optional
+arguments and block bodies with local named expressions:
+
+```sql
+($x, $y?) -> ($x + COALESCE($y, 0));
+($y) -> { $prefix = "x"; RETURN $prefix || $y; };
+```
+
 ### Conditional / math
 
 | Input | YQL output |

{ydb_sqlglot_plugin-0.2.1 → ydb_sqlglot_plugin-0.2.3}/README.md

@@ -75,6 +75,31 @@ LEFT JOIN (
 
 The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subqueries.
 
+#### GROUP BY aliases
+
+YDB accepts aliases directly inside `GROUP BY` items. The generator uses this
+form for grouped columns so later clauses and decorrelated subqueries can refer
+to a stable grouping name:
+
+```sql
+-- input
+SELECT user_id, COUNT(*) FROM events GROUP BY user_id
+
+-- output
+SELECT user_id, COUNT(*) FROM `events` GROUP BY user_id AS user_id
+```
+
+If a grouped column is selected under a generated alias, the `GROUP BY` item uses
+that alias as well:
+
+```sql
+SELECT user_id AS _u_1, COUNT(*) FROM `events` GROUP BY user_id AS _u_1
+```
+
+Positional `GROUP BY` references are expanded before generation. When a
+positional reference points to a constant expression, the grouping item is
+removed because YDB rejects grouping by constants.
+
 ---
 
 ### YDB → any SQL
@@ -92,7 +117,10 @@ The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-
 | `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
 | Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
 | `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| `GROUP BY expr AS alias` | `SELECT v, COUNT(*) FROM t GROUP BY v AS v` |
 | Named expressions | `$t = (SELECT 1 AS x)` |
+| Lambda expressions | `($x, $y?) -> ($x + COALESCE($y, 0))`, `($y) -> { $p = "x"; RETURN $p \|\| $y }` |
+| `IN COMPACT` | `WHERE key IN COMPACT $values` |
 | `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
 
 Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
@@ -177,11 +205,30 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 |---|---|
 | `ARRAY(v1, v2, ...)` | `AsList(v1, v2, ...)` |
 | `ARRAY_LENGTH(x)` / `ARRAY_SIZE(x)` | `ListLength(x)` |
-| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> {RETURN cond})` |
-| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> {RETURN cond}))` |
+| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> (cond))` |
+| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> (cond)))` |
 | `ARRAY_AGG(x)` | `AGGREGATE_LIST(x)` |
 | `UNNEST(x)` | `FLATTEN BY x` |
 
+Lambda expressions are represented with sqlglot's standard `exp.Lambda` AST node.
+When a source dialect parses lambdas, the YDB generator emits YQL lambda syntax:
+
+```sql
+-- DuckDB input
+SELECT list_filter(arr, x -> x > 0) FROM t
+
+-- YDB output
+SELECT ListFilter(arr, ($x) -> ($x > 0)) FROM `t`
+```
+
+YDB input also supports documented YQL lambda forms, including optional
+arguments and block bodies with local named expressions:
+
+```sql
+($x, $y?) -> ($x + COALESCE($y, 0));
+($y) -> { $prefix = "x"; RETURN $prefix || $y; };
+```
+
 ### Conditional / math
 
 | Input | YQL output |

{ydb_sqlglot_plugin-0.2.1 → ydb_sqlglot_plugin-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ydb-sqlglot-plugin"
-version = "0.2.1"  # AUTOVERSION
+version = "0.2.3"  # AUTOVERSION
 description = "YDB dialect plugin for sqlglot"
 readme = "README.md"
 license = {text = "Apache-2.0"}

ydb_sqlglot_plugin-0.2.3/ydb_sqlglot/version.py

@@ -0,0 +1 @@
+VERSION = "0.2.3"

{ydb_sqlglot_plugin-0.2.1 → ydb_sqlglot_plugin-0.2.3}/ydb_sqlglot/ydb.py

@@ -1,5 +1,7 @@
+import inspect as _inspect
 import re
 import typing as t
+from collections import defaultdict
 
 from sqlglot import Generator, TokenType, exp, generator, parser, tokens, transforms
 from sqlglot.dialects.dialect import Dialect, NormalizationStrategy, concat_to_dpipe_sql, unit_to_var
@@ -175,6 +177,8 @@ def _wrap_udf_group_by(expression: exp.Expression) -> None:
         # Resolve each GROUP BY expression to its "effective" form
         # (expand alias references to the aliased expression).
         def _effective(e: exp.Expression) -> exp.Expression:
+            if isinstance(e, exp.Alias):
+                return e.this
             if isinstance(e, exp.Column):
                 name = e.name
                 if name in alias_to_expr:
@@ -418,6 +422,21 @@ class YdbTuple(exp.Expression):
     arg_types = {"expressions": True, "nullable": False}
 
 
+class YdbAtString(exp.Expression):
+    """YDB @@...@@ string literal."""
+    arg_types = {"this": True}
+
+
+class YdbPostfixCall(exp.Expression):
+    """YDB call of an expression result, e.g. $grep(x) or DateTime::Format(fmt)(ts)."""
+    arg_types = {"this": True, "expressions": False}
+
+
+class YdbLambdaBlock(exp.Expression):
+    """YDB lambda body with local named expressions followed by RETURN."""
+    arg_types = {"this": True, "expressions": False}
+
+
 # Container types that use Generic<T, ...> syntax in YDB
 _YDB_GENERIC_TYPES = {
     "List": exp.DataType.Type.LIST,
@@ -426,6 +445,13 @@ _YDB_GENERIC_TYPES = {
 }
 
 
+# sqlglot >= 30.0.0 changed Parser.expression() to take a pre-built instance instead of
+# (cls, **kwargs). Detect once so the YDB parser override below can support both APIs.
+_EXPRESSION_TAKES_INSTANCE = (
+    "instance" in _inspect.signature(parser.Parser.expression).parameters
+)
+
+
 def _reassemble_ctes(
     statements: t.List[t.Optional[exp.Expression]],
 ) -> t.List[t.Optional[exp.Expression]]:
@@ -449,6 +475,9 @@ def _reassemble_ctes(
         pending_names.clear()
 
     for stmt in statements:
+        if stmt is None:
+            continue
+
         if (
             isinstance(stmt, exp.Alias)
             and isinstance(stmt.args.get("alias"), exp.Identifier)
@@ -459,6 +488,10 @@ def _reassemble_ctes(
             # Unwrap Subquery — CTE.this must be Select, not Subquery
             if isinstance(inner, exp.Subquery):
                 inner = inner.this
+            if not isinstance(inner, exp.Query):
+                _flush_as_aliases()
+                result.append(stmt)
+                continue
             # Replace any $prev_cte refs inside this CTE body
             inner = _replace_param_table_refs(inner, pending_names)
             pending_aliases.append(stmt)
@@ -572,16 +605,106 @@ class YDB(Dialect):
             ),
         }
 
+        LAMBDAS = {
+            **parser.Parser.LAMBDAS,
+            TokenType.ARROW: lambda self, expressions: self._parse_ydb_lambda(expressions),
+        }
+
         STATEMENT_PARSERS = {
             **parser.Parser.STATEMENT_PARSERS,
             TokenType.DECLARE: lambda self: self._parse_ydb_declare(),
             TokenType.PARAMETER: lambda self: self._parse_ydb_named_expr(),
         }
 
+        PRIMARY_PARSERS = {
+            **parser.Parser.PRIMARY_PARSERS,
+            TokenType.STRING: lambda self, token: self._parse_ydb_string(token),
+        }
+
         def parse(self, raw_tokens, sql=None):
-            statements = super().parse(raw_tokens, sql)
+            self.reset()
+            self.sql = sql or ""
+
+            chunks: t.List[t.List[tokens.Token]] = [[]]
+            brace_depth = 0
+            total = len(raw_tokens)
+
+            for i, token in enumerate(raw_tokens):
+                if token.token_type == TokenType.L_BRACE:
+                    brace_depth += 1
+                elif token.token_type == TokenType.R_BRACE and brace_depth:
+                    brace_depth -= 1
+
+                if token.token_type == TokenType.SEMICOLON and not brace_depth:
+                    if token.comments:
+                        chunks.append([token])
+                    if i < total - 1:
+                        chunks.append([])
+                else:
+                    chunks[-1].append(token)
+
+            self._chunks = chunks
+            statements = self._parse_ydb_chunks(self.__class__._parse_statement, chunks)
             return _reassemble_ctes(statements)
 
+        def _parse_ydb_chunks(self, parse_method, chunks):
+            expressions = []
+
+            # sqlglot >= 29 has this helper; 28.x only exposes the lower-level
+            # parser state, so keep a local fallback for the supported range.
+            if hasattr(self, "_parse_batch_statements"):
+                return self._parse_batch_statements(
+                    parse_method=parse_method,
+                    sep_first_statement=False,
+                )
+
+            for chunk in chunks:
+                self._index = -1
+                self._tokens = chunk
+                self._tokens_size = len(chunk)
+                self._advance()
+
+                expressions.append(parse_method(self))
+
+                if self._index < len(self._tokens):
+                    self.raise_error("Invalid expression / Unexpected token")
+
+                self.check_errors()
+
+            return expressions
+
+        def expression(self, exp_class_or_instance, token=None, comments=None, **kwargs):
+            """Bridge sqlglot's two `Parser.expression()` calling conventions.
+
+            sqlglot < 30 expects ``expression(cls, **kwargs)`` and instantiates internally.
+            sqlglot >= 30 expects a pre-built ``expression(instance)`` and rejects kwargs.
+            Several call sites in this dialect (and a few in upstream code paths we exercise)
+            mix both styles, so normalise here before delegating.
+            """
+            if _EXPRESSION_TAKES_INSTANCE:
+                if not isinstance(exp_class_or_instance, exp.Expression):
+                    exp_class_or_instance = exp_class_or_instance(**kwargs)
+                return super().expression(
+                    exp_class_or_instance, token=token, comments=comments
+                )
+
+            if isinstance(exp_class_or_instance, exp.Expression):
+                # Old super() would attempt instance(**kwargs) -> "object is not callable".
+                instance = exp_class_or_instance
+                if token is not None:
+                    update_positions = getattr(instance, "update_positions", None)
+                    if update_positions is not None:
+                        update_positions(token)
+                if comments:
+                    instance.add_comments(comments)
+                else:
+                    self._add_comments(instance)
+                return self.validate_expression(instance)
+
+            return super().expression(
+                exp_class_or_instance, token=token, comments=comments, **kwargs
+            )
+
         def _parse_dcolon(self) -> t.Optional[exp.Expression]:
             return self._parse_function(anonymous=True) or self._parse_var(any_token=True)
 
@@ -602,9 +725,40 @@ class YDB(Dialect):
                 )
             )
 
+        def _parse_lambda_arg(self) -> t.Optional[exp.Expression]:
+            if self._match(TokenType.PARAMETER):
+                name = self._parse_var(any_token=True)
+                if not name:
+                    return None
+
+                parameter = exp.Parameter(this=name)
+                if self._match(TokenType.PLACEHOLDER):
+                    parameter.meta["optional"] = True
+                return self.expression(parameter)
+
+            arg = super()._parse_lambda_arg()
+            if arg and self._match(TokenType.PLACEHOLDER):
+                arg.meta["optional"] = True
+            return arg
+
         def _parse_ydb_declare(self) -> exp.Declare:
             items = self._parse_csv(self._parse_ydb_declareitem)
-            return self.expression(exp.Declare(expressions=items))
+            declare = self.expression(exp.Declare(expressions=items))
+
+            item_comments = {
+                comment
+                for item in items
+                if item is not None
+                for comment in (item.comments or [])
+            }
+            if item_comments and declare.comments:
+                declare.comments = [
+                    comment
+                    for comment in declare.comments
+                    if comment.strip().startswith("!") or comment not in item_comments
+                ]
+
+            return declare
 
         def _parse_ydb_declareitem(self) -> t.Optional[exp.DeclareItem]:
             if not self._match(TokenType.PARAMETER):
@@ -614,7 +768,61 @@ class YDB(Dialect):
                 return None
             self._match(TokenType.ALIAS)
             kind = self._parse_types()
-            return self.expression(exp.DeclareItem(this=name, kind=kind))
+            comments = self._prev.comments if self._prev else None
+            return self.expression(
+                exp.DeclareItem(this=name, kind=kind),
+                comments=comments,
+            )
+
+        def _parse_ydb_string(self, token: tokens.Token) -> exp.Literal:
+            literal = self.expression(exp.Literal(this=token.text, is_string=True), token)
+            if (
+                self._curr
+                and self._curr.token_type == TokenType.VAR
+                and self._curr.text.lower() == "u"
+                and token.end + 1 == self._curr.start
+            ):
+                literal.meta["ydb_string_suffix"] = self._curr.text
+                self._advance()
+            return literal
+
+        def _parse_unary(self) -> t.Optional[exp.Expression]:
+            return self._parse_ydb_postfix_calls(super()._parse_unary())
+
+        def _parse_type(
+            self,
+            parse_interval: bool = True,
+            fallback_to_identifier: bool = False,
+        ) -> t.Optional[exp.Expression]:
+            if (
+                self._curr
+                and self._next
+                and self._curr.token_type == TokenType.STRING
+                and self._next.token_type == TokenType.VAR
+                and self._next.text.lower() == "u"
+                and self._curr.end + 1 == self._next.start
+            ):
+                token = self._curr
+                suffix = self._next.text
+                self._advance(2)
+                literal = self.expression(exp.Literal(this=token.text, is_string=True), token)
+                literal.meta["ydb_string_suffix"] = suffix
+                return literal
+
+            return super()._parse_type(
+                parse_interval=parse_interval,
+                fallback_to_identifier=fallback_to_identifier,
+            )
+
+        def _parse_ydb_postfix_calls(
+            self,
+            expression: t.Optional[exp.Expression],
+        ) -> t.Optional[exp.Expression]:
+            while expression is not None and self._match(TokenType.L_PAREN):
+                args = self._parse_function_args(alias=True)
+                self._match_r_paren(expression)
+                expression = self.expression(YdbPostfixCall(this=expression, expressions=args))
+            return expression
 
         def _parse_types(self, *args, **kwargs) -> t.Optional[exp.Expression]:
             # YDB generic types use Name<...> syntax; token type varies by keyword status
@@ -643,6 +851,17 @@ class YDB(Dialect):
                         nested=True,
                     )
 
+                if name == "Struct":
+                    self._advance()  # consume 'Struct'
+                    self._advance()  # consume '<'
+                    fields = self._parse_csv(self._parse_ydb_struct_field)
+                    self._match(TokenType.GT)
+                    return exp.DataType(
+                        this=exp.DataType.Type.STRUCT,
+                        expressions=[field for field in fields if field],
+                        nested=True,
+                    )
+
                 if name == "Tuple":
                     self._advance()  # consume 'Tuple'
                     self._advance()  # consume '<'
@@ -667,13 +886,28 @@ class YDB(Dialect):
                 dtype.set("nullable", True)
             return dtype
 
+        def _parse_ydb_struct_field(self) -> t.Optional[exp.ColumnDef]:
+            if self._curr and self._curr.token_type == TokenType.GT:
+                return None
+
+            field = self._parse_id_var(any_token=True)
+            if not field:
+                return None
+
+            if not self._match(TokenType.COLON):
+                self.raise_error("Expected colon after struct field name")
+
+            kind = self._parse_types()
+            return self.expression(exp.ColumnDef(this=field, kind=kind))
+
         def _parse_table_alias(self, alias_tokens=None):
             # Prevent YDB-specific keywords from being consumed as table aliases
-            if self._curr and self._curr.text.upper() in ("FLATTEN", "ASSUME"):
+            if self._curr and self._curr.text.upper() in ("FLATTEN", "ASSUME", "VIEW"):
                 # Also check that what follows is a YDB construct, not a regular alias
                 if self._next and (
                     self._next.text.upper() in ("BY", "LIST", "DICT")
                     or self._next.token_type == TokenType.ORDER_BY
+                    or self._curr.text.upper() == "VIEW"
                 ):
                     return None
             return super()._parse_table_alias(alias_tokens=alias_tokens)
@@ -691,8 +925,70 @@ class YDB(Dialect):
                     this.set("order", self.expression(AssumeOrderBy(this=order)))
             return super()._parse_query_modifiers(this)
 
+        def _parse_group(self, skip_group_by_token: bool = False) -> t.Optional[exp.Group]:
+            # This mirrors sqlglot's Parser._parse_group control flow so YDB keeps
+            # support for GROUP BY modifiers (WITH ROLLUP, CUBE, GROUPING SETS,
+            # TOTALS). The YDB-specific change is parsing each group item through
+            # _parse_alias(..., explicit=True), enabling `GROUP BY expr AS alias`.
+            if not skip_group_by_token and not self._match(TokenType.GROUP_BY):
+                return None
+            comments = self._prev_comments
+
+            elements: dict[str, t.Any] = defaultdict(list)
+
+            if self._match(TokenType.ALL):
+                elements["all"] = True
+            elif self._match(TokenType.DISTINCT):
+                elements["all"] = False
+
+            if self._match_set(self.QUERY_MODIFIER_TOKENS, advance=False):
+                return self.expression(exp.Group(**elements), comments=comments)
+
+            def _parse_group_expression():
+                if self._match_set((TokenType.CUBE, TokenType.ROLLUP), advance=False):
+                    return None
+                return self._parse_alias(self._parse_disjunction(), explicit=True)
+
+            while True:
+                index = self._index
+
+                elements["expressions"].extend(self._parse_csv(_parse_group_expression))
+
+                before_with_index = self._index
+                with_prefix = self._match(TokenType.WITH)
+
+                if cube_or_rollup := self._parse_cube_or_rollup(with_prefix=with_prefix):
+                    key = "rollup" if isinstance(cube_or_rollup, exp.Rollup) else "cube"
+                    elements[key].append(cube_or_rollup)
+                elif grouping_sets := self._parse_grouping_sets():
+                    elements["grouping_sets"].append(grouping_sets)
+                elif self._match_text_seq("TOTALS"):
+                    elements["totals"] = True
+
+                if before_with_index <= self._index <= before_with_index + 1:
+                    self._retreat(before_with_index)
+                    break
+
+                if index == self._index:
+                    break
+
+            return self.expression(exp.Group(**elements), comments=comments)
+
         def _parse_table(self, *args, **kwargs) -> t.Optional[exp.Expression]:
+            if self._match(TokenType.L_BRACKET):
+                parts = []
+                while self._curr and not self._match(TokenType.R_BRACKET, advance=False):
+                    parts.append(self._curr.text)
+                    self._advance()
+                self._match(TokenType.R_BRACKET)
+                table = self.expression(exp.Table(this=exp.to_identifier("".join(parts))))
+                table.set("alias", self._parse_table_alias())
+                return table
+
             table = super()._parse_table(*args, **kwargs)
+            if table and self._match(TokenType.VIEW):
+                table.set("ydb_index_view", self._parse_id_var(any_token=True))
+                table.set("alias", self._parse_table_alias())
             if table and self._curr and self._curr.text.upper() == "FLATTEN":
                 self._advance()
                 kind: t.Optional[str] = None
@@ -722,8 +1018,22 @@ class YDB(Dialect):
             return exp.EQ(this=key, expression=value)
 
         def _parse_primary(self) -> t.Optional[exp.Expression]:
+            if (
+                self._curr
+                and self._curr.token_type == TokenType.PARAMETER
+                and self._next
+                and self._next.token_type == TokenType.PARAMETER
+            ):
+                return self._parse_at_raw_string()
+
             if self._match(TokenType.L_PAREN):
                 comments = self._prev_comments
+                if self._next_matching_rparen_is_arrow():
+                    expressions = self._parse_csv(self._parse_lambda_arg)
+                    self._match_r_paren()
+                    self._match(TokenType.ARROW)
+                    return self._parse_ydb_lambda(expressions)
+
                 query = self._parse_select()
 
                 if query:
@@ -757,25 +1067,88 @@ class YDB(Dialect):
                 return this
             return super()._parse_primary()
 
+        def _next_matching_rparen_is_arrow(self) -> bool:
+            depth = 1
+            # _tokens_size is not available in all supported sqlglot versions.
+            for i in range(self._index, len(self._tokens)):
+                token = self._tokens[i]
+                if token.token_type == TokenType.L_PAREN:
+                    depth += 1
+                elif token.token_type == TokenType.R_PAREN:
+                    depth -= 1
+                    if depth == 0:
+                        return (
+                            i + 1 < len(self._tokens)
+                            and self._tokens[i + 1].token_type == TokenType.ARROW
+                        )
+            return False
+
+        def _parse_at_raw_string(self) -> YdbAtString:
+            self._advance()
+            self._advance()
+
+            parts = []
+            while self._curr:
+                if (
+                    self._curr.token_type == TokenType.PARAMETER
+                    and self._next
+                    and self._next.token_type == TokenType.PARAMETER
+                ):
+                    self._advance()
+                    self._advance()
+                    break
+
+                parts.append(self._curr.text)
+                self._advance()
+
+            return self.expression(YdbAtString(this="".join(parts)))
+
         def _parse_lambda_body(self, params):
             if (
-                    self._curr.token_type != TokenType.R_PAREN
+                    self._curr is None
+                    or self._curr.token_type != TokenType.R_PAREN
+                    or self._next is None
                     or self._next.token_type != TokenType.ARROW
             ):
                 return None
             self._advance()
             self._advance()
-            self._match(TokenType.L_PAREN)
+            return self._parse_ydb_lambda(params)
+
+        def _parse_ydb_lambda(self, params):
+            has_brace = self._match(TokenType.L_BRACE)
+            assignments = []
+
+            if has_brace:
+                while self._curr and self._curr.text != "RETURN":
+                    assignment = self._parse_ydb_named_expr()
+                    if not assignment:
+                        self.raise_error("Expected lambda body expression after '->'")
+                    assignments.append(assignment)
+                    self._match(TokenType.SEMICOLON)
+
+                if not self._match_text_seq("RETURN"):
+                    self.raise_error("Expected lambda body RETURN after '->'")
+
+            body = self._parse_disjunction()
 
-            if not (self._curr.text == "RETURN"):
-                self.raise_error("Expected lambda body expression after '->'")
-            self._advance()
-            body = self._parse_conjunction()
             if not body:
                 self.raise_error("Expected lambda body expression after '->'")
 
-            self._match(TokenType.R_BRACE)
-            return exp.Lambda(this=body, expressions=params)
+            self._match(TokenType.SEMICOLON)
+            if has_brace:
+                self._match(TokenType.R_BRACE, expression=body)
+                if assignments:
+                    body = self.expression(YdbLambdaBlock(this=body, expressions=assignments))
+
+            return self.expression(exp.Lambda(this=body, expressions=params))
+
+        def _parse_in(self, this: t.Optional[exp.Expression], alias: bool = False) -> exp.In:
+            if self._match_text_seq("COMPACT"):
+                expression = self.expression(exp.In(this=this, field=self._parse_column()))
+                expression.meta["compact"] = True
+                return expression
+            return super()._parse_in(this, alias=alias)
 
     class Generator(generator.Generator):
         """
@@ -870,6 +1243,10 @@ class YDB(Dialect):
             prefix = f"{expression.db}/" if expression.db else ""
             sql = f"`{prefix}{expression.name}`"
 
+            ydb_index_view = self.sql(expression, "ydb_index_view")
+            if ydb_index_view:
+                sql += f" VIEW {ydb_index_view}"
+
             if expression.alias:
                 sql += f" AS {expression.alias}"
 
@@ -898,6 +1275,11 @@ class YDB(Dialect):
             expr = self.sql(expression, "expression")
             return f"{this}::{expr}"
 
+        def literal_sql(self, expression: exp.Literal) -> str:
+            sql = super().literal_sql(expression)
+            suffix = expression.meta.get("ydb_string_suffix")
+            return f"{sql}{suffix}" if suffix else sql
+
         def declareitem_sql(self, expression: exp.DeclareItem) -> str:
             name = self.sql(expression, "this")
             kind = self.sql(expression, "kind")
@@ -919,6 +1301,48 @@ class YDB(Dialect):
             sql = f"Tuple<{inner}>"
             return f"Optional<{sql}>" if expression.args.get("nullable") else sql
 
+        def ydbatstring_sql(self, expression: YdbAtString) -> str:
+            return f"@@{expression.this}@@"
+
+        def ydbpostfixcall_sql(self, expression: YdbPostfixCall) -> str:
+            this = self.sql(expression, "this")
+            args = self.expressions(expression, flat=True)
+            return f"{this}({args})"
+
+        def ydblambdablock_sql(self, expression: YdbLambdaBlock) -> str:
+            assignments = [self.sql(assignment) for assignment in expression.expressions]
+            statements = [*assignments, f"RETURN {self.sql(expression, 'this')}"]
+            return "{ " + "; ".join(statements) + " }"
+
+        def in_sql(self, expression: exp.In) -> str:
+            if expression.meta.get("compact"):
+                return f"{self.sql(expression, 'this')} IN COMPACT {self.sql(expression, 'field')}"
+            return super().in_sql(expression)
+
+        def maybe_comment(
+            self,
+            sql: str,
+            expression: t.Optional[exp.Expression] = None,
+            comments: t.Optional[t.List[str]] = None,
+            separated: bool = False,
+        ) -> str:
+            comments = (
+                ((expression and expression.comments) if comments is None else comments)
+                if self.comments
+                else None
+            )
+            if not comments:
+                return sql
+
+            line_directives = [comment.strip() for comment in comments if comment.strip().startswith("!")]
+            if not line_directives:
+                return super().maybe_comment(sql, expression, comments=comments, separated=separated)
+
+            remaining_comments = [comment for comment in comments if not comment.strip().startswith("!")]
+            sql = super().maybe_comment(sql, expression, comments=remaining_comments, separated=separated)
+            prefix = "\n".join(f"--{directive}" for directive in line_directives)
+            return f"{prefix}\n{sql}" if sql else prefix
+
         def alias_sql(self, expression: exp.Alias) -> str:
             alias = expression.args.get("alias")
             if alias and alias.name.startswith("$"):
@@ -1063,6 +1487,12 @@ class YDB(Dialect):
             """
             nullable = expression.args.get("nullable")
 
+            def _struct_field_sql(field: exp.Expression) -> str:
+                if isinstance(field, exp.Identifier) and field.args.get("quoted"):
+                    name = field.name.replace("'", "\\'")
+                    return f"'{name}'"
+                return self.sql(field)
+
             # YDB generic container types rendered with <> syntax and correct casing
             if expression.args.get("nested"):
                 type_value = expression.this
@@ -1080,6 +1510,15 @@ class YDB(Dialect):
                     sql = f"Tuple<{inner}>"
                     return f"Optional<{sql}>" if nullable else sql
 
+                if type_value == exp.DataType.Type.STRUCT:
+                    inner = ", ".join(
+                        f"{_struct_field_sql(col.this)}: {self.sql(col, 'kind')}"
+                        for col in expression.expressions
+                        if isinstance(col, exp.ColumnDef)
+                    )
+                    sql = f"Struct<{inner}>"
+                    return f"Optional<{sql}>" if nullable else sql
+
                 inner = ", ".join(self.sql(e) for e in expression.expressions)
                 name = {
                     exp.DataType.Type.LIST: "List",
@@ -1405,7 +1844,7 @@ class YDB(Dialect):
                                 f"decorrelated in YDB — rewrite manually using a $variable subquery"
                             )
                     continue
-                if scope.external_columns:
+                if scope.external_columns and scope.scope_type != ScopeType.CTE:
                     self.decorrelate(select, parent, scope.external_columns, next_alias_name)
                 if scope.scope_type == ScopeType.SUBQUERY:
                     self.unnest(select, parent, next_alias_name)
@@ -1459,6 +1898,9 @@ class YDB(Dialect):
             """
             Unnests a subquery by transforming it into a join
             """
+            if isinstance(select.parent, exp.CTE):
+                return
+
             if len(select.selects) > 1:
                 return
             self.ensure_select_aliases(select)
@@ -2031,14 +2473,30 @@ class YDB(Dialect):
             Returns:
                 YDB-specific SQL for lambda functions
             """
-            for ident in expression.find_all(exp.Identifier):
-                new_ident = exp.to_identifier("$" + ident.alias_or_name)
-                new_ident.set("quoted", False)
-                ident.replace(new_ident)
+            def _arg_name(arg: exp.Expression) -> str:
+                if isinstance(arg, exp.Parameter):
+                    return arg.name
+                return arg.name if hasattr(arg, "name") else self.sql(arg).lstrip("$")
 
-            args = self.expressions(expression, flat=True)
-            args = f"({args})" if len(args.split(",")) > 1 else args
-            return f"({args}) {arrow_sep} {{RETURN {self.sql(expression, 'this')}}}"
+            def _arg_sql(arg: exp.Expression) -> str:
+                name = _arg_name(arg)
+                sql = f"${name}" if name and not name.startswith("$") else self.sql(arg)
+                return f"{sql}?" if arg.meta.get("optional") else sql
+
+            def _prefix_lambda_refs(node: exp.Expression) -> exp.Expression:
+                if (
+                    isinstance(node, exp.Identifier)
+                    and not node.name.startswith("$")
+                ):
+                    return exp.Identifier(this=f"${node.name}", quoted=False)
+                return node
+
+            args = ", ".join(_arg_sql(arg) for arg in expression.expressions)
+            body = expression.this.copy().transform(_prefix_lambda_refs)
+            if isinstance(body, YdbLambdaBlock):
+                return f"({args}) {arrow_sep} {self.sql(body)}"
+            body_sql = self.sql(body)
+            return f"({args}) {arrow_sep} {body_sql if isinstance(body, exp.Paren) else f'({body_sql})'}"
 
         def _is_simple_expression(self, expr: exp.Expression) -> bool:
             """
@@ -2536,7 +2994,7 @@ class YDB(Dialect):
                     transformed.append(alias_map[gb_sql].copy())
                 elif isinstance(gb_expr, (exp.Column, exp.Identifier)):
                     # Add column AS alias so YDB resolves unambiguously.
-                    # Strip any table qualifier from the column (e.g. y.a → a).
+                    # Strip any table qualifier from the column (e.g. y.a -> a).
                     # Use the SELECT-level alias if the column is aliased there
                     # (e.g. `a_id AS _u_1` in SELECT means GROUP BY `a_id AS _u_1`).
                     column_name = gb_expr.alias_or_name
@@ -2557,6 +3015,8 @@ class YDB(Dialect):
             # Build the GROUP BY clause
             if group_by_items:
                 result = f" GROUP BY {group_by_items}"
+            elif not (rollup or cube or grouping_sets):
+                return ""
             else:
                 result = " GROUP BY"
 
@@ -2657,6 +3117,9 @@ class YDB(Dialect):
             FlattenBy: lambda self, e: self.flattenby_sql(e),
             AssumeOrderBy: lambda self, e: self.assumeorderby_sql(e),
             YdbTuple: lambda self, e: self.ydbtuple_sql(e),
+            YdbAtString: lambda self, e: self.ydbatstring_sql(e),
+            YdbPostfixCall: lambda self, e: self.ydbpostfixcall_sql(e),
+            YdbLambdaBlock: lambda self, e: self.ydblambdablock_sql(e),
             exp.Create: create_sql,
             exp.DefaultColumnConstraint: lambda self, e: "",
             exp.DateTrunc: _date_trunc_sql,

{ydb_sqlglot_plugin-0.2.1 → ydb_sqlglot_plugin-0.2.3}/ydb_sqlglot_plugin.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ydb-sqlglot-plugin
-Version: 0.2.1
+Version: 0.2.3
 Summary: YDB dialect plugin for sqlglot
 Author: YDB Team
 License: Apache-2.0
@@ -102,6 +102,31 @@ LEFT JOIN (
 
 The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subqueries.
 
+#### GROUP BY aliases
+
+YDB accepts aliases directly inside `GROUP BY` items. The generator uses this
+form for grouped columns so later clauses and decorrelated subqueries can refer
+to a stable grouping name:
+
+```sql
+-- input
+SELECT user_id, COUNT(*) FROM events GROUP BY user_id
+
+-- output
+SELECT user_id, COUNT(*) FROM `events` GROUP BY user_id AS user_id
+```
+
+If a grouped column is selected under a generated alias, the `GROUP BY` item uses
+that alias as well:
+
+```sql
+SELECT user_id AS _u_1, COUNT(*) FROM `events` GROUP BY user_id AS _u_1
+```
+
+Positional `GROUP BY` references are expanded before generation. When a
+positional reference points to a constant expression, the grouping item is
+removed because YDB rejects grouping by constants.
+
 ---
 
 ### YDB → any SQL
@@ -119,7 +144,10 @@ The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-
 | `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
 | Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
 | `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| `GROUP BY expr AS alias` | `SELECT v, COUNT(*) FROM t GROUP BY v AS v` |
 | Named expressions | `$t = (SELECT 1 AS x)` |
+| Lambda expressions | `($x, $y?) -> ($x + COALESCE($y, 0))`, `($y) -> { $p = "x"; RETURN $p \|\| $y }` |
+| `IN COMPACT` | `WHERE key IN COMPACT $values` |
 | `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
 
 Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
@@ -204,11 +232,30 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 |---|---|
 | `ARRAY(v1, v2, ...)` | `AsList(v1, v2, ...)` |
 | `ARRAY_LENGTH(x)` / `ARRAY_SIZE(x)` | `ListLength(x)` |
-| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> {RETURN cond})` |
-| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> {RETURN cond}))` |
+| `ARRAY_FILTER(arr, x -> cond)` | `ListFilter(arr, ($x) -> (cond))` |
+| `ARRAY_ANY(arr, x -> cond)` | `ListHasItems(ListFilter(arr, ($x) -> (cond)))` |
 | `ARRAY_AGG(x)` | `AGGREGATE_LIST(x)` |
 | `UNNEST(x)` | `FLATTEN BY x` |
 
+Lambda expressions are represented with sqlglot's standard `exp.Lambda` AST node.
+When a source dialect parses lambdas, the YDB generator emits YQL lambda syntax:
+
+```sql
+-- DuckDB input
+SELECT list_filter(arr, x -> x > 0) FROM t
+
+-- YDB output
+SELECT ListFilter(arr, ($x) -> ($x > 0)) FROM `t`
+```
+
+YDB input also supports documented YQL lambda forms, including optional
+arguments and block bodies with local named expressions:
+
+```sql
+($x, $y?) -> ($x + COALESCE($y, 0));
+($y) -> { $prefix = "x"; RETURN $prefix || $y; };
+```
+
 ### Conditional / math
 
 | Input | YQL output |

ydb_sqlglot_plugin-0.2.1/ydb_sqlglot/version.py

@@ -1 +0,0 @@
-VERSION = "0.2.1"