sqlglot.executor.python

  1import ast
  2import collections
  3import itertools
  4import math
  5
  6from sqlglot import exp, generator, planner, tokens
  7from sqlglot.dialects.dialect import Dialect, inline_array_sql
  8from sqlglot.errors import ExecuteError
  9from sqlglot.executor.context import Context
 10from sqlglot.executor.env import ENV
 11from sqlglot.executor.table import RowReader, Table
 12from sqlglot.helper import csv_reader, subclasses
 13
 14
 15class PythonExecutor:
 16    def __init__(self, env=None, tables=None):
 17        self.generator = Python().generator(identify=True, comments=False)
 18        self.env = {**ENV, **(env or {})}
 19        self.tables = tables or {}
 20
 21    def execute(self, plan):
 22        finished = set()
 23        queue = set(plan.leaves)
 24        contexts = {}
 25
 26        while queue:
 27            node = queue.pop()
 28            try:
 29                context = self.context(
 30                    {
 31                        name: table
 32                        for dep in node.dependencies
 33                        for name, table in contexts[dep].tables.items()
 34                    }
 35                )
 36
 37                if isinstance(node, planner.Scan):
 38                    contexts[node] = self.scan(node, context)
 39                elif isinstance(node, planner.Aggregate):
 40                    contexts[node] = self.aggregate(node, context)
 41                elif isinstance(node, planner.Join):
 42                    contexts[node] = self.join(node, context)
 43                elif isinstance(node, planner.Sort):
 44                    contexts[node] = self.sort(node, context)
 45                elif isinstance(node, planner.SetOperation):
 46                    contexts[node] = self.set_operation(node, context)
 47                else:
 48                    raise NotImplementedError
 49
 50                finished.add(node)
 51
 52                for dep in node.dependents:
 53                    if all(d in contexts for d in dep.dependencies):
 54                        queue.add(dep)
 55
 56                for dep in node.dependencies:
 57                    if all(d in finished for d in dep.dependents):
 58                        contexts.pop(dep)
 59            except Exception as e:
 60                raise ExecuteError(f"Step '{node.id}' failed: {e}") from e
 61
 62        root = plan.root
 63        return contexts[root].tables[root.name]
 64
 65    def generate(self, expression):
 66        """Convert a SQL expression into literal Python code and compile it into bytecode."""
 67        if not expression:
 68            return None
 69
 70        sql = self.generator.generate(expression)
 71        return compile(sql, sql, "eval", optimize=2)
 72
 73    def generate_tuple(self, expressions):
 74        """Convert an array of SQL expressions into tuple of Python byte code."""
 75        if not expressions:
 76            return tuple()
 77        return tuple(self.generate(expression) for expression in expressions)
 78
 79    def context(self, tables):
 80        return Context(tables, env=self.env)
 81
 82    def table(self, expressions):
 83        return Table(
 84            expression.alias_or_name if isinstance(expression, exp.Expression) else expression
 85            for expression in expressions
 86        )
 87
 88    def scan(self, step, context):
 89        source = step.source
 90
 91        if source and isinstance(source, exp.Expression):
 92            source = source.name or source.alias
 93
 94        if source is None:
 95            context, table_iter = self.static()
 96        elif source in context:
 97            if not step.projections and not step.condition:
 98                return self.context({step.name: context.tables[source]})
 99            table_iter = context.table_iter(source)
