sqlglot.generators.starrocks

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, transforms
  6from sqlglot.dialects.dialect import (
  7    approx_count_distinct_sql,
  8    arrow_json_extract_sql,
  9    no_pivot_sql,
 10    rename_func,
 11    weekstart_unit_to_str,
 12    inline_array_sql,
 13    property_sql,
 14    var_map_sql,
 15)
 16from sqlglot.generator import Generator
 17from sqlglot.generators.mysql import MySQLGenerator
 18
 19
 20def _eliminate_between_in_delete(expression: exp.Expr) -> exp.Expr:
 21    """
 22    StarRocks doesn't support BETWEEN in DELETE statements, so we convert
 23    BETWEEN expressions to explicit comparisons.
 24
 25    https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/DELETE/#parameters
 26
 27    Example:
 28        >>> from sqlglot import parse_one
 29        >>> expr = parse_one("DELETE FROM t WHERE x BETWEEN 1 AND 10")
 30        >>> print(_eliminate_between_in_delete(expr).sql(dialect="starrocks"))
 31        DELETE FROM t WHERE x >= 1 AND x <= 10
 32    """
 33    if where := expression.args.get("where"):
 34        for between in where.find_all(exp.Between):
 35            between.replace(
 36                exp.and_(
 37                    exp.GTE(this=between.this.copy(), expression=between.args["low"]),
 38                    exp.LTE(this=between.this.copy(), expression=between.args["high"]),
 39                    copy=False,
 40                )
 41            )
 42    return expression
 43
 44
 45# https://docs.starrocks.io/docs/sql-reference/sql-functions/spatial-functions/st_distance_sphere/
 46def st_distance_sphere(self, expression: exp.StDistance) -> str:
 47    point1 = expression.this
 48    point2 = expression.expression
 49
 50    point1_x = self.func("ST_X", point1)
 51    point1_y = self.func("ST_Y", point1)
 52    point2_x = self.func("ST_X", point2)
 53    point2_y = self.func("ST_Y", point2)
 54
 55    return self.func("ST_Distance_Sphere", point1_x, point1_y, point2_x, point2_y)
 56
 57
 58class StarRocksGenerator(MySQLGenerator):
 59    EXCEPT_INTERSECT_SUPPORT_ALL_CLAUSE = False
 60    JSON_TYPE_REQUIRED_FOR_EXTRACTION = False
 61    VARCHAR_REQUIRES_SIZE = False
 62    PARSE_JSON_NAME: str | None = "PARSE_JSON"
 63    WITH_PROPERTIES_PREFIX = "PROPERTIES"
 64    UPDATE_STATEMENT_SUPPORTS_FROM = True
 65    INSERT_OVERWRITE = " OVERWRITE"
 66    CAST_TO_INT_ROUNDS = False
 67
 68    # StarRocks doesn't support "IS TRUE/FALSE" syntax.
 69    IS_BOOL_ALLOWED = False
 70    # StarRocks doesn't support renaming a table with a database.
 71    RENAME_TABLE_WITH_DB = False
 72
 73    ALTER_PARTITION_OPTIONS = (
 74        *MySQLGenerator.ALTER_PARTITION_OPTIONS,
 75        exp.DistributedByProperty,
 76        exp.PartitionedByProperty,
 77    )
 78
 79    CAST_MAPPING: t.ClassVar[dict[exp.DType, str]] = {}
 80
 81    TYPE_MAPPING = {
 82        **MySQLGenerator.TYPE_MAPPING,
 83        exp.DType.HLLSKETCH: "HLL",
 84        exp.DType.INT128: "LARGEINT",
 85        exp.DType.TEXT: "STRING",
 86        exp.DType.TIMESTAMP: "DATETIME",
 87        exp.DType.TIMESTAMPTZ: "DATETIME",
 88    }
 89
 90    SQL_SECURITY_VIEW_LOCATION: t.ClassVar = exp.Properties.Location.POST_SCHEMA
 91
 92    PROPERTIES_LOCATION = {
 93        **MySQLGenerator.PROPERTIES_LOCATION,
 94        exp.PrimaryKey: exp.Properties.Location.POST_SCHEMA,
 95        exp.UniqueKeyProperty: exp.Properties.Location.POST_SCHEMA,
 96        exp.RollupProperty: exp.Properties.Location.POST_SCHEMA,
 97        exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
 98        exp.Property: exp.Properties.Location.POST_WITH,
 99    }
