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.
table: sqlglot.executor.table.Table