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
 12
 13
 14class PythonExecutor:
 15    def __init__(self, env=None, tables=None):
 16        self.generator = Python().generator(identify=True, comments=False)
 17        self.env = {**ENV, **(env or {})}
 18        self.tables = tables or {}
 19
 20    def execute(self, plan):
 21        finished = set()
 22        queue = set(plan.leaves)
 23        contexts = {}
 24
 25        while queue:
 26            node = queue.pop()
 27            try:
 28                context = self.context(
 29                    {
 30                        name: table
 31                        for dep in node.dependencies
 32                        for name, table in contexts[dep].tables.items()
 33                    }
 34                )
 35
 36                if isinstance(node, planner.Scan):
 37                    contexts[node] = self.scan(node, context)
 38                elif isinstance(node, planner.Aggregate):
 39                    contexts[node] = self.aggregate(node, context)
 40                elif isinstance(node, planner.Join):
 41                    contexts[node] = self.join(node, context)
 42                elif isinstance(node, planner.Sort):
 43                    contexts[node] = self.sort(node, context)
 44                elif isinstance(node, planner.SetOperation):
 45                    contexts[node] = self.set_operation(node, context)
 46                else:
 47                    raise NotImplementedError
 48
 49                finished.add(node)
 50
 51                for dep in node.dependents:
 52                    if all(d in contexts for d in dep.dependencies):
 53                        queue.add(dep)
 54
 55                for dep in node.dependencies:
 56                    if all(d in finished for d in dep.dependents):
 57                        contexts.pop(dep)
 58            except Exception as e:
 59                raise ExecuteError(f"Step '{node.id}' failed: {e}") from e
 60
 61        root = plan.root
 62        return contexts[root].tables[root.name]
 63
 64    def generate(self, expression):
 65        """Convert a SQL expression into literal Python code and compile it into bytecode."""
 66        if not expression:
 67            return None
 68
 69        sql = self.generator.generate(expression)
 70        return compile(sql, sql, "eval", optimize=2)
 71
 72    def generate_tuple(self, expressions):
 73        """Convert an array of SQL expressions into tuple of Python byte code."""
 74        if not expressions:
 75            return tuple()
 76        return tuple(self.generate(expression) for expression in expressions)
 77
 78    def context(self, tables):
 79        return Context(tables, env=self.env)
 80
 81    def table(self, expressions):
 82        return Table(
 83            expression.alias_or_name if isinstance(expression, exp.Expr) else expression
 84            for expression in expressions
 85        )
 86
 87    def scan(self, step, context):
 88        source = step.source
 89
 90        if source and isinstance(source, exp.Expr):
 91            source = source.name or source.alias
 92
 93        if source is None:
 94            context, table_iter = self.static()
 95        elif source in context:
 96            if not step.projections and not step.condition:
 97                return self.context({step.name: context.tables[source]})
 98            table_iter = context.table_iter(source)
 99        else:
