sqlglot.executor.env

  1import datetime
  2import inspect
  3import re
  4import statistics
  5from functools import wraps
  6
  7from sqlglot import exp
  8from sqlglot.helper import PYTHON_VERSION
  9
 10
 11class reverse_key:
 12    def __init__(self, obj):
 13        self.obj = obj
 14
 15    def __eq__(self, other):
 16        return other.obj == self.obj
 17
 18    def __lt__(self, other):
 19        return other.obj < self.obj
 20
 21
 22def filter_nulls(func, empty_null=True):
 23    @wraps(func)
 24    def _func(values):
 25        filtered = tuple(v for v in values if v is not None)
 26        if not filtered and empty_null:
 27            return None
 28        return func(filtered)
 29
 30    return _func
 31
 32
 33def null_if_any(*required):
 34    """
 35    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
 36
 37    This also supports decoration with no arguments, e.g.:
 38
 39        @null_if_any
 40        def foo(a, b): ...
 41
 42    In which case all arguments are required.
 43    """
 44    f = None
 45    if len(required) == 1 and callable(required[0]):
 46        f = required[0]
 47        required = ()
 48
 49    def decorator(func):
 50        if required:
 51            required_indices = [
 52                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
 53            ]
 54
 55            def predicate(*args):
 56                return any(args[i] is None for i in required_indices)
 57
 58        else:
 59
 60            def predicate(*args):
 61                return any(a is None for a in args)
 62
 63        @wraps(func)
 64        def _func(*args):
 65            if predicate(*args):
 66                return None
 67            return func(*args)
 68
 69        return _func
 70
 71    if f:
 72        return decorator(f)
 73
 74    return decorator
 75
 76
 77@null_if_any("substr", "this")
 78def str_position(substr, this, position=None):
 79    position = position - 1 if position is not None else position
 80    return this.find(substr, position) + 1
 81
 82
 83@null_if_any("this")
 84def substring(this, start=None, length=None):
 85    if start is None:
 86        return this
 87    elif start == 0:
 88        return ""
 89    elif start < 0:
 90        start = len(this) + start
 91    else:
 92        start -= 1
 93
 94    end = None if length is None else start + length
 95
 96    return this[start:end]
 97
 98
 99@null_if_any
100def cast(this, to):
101    if to == exp.DataType.Type.DATE:
102        return datetime.date.fromisoformat(this)
103    if to == exp.DataType.Type.DATETIME:
104        return datetime.datetime.fromisoformat(this)
105    if to in exp.DataType.TEXT_TYPES:
106        return str(this)
107    if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}:
108        return float(this)
109    if to in exp.DataType.NUMERIC_TYPES:
110        return int(this)
111    raise NotImplementedError(f"Casting to '{to}' not implemented.")
112
113
114def ordered(this, desc, nulls_first):
115    if desc:
116        return reverse_key(this)
117    return this
118
119
120@null_if_any
121def interval(this, unit):
122    if unit == "DAY":
123        return datetime.timedelta(days=float(this))
124    raise NotImplementedError
125
126
127ENV = {
128    "exp": exp,
129    # aggs
130    "ARRAYAGG": list,
131    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
132    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
133    "MAX": filter_nulls(max),
134    "MIN": filter_nulls(min),
135    "SUM": filter_nulls(sum),
136    # scalar functions
137    "ABS": null_if_any(lambda this: abs(this)),
138    "ADD": null_if_any(lambda e, this: e + this),
139    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
140    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
141    "BITWISEAND": null_if_any(lambda this, e: this & e),
142    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
143    "BITWISEOR": null_if_any(lambda this, e: this | e),
144    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
145    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
146    "CAST": cast,
147    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
148    "CONCAT": null_if_any(lambda *args: "".join(args)),
149    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
150    "DIV": null_if_any(lambda e, this: e / this),
151    "EQ": null_if_any(lambda this, e: this == e),
152    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
153    "GT": null_if_any(lambda this, e: this > e),
154    "GTE": null_if_any(lambda this, e: this >= e),
155    "IFNULL": lambda e, alt: alt if e is None else e,
156    "IF": lambda predicate, true, false: true if predicate else false,
157    "INTDIV": null_if_any(lambda e, this: e // this),
158    "INTERVAL": interval,
159    "LIKE": null_if_any(
160        lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this))
161    ),
162    "LOWER": null_if_any(lambda arg: arg.lower()),
163    "LT": null_if_any(lambda this, e: this < e),
164    "LTE": null_if_any(lambda this, e: this <= e),
165    "MOD": null_if_any(lambda e, this: e % this),
166    "MUL": null_if_any(lambda e, this: e * this),
167    "NEQ": null_if_any(lambda this, e: this != e),
168    "ORD": null_if_any(ord),
169    "ORDERED": ordered,
170    "POW": pow,
171    "STRPOSITION": str_position,
172    "SUB": null_if_any(lambda e, this: e - this),
173    "SUBSTRING": substring,
174    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
175    "UPPER": null_if_any(lambda arg: arg.upper()),
176}