Edit on GitHub

# sqlglot.executor.context

View Source

      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) -> t.Iterator[t.Tuple[TableIter, Context]]:
     67        self.env["scope"] = self.row_readers
     68
     69        for reader in self.tables[table]:
     70            yield reader, self
     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 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

class Context: View Source

     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) -> t.Iterator[t.Tuple[TableIter, Context]]:
     68        self.env["scope"] = self.row_readers
     69
     70        for reader in self.tables[table]:
     71            yield reader, self
     72
     73    def filter(self, condition) -> None:
     74        rows = [reader.row for reader, _ in self if self.eval(condition)]
     75
     76        for table in self.tables.values():
     77            table.rows = rows
     78
     79    def sort(self, key) -> None:
     80        def sort_key(row: t.Tuple) -> t.Tuple:
     81            self.set_row(row)
     82            return self.eval_tuple(key)
     83
     84        self.table.rows.sort(key=sort_key)
     85
     86    def set_row(self, row: t.Tuple) -> None:
     87        for table in self.tables.values():
     88            table.reader.row = row
     89        self.env["scope"] = self.row_readers
     90
     91    def set_index(self, index: int) -> None:
     92        for table in self.tables.values():
     93            table[index]
     94        self.env["scope"] = self.row_readers
     95
     96    def set_range(self, start: int, end: int) -> None:
     97        for name in self.tables:
     98            self.range_readers[name].range = range(start, end)
     99        self.env["scope"] = self.range_readers
    100
    101    def __contains__(self, table: str) -> bool:
    102        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) View Source

    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.

def eval(self, code): View Source

    35    def eval(self, code):
    36        return eval(code, self.env)

def eval_tuple(self, codes): View Source

    38    def eval_tuple(self, codes):
    39        return tuple(self.eval(code) for code in codes)

def add_columns(self, *columns: str) -> None: View Source

    52    def add_columns(self, *columns: str) -> None:
    53        for table in self.tables.values():
    54            table.add_columns(*columns)

def table_iter( self, table: str) -> Iterator[Tuple[sqlglot.executor.table.TableIter, sqlglot.executor.context.Context]]: View Source

    67    def table_iter(self, table: str) -> t.Iterator[t.Tuple[TableIter, Context]]:
    68        self.env["scope"] = self.row_readers
    69
    70        for reader in self.tables[table]:
    71            yield reader, self

def filter(self, condition) -> None: View Source

    73    def filter(self, condition) -> None:
    74        rows = [reader.row for reader, _ in self if self.eval(condition)]
    75
    76        for table in self.tables.values():
    77            table.rows = rows

def sort(self, key) -> None: View Source

    79    def sort(self, key) -> None:
    80        def sort_key(row: t.Tuple) -> t.Tuple:
    81            self.set_row(row)
    82            return self.eval_tuple(key)
    83
    84        self.table.rows.sort(key=sort_key)

def set_row(self, row: Tuple) -> None: View Source

    86    def set_row(self, row: t.Tuple) -> None:
    87        for table in self.tables.values():
    88            table.reader.row = row
    89        self.env["scope"] = self.row_readers

def set_index(self, index: int) -> None: View Source

    91    def set_index(self, index: int) -> None:
    92        for table in self.tables.values():
    93            table[index]
    94        self.env["scope"] = self.row_readers

def set_range(self, start: int, end: int) -> None: View Source

    96    def set_range(self, start: int, end: int) -> None:
    97        for name in self.tables:
    98            self.range_readers[name].range = range(start, end)
    99        self.env["scope"] = self.range_readers