altimate-code 0.5.2 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dialects/spark.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+from sqlglot.dialects.dialect import rename_func, unit_to_var, timestampdiff_sql, build_date_delta
+from sqlglot.dialects.hive import _build_with_ignore_nulls
+from sqlglot.dialects.spark2 import Spark2, temporary_storage_provider, _build_as_cast
+from sqlglot.helper import ensure_list, seq_get
+from sqlglot.transforms import (
+    ctas_with_tmp_tables_to_create_tmp_view,
+    remove_unique_constraints,
+    preprocess,
+    move_partitioned_by_to_schema_columns,
+)
+
+
+def _build_datediff(args: t.List) -> exp.Expression:
+    """
+    Although Spark docs don't mention the "unit" argument, Spark3 added support for
+    it at some point. Databricks also supports this variant (see below).
+
+    For example, in spark-sql (v3.3.1):
+    - SELECT DATEDIFF('2020-01-01', '2020-01-05') results in -4
+    - SELECT DATEDIFF(day, '2020-01-01', '2020-01-05') results in 4
+
+    See also:
+    - https://docs.databricks.com/sql/language-manual/functions/datediff3.html
+    - https://docs.databricks.com/sql/language-manual/functions/datediff.html
+    """
+    unit = None
+    this = seq_get(args, 0)
+    expression = seq_get(args, 1)
+
+    if len(args) == 3:
+        unit = exp.var(t.cast(exp.Expression, this).name)
+        this = args[2]
+
+    return exp.DateDiff(
+        this=exp.TsOrDsToDate(this=this), expression=exp.TsOrDsToDate(this=expression), unit=unit
+    )
+
+
+def _build_dateadd(args: t.List) -> exp.Expression:
+    expression = seq_get(args, 1)
+
+    if len(args) == 2:
+        # DATE_ADD(startDate, numDays INTEGER)
+        # https://docs.databricks.com/en/sql/language-manual/functions/date_add.html
+        return exp.TsOrDsAdd(
+            this=seq_get(args, 0), expression=expression, unit=exp.Literal.string("DAY")
+        )
+
+    # DATE_ADD / DATEADD / TIMESTAMPADD(unit, value integer, expr)
+    # https://docs.databricks.com/en/sql/language-manual/functions/date_add3.html
+    return exp.TimestampAdd(this=seq_get(args, 2), expression=expression, unit=seq_get(args, 0))
+
+
+def _normalize_partition(e: exp.Expression) -> exp.Expression:
+    """Normalize the expressions in PARTITION BY (<expression>, <expression>, ...)"""
+    if isinstance(e, str):
+        return exp.to_identifier(e)
+    if isinstance(e, exp.Literal):
+        return exp.to_identifier(e.name)
+    return e
+
+
+def _dateadd_sql(self: Spark.Generator, expression: exp.TsOrDsAdd | exp.TimestampAdd) -> str:
+    if not expression.unit or (
+        isinstance(expression, exp.TsOrDsAdd) and expression.text("unit").upper() == "DAY"
+    ):
+        # Coming from Hive/Spark2 DATE_ADD or roundtripping the 2-arg version of Spark3/DB
+        return self.func("DATE_ADD", expression.this, expression.expression)
+
+    this = self.func(
+        "DATE_ADD",
+        unit_to_var(expression),
+        expression.expression,
+        expression.this,
+    )
+
+    if isinstance(expression, exp.TsOrDsAdd):
+        # The 3 arg version of DATE_ADD produces a timestamp in Spark3/DB but possibly not
+        # in other dialects
+        return_type = expression.return_type
+        if not return_type.is_type(exp.DataType.Type.TIMESTAMP, exp.DataType.Type.DATETIME):
+            this = f"CAST({this} AS {return_type})"
+
+    return this
+
+
+class Spark(Spark2):
+    SUPPORTS_ORDER_BY_ALL = True
+
+    class Tokenizer(Spark2.Tokenizer):
+        STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS = False
+
+        RAW_STRINGS = [
+            (prefix + q, q)
+            for q in t.cast(t.List[str], Spark2.Tokenizer.QUOTES)
+            for prefix in ("r", "R")
+        ]
+
+    class Parser(Spark2.Parser):
+        FUNCTIONS = {
+            **Spark2.Parser.FUNCTIONS,
+            "ANY_VALUE": _build_with_ignore_nulls(exp.AnyValue),
+            "DATE_ADD": _build_dateadd,
+            "DATEADD": _build_dateadd,
+            "TIMESTAMPADD": _build_dateadd,
+            "TIMESTAMPDIFF": build_date_delta(exp.TimestampDiff),
+            "DATEDIFF": _build_datediff,
+            "DATE_DIFF": _build_datediff,
+            "TIMESTAMP_LTZ": _build_as_cast("TIMESTAMP_LTZ"),
+            "TIMESTAMP_NTZ": _build_as_cast("TIMESTAMP_NTZ"),
+            "TRY_ELEMENT_AT": lambda args: exp.Bracket(
+                this=seq_get(args, 0),
+                expressions=ensure_list(seq_get(args, 1)),
+                offset=1,
+                safe=True,
+            ),
+        }
+
+        def _parse_generated_as_identity(
+            self,
+        ) -> (
+            exp.GeneratedAsIdentityColumnConstraint
+            | exp.ComputedColumnConstraint
+            | exp.GeneratedAsRowColumnConstraint
+        ):
+            this = super()._parse_generated_as_identity()
+            if this.expression:
+                return self.expression(exp.ComputedColumnConstraint, this=this.expression)
+            return this
+
+    class Generator(Spark2.Generator):
+        SUPPORTS_TO_NUMBER = True
+        PAD_FILL_PATTERN_IS_REQUIRED = False
+        SUPPORTS_CONVERT_TIMEZONE = True
+        SUPPORTS_MEDIAN = True
+        SUPPORTS_UNIX_SECONDS = True
+
+        TYPE_MAPPING = {
+            **Spark2.Generator.TYPE_MAPPING,
+            exp.DataType.Type.MONEY: "DECIMAL(15, 4)",
+            exp.DataType.Type.SMALLMONEY: "DECIMAL(6, 4)",
+            exp.DataType.Type.UUID: "STRING",
+            exp.DataType.Type.TIMESTAMPLTZ: "TIMESTAMP_LTZ",
+            exp.DataType.Type.TIMESTAMPNTZ: "TIMESTAMP_NTZ",
+        }
+
+        TRANSFORMS = {
+            **Spark2.Generator.TRANSFORMS,
+            exp.ArrayConstructCompact: lambda self, e: self.func(
+                "ARRAY_COMPACT", self.func("ARRAY", *e.expressions)
+            ),
+            exp.Create: preprocess(
+                [
+                    remove_unique_constraints,
+                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
+                        e, temporary_storage_provider
+                    ),
+                    move_partitioned_by_to_schema_columns,
+                ]
+            ),
+            exp.EndsWith: rename_func("ENDSWITH"),
+            exp.PartitionedByProperty: lambda self,
+            e: f"PARTITIONED BY {self.wrap(self.expressions(sqls=[_normalize_partition(e) for e in e.this.expressions], skip_first=True))}",
+            exp.StartsWith: rename_func("STARTSWITH"),
+            exp.TsOrDsAdd: _dateadd_sql,
+            exp.TimestampAdd: _dateadd_sql,
+            exp.DatetimeDiff: timestampdiff_sql,
+            exp.TimestampDiff: timestampdiff_sql,
+            exp.TryCast: lambda self, e: (
+                self.trycast_sql(e) if e.args.get("safe") else self.cast_sql(e)
+            ),
+        }
+        TRANSFORMS.pop(exp.AnyValue)
+        TRANSFORMS.pop(exp.DateDiff)
+        TRANSFORMS.pop(exp.Group)
+
+        def bracket_sql(self, expression: exp.Bracket) -> str:
+            if expression.args.get("safe"):
+                key = seq_get(self.bracket_offset_expressions(expression, index_offset=1), 0)
+                return self.func("TRY_ELEMENT_AT", expression.this, key)
+
+            return super().bracket_sql(expression)
+
+        def computedcolumnconstraint_sql(self, expression: exp.ComputedColumnConstraint) -> str:
+            return f"GENERATED ALWAYS AS ({self.sql(expression, 'this')})"
+
+        def anyvalue_sql(self, expression: exp.AnyValue) -> str:
+            return self.function_fallback_sql(expression)
+
+        def datediff_sql(self, expression: exp.DateDiff) -> str:
+            end = self.sql(expression, "this")
+            start = self.sql(expression, "expression")
+
+            if expression.unit:
+                return self.func("DATEDIFF", unit_to_var(expression), start, end)
+
+            return self.func("DATEDIFF", end, start)

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dialects/spark2.py

