sqlglot.dialects.databricks

 29class Databricks(Spark):
 30    SAFE_DIVISION = False
 31    COPY_PARAMS_ARE_CSV = False
 32
 33    COERCES_TO = defaultdict(set, deepcopy(TypeAnnotator.COERCES_TO))
 34    for text_type in exp.DataType.TEXT_TYPES:
 35        COERCES_TO[text_type] |= {
 36            *exp.DataType.NUMERIC_TYPES,
 37            *exp.DataType.TEMPORAL_TYPES,
 38            exp.DataType.Type.BINARY,
 39            exp.DataType.Type.BOOLEAN,
 40            exp.DataType.Type.INTERVAL,
 41        }
 42
 43    class JSONPathTokenizer(jsonpath.JSONPathTokenizer):
 44        IDENTIFIERS = ["`", '"']
 45
 46    class Tokenizer(Spark.Tokenizer):
 47        KEYWORDS = {
 48            **Spark.Tokenizer.KEYWORDS,
 49            "VOID": TokenType.VOID,
 50        }
 51
 52    class Parser(Spark.Parser):
 53        LOG_DEFAULTS_TO_LN = True
 54        STRICT_CAST = True
 55        COLON_IS_VARIANT_EXTRACT = True
 56
 57        FUNCTIONS = {
 58            **Spark.Parser.FUNCTIONS,
 59            "GETDATE": exp.CurrentTimestamp.from_arg_list,
 60            "DATEADD": build_date_delta(exp.DateAdd),
 61            "DATE_ADD": build_date_delta(exp.DateAdd),
 62            "DATEDIFF": build_date_delta(exp.DateDiff),
 63            "DATE_DIFF": build_date_delta(exp.DateDiff),
 64            "NOW": exp.CurrentTimestamp.from_arg_list,
 65            "TO_DATE": build_formatted_time(exp.TsOrDsToDate, "databricks"),
 66            "UNIFORM": lambda args: exp.Uniform(
 67                this=seq_get(args, 0), expression=seq_get(args, 1), seed=seq_get(args, 2)
 68            ),
 69        }
 70
 71        NO_PAREN_FUNCTION_PARSERS = {
 72            **Spark.Parser.NO_PAREN_FUNCTION_PARSERS,
 73            "CURDATE": lambda self: self._parse_curdate(),
 74        }
 75
 76        FACTOR = {
 77            **Spark.Parser.FACTOR,
 78            TokenType.COLON: exp.JSONExtract,
 79        }
 80
 81        COLUMN_OPERATORS = {
 82            **parser.Parser.COLUMN_OPERATORS,
 83            TokenType.QDCOLON: lambda self, this, to: self.expression(
 84                exp.TryCast,
 85                this=this,
 86                to=to,
 87            ),
 88        }
 89
 90        def _parse_curdate(self) -> exp.CurrentDate:
 91            # CURDATE, an alias for CURRENT_DATE, has optional parentheses
 92            if self._match(TokenType.L_PAREN):
 93                self._match_r_paren()
 94            return self.expression(exp.CurrentDate)
 95
 96    class Generator(Spark.Generator):
 97        TABLESAMPLE_SEED_KEYWORD = "REPEATABLE"
 98        COPY_PARAMS_ARE_WRAPPED = False
 99        COPY_PARAMS_EQ_REQUIRED = True