100
101    TRANSFORMS = {
102        # StarRocks uses the native TRIM(str, chars)/LTRIM/RTRIM function form, not
103        # MySQL's TRIM(chars FROM str) syntax, so it falls back to the base generator.
104        **{
105            k: v
106            for k, v in MySQLGenerator.TRANSFORMS.items()
107            if k not in (exp.DateDiff, exp.DateTrunc, exp.Pivot, exp.ToMap, exp.Trim)
108        },
109        exp.ArgMax: rename_func("MAX_BY"),
110        exp.ArgMin: rename_func("MIN_BY"),
111        exp.Array: inline_array_sql,
112        exp.ArrayAgg: rename_func("ARRAY_AGG"),
113        # a <@ b (ArrayContainedBy) is equivalent to ARRAY_CONTAINS_ALL(b, a)
114        exp.ArrayContainedBy: lambda self, e: self.func("ARRAY_CONTAINS_ALL", e.expression, e.this),
115        exp.ArrayContainsAll: rename_func("ARRAY_CONTAINS_ALL"),
116        exp.ArrayFilter: rename_func("ARRAY_FILTER"),
117        exp.ArrayToString: rename_func("ARRAY_JOIN"),
118        exp.ApproxDistinct: approx_count_distinct_sql,
119        exp.CurrentVersion: lambda *_: "CURRENT_VERSION()",
120        exp.Delete: transforms.preprocess([_eliminate_between_in_delete]),
121        exp.Flatten: rename_func("ARRAY_FLATTEN"),
122        exp.JSONExtractScalar: arrow_json_extract_sql,
123        exp.JSONExtract: arrow_json_extract_sql,
124        # Both MAP forms (two-array MAP([keys], [values]) and variadic MAP(k1, v1, ...))
125        # generate StarRocks' variadic MAP(k1, v1, k2, v2, ...) constructor
126        exp.Map: lambda self, e: var_map_sql(self, e, "MAP"),
127        exp.Property: property_sql,
128        exp.RegexpLike: rename_func("REGEXP"),
129        # Inherited from MySQL, minus operations StarRocks supports natively
130        # (QUALIFY, FULL OUTER JOIN, SEMI/ANTI JOIN)
131        exp.Select: transforms.preprocess(
132            [
133                transforms.eliminate_distinct_on,
134                transforms.unnest_generate_date_array_using_recursive_cte,
135            ]
136        ),
137        exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
138        exp.SqlSecurityProperty: lambda self, e: f"SECURITY {self.sql(e.this)}",
139        exp.StDistance: st_distance_sphere,
140        exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
141        exp.TimestampTrunc: lambda self, e: self.func(
142            "DATE_TRUNC", weekstart_unit_to_str(self, e), e.this
143        ),
144        exp.TimeStrToDate: rename_func("TO_DATE"),
145        exp.UnixToStr: lambda self, e: self.func("FROM_UNIXTIME", e.this, self.format_time(e)),
146        exp.UnixToTime: rename_func("FROM_UNIXTIME"),
147        exp.VarMap: lambda self, e: var_map_sql(self, e, "MAP"),
148    }
149
150    # https://docs.starrocks.io/docs/sql-reference/sql-statements/keywords/#reserved-keywords
151    RESERVED_KEYWORDS = {
152        "add",
153        "all",
154        "alter",
155        "analyze",
156        "and",
157        "array",
158        "as",
159        "asc",
160        "between",
161        "bigint",
162        "bitmap",
163        "both",
164        "by",
165        "case",
166        "char",
167        "character",
168        "check",
169        "collate",
170        "column",
171        "compaction",
172        "convert",
173        "create",
174        "cross",
175        "cube",
176        "current_date",
177        "current_role",
178        "current_time",
179        "current_timestamp",
180        "current_user",
181        "database",
182        "databases",
183        "decimal",
184        "decimalv2",
185        "decimal32",
186        "decimal64",
187        "decimal128",
188        "default",
189        "deferred",
190        "delete",
191        "dense_rank",
192        "desc",
193        "describe",
194        "distinct",
195        "double",
196        "drop",
197        "dual",
198        "else",
199        "except",
200        "exists",
201        "explain",
202        "false",
203        "first_value",
204        "float",
205        "for",
206        "force",
207        "from",
208        "full",
209        "function",
210        "grant",
211        "group",
212        "grouping",
213        "grouping_id",
214        "groups",
215        "having",
216        "hll",
217        "host",
218        "if",
219        "ignore",
220        "immediate",
221        "in",
222        "index",
223        "infile",
224        "inner",
225        "insert",
226        "int",
227        "integer",
228        "intersect",
229        "into",
230        "is",
231        "join",
232        "json",
233        "key",
234        "keys",
235        "kill",
236        "lag",
237        "largeint",
238        "last_value",
239        "lateral",
240        "lead",
241        "left",
242        "like",
243        "limit",
244        "load",
245        "localtime",
246        "localtimestamp",
247        "maxvalue",
248        "minus",
249        "mod",
250        "not",
251        "ntile",
252        "null",
253        "on",
254        "or",
255        "order",
256        "outer",
257        "outfile",
258        "over",
259        "partition",
260        "percentile",
261        "primary",
262        "procedure",
263        "qualify",
264        "range",
265        "rank",
266        "read",
267        "regexp",
268        "release",
269        "rename",
270        "replace",
271        "revoke",
272        "right",
273        "rlike",
274        "row",
275        "row_number",
276        "rows",
277        "schema",
278        "schemas",
279        "select",
280        "set",
281        "set_var",
282        "show",
283        "smallint",
284        "system",
285        "table",
286        "terminated",
287        "text",
288        "then",
289        "tinyint",
290        "to",
291        "true",
292        "union",
293        "unique",
294        "unsigned",
295        "update",
296        "use",
297        "using",
298        "values",
299        "varchar",
300        "when",
301        "where",
302        "with",
303    }
304
305    def ignorenulls_sql(self, expression: exp.IgnoreNulls) -> str:
306        # IGNORE NULLS follows the first argument, e.g. LAG(x IGNORE NULLS, 1, 0)
307        # https://docs.starrocks.io/docs/sql-reference/sql-functions/Window_function/
308        func = expression.this
309        if isinstance(func, (exp.FirstValue, exp.LastValue, exp.Lag, exp.Lead)):
310            return self.func(
311                func.sql_name(),
312                Generator.ignorenulls_sql(self, exp.IgnoreNulls(this=func.this)),
313                func.args.get("offset"),
314                func.args.get("default"),
315            )
316
317        return super().ignorenulls_sql(expression)
318
319    def respectnulls_sql(self, expression: exp.RespectNulls) -> str:
320        # RESPECT NULLS is the default behavior, and StarRocks rejects the keyword
321        return self.sql(expression, "this")
322
323    def create_sql(self, expression: exp.Create) -> str:
324        # Starrocks' primary key is defined outside of the schema, so we need to move it there
325        schema = expression.this
326        if isinstance(schema, exp.Schema):
327            primary_key = schema.find(exp.PrimaryKey)
328
329            if primary_key:
330                props = expression.args.get("properties")
331
332                if not props:
333                    props = exp.Properties(expressions=[])
334                    expression.set("properties", props)
335
336                # Verify if the first one is an engine property. Is true then insert it after the engine,
337                # otherwise insert it at the beginning
338                engine = props.find(exp.EngineProperty)
339                engine_index = (engine.index or 0) if engine else -1
340                props.set("expressions", primary_key.pop(), engine_index + 1, overwrite=False)
341
342        return super().create_sql(expression)
343
344    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
345        this = expression.this
346        if isinstance(this, exp.Schema):
347            # For MVs, StarRocks needs outer parentheses.
348            create = expression.find_ancestor(exp.Create)
349
350            sql = self.expressions(this, flat=True)
351            if (create and create.kind == "VIEW") or all(
352                isinstance(col, (exp.Column, exp.Identifier)) for col in this.expressions
353            ):
354                sql = f"({sql})"
355
356            return f"PARTITION BY {sql}"
357
358        return f"PARTITION BY {self.sql(this)}"
359
360    def clusterproperty_sql(self, expression: exp.ClusterProperty) -> str:
361        """Generate StarRocks ORDER BY clause for clustering."""
362        if expression.this:
363            self.unsupported(f"Unsupported CLUSTER BY {self.sql(expression, 'this')}")
364            return ""
365        expressions = self.expressions(expression, flat=True)
366        return f"ORDER BY ({expressions})"
367
368    def refreshtriggerproperty_sql(self, expression: exp.RefreshTriggerProperty) -> str:
369        """Generate StarRocks REFRESH clause for materialized views.
370        There is a little difference of the syntax between StarRocks and Doris.
371        """
372        method = self.sql(expression, "method")
373        method = f" {method}" if method else ""
374        kind = self.sql(expression, "kind")
375        kind = f" {kind}" if kind else ""
376        starts = self.sql(expression, "starts")
377        starts = f" START ({starts})" if starts else ""
378        every = self.sql(expression, "every")
379        unit = self.sql(expression, "unit")
380        every = f" EVERY (INTERVAL {every} {unit})" if every and unit else ""
381
382        return f"REFRESH{method}{kind}{starts}{every}"
383
384    def refresh_sql(self, expression: exp.Refresh) -> str:
385        this = self.sql(expression, "this")
386        kind = "" if isinstance(expression.this, exp.Literal) else f"{expression.text('kind')} "
387        partition_start = self.sql(expression, "partition_start")
388        partition_end = self.sql(expression, "partition_end")
389        partition = (
390            f" PARTITION START ({partition_start}) END ({partition_end})" if partition_start else ""
391        )
392        force = " FORCE" if expression.args.get("force") else ""
393        mode = self.sql(expression, "mode")
394        mode = f" WITH {mode} MODE" if mode else ""
395
396        return f"REFRESH {kind}{this}{partition}{force}{mode}"
397
398    def datediff_sql(self, expression: exp.DateDiff) -> str:
399        # DATEDIFF compares only the date parts, whereas DATE_DIFF counts elapsed units
400        # https://docs.starrocks.io/docs/sql-reference/sql-functions/date-time-functions/datediff/
401        unit = expression.args.get("unit")
402        if expression.args.get("date_part_boundary") and (not unit or unit.name.upper() == "DAY"):
403            return self.func("DATEDIFF", expression.this, expression.expression)
404
405        return self.func(
406            "DATE_DIFF",
407            weekstart_unit_to_str(self, expression),
408            expression.this,
409            expression.expression,
410        )
411
412    def pivot_sql(self, expression: exp.Pivot) -> str:
413        # Only PIVOT (...) is supported, not UNPIVOT or DuckDB's PIVOT statement
414        # https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/SELECT/SELECT_PIVOT/
415        if expression.unpivot or expression.this:
416            return no_pivot_sql(self, expression)
417        return super().pivot_sql(expression)
418
419    def tomap_sql(self, expression: exp.ToMap) -> str:
420        struct = expression.this
421        if not isinstance(struct, exp.Struct):
422            return f"MAP {self.sql(struct)}"
423
424        entries = ", ".join(
425            f"{self.sql(e, 'this')}:{self.sql(e, 'expression')}" for e in struct.expressions
426        )
427        return f"MAP{{{entries}}}"
428
429    def timeslice_sql(self, expression: exp.TimeSlice) -> str:
430        interval = exp.Interval(this=expression.expression, unit=expression.unit)
431        kind = expression.text("kind").upper()
432        boundary = exp.var("CEIL" if kind == "END" else "FLOOR") if kind else None
433        return self.func("TIME_SLICE", expression.this, interval, boundary)