sqlglot.executor.python

  1import collections
  2import itertools
  3import math
  4
  5from sqlglot import exp, planner, tokens
  6from sqlglot.dialects.dialect import Dialect
  7from sqlglot.errors import ExecuteError
  8from sqlglot.executor.context import Context
  9from sqlglot.executor.env import ENV
 10from sqlglot.executor.table import RowReader, Table
 11from sqlglot.generators.python import PythonGenerator
 12from sqlglot.optimizer.scope import build_scope
 13
 14SUBQUERY_NODES = (exp.Subquery, exp.Exists, exp.All, exp.Any)
 15
 16
 17class PythonExecutor:
 18    def __init__(self, env=None, tables=None):
 19        self.generator = Python().generator(identify=True, comments=False)
 20        self.env = {**ENV, **(env or {})}
 21        self.tables = tables or {}
 22        self._subquery_plans = {}
 23        self._plan_names_by_sql = {}
 24        self._ctes = None
 25        self._outer_scope = None
 26        self.env.update(
 27            SUBQUERY_COMPARISON=self._subquery_comparison,
 28            SUBQUERY_EXISTS=self._subquery_exists,
 29            SUBQUERY_SCALAR=self._subquery_scalar,
 30        )
 31
 32    def execute(self, plan, outer_scope=None):
 33        ctes, scope = self._ctes, self._outer_scope
 34        self._ctes, self._outer_scope = plan.ctes, outer_scope
 35        try:
 36            return self._execute(plan)
 37        finally:
 38            self._ctes, self._outer_scope = ctes, scope
 39
 40    def _execute(self, plan):
 41        finished = set()
 42        queue = set(plan.leaves)
 43        contexts = {}
 44
 45        while queue:
 46            node = queue.pop()
 47            try:
 48                context = self.context(
 49                    {
 50                        name: table
 51                        for dep in node.dependencies
 52                        for name, table in contexts[dep].tables.items()
 53                    }
 54                )
 55
 56                if isinstance(node, planner.Scan):
 57                    contexts[node] = self.scan(node, context)
 58                elif isinstance(node, planner.Aggregate):
 59                    contexts[node] = self.aggregate(node, context)
 60                elif isinstance(node, planner.Join):
 61                    contexts[node] = self.join(node, context)
 62                elif isinstance(node, planner.Sort):
 63                    contexts[node] = self.sort(node, context)
 64                elif isinstance(node, planner.SetOperation):
 65                    contexts[node] = self.set_operation(node, context)
 66                else:
 67                    raise NotImplementedError
 68
 69                if node.offset:
 70                    table = contexts[node].tables[node.name]
 71                    table.rows = table.rows[node.offset :]
 72
 73                finished.add(node)
 74
 75                for dep in node.dependents:
 76                    if all(d in contexts for d in dep.dependencies):
 77                        queue.add(dep)
 78
 79                for dep in node.dependencies:
 80                    if all(d in finished for d in dep.dependents):
 81                        contexts.pop(dep)
 82            except Exception as e:
 83                raise ExecuteError(f"Step '{node.id}' failed: {e}") from e
 84
 85        root = plan.root
 86        return contexts[root].tables[root.name]
 87
 88    def generate(self, expression):
 89        """Convert a SQL expression into literal Python code and compile it into bytecode."""
 90        if not expression:
 91            return None
 92
 93        expression = self._replace_subqueries(expression)
 94        sql = self.generator.generate(expression)
 95        return compile(sql, sql, "eval", optimize=2)
 96
 97    def _replace_subqueries(self, expression):
 98        if not expression.find(*SUBQUERY_NODES):
 99            return expression