100        JSON_PATH_SINGLE_QUOTE_ESCAPE = False
101        QUOTE_JSON_PATH = False
102        PARSE_JSON_NAME = "PARSE_JSON"
103
104        TRANSFORMS = {
105            **Spark.Generator.TRANSFORMS,
106            exp.CurrentVersion: lambda *_: "CURRENT_VERSION()",
107            exp.DateAdd: date_delta_sql("DATEADD"),
108            exp.DateDiff: date_delta_sql("DATEDIFF"),
109            exp.DatetimeAdd: lambda self, e: self.func(
110                "TIMESTAMPADD", e.unit, e.expression, e.this
111            ),
112            exp.DatetimeSub: lambda self, e: self.func(
113                "TIMESTAMPADD",
114                e.unit,
115                exp.Mul(this=e.expression, expression=exp.Literal.number(-1)),
116                e.this,
117            ),
118            exp.DatetimeTrunc: timestamptrunc_sql(),
119            exp.GroupConcat: groupconcat_sql,
120            exp.Select: transforms.preprocess(
121                [
122                    transforms.eliminate_distinct_on,
123                    transforms.unnest_to_explode,
124                    transforms.any_to_exists,
125                ]
126            ),
127            exp.JSONExtract: _jsonextract_sql,
128            exp.JSONExtractScalar: _jsonextract_sql,
129            exp.JSONPathRoot: lambda *_: "",
130            exp.ToChar: lambda self, e: (
131                self.cast_sql(exp.Cast(this=e.this, to=exp.DataType(this="STRING")))
132                if e.args.get("is_numeric")
133                else self.function_fallback_sql(e)
134            ),
135            exp.CurrentCatalog: lambda *_: "CURRENT_CATALOG()",
136        }
137
138        TRANSFORMS.pop(exp.RegexpLike)
139        TRANSFORMS.pop(exp.TryCast)
140
141        TYPE_MAPPING = {
142            **Spark.Generator.TYPE_MAPPING,
143            exp.DataType.Type.NULL: "VOID",
144        }
145
146        def columndef_sql(self, expression: exp.ColumnDef, sep: str = " ") -> str:
147            constraint = expression.find(exp.GeneratedAsIdentityColumnConstraint)
148            kind = expression.kind
149            if (
150                constraint
151                and isinstance(kind, exp.DataType)
152                and kind.this in exp.DataType.INTEGER_TYPES
153            ):
154                # only BIGINT generated identity constraints are supported
155                expression.set("kind", exp.DataType.build("bigint"))
156
157            return super().columndef_sql(expression, sep)
158
159        def generatedasidentitycolumnconstraint_sql(
160            self, expression: exp.GeneratedAsIdentityColumnConstraint
161        ) -> str:
162            expression.set("this", True)  # trigger ALWAYS in super class
163            return super().generatedasidentitycolumnconstraint_sql(expression)
164
165        def jsonpath_sql(self, expression: exp.JSONPath) -> str:
166            expression.set("escape", None)
167            return super().jsonpath_sql(expression)
168
169        def uniform_sql(self, expression: exp.Uniform) -> str:
170            gen = expression.args.get("gen")
171            seed = expression.args.get("seed")
172
173            # From Snowflake UNIFORM(min, max, gen) as RANDOM(), RANDOM(seed), or constant value -> Extract seed
174            if gen:
175                seed = gen.this
176
177            return self.func("UNIFORM", expression.this, expression.expression, seed)

43    class JSONPathTokenizer(jsonpath.JSONPathTokenizer):
44        IDENTIFIERS = ["`", '"']

46    class Tokenizer(Spark.Tokenizer):
47        KEYWORDS = {
48            **Spark.Tokenizer.KEYWORDS,
49            "VOID": TokenType.VOID,
50        }

52    class Parser(Spark.Parser):
53        LOG_DEFAULTS_TO_LN = True
54        STRICT_CAST = True
55        COLON_IS_VARIANT_EXTRACT = True
56
57        FUNCTIONS = {
58            **Spark.Parser.FUNCTIONS,
59            "GETDATE": exp.CurrentTimestamp.from_arg_list,
60            "DATEADD": build_date_delta(exp.DateAdd),
61            "DATE_ADD": build_date_delta(exp.DateAdd),
62            "DATEDIFF": build_date_delta(exp.DateDiff),
63            "DATE_DIFF": build_date_delta(exp.DateDiff),
64            "NOW": exp.CurrentTimestamp.from_arg_list,
65            "TO_DATE": build_formatted_time(exp.TsOrDsToDate, "databricks"),
66            "UNIFORM": lambda args: exp.Uniform(
67                this=seq_get(args, 0), expression=seq_get(args, 1), seed=seq_get(args, 2)
68            ),
69        }
70
71        NO_PAREN_FUNCTION_PARSERS = {
72            **Spark.Parser.NO_PAREN_FUNCTION_PARSERS,
73            "CURDATE": lambda self: self._parse_curdate(),
74        }
75
76        FACTOR = {
77            **Spark.Parser.FACTOR,
78            TokenType.COLON: exp.JSONExtract,
79        }
80
81        COLUMN_OPERATORS = {
82            **parser.Parser.COLUMN_OPERATORS,
83            TokenType.QDCOLON: lambda self, this, to: self.expression(
84                exp.TryCast,
85                this=this,
86                to=to,
87            ),
88        }
89
90        def _parse_curdate(self) -> exp.CurrentDate:
91            # CURDATE, an alias for CURRENT_DATE, has optional parentheses
92            if self._match(TokenType.L_PAREN):
93                self._match_r_paren()
94            return self.expression(exp.CurrentDate)

 96    class Generator(Spark.Generator):
 97        TABLESAMPLE_SEED_KEYWORD = "REPEATABLE"
 98        COPY_PARAMS_ARE_WRAPPED = False
 99        COPY_PARAMS_EQ_REQUIRED = True
