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