class Count(Reduction):
# TODO: count(col) takes down Impala, must always do count(*) in generated
# SQL
input_type = [rules.collection, boolean(name='where', optional=True)]
# TODO: counts are actually table-level operations. Let's address
# during the SQL generation exercise
def output_type(self):
return ir.Int64Scalar
class Where(ValueOp):
"""
Ternary case expression, equivalent to
bool_expr.case()
.when(True, true_expr)
.else_(false_or_null_expr)
"""
input_type = [
boolean(name='bool_expr'),
value(name='true_expr'),
value(name='false_null_expr')
]
def output_type(self):
return rules.shape_like(self.args[0], self.args[1].type())
class Reduction(ValueOp):
input_type = [rules.array, boolean(name='where', optional=True)]
_reduction = True