sqlglot.executor.context API documentation

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

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

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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: View Source

    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