sqlglot.dialects.spark

 54class Spark(Spark2):
 55    class Tokenizer(Spark2.Tokenizer):
 56        RAW_STRINGS = [
 57            (prefix + q, q)
 58            for q in t.cast(t.List[str], Spark2.Tokenizer.QUOTES)
 59            for prefix in ("r", "R")
 60        ]
 61
 62    class Parser(Spark2.Parser):
 63        FUNCTIONS = {
 64            **Spark2.Parser.FUNCTIONS,
 65            "ANY_VALUE": _build_with_ignore_nulls(exp.AnyValue),
 66            "DATEDIFF": _build_datediff,
 67            "TIMESTAMP_LTZ": _build_as_cast("TIMESTAMP_LTZ"),
 68            "TIMESTAMP_NTZ": _build_as_cast("TIMESTAMP_NTZ"),
 69            "TRY_ELEMENT_AT": lambda args: exp.Bracket(
 70                this=seq_get(args, 0), expressions=ensure_list(seq_get(args, 1)), safe=True
 71            ),
 72        }
 73
 74        def _parse_generated_as_identity(
 75            self,
 76        ) -> (
 77            exp.GeneratedAsIdentityColumnConstraint
 78            | exp.ComputedColumnConstraint
 79            | exp.GeneratedAsRowColumnConstraint
 80        ):
 81            this = super()._parse_generated_as_identity()
 82            if this.expression:
 83                return self.expression(exp.ComputedColumnConstraint, this=this.expression)
 84            return this
 85
 86    class Generator(Spark2.Generator):
 87        SUPPORTS_TO_NUMBER = True
 88
 89        TYPE_MAPPING = {
 90            **Spark2.Generator.TYPE_MAPPING,
 91            exp.DataType.Type.MONEY: "DECIMAL(15, 4)",
 92            exp.DataType.Type.SMALLMONEY: "DECIMAL(6, 4)",
 93            exp.DataType.Type.UNIQUEIDENTIFIER: "STRING",
 94            exp.DataType.Type.TIMESTAMPLTZ: "TIMESTAMP_LTZ",
 95            exp.DataType.Type.TIMESTAMPNTZ: "TIMESTAMP_NTZ",
 96        }
 97
 98        TRANSFORMS = {
 99            **Spark2.Generator.TRANSFORMS,
100            exp.ArrayConstructCompact: lambda self, e: self.func(
101                "ARRAY_COMPACT", self.func("ARRAY", *e.expressions)
102            ),
103            exp.Create: preprocess(
104                [
105                    remove_unique_constraints,
106                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
107                        e, temporary_storage_provider
108                    ),
109                    move_partitioned_by_to_schema_columns,
110                ]
111            ),
112            exp.PartitionedByProperty: lambda self,
113            e: f"PARTITIONED BY {self.wrap(self.expressions(sqls=[_normalize_partition(e) for e in e.this.expressions], skip_first=True))}",
114            exp.StartsWith: rename_func("STARTSWITH"),
115            exp.TimestampAdd: lambda self, e: self.func(
116                "DATEADD", unit_to_var(e), e.expression, e.this
117            ),
118            exp.TryCast: lambda self, e: (
119                self.trycast_sql(e) if e.args.get("safe") else self.cast_sql(e)
120            ),
121        }
122        TRANSFORMS.pop(exp.AnyValue)
123        TRANSFORMS.pop(exp.DateDiff)
124        TRANSFORMS.pop(exp.Group)
125
126        def bracket_sql(self, expression: exp.Bracket) -> str:
127            if expression.args.get("safe"):
128                key = seq_get(self.bracket_offset_expressions(expression), 0)
129                return self.func("TRY_ELEMENT_AT", expression.this, key)
130
131            return super().bracket_sql(expression)
132
133        def computedcolumnconstraint_sql(self, expression: exp.ComputedColumnConstraint) -> str:
134            return f"GENERATED ALWAYS AS ({self.sql(expression, 'this')})"
135
136        def anyvalue_sql(self, expression: exp.AnyValue) -> str:
137            return self.function_fallback_sql(expression)
138
139        def datediff_sql(self, expression: exp.DateDiff) -> str:
140            end = self.sql(expression, "this")
141            start = self.sql(expression, "expression")
142
143            if expression.unit:
144                return self.func("DATEDIFF", unit_to_var(expression), start, end)
145
146            return self.func("DATEDIFF", end, start)

