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
     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

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
     46            for other in self.tables.values():
     47                if self._table.columns != other.columns:
     48                    raise Exception("Columns are different.")
     49                if len(self._table.rows) != len(other.rows):
     50                    raise Exception("Rows are different.")
     51
     52        return self._table
     53
     54    def add_columns(self, *columns: str) -> None:
     55        for table in self.tables.values():
     56            table.add_columns(*columns)
     57
     58    @property
     59    def columns(self) -> t.Tuple:
     60        return self.table.columns
     61
     62    def __iter__(self):
     63        self.env["scope"] = self.row_readers
     64        for i in range(len(self.table.rows)):
     65            for table in self.tables.values():
     66                reader = table[i]
     67            yield reader, self
     68
     69    def table_iter(self, table: str) -> TableIter:
     70        self.env["scope"] = self.row_readers
     71        return iter(self.tables[table])
     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 tuple((t is None, t) for t in 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.

tables
range_readers
row_readers
env
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)

table: sqlglot.executor.table.Table View Source

    41    @property
    42    def table(self) -> Table:
    43        if self._table is None:
    44            self._table = list(self.tables.values())[0]
    45
    46            for other in self.tables.values():
    47                if self._table.columns != other.columns:
    48                    raise Exception("Columns are different.")
    49                if len(self._table.rows) != len(other.rows):
    50                    raise Exception("Rows are different.")
    51
    52        return self._table

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

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

columns: Tuple View Source

    58    @property
    59    def columns(self) -> t.Tuple:
    60        return self.table.columns

def table_iter(self, table: str) -> sqlglot.executor.table.TableIter: View Source

    69    def table_iter(self, table: str) -> TableIter:
    70        self.env["scope"] = self.row_readers
    71        return iter(self.tables[table])

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 tuple((t is None, t) for t in 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