def __call__(self, expr: ir.TableExpr) -> ir.TableExpr: op = expr.op() if isinstance(op, ops.Join): scope = {X: op.left, Y: op.right} else: scope = {X: expr} return expr.projection( list( itertools.chain.from_iterable( ibis.util.promote_list(column.resolve(expr, scope)) for column in self.columns)))
def flatten(table: ir.TableExpr): """Extract all intersection or difference queries from `table`. Parameters ---------- table : TableExpr Returns ------- Iterable[Union[TableExpr]] """ op = table.op() return list(toolz.concatv(flatten_union(op.left), flatten_union(op.right)))
def flatten_difference(table: ir.TableExpr): """Extract all intersection queries from `table`. Parameters ---------- table : TableExpr Returns ------- Iterable[Union[TableExpr]] """ op = table.op() if isinstance(op, ops.Difference): return toolz.concatv(flatten_union(op.left), flatten_union(op.right)) return [table]
def flatten_union(table: ir.TableExpr): """Extract all union queries from `table`. Parameters ---------- table : TableExpr Returns ------- Iterable[Union[TableExpr, bool]] """ op = table.op() if isinstance(op, ops.Union): return toolz.concatv(flatten_union(op.left), [op.distinct], flatten_union(op.right)) return [table]
def flatten_union(table: ir.TableExpr): """Extract all union queries from `table`. Parameters ---------- table : TableExpr Returns ------- Iterable[Union[TableExpr, bool]] """ op = table.op() if isinstance(op, ops.Union): # For some reason mypy considers `op.left` and `op.right` # of `Argument` type, and fails the validation. While in # `flatten` types are the same, and it works return toolz.concatv( flatten_union(op.left), # type: ignore [op.distinct], flatten_union(op.right), # type: ignore ) return [table]
def _fmt_selection_column_table_expr(expr: ir.TableExpr, *, aliases: Aliases, **_: Any) -> str: return str(aliases[expr.op()])