100            context, table_iter = self.scan_table(step)
101
102        return self.context({step.name: self._project_and_filter(context, step, table_iter)})
103
104    def _project_and_filter(self, context, step, table_iter):
105        sink = self.table(step.projections if step.projections else context.columns)
106        condition = self.generate(step.condition)
107        projections = self.generate_tuple(step.projections)
108
109        for reader in table_iter:
110            if len(sink) >= step.limit:
111                break
112
113            if condition and not context.eval(condition):
114                continue
115
116            if projections:
117                sink.append(context.eval_tuple(projections))
118            else:
119                sink.append(reader.row)
120
121        return sink
122
123    def static(self):
124        return self.context({}), [RowReader(())]
125
126    def scan_table(self, step):
127        table = self.tables.find(step.source)
128        context = self.context({step.source.alias_or_name: table})
129        return context, iter(table)
130
131    def join(self, step, context):
132        source = step.source_name
133
134        source_table = context.tables[source]
135        source_context = self.context({source: source_table})
136        column_ranges = {source: range(0, len(source_table.columns))}
137
138        for name, join in step.joins.items():
139            table = context.tables[name]
140            start = max(r.stop for r in column_ranges.values())
141            column_ranges[name] = range(start, len(table.columns) + start)
142            join_context = self.context({name: table})
143            condition = self.generate(join["condition"])
144            condition_context = (
145                self.context(
146                    {
147                        name: Table(
148                            source_context.columns + join_context.columns,
149                            column_range=column_range,
150                        )
151                        for name, column_range in column_ranges.items()
152                    }
153                )
154                if condition
155                else None
156            )
157
158            if join.get("source_key"):
159                table = self.hash_join(
160                    join, source_context, join_context, condition, condition_context
161                )
162            else:
163                table = self.nested_loop_join(
164                    join, source_context, join_context, condition, condition_context
165                )
166
167            source_context = self.context(
168                {
169                    name: Table(table.columns, table.rows, column_range)
170                    for name, column_range in column_ranges.items()
171                }
172            )
173        if not step.condition and not step.projections:
174            return source_context
175
176        sink = self._project_and_filter(
177            source_context,
178            step,
179            (reader for reader, _ in iter(source_context)),
180        )
181
182        if step.projections:
183            return self.context({step.name: sink})
184        else:
185            return self.context(
186                {
187                    name: Table(table.columns, sink.rows, table.column_range)
188                    for name, table in source_context.tables.items()
189                }
190            )
191
192    @staticmethod
193    def _join_matches(row, condition, condition_context):
194        if not condition:
195            return True
196
197        condition_context.set_row(row)
198        return condition_context.eval(condition) is True
199
200    def nested_loop_join(self, join, source_context, join_context, condition, condition_context):
201        table = Table(source_context.columns + join_context.columns)
202        source_rows = source_context.table.rows
203        join_rows = join_context.table.rows
204        matched_source = set()
205        matched_join = set()
206
207        for source_index, source_row in enumerate(source_rows):
208            for join_index, join_row in enumerate(join_rows):
209                row = source_row + join_row
210                if self._join_matches(row, condition, condition_context):
211                    table.append(row)
212                    matched_source.add(source_index)
213                    matched_join.add(join_index)
214
215        self._append_unmatched_join_rows(
216            table, join, source_rows, join_rows, matched_source, matched_join
217        )
218
219        return table
220
221    def hash_join(self, join, source_context, join_context, condition, condition_context):
222        source_key = self.generate_tuple(join["source_key"])
223        join_key = self.generate_tuple(join["join_key"])
224        results = collections.defaultdict(lambda: ([], []))
225
226        for index, (reader, ctx) in enumerate(source_context):
227            key = ctx.eval_tuple(source_key)
228            if all(value is not None for value in key):
229                results[key][0].append((index, reader.row))
230        for index, (reader, ctx) in enumerate(join_context):
231            key = ctx.eval_tuple(join_key)
232            if all(value is not None for value in key):
233                results[key][1].append((index, reader.row))
234
235        table = Table(source_context.columns + join_context.columns)
236        matched_source = set()
237        matched_join = set()
238
239        for source_group, join_group in results.values():
240            for (source_index, source_row), (join_index, join_row) in itertools.product(
241                source_group, join_group
242            ):
243                row = source_row + join_row
244                if self._join_matches(row, condition, condition_context):
245                    table.append(row)
246                    matched_source.add(source_index)
247                    matched_join.add(join_index)
248
249        self._append_unmatched_join_rows(
250            table,
251            join,
252            source_context.table.rows,
253            join_context.table.rows,
254            matched_source,
255            matched_join,
256        )
257
258        return table
259
260    @staticmethod
261    def _append_unmatched_join_rows(
262        table, join, source_rows, join_rows, matched_source, matched_join
263    ):
264        side = join.get("side")
265        if side in ("LEFT", "FULL"):
266            join_nulls = (None,) * (len(table.columns) - len(source_rows[0]) if source_rows else 0)
267            for index, row in enumerate(source_rows):
268                if index not in matched_source:
269                    table.append(row + join_nulls)
270
271        if side in ("RIGHT", "FULL"):
272            source_width = len(table.columns) - (len(join_rows[0]) if join_rows else 0)
273            source_nulls = (None,) * source_width
274            for index, row in enumerate(join_rows):
275                if index not in matched_join:
276                    table.append(source_nulls + row)
277
278    def aggregate(self, step, context):
279        group_by = self.generate_tuple(step.group.values())
280        aggregations = self.generate_tuple(step.aggregations)
281        operands = self.generate_tuple(step.operands)
282
283        if operands:
284            operand_table = Table(self.table(step.operands).columns)
285
286            for reader, ctx in context:
287                operand_table.append(ctx.eval_tuple(operands))
288
289            for i, (a, b) in enumerate(zip(context.table.rows, operand_table.rows)):
290                context.table.rows[i] = a + b
291
292            width = len(context.columns)
293            context.add_columns(*operand_table.columns)
294
295            operand_table = Table(
296                context.columns,
297                context.table.rows,
298                range(width, width + len(operand_table.columns)),
299            )
300
301            context = self.context(
302                {
303                    None: operand_table,
304                    **context.tables,
305                }
306            )
307
308        context.sort(group_by)
309
310        group = None
311        start = 0
312        end = 1
313        length = len(context.table)
314        table = self.table(list(step.group) + step.aggregations)
315
316        def add_row():
317            table.append(group + context.eval_tuple(aggregations))
318
319        if length:
320            for i in range(length):
321                context.set_index(i)
322                key = context.eval_tuple(group_by)
323                group = key if group is None else group
324                end += 1
325                if key != group:
326                    context.set_range(start, end - 2)
327                    add_row()
328                    group = key
329                    start = end - 2
330                if len(table.rows) >= step.limit:
331                    break
332                if i == length - 1:
333                    context.set_range(start, end - 1)
334                    add_row()
335        elif step.limit > 0 and not group_by:
336            context.set_range(0, 0)
337            table.append(context.eval_tuple(aggregations))
338
339        context = self.context({step.name: table, **{name: table for name in context.tables}})
340
341        if step.projections or step.condition:
342            return self.scan(step, context)
343        return context
344
345    def sort(self, step, context):
346        projections = self.generate_tuple(step.projections)
347        projection_columns = [p.alias_or_name for p in step.projections]
348        all_columns = list(context.columns) + projection_columns
349        sink = self.table(all_columns)
350        for reader, ctx in context:
351            sink.append(reader.row + ctx.eval_tuple(projections))
352
353        sort_ctx = self.context(
354            {
355                None: sink,
356                **{table: sink for table in context.tables},
357            }
358        )
359        sort_ctx.sort(self.generate_tuple(step.key))
360
361        if not math.isinf(step.limit):
362            sort_ctx.table.rows = sort_ctx.table.rows[0 : step.limit]
363
364        output = Table(
365            projection_columns,
366            rows=[r[len(context.columns) : len(all_columns)] for r in sort_ctx.table.rows],
367        )
368        return self.context({step.name: output})
369
370    def set_operation(self, step, context):
371        left = context.tables[step.left]
372        right = context.tables[step.right]
373
374        sink = self.table(left.columns)
375
376        if issubclass(step.op, exp.Intersect):
377            right_counts = collections.Counter(right.rows)
378            seen = set()
379            for row in left.rows:
380                if right_counts[row] and (not step.distinct or row not in seen):
381                    sink.append(row)
382                    seen.add(row)
383                    if not step.distinct:
384                        right_counts[row] -= 1
385        elif issubclass(step.op, exp.Except):
386            right_counts = collections.Counter(right.rows)
387            seen = set()
388            for row in left.rows:
389                if right_counts[row] and not step.distinct:
390                    right_counts[row] -= 1
391                elif not right_counts[row] and (not step.distinct or row not in seen):
392                    sink.append(row)
393                    seen.add(row)
394        elif issubclass(step.op, exp.Union) and step.distinct:
395            sink.rows = list(set(left.rows).union(set(right.rows)))
396        else:
397            sink.rows = left.rows + right.rows
398
399        if not math.isinf(step.limit):
400            sink.rows = sink.rows[0 : step.limit]
401
402        return self.context({step.name: sink})
403
404
405class Python(Dialect):
406    class Tokenizer(tokens.Tokenizer):
407        STRING_ESCAPES = ["\\"]
408
409    Generator = PythonGenerator