@@ -0,0 +1,349 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp, transforms
+from sqlglot.dialects.dialect import (
+    binary_from_function,
+    build_formatted_time,
+    is_parse_json,
+    pivot_column_names,
+    rename_func,
+    trim_sql,
+    unit_to_str,
+)
+from sqlglot.dialects.hive import Hive
+from sqlglot.helper import seq_get, ensure_list
+from sqlglot.tokens import TokenType
+from sqlglot.transforms import (
+    preprocess,
+    remove_unique_constraints,
+    ctas_with_tmp_tables_to_create_tmp_view,
+    move_schema_columns_to_partitioned_by,
+)
+
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+
+    from sqlglot.optimizer.annotate_types import TypeAnnotator
+
+
+def _map_sql(self: Spark2.Generator, expression: exp.Map) -> str:
+    keys = expression.args.get("keys")
+    values = expression.args.get("values")
+
+    if not keys or not values:
+        return self.func("MAP")
+
+    return self.func("MAP_FROM_ARRAYS", keys, values)
+
+
+def _build_as_cast(to_type: str) -> t.Callable[[t.List], exp.Expression]:
+    return lambda args: exp.Cast(this=seq_get(args, 0), to=exp.DataType.build(to_type))
+
+
+def _str_to_date(self: Spark2.Generator, expression: exp.StrToDate) -> str:
+    time_format = self.format_time(expression)
+    if time_format == Hive.DATE_FORMAT:
+        return self.func("TO_DATE", expression.this)
+    return self.func("TO_DATE", expression.this, time_format)
+
+
+def _unix_to_time_sql(self: Spark2.Generator, expression: exp.UnixToTime) -> str:
+    scale = expression.args.get("scale")
+    timestamp = expression.this
+
+    if scale is None:
+        return self.sql(exp.cast(exp.func("from_unixtime", timestamp), exp.DataType.Type.TIMESTAMP))
+    if scale == exp.UnixToTime.SECONDS:
+        return self.func("TIMESTAMP_SECONDS", timestamp)
+    if scale == exp.UnixToTime.MILLIS:
+        return self.func("TIMESTAMP_MILLIS", timestamp)
+    if scale == exp.UnixToTime.MICROS:
+        return self.func("TIMESTAMP_MICROS", timestamp)
+
+    unix_seconds = exp.Div(this=timestamp, expression=exp.func("POW", 10, scale))
+    return self.func("TIMESTAMP_SECONDS", unix_seconds)
+
+
+def _unalias_pivot(expression: exp.Expression) -> exp.Expression:
+    """
+    Spark doesn't allow PIVOT aliases, so we need to remove them and possibly wrap a
+    pivoted source in a subquery with the same alias to preserve the query's semantics.
+
+    Example:
+        >>> from sqlglot import parse_one
+        >>> expr = parse_one("SELECT piv.x FROM tbl PIVOT (SUM(a) FOR b IN ('x')) piv")
+        >>> print(_unalias_pivot(expr).sql(dialect="spark"))
+        SELECT piv.x FROM (SELECT * FROM tbl PIVOT(SUM(a) FOR b IN ('x'))) AS piv
+    """
+    if isinstance(expression, exp.From) and expression.this.args.get("pivots"):
+        pivot = expression.this.args["pivots"][0]
+        if pivot.alias:
+            alias = pivot.args["alias"].pop()
+            return exp.From(
+                this=expression.this.replace(
+                    exp.select("*")
+                    .from_(expression.this.copy(), copy=False)
+                    .subquery(alias=alias, copy=False)
+                )
+            )
+
+    return expression
+
+
+def _unqualify_pivot_columns(expression: exp.Expression) -> exp.Expression:
+    """
+    Spark doesn't allow the column referenced in the PIVOT's field to be qualified,
+    so we need to unqualify it.
+
+    Example:
+        >>> from sqlglot import parse_one
+        >>> expr = parse_one("SELECT * FROM tbl PIVOT (SUM(tbl.sales) FOR tbl.quarter IN ('Q1', 'Q2'))")
+        >>> print(_unqualify_pivot_columns(expr).sql(dialect="spark"))
+        SELECT * FROM tbl PIVOT(SUM(tbl.sales) FOR quarter IN ('Q1', 'Q1'))
+    """
+    if isinstance(expression, exp.Pivot):
+        expression.set(
+            "fields", [transforms.unqualify_columns(field) for field in expression.fields]
+        )
+
+    return expression
+
+
+def temporary_storage_provider(expression: exp.Expression) -> exp.Expression:
+    # spark2, spark, Databricks require a storage provider for temporary tables
+    provider = exp.FileFormatProperty(this=exp.Literal.string("parquet"))
+    expression.args["properties"].append("expressions", provider)
+    return expression
+
+
+def _annotate_by_similar_args(
+    self: TypeAnnotator, expression: E, *args: str, target_type: exp.DataType | exp.DataType.Type
+) -> E:
+    """
+    Infers the type of the expression according to the following rules:
+    - If all args are of the same type OR any arg is of target_type, the expr is inferred as such
+    - If any arg is of UNKNOWN type and none of target_type, the expr is inferred as UNKNOWN
+    """
+    self._annotate_args(expression)
+
+    expressions: t.List[exp.Expression] = []
+    for arg in args:
+        arg_expr = expression.args.get(arg)
+        expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
+
+    last_datatype = None
+
+    has_unknown = False
+    for expr in expressions:
+        if expr.is_type(exp.DataType.Type.UNKNOWN):
+            has_unknown = True
+        elif expr.is_type(target_type):
+            has_unknown = False
+            last_datatype = target_type
+            break
+        else:
+            last_datatype = expr.type
+
+    self._set_type(expression, exp.DataType.Type.UNKNOWN if has_unknown else last_datatype)
+    return expression
+
+
+class Spark2(Hive):
+    ANNOTATORS = {
+        **Hive.ANNOTATORS,
+        exp.Substring: lambda self, e: self._annotate_by_args(e, "this"),
+        exp.Concat: lambda self, e: _annotate_by_similar_args(
+            self, e, "expressions", target_type=exp.DataType.Type.TEXT
+        ),
+        exp.Pad: lambda self, e: _annotate_by_similar_args(
+            self, e, "this", "fill_pattern", target_type=exp.DataType.Type.TEXT
+        ),
+    }
+
+    class Tokenizer(Hive.Tokenizer):
+        HEX_STRINGS = [("X'", "'"), ("x'", "'")]
+
+        KEYWORDS = {
+            **Hive.Tokenizer.KEYWORDS,
+            "TIMESTAMP": TokenType.TIMESTAMPTZ,
+        }
+
+    class Parser(Hive.Parser):
+        TRIM_PATTERN_FIRST = True
+
+        FUNCTIONS = {
+            **Hive.Parser.FUNCTIONS,
+            "AGGREGATE": exp.Reduce.from_arg_list,
+            "APPROX_PERCENTILE": exp.ApproxQuantile.from_arg_list,
+            "BOOLEAN": _build_as_cast("boolean"),
+            "DATE": _build_as_cast("date"),
+            "DATE_TRUNC": lambda args: exp.TimestampTrunc(
+                this=seq_get(args, 1), unit=exp.var(seq_get(args, 0))
+            ),
+            "DAYOFMONTH": lambda args: exp.DayOfMonth(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
+            "DAYOFWEEK": lambda args: exp.DayOfWeek(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
+            "DAYOFYEAR": lambda args: exp.DayOfYear(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
+            "DOUBLE": _build_as_cast("double"),
+            "FLOAT": _build_as_cast("float"),
+            "FROM_UTC_TIMESTAMP": lambda args, dialect: exp.AtTimeZone(
+                this=exp.cast(
+                    seq_get(args, 0) or exp.Var(this=""),
+                    exp.DataType.Type.TIMESTAMP,
+                    dialect=dialect,
+                ),
+                zone=seq_get(args, 1),
+            ),
+            "INT": _build_as_cast("int"),
+            "MAP_FROM_ARRAYS": exp.Map.from_arg_list,
+            "RLIKE": exp.RegexpLike.from_arg_list,
+            "SHIFTLEFT": binary_from_function(exp.BitwiseLeftShift),
+            "SHIFTRIGHT": binary_from_function(exp.BitwiseRightShift),
+            "STRING": _build_as_cast("string"),
+            "TIMESTAMP": _build_as_cast("timestamp"),
+            "TO_TIMESTAMP": lambda args: (
+                _build_as_cast("timestamp")(args)
+                if len(args) == 1
+                else build_formatted_time(exp.StrToTime, "spark")(args)
+            ),
+            "TO_UNIX_TIMESTAMP": exp.StrToUnix.from_arg_list,
+            "TO_UTC_TIMESTAMP": lambda args, dialect: exp.FromTimeZone(
+                this=exp.cast(
+                    seq_get(args, 0) or exp.Var(this=""),
+                    exp.DataType.Type.TIMESTAMP,
+                    dialect=dialect,
+                ),
+                zone=seq_get(args, 1),
+            ),
+            "TRUNC": lambda args: exp.DateTrunc(unit=seq_get(args, 1), this=seq_get(args, 0)),
+            "WEEKOFYEAR": lambda args: exp.WeekOfYear(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
+        }
+
+        FUNCTION_PARSERS = {
+            **Hive.Parser.FUNCTION_PARSERS,
+            "BROADCAST": lambda self: self._parse_join_hint("BROADCAST"),
+            "BROADCASTJOIN": lambda self: self._parse_join_hint("BROADCASTJOIN"),
+            "MAPJOIN": lambda self: self._parse_join_hint("MAPJOIN"),
+            "MERGE": lambda self: self._parse_join_hint("MERGE"),
+            "SHUFFLEMERGE": lambda self: self._parse_join_hint("SHUFFLEMERGE"),
+            "MERGEJOIN": lambda self: self._parse_join_hint("MERGEJOIN"),
+            "SHUFFLE_HASH": lambda self: self._parse_join_hint("SHUFFLE_HASH"),
+            "SHUFFLE_REPLICATE_NL": lambda self: self._parse_join_hint("SHUFFLE_REPLICATE_NL"),
+        }
+
+        def _parse_drop_column(self) -> t.Optional[exp.Drop | exp.Command]:
+            return self._match_text_seq("DROP", "COLUMNS") and self.expression(
+                exp.Drop, this=self._parse_schema(), kind="COLUMNS"
+            )
+
+        def _pivot_column_names(self, aggregations: t.List[exp.Expression]) -> t.List[str]:
+            if len(aggregations) == 1:
+                return []
+            return pivot_column_names(aggregations, dialect="spark")
+
+    class Generator(Hive.Generator):
+        QUERY_HINTS = True
+        NVL2_SUPPORTED = True
+        CAN_IMPLEMENT_ARRAY_ANY = True
+
+        PROPERTIES_LOCATION = {
+            **Hive.Generator.PROPERTIES_LOCATION,
+            exp.EngineProperty: exp.Properties.Location.UNSUPPORTED,
+            exp.AutoIncrementProperty: exp.Properties.Location.UNSUPPORTED,
+            exp.CharacterSetProperty: exp.Properties.Location.UNSUPPORTED,
+            exp.CollateProperty: exp.Properties.Location.UNSUPPORTED,
+        }
+
+        TRANSFORMS = {
+            **Hive.Generator.TRANSFORMS,
+            exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
+            exp.ArraySum: lambda self,
+            e: f"AGGREGATE({self.sql(e, 'this')}, 0, (acc, x) -> acc + x, acc -> acc)",
+            exp.ArrayToString: rename_func("ARRAY_JOIN"),
+            exp.AtTimeZone: lambda self, e: self.func(
+                "FROM_UTC_TIMESTAMP", e.this, e.args.get("zone")
+            ),
+            exp.BitwiseLeftShift: rename_func("SHIFTLEFT"),
+            exp.BitwiseRightShift: rename_func("SHIFTRIGHT"),
+            exp.Create: preprocess(
+                [
+                    remove_unique_constraints,
+                    lambda e: ctas_with_tmp_tables_to_create_tmp_view(
+                        e, temporary_storage_provider
+                    ),
+                    move_schema_columns_to_partitioned_by,
+                ]
+            ),
+            exp.DateFromParts: rename_func("MAKE_DATE"),
+            exp.DateTrunc: lambda self, e: self.func("TRUNC", e.this, unit_to_str(e)),
+            exp.DayOfMonth: rename_func("DAYOFMONTH"),
+            exp.DayOfWeek: rename_func("DAYOFWEEK"),
+            # (DAY_OF_WEEK(datetime) % 7) + 1 is equivalent to DAYOFWEEK_ISO(datetime)
+            exp.DayOfWeekIso: lambda self, e: f"(({self.func('DAYOFWEEK', e.this)} % 7) + 1)",
+            exp.DayOfYear: rename_func("DAYOFYEAR"),
+            exp.FileFormatProperty: lambda self, e: f"USING {e.name.upper()}",
+            exp.From: transforms.preprocess([_unalias_pivot]),
+            exp.FromTimeZone: lambda self, e: self.func(
+                "TO_UTC_TIMESTAMP", e.this, e.args.get("zone")
+            ),
+            exp.LogicalAnd: rename_func("BOOL_AND"),
+            exp.LogicalOr: rename_func("BOOL_OR"),
+            exp.Map: _map_sql,
+            exp.Pivot: transforms.preprocess([_unqualify_pivot_columns]),
+            exp.Reduce: rename_func("AGGREGATE"),
+            exp.RegexpReplace: lambda self, e: self.func(
+                "REGEXP_REPLACE",
+                e.this,
+                e.expression,
+                e.args["replacement"],
+                e.args.get("position"),
+            ),
+            exp.Select: transforms.preprocess(
+                [
+                    transforms.eliminate_qualify,
+                    transforms.eliminate_distinct_on,
+                    transforms.unnest_to_explode,
+                    transforms.any_to_exists,
+                ]
+            ),
+            exp.StrToDate: _str_to_date,
+            exp.StrToTime: lambda self, e: self.func("TO_TIMESTAMP", e.this, self.format_time(e)),
+            exp.TimestampTrunc: lambda self, e: self.func("DATE_TRUNC", unit_to_str(e), e.this),
+            exp.Trim: trim_sql,
+            exp.UnixToTime: _unix_to_time_sql,
+            exp.VariancePop: rename_func("VAR_POP"),
+            exp.WeekOfYear: rename_func("WEEKOFYEAR"),
+            exp.WithinGroup: transforms.preprocess(
+                [transforms.remove_within_group_for_percentiles]
+            ),
+        }
+        TRANSFORMS.pop(exp.ArraySort)
+        TRANSFORMS.pop(exp.ILike)
+        TRANSFORMS.pop(exp.Left)
+        TRANSFORMS.pop(exp.MonthsBetween)
+        TRANSFORMS.pop(exp.Right)
+
+        WRAP_DERIVED_VALUES = False
+        CREATE_FUNCTION_RETURN_AS = False
+
+        def struct_sql(self, expression: exp.Struct) -> str:
+            from sqlglot.generator import Generator
+
+            return Generator.struct_sql(self, expression)
+
+        def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) -> str:
+            arg = expression.this
+            is_json_extract = isinstance(
+                arg, (exp.JSONExtract, exp.JSONExtractScalar)
+            ) and not arg.args.get("variant_extract")
+
+            # We can't use a non-nested type (eg. STRING) as a schema
+            if expression.to.args.get("nested") and (is_parse_json(arg) or is_json_extract):
+                schema = f"'{self.sql(expression, 'to')}'"
+                return self.func("FROM_JSON", arg if is_json_extract else arg.this, schema)
+
+            if is_parse_json(expression):
+                return self.func("TO_JSON", arg)
+
+            return super(Hive.Generator, self).cast_sql(expression, safe_prefix=safe_prefix)