55    class Tokenizer(Spark2.Tokenizer):
56        RAW_STRINGS = [
57            (prefix + q, q)
58            for q in t.cast(t.List[str], Spark2.Tokenizer.QUOTES)
59            for prefix in ("r", "R")
60        ]

62    class Parser(Spark2.Parser):
63        FUNCTIONS = {
64            **Spark2.Parser.FUNCTIONS,
65            "ANY_VALUE": _build_with_ignore_nulls(exp.AnyValue),
66            "DATEDIFF": _build_datediff,
67            "TIMESTAMP_LTZ": _build_as_cast("TIMESTAMP_LTZ"),
68            "TIMESTAMP_NTZ": _build_as_cast("TIMESTAMP_NTZ"),
69            "TRY_ELEMENT_AT": lambda args: exp.Bracket(
70                this=seq_get(args, 0), expressions=ensure_list(seq_get(args, 1)), safe=True
71            ),
72        }
73
74        def _parse_generated_as_identity(
75            self,
76        ) -> (
77            exp.GeneratedAsIdentityColumnConstraint
78            | exp.ComputedColumnConstraint
79            | exp.GeneratedAsRowColumnConstraint
80        ):
81            this = super()._parse_generated_as_identity()
82            if this.expression:
83                return self.expression(exp.ComputedColumnConstraint, this=this.expression)
84            return this

 86    class Generator(Spark2.Generator):
 87        SUPPORTS_TO_NUMBER = True
 88
 89        TYPE_MAPPING = {
 90            **Spark2.Generator.TYPE_MAPPING,
 91            exp.DataType.Type.MONEY: "DECIMAL(15, 4)",
 92            exp.DataType.Type.SMALLMONEY: "DECIMAL(6, 4)",
 93            exp.DataType.Type.UNIQUEIDENTIFIER: "STRING",
 94            exp.DataType.Type.TIMESTAMPLTZ: "TIMESTAMP_LTZ",
 95            exp.DataType.Type.TIMESTAMPNTZ: "TIMESTAMP_NTZ",
 96        }
 97
 98        TRANSFORMS = {
 99            **Spark2.Generator.TRANSFORMS,
100            exp.ArrayConstructCompact: lambda self, e: self.func(
101                "ARRAY_COMPACT", self.func("ARRAY", *e.expressions)
102            ),
103            exp.Create: preprocess(
104                [
105                    remove_unique_constraints,
106                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
107                        e, temporary_storage_provider
108                    ),
109                    move_partitioned_by_to_schema_columns,
110                ]
111            ),
112            exp.PartitionedByProperty: lambda self,
113            e: f"PARTITIONED BY {self.wrap(self.expressions(sqls=[_normalize_partition(e) for e in e.this.expressions], skip_first=True))}",
114            exp.StartsWith: rename_func("STARTSWITH"),
115            exp.TimestampAdd: lambda self, e: self.func(
116                "DATEADD", unit_to_var(e), e.expression, e.this
117            ),
118            exp.TryCast: lambda self, e: (
119                self.trycast_sql(e) if e.args.get("safe") else self.cast_sql(e)
120            ),
121        }
122        TRANSFORMS.pop(exp.AnyValue)
123        TRANSFORMS.pop(exp.DateDiff)
124        TRANSFORMS.pop(exp.Group)
125
126        def bracket_sql(self, expression: exp.Bracket) -> str:
127            if expression.args.get("safe"):
128                key = seq_get(self.bracket_offset_expressions(expression), 0)
129                return self.func("TRY_ELEMENT_AT", expression.this, key)
130
131            return super().bracket_sql(expression)
132
133        def computedcolumnconstraint_sql(self, expression: exp.ComputedColumnConstraint) -> str:
134            return f"GENERATED ALWAYS AS ({self.sql(expression, 'this')})"
135
136        def anyvalue_sql(self, expression: exp.AnyValue) -> str:
137            return self.function_fallback_sql(expression)
138
139        def datediff_sql(self, expression: exp.DateDiff) -> str:
140            end = self.sql(expression, "this")
141            start = self.sql(expression, "expression")
142
143            if expression.unit:
144                return self.func("DATEDIFF", unit_to_var(expression), start, end)
145
146            return self.func("DATEDIFF", end, start)