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()])