100        JSON_PATH_SINGLE_QUOTE_ESCAPE = False
101        QUOTE_JSON_PATH = False
102        PARSE_JSON_NAME = "PARSE_JSON"
103
104        TRANSFORMS = {
105            **Spark.Generator.TRANSFORMS,
106            exp.CurrentVersion: lambda *_: "CURRENT_VERSION()",
107            exp.DateAdd: date_delta_sql("DATEADD"),
108            exp.DateDiff: date_delta_sql("DATEDIFF"),
109            exp.DatetimeAdd: lambda self, e: self.func(
110                "TIMESTAMPADD", e.unit, e.expression, e.this
111            ),
112            exp.DatetimeSub: lambda self, e: self.func(
113                "TIMESTAMPADD",
114                e.unit,
115                exp.Mul(this=e.expression, expression=exp.Literal.number(-1)),
116                e.this,
117            ),
118            exp.DatetimeTrunc: timestamptrunc_sql(),
119            exp.GroupConcat: groupconcat_sql,
120            exp.Select: transforms.preprocess(
121                [
122                    transforms.eliminate_distinct_on,
123                    transforms.unnest_to_explode,
124                    transforms.any_to_exists,
125                ]
126            ),
127            exp.JSONExtract: _jsonextract_sql,
128            exp.JSONExtractScalar: _jsonextract_sql,
129            exp.JSONPathRoot: lambda *_: "",
130            exp.ToChar: lambda self, e: (
131                self.cast_sql(exp.Cast(this=e.this, to=exp.DataType(this="STRING")))
132                if e.args.get("is_numeric")
133                else self.function_fallback_sql(e)
134            ),
135            exp.CurrentCatalog: lambda *_: "CURRENT_CATALOG()",
136        }
137
138        TRANSFORMS.pop(exp.RegexpLike)
139        TRANSFORMS.pop(exp.TryCast)
140
141        TYPE_MAPPING = {
142            **Spark.Generator.TYPE_MAPPING,
143            exp.DataType.Type.NULL: "VOID",
144        }
145
146        def columndef_sql(self, expression: exp.ColumnDef, sep: str = " ") -> str:
147            constraint = expression.find(exp.GeneratedAsIdentityColumnConstraint)
148            kind = expression.kind
149            if (
150                constraint
151                and isinstance(kind, exp.DataType)
152                and kind.this in exp.DataType.INTEGER_TYPES
153            ):
154                # only BIGINT generated identity constraints are supported
155                expression.set("kind", exp.DataType.build("bigint"))
156
157            return super().columndef_sql(expression, sep)
158
159        def generatedasidentitycolumnconstraint_sql(
160            self, expression: exp.GeneratedAsIdentityColumnConstraint
161        ) -> str:
162            expression.set("this", True)  # trigger ALWAYS in super class
163            return super().generatedasidentitycolumnconstraint_sql(expression)
164
165        def jsonpath_sql(self, expression: exp.JSONPath) -> str:
166            expression.set("escape", None)
167            return super().jsonpath_sql(expression)
168
169        def uniform_sql(self, expression: exp.Uniform) -> str:
170            gen = expression.args.get("gen")
171            seed = expression.args.get("seed")
172
173            # From Snowflake UNIFORM(min, max, gen) as RANDOM(), RANDOM(seed), or constant value -> Extract seed
174            if gen:
175                seed = gen.this
176
177            return self.func("UNIFORM", expression.this, expression.expression, seed)