sqlglot.typing.spark2

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlglot import exp
 6from sqlglot.helper import ensure_list
 7from sqlglot.typing.hive import EXPRESSION_METADATA as HIVE_EXPRESSION_METADATA
 8
 9if t.TYPE_CHECKING:
10    from sqlglot._typing import E
11    from sqlglot.optimizer.annotate_types import TypeAnnotator
12    from sqlglot.typing import ExprMetadataType
13
14
15def _annotate_by_similar_args(self: TypeAnnotator, expression: E, *arg_keys: str) -> E:
16    """
17    Type inference for CONCAT-family expressions (CONCAT, LPAD, RPAD).
18
19    - All-BINARY → BINARY (the binary overload).
20    - Otherwise, if any arg has a known, non-array, non-binary type → STRING.
21      Spark coerces scalars (dates, ints, etc.) to string when mixed with a
22      string-resolving arg. The binary exclusion preserves the binary+unknown
23      case as UNKNOWN: Spark can't disambiguate the string vs. binary overload
24      there.
25    - Else → UNKNOWN. Covers all-unknown, binary+unknown, and anything
26      involving arrays (array handling is intentionally out of scope here).
27    """
28    arg_exprs: list[exp.Expression] = []
29    for key in arg_keys:
30        arg_exprs.extend(e for e in ensure_list(expression.args.get(key)) if e)
31
32    if arg_exprs and all(e.is_type(exp.DType.BINARY) for e in arg_exprs):
33        result: exp.DataType | exp.DType = exp.DType.BINARY
34    elif any(
35        e.type is not None and not e.is_type(exp.DType.UNKNOWN, exp.DType.ARRAY, exp.DType.BINARY)
36        for e in arg_exprs
37    ):
38        result = exp.DType.TEXT
39    else:
40        result = exp.DType.UNKNOWN
41
42    self._set_type(expression, result)
43    return expression
44
45
46EXPRESSION_METADATA: ExprMetadataType = {
47    **HIVE_EXPRESSION_METADATA,
48    **{
49        expr_type: {"returns": exp.DType.DOUBLE}
50        for expr_type in {
51            exp.Atan2,
52            exp.Randn,
53        }
54    },
55    **{
56        exp_type: {"returns": exp.DType.VARCHAR}
57        for exp_type in {
58            exp.Format,
59            exp.Right,
60        }
61    },
62    **{
63        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
64        for expr_type in {
65            exp.ArrayFilter,
66            exp.Substring,
67        }
68    },
69    exp.AddMonths: {"returns": exp.DType.DATE},
70    exp.ApproxQuantile: {
71        "annotator": lambda self, e: self._annotate_by_args(
72            e, "this", array=e.args["quantile"].is_type(exp.DType.ARRAY)
73        )
74    },
75    exp.AtTimeZone: {"returns": exp.DType.TIMESTAMP},
76    exp.Concat: {"annotator": lambda self, e: _annotate_by_similar_args(self, e, "expressions")},
77    exp.NextDay: {"returns": exp.DType.DATE},
78    exp.Pad: {
79        "annotator": lambda self, e: _annotate_by_similar_args(self, e, "this", "fill_pattern")
80    },
81}