sqlglot.executor.python

  1import collections
  2import itertools
  3import math
  4
  5from sqlglot import exp, generator, planner, tokens
  6from sqlglot.dialects.dialect import Dialect, inline_array_sql
  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.helper import subclasses
 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.Expression) 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.Expression):
 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
144            if join.get("source_key"):
145                table = self.hash_join(join, source_context, join_context)
146            else:
147                table = self.nested_loop_join(join, source_context, join_context)
148
149            source_context = self.context(
150                {
151                    name: Table(table.columns, table.rows, column_range)
152                    for name, column_range in column_ranges.items()
153                }
154            )
155            condition = self.generate(join["condition"])
156            if condition:
157                source_context.filter(condition)
158
159        if not step.condition and not step.projections:
160            return source_context
161
162        sink = self._project_and_filter(
163            source_context,
164            step,
165            (reader for reader, _ in iter(source_context)),
166        )
167
168        if step.projections:
169            return self.context({step.name: sink})
170        else:
171            return self.context(
172                {
173                    name: Table(table.columns, sink.rows, table.column_range)
174                    for name, table in source_context.tables.items()
175                }
176            )
177
178    def nested_loop_join(self, _join, source_context, join_context):
179        table = Table(source_context.columns + join_context.columns)
180
181        for reader_a, _ in source_context:
182            for reader_b, _ in join_context:
183                table.append(reader_a.row + reader_b.row)
184
185        return table
186
187    def hash_join(self, join, source_context, join_context):
188        source_key = self.generate_tuple(join["source_key"])
189        join_key = self.generate_tuple(join["join_key"])
190        left = join.get("side") == "LEFT"
191        right = join.get("side") == "RIGHT"
192
193        results = collections.defaultdict(lambda: ([], []))
194
195        for reader, ctx in source_context:
196            results[ctx.eval_tuple(source_key)][0].append(reader.row)
197        for reader, ctx in join_context:
198            results[ctx.eval_tuple(join_key)][1].append(reader.row)
199
200        table = Table(source_context.columns + join_context.columns)
201        nulls = [(None,) * len(join_context.columns if left else source_context.columns)]
202
203        for a_group, b_group in results.values():
204            if left:
205                b_group = b_group or nulls
206            elif right:
207                a_group = a_group or nulls
208
209            for a_row, b_row in itertools.product(a_group, b_group):
210                table.append(a_row + b_row)
211
212        return table
213
214    def aggregate(self, step, context):
215        group_by = self.generate_tuple(step.group.values())
216        aggregations = self.generate_tuple(step.aggregations)
217        operands = self.generate_tuple(step.operands)
218
219        if operands:
220            operand_table = Table(self.table(step.operands).columns)
221
222            for reader, ctx in context:
223                operand_table.append(ctx.eval_tuple(operands))
224
225            for i, (a, b) in enumerate(zip(context.table.rows, operand_table.rows)):
226                context.table.rows[i] = a + b
227
228            width = len(context.columns)
229            context.add_columns(*operand_table.columns)
230
231            operand_table = Table(
232                context.columns,
233                context.table.rows,
234                range(width, width + len(operand_table.columns)),
235            )
236
237            context = self.context(
238                {
239                    None: operand_table,
240                    **context.tables,
241                }
242            )
243
244        context.sort(group_by)
245
246        group = None
247        start = 0
248        end = 1
249        length = len(context.table)
250        table = self.table(list(step.group) + step.aggregations)
251
252        def add_row():
253            table.append(group + context.eval_tuple(aggregations))
254
255        if length:
256            for i in range(length):
257                context.set_index(i)
258                key = context.eval_tuple(group_by)
259                group = key if group is None else group
260                end += 1
261                if key != group:
262                    context.set_range(start, end - 2)
263                    add_row()
264                    group = key
265                    start = end - 2
266                if len(table.rows) >= step.limit:
267                    break
268                if i == length - 1:
269                    context.set_range(start, end - 1)
270                    add_row()
271        elif step.limit > 0 and not group_by:
272            context.set_range(0, 0)
273            table.append(context.eval_tuple(aggregations))
274
275        context = self.context({step.name: table, **{name: table for name in context.tables}})
276
277        if step.projections or step.condition:
278            return self.scan(step, context)
279        return context
280
281    def sort(self, step, context):
282        projections = self.generate_tuple(step.projections)
283        projection_columns = [p.alias_or_name for p in step.projections]
284        all_columns = list(context.columns) + projection_columns
285        sink = self.table(all_columns)
286        for reader, ctx in context:
287            sink.append(reader.row + ctx.eval_tuple(projections))
288
289        sort_ctx = self.context(
290            {
291                None: sink,
292                **{table: sink for table in context.tables},
293            }
294        )
295        sort_ctx.sort(self.generate_tuple(step.key))
296
297        if not math.isinf(step.limit):
298            sort_ctx.table.rows = sort_ctx.table.rows[0 : step.limit]
299
300        output = Table(
301            projection_columns,
302            rows=[r[len(context.columns) : len(all_columns)] for r in sort_ctx.table.rows],
303        )
304        return self.context({step.name: output})
305
306    def set_operation(self, step, context):
307        left = context.tables[step.left]
308        right = context.tables[step.right]
309
310        sink = self.table(left.columns)
311
312        if issubclass(step.op, exp.Intersect):
313            sink.rows = list(set(left.rows).intersection(set(right.rows)))
314        elif issubclass(step.op, exp.Except):
315            sink.rows = list(set(left.rows).difference(set(right.rows)))
316        elif issubclass(step.op, exp.Union) and step.distinct:
317            sink.rows = list(set(left.rows).union(set(right.rows)))
318        else:
319            sink.rows = left.rows + right.rows
320
321        if not math.isinf(step.limit):
322            sink.rows = sink.rows[0 : step.limit]
323
324        return self.context({step.name: sink})
325
326
327def _ordered_py(self, expression):
328    this = self.sql(expression, "this")
329    desc = "True" if expression.args.get("desc") else "False"
330    nulls_first = "True" if expression.args.get("nulls_first") else "False"
331    return f"ORDERED({this}, {desc}, {nulls_first})"
332
333
334def _rename(self, e):
335    try:
336        values = list(e.args.values())
337
338        if len(values) == 1:
339            values = values[0]
340            if not isinstance(values, list):
341                return self.func(e.key, values)
342            return self.func(e.key, *values)
343
344        if isinstance(e, exp.Func) and e.is_var_len_args:
345            args = itertools.chain.from_iterable(x if isinstance(x, list) else [x] for x in values)
346            return self.func(e.key, *args)
347
348        return self.func(e.key, *values)
349    except Exception as ex:
350        raise Exception(f"Could not rename {repr(e)}") from ex
351
352
353def _case_sql(self, expression):
354    this = self.sql(expression, "this")
355    chain = self.sql(expression, "default") or "None"
356
357    for e in reversed(expression.args["ifs"]):
358        true = self.sql(e, "true")
359        condition = self.sql(e, "this")
360        condition = f"{this} = ({condition})" if this else condition
361        chain = f"{true} if {condition} else ({chain})"
362
363    return chain
364
365
366def _lambda_sql(self, e: exp.Lambda) -> str:
367    names = {e.name.lower() for e in e.expressions}
368
369    e = e.transform(
370        lambda n: (
371            exp.var(n.name) if isinstance(n, exp.Identifier) and n.name.lower() in names else n
372        )
373    ).assert_is(exp.Lambda)
374
375    return f"lambda {self.expressions(e, flat=True)}: {self.sql(e, 'this')}"
376
377
378def _div_sql(self: generator.Generator, e: exp.Div) -> str:
379    denominator = self.sql(e, "expression")
380
381    if e.args.get("safe"):
382        denominator += " or None"
383
384    sql = f"DIV({self.sql(e, 'this')}, {denominator})"
385
386    if e.args.get("typed"):
387        sql = f"int({sql})"
388
389    return sql
390
391
392class Python(Dialect):
393    class Tokenizer(tokens.Tokenizer):
394        STRING_ESCAPES = ["\\"]
395
396    class Generator(generator.Generator):
397        TRANSFORMS = {
398            **{klass: _rename for klass in subclasses(exp.__name__, exp.Binary)},
399            **{klass: _rename for klass in exp.ALL_FUNCTIONS},
400            exp.Case: _case_sql,
401            exp.Alias: lambda self, e: self.sql(e.this),
402            exp.Array: inline_array_sql,
403            exp.And: lambda self, e: self.binary(e, "and"),
404            exp.Between: _rename,
405            exp.Boolean: lambda self, e: "True" if e.this else "False",
406            exp.Cast: lambda self, e: f"CAST({self.sql(e.this)}, exp.DataType.Type.{e.args['to']})",
407            exp.Column: lambda self,
408            e: f"scope[{self.sql(e, 'table') or None}][{self.sql(e.this)}]",
409            exp.Concat: lambda self, e: self.func(
410                "SAFECONCAT" if e.args.get("safe") else "CONCAT", *e.expressions
411            ),
412            exp.Distinct: lambda self, e: f"set({self.sql(e, 'this')})",
413            exp.Div: _div_sql,
414            exp.Extract: lambda self,
415            e: f"EXTRACT('{e.name.lower()}', {self.sql(e, 'expression')})",
416            exp.In: lambda self,
417            e: f"{self.sql(e, 'this')} in {{{self.expressions(e, flat=True)}}}",
418            exp.Interval: lambda self, e: f"INTERVAL({self.sql(e.this)}, '{self.sql(e.unit)}')",
419            exp.Is: lambda self, e: (
420                self.binary(e, "==") if isinstance(e.this, exp.Literal) else self.binary(e, "is")
421            ),
422            exp.JSONExtract: lambda self, e: self.func(e.key, e.this, e.expression, *e.expressions),
423            exp.JSONPath: lambda self, e: f"[{','.join(self.sql(p) for p in e.expressions[1:])}]",
424            exp.JSONPathKey: lambda self, e: f"'{self.sql(e.this)}'",
425            exp.JSONPathSubscript: lambda self, e: f"'{e.this}'",
426            exp.Lambda: _lambda_sql,
427            exp.Not: lambda self, e: f"not {self.sql(e.this)}",
428            exp.Null: lambda *_: "None",
429            exp.Or: lambda self, e: self.binary(e, "or"),
430            exp.Ordered: _ordered_py,
431            exp.Star: lambda *_: "1",
432        }