File: test_starrocks.py

package info (click to toggle)
sqlglot 30.0.3-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 15,644 kB
  • sloc: python: 96,110; sql: 23,456; makefile: 75
file content (307 lines) | stat: -rw-r--r-- 15,629 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
from sqlglot import exp
from sqlglot.errors import UnsupportedError
from tests.dialects.test_dialect import Validator


class TestStarrocks(Validator):
    dialect = "starrocks"

    def test_starrocks(self):
        self.assertEqual(self.validate_identity("arr[1]").expressions[0], exp.Literal.number(0))

        self.validate_identity("SELECT ARRAY_JOIN([1, 3, 5, NULL], '_', 'NULL')")
        self.validate_identity("SELECT ARRAY_JOIN([1, 3, 5, NULL], '_')")
        self.validate_identity("ALTER TABLE a SWAP WITH b")
        self.validate_identity("SELECT ARRAY_AGG(a) FROM x")
        self.validate_identity("SELECT ST_POINT(10, 20)")
        self.validate_identity("SELECT ST_DISTANCE_SPHERE(10.1, 20.2, 30.3, 40.4)")
        self.validate_identity("ARRAY_FLATTEN(arr)").assert_is(exp.Flatten)

        self.validate_all(
            "SELECT * FROM t WHERE cond",
            read={
                "": "SELECT * FROM t WHERE cond IS TRUE",
                "starrocks": "SELECT * FROM t WHERE cond",
            },
        )

        self.validate_identity("CURRENT_VERSION()")

    def test_ddl(self):
        self.validate_identity("INSERT OVERWRITE my_table SELECT * FROM other_table")
        self.validate_identity("CREATE TABLE t (c INT) COMMENT 'c'")

        ddl_sqls = [
            "PARTITION BY (col1, col2)",
            "PARTITION BY DATE_TRUNC('DAY', col2), col1",
            "PARTITION BY FROM_UNIXTIME(col2)",
            "DISTRIBUTED BY HASH (col1) BUCKETS 1",
            "DISTRIBUTED BY HASH (col1)",
            "DISTRIBUTED BY RANDOM BUCKETS 1",
            "DISTRIBUTED BY RANDOM",
            "DISTRIBUTED BY HASH (col1) ORDER BY (col1)",
            "DISTRIBUTED BY HASH (col1) PROPERTIES ('replication_num'='1')",
            "PRIMARY KEY (col1) DISTRIBUTED BY HASH (col1)",
            "DUPLICATE KEY (col1, col2) DISTRIBUTED BY HASH (col1)",
            "UNIQUE KEY (col1, col2) PARTITION BY RANGE (col1) (START ('2024-01-01') END ('2024-01-31') EVERY (INTERVAL 1 DAY)) DISTRIBUTED BY HASH (col1)",
            "UNIQUE KEY (col1, col2) PARTITION BY RANGE (col1, col2) (START ('1') END ('10') EVERY (1), START ('10') END ('100') EVERY (10)) DISTRIBUTED BY HASH (col1)",
            "ORDER BY (col1, col2)",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(event_day, siteid), r2(event_day, citycode), r3(event_day))",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(col2))",
            "DISTRIBUTED BY HASH (col1) ROLLUP (`r1`(`col2`))",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(col2) FROM base_index)",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(col2) PROPERTIES ('storage_type'='column'))",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(col2) FROM base_index PROPERTIES ('k'='v'))",
            "DISTRIBUTED BY HASH (col1) ROLLUP (r1(col2) PROPERTIES ('k1'='v1', 'k2'='v2'))",
        ]

        for properties in ddl_sqls:
            with self.subTest(f"Testing create scheme: {properties}"):
                self.validate_identity(f"CREATE TABLE foo (col1 BIGINT, col2 BIGINT) {properties}")
                self.validate_identity(
                    f"CREATE TABLE foo (col1 BIGINT, col2 BIGINT) ENGINE=OLAP {properties}"
                )

        # Test the different wider DECIMAL types
        self.validate_identity(
            "CREATE TABLE foo (col0 DECIMAL(9, 1), col1 DECIMAL32(9, 1), col2 DECIMAL64(18, 10), col3 DECIMAL128(38, 10)) DISTRIBUTED BY HASH (col1) BUCKETS 1"
        )
        self.validate_identity(
            "CREATE TABLE foo (col1 LARGEINT) DISTRIBUTED BY HASH (col1) BUCKETS 1"
        )
        self.validate_identity(
            "CREATE VIEW foo (foo_col1) SECURITY NONE AS SELECT bar_col1 FROM bar"
        )

        # Test ROLLUP property
        self.validate_all(
            "CREATE TABLE foo (col1 BIGINT, col2 BIGINT) ROLLUP (r1(col1, col2), r2(col1))",
            write={
                "starrocks": "CREATE TABLE foo (col1 BIGINT, col2 BIGINT) ROLLUP (r1(col1, col2), r2(col1))",
                "spark": "CREATE TABLE foo (col1 BIGINT, col2 BIGINT)",
                "duckdb": "CREATE TABLE foo (col1 BIGINT, col2 BIGINT)",
                "postgres": "CREATE TABLE foo (col1 BIGINT, col2 BIGINT)",
            },
        )

        multi_column_cluster = exp.Cluster(
            expressions=[
                exp.column("c"),
                exp.column("d"),
            ]
        )
        self.assertEqual(multi_column_cluster.sql(dialect="starrocks"), "ORDER BY (c, d)")

        single_column_cluster = exp.Cluster(expressions=[exp.column("c")])
        self.assertEqual(single_column_cluster.sql(dialect="starrocks"), "ORDER BY (c)")

        mv_properties = [
            # partitioning in MV
            "PARTITION BY (DATE_FUNC(ts), region) REFRESH ASYNC",
            "PARTITION BY (DATE_TRUNC('DAY', ts)) REFRESH ASYNC",
            "PARTITION BY (col1, col2) REFRESH ASYNC",
            # MV: Refresh trigger property
            "REFRESH ASYNC",
            "REFRESH IMMEDIATE",
            "REFRESH DEFERRED",
            "REFRESH DEFERRED ASYNC",
            "REFRESH IMMEDIATE ASYNC",
            "REFRESH DEFERRED MANUAL",
            "REFRESH IMMEDIATE MANUAL",
            "REFRESH IMMEDIATE START ('2025-01-01 00:00:00') EVERY (INTERVAL 5 MINUTE)",
            "REFRESH IMMEDIATE ASYNC EVERY (INTERVAL 5 MINUTE)",
            "REFRESH DEFERRED START ('2025-01-01 00:00:00') EVERY (INTERVAL 5 MINUTE)",
            "REFRESH DEFERRED ASYNC EVERY (INTERVAL 5 MINUTE)",
            "REFRESH ASYNC START ('2025-01-01 00:00:00') EVERY (INTERVAL 5 MINUTE)",
            "REFRESH ASYNC EVERY (INTERVAL 5 MINUTE)",
        ]
        for properties in mv_properties:
            with self.subTest(f"Testing refresh clause: {properties}"):
                self.validate_identity(f"CREATE MATERIALIZED VIEW mv {properties} AS SELECT 1")

        # RENAME table without TO keyword
        self.validate_identity("ALTER TABLE t1 RENAME t2")

    def test_identity(self):
        self.validate_identity("SELECT CAST(`a`.`b` AS INT) FROM foo")
        self.validate_identity("SELECT APPROX_COUNT_DISTINCT(a) FROM x")
        self.validate_identity("SELECT [1, 2, 3]")
        self.validate_identity(
            """SELECT CAST(PARSE_JSON(fieldvalue) -> '00000000-0000-0000-0000-00000000' AS VARCHAR) AS `code` FROM (SELECT '{"00000000-0000-0000-0000-00000000":"code01"}') AS t(fieldvalue)"""
        )
        self.validate_identity(
            "SELECT text FROM example_table", write_sql="SELECT `text` FROM example_table"
        )

    def test_time(self):
        self.validate_identity("TIMESTAMP('2022-01-01')")
        self.validate_identity(
            "SELECT DATE_DIFF('SECOND', '2010-11-30 23:59:59', '2010-11-30 20:58:59')"
        )
        self.validate_identity(
            "SELECT DATE_DIFF('MINUTE', '2010-11-30 23:59:59', '2010-11-30 20:58:59')"
        )

    def test_regex(self):
        self.validate_all(
            "SELECT REGEXP(abc, '%foo%')",
            read={
                "mysql": "SELECT REGEXP_LIKE(abc, '%foo%')",
                "starrocks": "SELECT REGEXP(abc, '%foo%')",
            },
            write={
                "mysql": "SELECT REGEXP_LIKE(abc, '%foo%')",
            },
        )

    def test_unnest(self):
        self.validate_identity(
            "SELECT student, score, t.unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS t",
            "SELECT student, score, t.unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS t(unnest)",
        )
        self.validate_all(
            "SELECT student, score, unnest FROM tests CROSS JOIN LATERAL UNNEST(scores)",
            write={
                "spark": "SELECT student, score, unnest FROM tests LATERAL VIEW EXPLODE(scores) unnest AS unnest",
                "starrocks": "SELECT student, score, unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS unnest(unnest)",
            },
        )
        self.validate_all(
            r"""SELECT * FROM UNNEST(array['John','Jane','Jim','Jamie'], array[24,25,26,27]) AS t(name, age)""",
            write={
                "postgres": "SELECT * FROM UNNEST(ARRAY['John', 'Jane', 'Jim', 'Jamie'], ARRAY[24, 25, 26, 27]) AS t(name, age)",
                "spark": "SELECT * FROM INLINE(ARRAYS_ZIP(ARRAY('John', 'Jane', 'Jim', 'Jamie'), ARRAY(24, 25, 26, 27))) AS t(name, age)",
                "starrocks": "SELECT * FROM UNNEST(['John', 'Jane', 'Jim', 'Jamie'], [24, 25, 26, 27]) AS t(name, age)",
            },
        )

        # Use UNNEST to convert into multiple columns
        # see: https://docs.starrocks.io/docs/sql-reference/sql-functions/array-functions/unnest/
        self.validate_all(
            r"""SELECT id, t.type, t.scores FROM example_table, unnest(split(type, ";"), scores) AS t(type,scores)""",
            write={
                "postgres": "SELECT id, t.type, t.scores FROM example_table, UNNEST(SPLIT(type, ';'), scores) AS t(type, scores)",
                "spark": r"""SELECT id, t.type, t.scores FROM example_table LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';', '\\E')), scores)) t AS type, scores""",
                "databricks": r"""SELECT id, t.type, t.scores FROM example_table LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';', '\\E')), scores)) t AS type, scores""",
                "starrocks": r"""SELECT id, t.type, t.scores FROM example_table, UNNEST(SPLIT(type, ';'), scores) AS t(type, scores)""",
                "hive": UnsupportedError,
            },
        )

        self.validate_all(
            r"""SELECT id, t.type, t.scores FROM example_table_2 CROSS JOIN LATERAL unnest(split(type, ";"), scores) AS t(type,scores)""",
            write={
                "spark": r"""SELECT id, t.type, t.scores FROM example_table_2 LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';', '\\E')), scores)) t AS type, scores""",
                "starrocks": r"""SELECT id, t.type, t.scores FROM example_table_2 CROSS JOIN LATERAL UNNEST(SPLIT(type, ';'), scores) AS t(type, scores)""",
                "hive": UnsupportedError,
            },
        )

        lateral_explode_sqls = [
            "SELECT id, t.col FROM tbl, UNNEST(scores) AS t(col)",
            "SELECT id, t.col FROM tbl CROSS JOIN LATERAL UNNEST(scores) AS t(col)",
        ]

        for sql in lateral_explode_sqls:
            with self.subTest(f"Testing Starrocks roundtrip & transpilation of: {sql}"):
                self.validate_all(
                    sql,
                    write={
                        "starrocks": sql,
                        "spark": "SELECT id, t.col FROM tbl LATERAL VIEW EXPLODE(scores) t AS col",
                    },
                )

    def test_analyze(self):
        self.validate_identity("ANALYZE TABLE TBL(c1, c2) PROPERTIES ('prop1'=val1)")
        self.validate_identity("ANALYZE FULL TABLE TBL(c1, c2) PROPERTIES ('prop1'=val1)")
        self.validate_identity("ANALYZE SAMPLE TABLE TBL(c1, c2) PROPERTIES ('prop1'=val1)")
        self.validate_identity("ANALYZE TABLE TBL(c1, c2) WITH SYNC MODE PROPERTIES ('prop1'=val1)")
        self.validate_identity(
            "ANALYZE TABLE TBL(c1, c2) WITH ASYNC MODE PROPERTIES ('prop1'=val1)"
        )
        self.validate_identity(
            "ANALYZE TABLE TBL UPDATE HISTOGRAM ON c1, c2 PROPERTIES ('prop1'=val1)"
        )
        self.validate_identity(
            "ANALYZE TABLE TBL UPDATE HISTOGRAM ON c1, c2 WITH 5 BUCKETS PROPERTIES ('prop1'=val1)"
        )
        self.validate_identity(
            "ANALYZE TABLE TBL UPDATE HISTOGRAM ON c1, c2 WITH SYNC MODE WITH 5 BUCKETS PROPERTIES ('prop1'=val1)"
        )
        self.validate_identity(
            "ANALYZE TABLE TBL UPDATE HISTOGRAM ON c1, c2 WITH ASYNC MODE WITH 5 BUCKETS PROPERTIES ('prop1'=val1)"
        )

    def test_between(self):
        self.validate_all(
            "SELECT * FROM t WHERE a BETWEEN 1 AND 5",
            write={
                "starrocks": "SELECT * FROM t WHERE a BETWEEN 1 AND 5",
                "mysql": "SELECT * FROM t WHERE a BETWEEN 1 AND 5",
            },
        )
        self.validate_identity("SELECT a BETWEEN 1 AND 5 FROM t")
        self.validate_identity(
            "DELETE FROM t WHERE a BETWEEN b AND c",
            "DELETE FROM t WHERE a >= b AND a <= c",
        )
        self.validate_identity(
            "DELETE FROM t WHERE a BETWEEN 1 AND 10 AND b BETWEEN 20 AND 30 OR c BETWEEN 'x' AND 'z'",
            "DELETE FROM t WHERE a >= 1 AND a <= 10 AND b >= 20 AND b <= 30 OR c >= 'x' AND c <= 'z'",
        )

    def test_partition(self):
        # Column-based partitioning
        for cols in "col1", "col1, col2":
            with self.subTest(f"Testing PARTITION BY with {cols}"):
                self.validate_identity(
                    f"CREATE TABLE test_table (col1 INT, col2 DATE) PARTITION BY ({cols})"
                )
                self.validate_identity(
                    f"CREATE TABLE test_table (col1 INT, col2 DATE) PARTITION BY {cols}",
                    f"CREATE TABLE test_table (col1 INT, col2 DATE) PARTITION BY ({cols})",
                )

        # Expr-based partitioning
        self.validate_identity(
            "CREATE TABLE test_table (col2 DATE) PARTITION BY DATE_TRUNC('DAY', col2)"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col2 BIGINT) PARTITION BY FROM_UNIXTIME(col2, '%Y%m%d')"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col1 STRING, col2 BIGINT) PARTITION BY FROM_UNIXTIME(col2, '%Y%m%d'), col1"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col1 BIGINT, col2 DATE) PARTITION BY FROM_UNIXTIME(col2, '%Y%m%d'), DATE_TRUNC('DAY', col1)"
        )

        # LIST partitioning
        self.validate_identity(
            "CREATE TABLE test_table (col1 STRING) PARTITION BY LIST (col1) (PARTITION pLos_Angeles VALUES IN ('Los Angeles'), PARTITION pSan_Francisco VALUES IN ('San Francisco'))"
        )

        # Multi-column LIST partitioning
        self.validate_identity(
            "CREATE TABLE test_table (col1 DATE, col2 STRING) PARTITION BY LIST (col1, col2) (PARTITION p1 VALUES IN (('2022-04-01', 'LA'), ('2022-04-01', 'SF')))"
        )

        # RANGE partitioning with explicit values
        self.validate_identity(
            "CREATE TABLE test_table (col1 DATE) PARTITION BY RANGE (col1) (PARTITION p1 VALUES LESS THAN ('2020-01-31'), PARTITION p2 VALUES LESS THAN ('2020-02-29'), PARTITION p3 VALUES LESS THAN ('2020-03-31'))"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col1 STRING) PARTITION BY RANGE (STR2DATE(col1, '%Y-%m-%d')) (PARTITION p1 VALUES LESS THAN ('2021-01-01'), PARTITION p2 VALUES LESS THAN ('2021-01-02'), PARTITION p3 VALUES LESS THAN ('2021-01-03'))"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col1 DATE) PARTITION BY RANGE (col1) (PARTITION p1 VALUES LESS THAN ('2020-01-31'), PARTITION p_max VALUES LESS THAN (MAXVALUE))"
        )

        # RANGE partitioning with START/END/EVERY
        self.validate_identity(
            "CREATE TABLE test_table (col1 BIGINT) PARTITION BY RANGE (col1) (START ('1') END ('10') EVERY (1), START ('10') END ('100') EVERY (10))"
        )
        self.validate_identity(
            "CREATE TABLE test_table (col1 DATE) PARTITION BY RANGE (col1) (START ('2019-01-01') END ('2021-01-01') EVERY (INTERVAL 1 YEAR), START ('2021-01-01') END ('2021-05-01') EVERY (INTERVAL 1 MONTH), START ('2021-05-01') END ('2021-05-04') EVERY (INTERVAL 1 DAY))"
        )