Edit on GitHub

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

Execution context for sql expressions.

Context is used to hold relevant data tables which can then be queried on with eval.

References to columns can either be scalar or vectors. When set_row is used, column references evaluate to scalars while set_range evaluates to vectors. This allows convenient and efficient evaluation of aggregation functions.

Context( tables: Dict[str, sqlglot.executor.table.Table], env: Optional[Dict] = None)
23    def __init__(self, tables: t.Dict[str, Table], env: t.Optional[t.Dict] = None) -> None:
24        """
25        Args
26            tables: representing the scope of the current execution context.
27            env: dictionary of functions within the execution context.
28        """
29        self.tables = tables
30        self._table: t.Optional[Table] = None
31        self.range_readers = {name: table.range_reader for name, table in self.tables.items()}
32        self.row_readers = {name: table.reader for name, table in tables.items()}
33        self.env = {**ENV, **(env or {}), "scope": self.row_readers}

Args tables: representing the scope of the current execution context. env: dictionary of functions within the execution context.

tables
range_readers
row_readers
env
def eval(self, code):
35    def eval(self, code):
36        return eval(code, self.env)
def eval_tuple(self, codes):
38    def eval_tuple(self, codes):
39        return tuple(self.eval(code) for code in codes)
def add_columns(self, *columns: str) -> None:
52    def add_columns(self, *columns: str) -> None:
53        for table in self.tables.values():
54            table.add_columns(*columns)
columns: Tuple
def table_iter(self, table: str) -> sqlglot.executor.table.TableIter:
67    def table_iter(self, table: str) -> TableIter:
68        self.env["scope"] = self.row_readers
69        return iter(self.tables[table])
def filter(self, condition) -> None:
71    def filter(self, condition) -> None:
72        rows = [reader.row for reader, _ in self if self.eval(condition)]
73
74        for table in self.tables.values():
75            table.rows = rows
def sort(self, key) -> None:
77    def sort(self, key) -> None:
78        def sort_key(row: t.Tuple) -> t.Tuple:
79            self.set_row(row)
80            return self.eval_tuple(key)
81
82        self.table.rows.sort(key=sort_key)
def set_row(self, row: Tuple) -> None:
84    def set_row(self, row: t.Tuple) -> None:
85        for table in self.tables.values():
86            table.reader.row = row
87        self.env["scope"] = self.row_readers
def set_index(self, index: int) -> None:
89    def set_index(self, index: int) -> None:
90        for table in self.tables.values():
91            table[index]
92        self.env["scope"] = self.row_readers
def set_range(self, start: int, end: int) -> None:
94    def set_range(self, start: int, end: int) -> None:
95        for name in self.tables:
96            self.range_readers[name].range = range(start, end)
97        self.env["scope"] = self.range_readers