100
101        expression = expression.copy()
102
103        while True:
104            subquery = expression.find(*SUBQUERY_NODES)
105
106            if subquery is None:
107                return expression
108
109            target, replacement = self._compile_subquery(subquery)
110
111            if target is expression:
112                expression = replacement
113            else:
114                target.replace(replacement)
115
116    def _compile_subquery(self, subquery):
117        query = subquery.this.unnest()
118        scope = build_scope(query)
119
120        outer_columns = list(scope.external_columns if scope else [])
121
122        plan = self._register_subquery(query)
123        parent = subquery.parent
124
125        if isinstance(subquery, exp.Exists):
126            return subquery, exp.func(
127                "SUBQUERY_EXISTS", plan, exp.var("scope"), *outer_columns, copy=False
128            )
129
130        if len(query.selects) != 1:
131            raise ExecuteError(
132                f"Subquery used as an expression returned {len(query.selects)} columns"
133            )
134
135        if isinstance(subquery, (exp.All, exp.Any)):
136            return self._compile_quantified(parent, subquery.key.upper(), plan, outer_columns)
137
138        if isinstance(parent, exp.In) and subquery is parent.args.get("query"):
139            return self._compile_quantified(parent, "ANY", plan, outer_columns, op="EQ")
140
141        return subquery, exp.func(
142            "SUBQUERY_SCALAR", plan, exp.var("scope"), *outer_columns, copy=False
143        )
144
145    def _compile_quantified(self, comparison, quantifier, plan, outer_columns, op=None):
146        if not isinstance(comparison, (exp.Binary, exp.In)):
147            raise ExecuteError(f"Unsupported {quantifier} subquery: expected a comparison")
148
149        return comparison, exp.func(
150            "SUBQUERY_COMPARISON",
151            comparison.this,
152            plan,
153            exp.var("scope"),
154            exp.Literal.string(op or comparison.key.upper()),
155            exp.Literal.string(quantifier),
156            *outer_columns,
157            copy=False,
158        )
159
160    def _register_subquery(self, query):
161        if self._ctes is not None and not query.args.get("with_"):
162            query.set("with_", self._ctes)
163
164        sql = query.sql()
165        name = self._plan_names_by_sql.get(sql)
166
167        if name is None:
168            name = self._plan_names_by_sql[sql] = f"_sq_{len(self._subquery_plans)}"
169            self._subquery_plans[name] = (planner.Plan(query), {})
170
171        return exp.Literal.string(name)
172
173    def _subquery_table(self, plan_name, scope, args):
174        plan, cache = self._subquery_plans[plan_name]
175
176        try:
177            return cache[args]
178        except KeyError:
179            pass
180        except TypeError:  # an unhashable correlated value can't be memoized
181            cache = None
182
183        table = self.execute(plan, scope)
184
185        if cache is not None:
186            cache[args] = table
187
188        return table
189
190    def _subquery_exists(self, plan_name, scope, *args):
191        return bool(self._subquery_table(plan_name, scope, args).rows)
192
193    def _subquery_scalar(self, plan_name, scope, *args):
194        rows = self._subquery_table(plan_name, scope, args).rows
195
196        if len(rows) > 1:
197            raise ExecuteError("More than one row returned by a subquery used as an expression")
198
199        return rows[0][0] if rows else None
200
201    def _subquery_comparison(self, value, plan_name, scope, op, quantifier, *args):
202        compare = self.env[op]
203        is_any = quantifier == "ANY"
204        saw_null = False
205
206        for row in self._subquery_table(plan_name, scope, args).rows:
207            result = compare(value, row[0])
208
209            if result is None:
210                saw_null = True
211            elif bool(result) is is_any:
212                return is_any
213
214        return None if saw_null else not is_any
215
216    def generate_tuple(self, expressions):
217        """Convert an array of SQL expressions into tuple of Python byte code."""
218        if not expressions:
219            return tuple()
220        return tuple(self.generate(expression) for expression in expressions)
221
222    def context(self, tables):
223        return Context(tables, env=self.env, outer=self._outer_scope)
224
225    def table(self, expressions):
226        return Table(
227            expression.alias_or_name if isinstance(expression, exp.Expr) else expression
228            for expression in expressions
229        )
230
231    def scan(self, step, context):
232        source = step.source
233
234        if source and isinstance(source, exp.Expr):
235            source = source.name or source.alias
236
237        if source is None:
238            context, table_iter = self.static()
239        elif source in context:
240            if not step.projections and not step.condition:
241                return self.context({step.name: context.tables[source]})
242            table_iter = context.table_iter(source)
243        else:
244            context, table_iter = self.scan_table(step)
245
246        return self.context({step.name: self._project_and_filter(context, step, table_iter)})
247
248    def _project_and_filter(self, context, step, table_iter):
249        sink = self.table(step.projections if step.projections else context.columns)
250        condition = self.generate(step.condition)
251        projections = self.generate_tuple(step.projections)
252
253        for reader in table_iter:
254            if len(sink) >= step.offset + step.limit:
255                break
256
257            if condition and not context.eval(condition):
258                continue
259
260            if projections:
261                sink.append(context.eval_tuple(projections))
262            else:
263                sink.append(reader.row)
264
265        return sink
266
267    def static(self):
268        return self.context({}), [RowReader(())]
269
270    def scan_table(self, step):
271        table = self.tables.find(step.source)
272        context = self.context({step.source.alias_or_name: table})
273        return context, iter(table)
274
275    def join(self, step, context):
276        source = step.source_name
277
278        source_table = context.tables[source]
279        source_context = self.context({source: source_table})
280        column_ranges = {source: range(0, len(source_table.columns))}
281
282        for name, join in step.joins.items():
283            table = context.tables[name]
284            start = max(r.stop for r in column_ranges.values())
285            column_ranges[name] = range(start, len(table.columns) + start)
286            join_context = self.context({name: table})
287            condition = self.generate(join["condition"])
288            condition_context = (
289                self.context(
290                    {
291                        name: Table(
292                            source_context.columns + join_context.columns,
293                            column_range=column_range,
294                        )
295                        for name, column_range in column_ranges.items()
296                    }
297                )
298                if condition
299                else None
300            )
301
302            if join.get("source_key"):
303                table = self.hash_join(
304                    join, source_context, join_context, condition, condition_context
305                )
306            else:
307                table = self.nested_loop_join(
308                    join, source_context, join_context, condition, condition_context
309                )
310
311            source_context = self.context(
312                {
313                    name: Table(table.columns, table.rows, column_range)
314                    for name, column_range in column_ranges.items()
315                }
316            )
317        if not step.condition and not step.projections:
318            return source_context
319
320        sink = self._project_and_filter(
321            source_context,
322            step,
323            (reader for reader, _ in iter(source_context)),
324        )
325
326        if step.projections:
327            return self.context({step.name: sink})
328        else:
329            return self.context(
330                {
331                    name: Table(table.columns, sink.rows, table.column_range)
332                    for name, table in source_context.tables.items()
333                }
334            )
335
336    @staticmethod
337    def _join_matches(row, condition, condition_context):
338        if not condition:
339            return True
340
341        condition_context.set_row(row)
342        return condition_context.eval(condition) is True
343
344    def nested_loop_join(self, join, source_context, join_context, condition, condition_context):
345        table = Table(source_context.columns + join_context.columns)
346        source_rows = source_context.table.rows
347        join_rows = join_context.table.rows
348        matched_source = set()
349        matched_join = set()
350
351        for source_index, source_row in enumerate(source_rows):
352            for join_index, join_row in enumerate(join_rows):
353                row = source_row + join_row
354                if self._join_matches(row, condition, condition_context):
355                    table.append(row)
356                    matched_source.add(source_index)
357                    matched_join.add(join_index)
358
359        self._append_unmatched_join_rows(
360            table, join, source_rows, join_rows, matched_source, matched_join
361        )
362
363        return table
364
365    def hash_join(self, join, source_context, join_context, condition, condition_context):
366        source_key = self.generate_tuple(join["source_key"])
367        join_key = self.generate_tuple(join["join_key"])
368        results = collections.defaultdict(lambda: ([], []))
369
370        for index, (reader, ctx) in enumerate(source_context):
371            key = ctx.eval_tuple(source_key)
372            if all(value is not None for value in key):
373                results[key][0].append((index, reader.row))
374        for index, (reader, ctx) in enumerate(join_context):
375            key = ctx.eval_tuple(join_key)
376            if all(value is not None for value in key):
377                results[key][1].append((index, reader.row))
378
379        table = Table(source_context.columns + join_context.columns)
380        matched_source = set()
381        matched_join = set()
382
383        for source_group, join_group in results.values():
384            for (source_index, source_row), (join_index, join_row) in itertools.product(
385                source_group, join_group
386            ):
387                row = source_row + join_row
388                if self._join_matches(row, condition, condition_context):
389                    table.append(row)
390                    matched_source.add(source_index)
391                    matched_join.add(join_index)
392
393        self._append_unmatched_join_rows(
394            table,
395            join,
396            source_context.table.rows,
397            join_context.table.rows,
398            matched_source,
399            matched_join,
400        )
401
402        return table
403
404    @staticmethod
405    def _append_unmatched_join_rows(
406        table, join, source_rows, join_rows, matched_source, matched_join
407    ):
408        side = join.get("side")
409        if side in ("LEFT", "FULL"):
410            join_nulls = (None,) * (len(table.columns) - len(source_rows[0]) if source_rows else 0)
411            for index, row in enumerate(source_rows):
412                if index not in matched_source:
413                    table.append(row + join_nulls)
414
415        if side in ("RIGHT", "FULL"):
416            source_width = len(table.columns) - (len(join_rows[0]) if join_rows else 0)
417            source_nulls = (None,) * source_width
418            for index, row in enumerate(join_rows):
419                if index not in matched_join:
420                    table.append(source_nulls + row)
421
422    def aggregate(self, step, context):
423        group_by = self.generate_tuple(step.group.values())
424        aggregations = self.generate_tuple(step.aggregations)
425        operands = self.generate_tuple(step.operands)
426
427        if operands:
428            operand_table = Table(self.table(step.operands).columns)
429
430            for reader, ctx in context:
431                operand_table.append(ctx.eval_tuple(operands))
432
433            for i, (a, b) in enumerate(zip(context.table.rows, operand_table.rows)):
434                context.table.rows[i] = a + b
435
436            width = len(context.columns)
437            context.add_columns(*operand_table.columns)
438
439            operand_table = Table(
440                context.columns,
441                context.table.rows,
442                range(width, width + len(operand_table.columns)),
443            )
444
445            context = self.context(
446                {
447                    None: operand_table,
448                    **context.tables,
449                }
450            )
451
452        context.sort(group_by)
453
454        group = None
455        start = 0
456        end = 1
457        length = len(context.table)
458        table = self.table(list(step.group) + step.aggregations)
459
460        def add_row():
461            table.append(group + context.eval_tuple(aggregations))
462
463        if length:
464            for i in range(length):
465                context.set_index(i)
466                key = context.eval_tuple(group_by)
467                group = key if group is None else group
468                end += 1
469                if key != group:
470                    context.set_range(start, end - 2)
471                    add_row()
472                    group = key
473                    start = end - 2
474                if not step.condition and len(table.rows) >= step.offset + step.limit:
475                    break
476                if i == length - 1:
477                    context.set_range(start, end - 1)
478                    add_row()
479        elif step.limit > 0 and not group_by:
480            context.set_range(0, 0)
481            table.append(context.eval_tuple(aggregations))
482
483        context = self.context({step.name: table, **{name: table for name in context.tables}})
484
485        if step.projections or step.condition:
486            return self.context(
487                {step.name: self._project_and_filter(context, step, context.table_iter(step.name))}
488            )
489        return context
490
491    def sort(self, step, context):
492        projections = self.generate_tuple(step.projections)
493        projection_columns = [p.alias_or_name for p in step.projections]
494        all_columns = list(context.columns) + projection_columns
495        sink = self.table(all_columns)
496        for reader, ctx in context:
497            sink.append(reader.row + ctx.eval_tuple(projections))
498
499        sort_ctx = self.context(
500            {
501                None: sink,
502                **{table: sink for table in context.tables},
503            }
504        )
505        sort_ctx.sort(self.generate_tuple(step.key))
506
507        if not math.isinf(step.limit):
508            sort_ctx.table.rows = sort_ctx.table.rows[0 : step.offset + step.limit]
509
510        rows = sort_ctx.table.rows
511
512        if projection_columns:
513            rows = [row[len(context.columns) : len(all_columns)] for row in rows]
514
515        output = Table(projection_columns or context.columns, rows=rows)
516        return self.context({step.name: output})
517
518    def set_operation(self, step, context):
519        left = context.tables[step.left]
520        right = context.tables[step.right]
521
522        sink = self.table(left.columns)
523
524        if issubclass(step.op, exp.Intersect):
525            right_counts = collections.Counter(right.rows)
526            seen = set()
527            for row in left.rows:
528                if right_counts[row] and (not step.distinct or row not in seen):
529                    sink.append(row)
530                    seen.add(row)
531                    if not step.distinct:
532                        right_counts[row] -= 1
533        elif issubclass(step.op, exp.Except):
534            right_counts = collections.Counter(right.rows)
535            seen = set()
536            for row in left.rows:
537                if right_counts[row] and not step.distinct:
538                    right_counts[row] -= 1
539                elif not right_counts[row] and (not step.distinct or row not in seen):
540                    sink.append(row)
541                    seen.add(row)
542        elif issubclass(step.op, exp.Union) and step.distinct:
543            sink.rows = list(set(left.rows).union(set(right.rows)))
544        else:
545            sink.rows = left.rows + right.rows
546
547        if not math.isinf(step.limit):
548            sink.rows = sink.rows[0 : step.offset + step.limit]
549
550        return self.context({step.name: sink})
551
552
553class Python(Dialect):
554    class Tokenizer(tokens.Tokenizer):
555        STRING_ESCAPES = ["\\"]
556
557    Generator = PythonGenerator