100        elif isinstance(step.source, exp.Table) and isinstance(step.source.this, exp.ReadCSV):
101            table_iter = self.scan_csv(step)
102            context = next(table_iter)
103        else:
104            context, table_iter = self.scan_table(step)
105
106        return self.context({step.name: self._project_and_filter(context, step, table_iter)})
107
108    def _project_and_filter(self, context, step, table_iter):
109        sink = self.table(step.projections if step.projections else context.columns)
110        condition = self.generate(step.condition)
111        projections = self.generate_tuple(step.projections)
112
113        for reader in table_iter:
114            if len(sink) >= step.limit:
115                break
116
117            if condition and not context.eval(condition):
118                continue
119
120            if projections:
121                sink.append(context.eval_tuple(projections))
122            else:
123                sink.append(reader.row)
124
125        return sink
126
127    def static(self):
128        return self.context({}), [RowReader(())]
129
130    def scan_table(self, step):
131        table = self.tables.find(step.source)
132        context = self.context({step.source.alias_or_name: table})
133        return context, iter(table)
134
135    def scan_csv(self, step):
136        alias = step.source.alias
137        source = step.source.this
138
139        with csv_reader(source) as reader:
140            columns = next(reader)
141            table = Table(columns)
142            context = self.context({alias: table})
143            yield context
144            types = []
145
146            for row in reader:
147                if not types:
148                    for v in row:
149                        try:
150                            types.append(type(ast.literal_eval(v)))
151                        except (ValueError, SyntaxError):
152                            types.append(str)
153                context.set_row(tuple(t(v) for t, v in zip(types, row)))
154                yield context.table.reader
155
156    def join(self, step, context):
157        source = step.name
158
159        source_table = context.tables[source]
160        source_context = self.context({source: source_table})
161        column_ranges = {source: range(0, len(source_table.columns))}
162
163        for name, join in step.joins.items():
164            table = context.tables[name]
165            start = max(r.stop for r in column_ranges.values())
166            column_ranges[name] = range(start, len(table.columns) + start)
167            join_context = self.context({name: table})
168
169            if join.get("source_key"):
170                table = self.hash_join(join, source_context, join_context)
171            else:
172                table = self.nested_loop_join(join, source_context, join_context)
173
174            source_context = self.context(
175                {
176                    name: Table(table.columns, table.rows, column_range)
177                    for name, column_range in column_ranges.items()
178                }
179            )
180            condition = self.generate(join["condition"])
181            if condition:
182                source_context.filter(condition)
183
184        if not step.condition and not step.projections:
185            return source_context
186
187        sink = self._project_and_filter(
188            source_context,
189            step,
190            (reader for reader, _ in iter(source_context)),
191        )
192
193        if step.projections:
194            return self.context({step.name: sink})
195        else:
196            return self.context(
197                {
198                    name: Table(table.columns, sink.rows, table.column_range)
199                    for name, table in source_context.tables.items()
200                }
201            )
202
203    def nested_loop_join(self, _join, source_context, join_context):
204        table = Table(source_context.columns + join_context.columns)
205
206        for reader_a, _ in source_context:
207            for reader_b, _ in join_context:
208                table.append(reader_a.row + reader_b.row)
209
210        return table
211
212    def hash_join(self, join, source_context, join_context):
213        source_key = self.generate_tuple(join["source_key"])
214        join_key = self.generate_tuple(join["join_key"])
215        left = join.get("side") == "LEFT"
216        right = join.get("side") == "RIGHT"
217
218        results = collections.defaultdict(lambda: ([], []))
219
220        for reader, ctx in source_context:
221            results[ctx.eval_tuple(source_key)][0].append(reader.row)
222        for reader, ctx in join_context:
223            results[ctx.eval_tuple(join_key)][1].append(reader.row)
224
225        table = Table(source_context.columns + join_context.columns)
226        nulls = [(None,) * len(join_context.columns if left else source_context.columns)]
227
228        for a_group, b_group in results.values():
229            if left:
230                b_group = b_group or nulls
231            elif right:
232                a_group = a_group or nulls
233
234            for a_row, b_row in itertools.product(a_group, b_group):
235                table.append(a_row + b_row)
236
237        return table
238
239    def aggregate(self, step, context):
240        group_by = self.generate_tuple(step.group.values())
241        aggregations = self.generate_tuple(step.aggregations)
242        operands = self.generate_tuple(step.operands)
243
244        if operands:
245            operand_table = Table(self.table(step.operands).columns)
246
247            for reader, ctx in context:
248                operand_table.append(ctx.eval_tuple(operands))
249
250            for i, (a, b) in enumerate(zip(context.table.rows, operand_table.rows)):
251                context.table.rows[i] = a + b
252
253            width = len(context.columns)
254            context.add_columns(*operand_table.columns)
255
256            operand_table = Table(
257                context.columns,
258                context.table.rows,
259                range(width, width + len(operand_table.columns)),
260            )
261
262            context = self.context(
263                {
264                    None: operand_table,
265                    **context.tables,
266                }
267            )
268
269        context.sort(group_by)
270
271        group = None
272        start = 0
273        end = 1
274        length = len(context.table)
275        table = self.table(list(step.group) + step.aggregations)
276
277        def add_row():
278            table.append(group + context.eval_tuple(aggregations))
279
280        if length:
281            for i in range(length):
282                context.set_index(i)
283                key = context.eval_tuple(group_by)
284                group = key if group is None else group
285                end += 1
286                if key != group:
287                    context.set_range(start, end - 2)
288                    add_row()
289                    group = key
290                    start = end - 2
291                if len(table.rows) >= step.limit:
292                    break
293                if i == length - 1:
294                    context.set_range(start, end - 1)
295                    add_row()
296        elif step.limit > 0 and not group_by:
297            context.set_range(0, 0)
298            table.append(context.eval_tuple(aggregations))
299
300        context = self.context({step.name: table, **{name: table for name in context.tables}})
301
302        if step.projections or step.condition:
303            return self.scan(step, context)
304        return context
305
306    def sort(self, step, context):
307        projections = self.generate_tuple(step.projections)
308        projection_columns = [p.alias_or_name for p in step.projections]
309        all_columns = list(context.columns) + projection_columns
310        sink = self.table(all_columns)
311        for reader, ctx in context:
312            sink.append(reader.row + ctx.eval_tuple(projections))
313
314        sort_ctx = self.context(
315            {
316                None: sink,
317                **{table: sink for table in context.tables},
318            }
319        )
320        sort_ctx.sort(self.generate_tuple(step.key))
321
322        if not math.isinf(step.limit):
323            sort_ctx.table.rows = sort_ctx.table.rows[0 : step.limit]
324
325        output = Table(
326            projection_columns,
327            rows=[r[len(context.columns) : len(all_columns)] for r in sort_ctx.table.rows],
328        )
329        return self.context({step.name: output})
330
331    def set_operation(self, step, context):
332        left = context.tables[step.left]
333        right = context.tables[step.right]
334
335        sink = self.table(left.columns)
336
337        if issubclass(step.op, exp.Intersect):
338            sink.rows = list(set(left.rows).intersection(set(right.rows)))
339        elif issubclass(step.op, exp.Except):
340            sink.rows = list(set(left.rows).difference(set(right.rows)))
341        elif issubclass(step.op, exp.Union) and step.distinct:
342            sink.rows = list(set(left.rows).union(set(right.rows)))
343        else:
344            sink.rows = left.rows + right.rows
345
346        return self.context({step.name: sink})
347
348
349def _ordered_py(self, expression):
350    this = self.sql(expression, "this")
351    desc = "True" if expression.args.get("desc") else "False"
352    nulls_first = "True" if expression.args.get("nulls_first") else "False"
353    return f"ORDERED({this}, {desc}, {nulls_first})"
354
355
356def _rename(self, e):
357    try:
358        values = list(e.args.values())
359
360        if len(values) == 1:
361            values = values[0]
362            if not isinstance(values, list):
363                return self.func(e.key, values)
364            return self.func(e.key, *values)
365
366        if isinstance(e, exp.Func) and e.is_var_len_args:
367            *head, tail = values
368            return self.func(e.key, *head, *tail)
369
370        return self.func(e.key, *values)
371    except Exception as ex:
372        raise Exception(f"Could not rename {repr(e)}") from ex
373
374
375def _case_sql(self, expression):
376    this = self.sql(expression, "this")
377    chain = self.sql(expression, "default") or "None"
378
379    for e in reversed(expression.args["ifs"]):
380        true = self.sql(e, "true")
381        condition = self.sql(e, "this")
382        condition = f"{this} = ({condition})" if this else condition
383        chain = f"{true} if {condition} else ({chain})"
384
385    return chain
386
387
388def _lambda_sql(self, e: exp.Lambda) -> str:
389    names = {e.name.lower() for e in e.expressions}
390
391    e = e.transform(
392        lambda n: exp.var(n.name)
393        if isinstance(n, exp.Identifier) and n.name.lower() in names
394        else n
395    )
396
397    return f"lambda {self.expressions(e, flat=True)}: {self.sql(e, 'this')}"
398
399
400class Python(Dialect):
401    class Tokenizer(tokens.Tokenizer):
402        STRING_ESCAPES = ["\\"]
403
404    class Generator(generator.Generator):
405        TRANSFORMS = {
406            **{klass: _rename for klass in subclasses(exp.__name__, exp.Binary)},
407            **{klass: _rename for klass in exp.ALL_FUNCTIONS},
408            exp.Case: _case_sql,
409            exp.Alias: lambda self, e: self.sql(e.this),
410            exp.Array: inline_array_sql,
411            exp.And: lambda self, e: self.binary(e, "and"),
412            exp.Between: _rename,
413            exp.Boolean: lambda self, e: "True" if e.this else "False",
414            exp.Cast: lambda self, e: f"CAST({self.sql(e.this)}, exp.DataType.Type.{e.args['to']})",
415            exp.Column: lambda self, e: f"scope[{self.sql(e, 'table') or None}][{self.sql(e.this)}]",
416            exp.Distinct: lambda self, e: f"set({self.sql(e, 'this')})",
417            exp.Extract: lambda self, e: f"EXTRACT('{e.name.lower()}', {self.sql(e, 'expression')})",
418            exp.In: lambda self, e: f"{self.sql(e, 'this')} in {{{self.expressions(e, flat=True)}}}",
419            exp.Interval: lambda self, e: f"INTERVAL({self.sql(e.this)}, '{self.sql(e.unit)}')",
420            exp.Is: lambda self, e: self.binary(e, "==")
421            if isinstance(e.this, exp.Literal)
422            else self.binary(e, "is"),
423            exp.Lambda: _lambda_sql,
424            exp.Not: lambda self, e: f"not {self.sql(e.this)}",
425            exp.Null: lambda *_: "None",
426            exp.Or: lambda self, e: self.binary(e, "or"),
427            exp.Ordered: _ordered_py,
428            exp.Star: lambda *_: "1",
429        }