diff options
Diffstat (limited to 'sqlglot/executor/python.py')
| -rw-r--r-- | sqlglot/executor/python.py | 46 |
1 files changed, 28 insertions, 18 deletions
diff --git a/sqlglot/executor/python.py b/sqlglot/executor/python.py index cb2543c..908b80a 100644 --- a/sqlglot/executor/python.py +++ b/sqlglot/executor/python.py @@ -115,6 +115,9 @@ class PythonExecutor: sink = self.table(context.columns) for reader in table_iter: + if len(sink) >= step.limit: + break + if condition and not context.eval(condition): continue @@ -123,9 +126,6 @@ class PythonExecutor: else: sink.append(reader.row) - if len(sink) >= step.limit: - break - return self.context({step.name: sink}) def static(self): @@ -288,21 +288,32 @@ class PythonExecutor: end = 1 length = len(context.table) table = self.table(list(step.group) + step.aggregations) + condition = self.generate(step.condition) - for i in range(length): - context.set_index(i) - key = context.eval_tuple(group_by) - group = key if group is None else group - end += 1 - if key != group: - context.set_range(start, end - 2) - table.append(group + context.eval_tuple(aggregations)) - group = key - start = end - 2 - if i == length - 1: - context.set_range(start, end - 1) + def add_row(): + if not condition or context.eval(condition): table.append(group + context.eval_tuple(aggregations)) + if length: + for i in range(length): + context.set_index(i) + key = context.eval_tuple(group_by) + group = key if group is None else group + end += 1 + if key != group: + context.set_range(start, end - 2) + add_row() + group = key + start = end - 2 + if len(table.rows) >= step.limit: + break + if i == length - 1: + context.set_range(start, end - 1) + add_row() + elif step.limit > 0: + context.set_range(0, 0) + table.append(context.eval_tuple(group_by) + context.eval_tuple(aggregations)) + context = self.context({step.name: table, **{name: table for name in context.tables}}) if step.projections: @@ -311,11 +322,9 @@ class PythonExecutor: def sort(self, step, context): projections = self.generate_tuple(step.projections) - projection_columns = [p.alias_or_name for p in step.projections] all_columns = list(context.columns) + projection_columns sink = self.table(all_columns) - for reader, ctx in context: sink.append(reader.row + ctx.eval_tuple(projections)) @@ -401,8 +410,9 @@ class Python(Dialect): exp.Boolean: lambda self, e: "True" if e.this else "False", exp.Cast: lambda self, e: f"CAST({self.sql(e.this)}, exp.DataType.Type.{e.args['to']})", exp.Column: lambda self, e: f"scope[{self.sql(e, 'table') or None}][{self.sql(e.this)}]", + exp.Distinct: lambda self, e: f"set({self.sql(e, 'this')})", exp.Extract: lambda self, e: f"EXTRACT('{e.name.lower()}', {self.sql(e, 'expression')})", - exp.In: lambda self, e: f"{self.sql(e, 'this')} in {self.expressions(e)}", + exp.In: lambda self, e: f"{self.sql(e, 'this')} in ({self.expressions(e, flat=True)})", exp.Is: lambda self, e: self.binary(e, "is"), exp.Not: lambda self, e: f"not {self.sql(e.this)}", exp.Null: lambda *_: "None", |