sqlglot.executor.context

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot.executor.env import ENV
  6
  7if t.TYPE_CHECKING:
  8    from sqlglot.executor.table import Table, TableIter
  9
 10
 11class Context:
 12    """
 13    Execution context for sql expressions.
 14
 15    Context is used to hold relevant data tables which can then be queried on with eval.
 16
 17    References to columns can either be scalar or vectors. When set_row is used, column references
 18    evaluate to scalars while set_range evaluates to vectors. This allows convenient and efficient
 19    evaluation of aggregation functions.
 20    """
 21
 22    def __init__(self, tables: t.Dict[str, Table], env: t.Optional[t.Dict] = None) -> None:
 23        """
 24        Args
 25            tables: representing the scope of the current execution context.
 26            env: dictionary of functions within the execution context.
 27        """
 28        self.tables = tables
 29        self._table: t.Optional[Table] = None
 30        self.range_readers = {name: table.range_reader for name, table in self.tables.items()}
 31        self.row_readers = {name: table.reader for name, table in tables.items()}
 32        self.env = {**ENV, **(env or {}), "scope": self.row_readers}
 33
 34    def eval(self, code):
 35        return eval(code, self.env)
 36
 37    def eval_tuple(self, codes):
 38        return tuple(self.eval(code) for code in codes)
 39
 40    @property
 41    def table(self) -> Table:
 42        if self._table is None:
 43            self._table = list(self.tables.values())[0]
 44
 45            for other in self.tables.values():
 46                if self._table.columns != other.columns:
 47                    raise Exception("Columns are different.")
 48                if len(self._table.rows) != len(other.rows):
 49                    raise Exception("Rows are different.")
 50
 51        return self._table
 52
 53    def add_columns(self, *columns: str) -> None:
 54        for table in self.tables.values():
 55            table.add_columns(*columns)
 56
 57    @property
 58    def columns(self) -> t.Tuple:
 59        return self.table.columns
 60
 61    def __iter__(self):
 62        self.env["scope"] = self.row_readers
 63        for i in range(len(self.table.rows)):
 64            for table in self.tables.values():
 65                reader = table[i]
 66            yield reader, self
 67
 68    def table_iter(self, table: str) -> TableIter:
 69        self.env["scope"] = self.row_readers
 70        return iter(self.tables[table])
 71
 72    def filter(self, condition) -> None:
 73        rows = [reader.row for reader, _ in self if self.eval(condition)]
 74
 75        for table in self.tables.values():
 76            table.rows = rows
 77
 78    def sort(self, key) -> None:
 79        def sort_key(row: t.Tuple) -> t.Tuple:
 80            self.set_row(row)
 81            return tuple((t is None, t) for t in self.eval_tuple(key))
 82
 83        self.table.rows.sort(key=sort_key)
 84
 85    def set_row(self, row: t.Tuple) -> None:
 86        for table in self.tables.values():
 87            table.reader.row = row
 88        self.env["scope"] = self.row_readers
 89
 90    def set_index(self, index: int) -> None:
 91        for table in self.tables.values():
 92            table[index]
 93        self.env["scope"] = self.row_readers
 94
 95    def set_range(self, start: int, end: int) -> None:
 96        for name in self.tables:
 97            self.range_readers[name].range = range(start, end)
 98        self.env["scope"] = self.range_readers
 99
100    def __contains__(self, table: str) -> bool:
101        return table in self.tables