Skip to content

Commit

Permalink
migrate bitwise functions to UDAF
Browse files Browse the repository at this point in the history
The functions now take a single expression instead of a Vec<_>.

Ref: apache/datafusion#10930
  • Loading branch information
Michael-J-Ward committed Jul 24, 2024
1 parent 27f5c73 commit bbe02ff
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 20 deletions.
25 changes: 10 additions & 15 deletions python/datafusion/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1482,31 +1482,26 @@ def last_value(
)


def bit_and(*args: Expr, distinct: bool = False) -> Expr:
def bit_and(arg: Expr, distinct: bool = False) -> Expr:
"""Computes the bitwise AND of the argument."""
args = [arg.expr for arg in args]
return Expr(f.bit_and(*args, distinct=distinct))
return Expr(f.bit_and(arg.expr, distinct=distinct))


def bit_or(*args: Expr, distinct: bool = False) -> Expr:
def bit_or(arg: Expr, distinct: bool = False) -> Expr:
"""Computes the bitwise OR of the argument."""
args = [arg.expr for arg in args]
return Expr(f.bit_or(*args, distinct=distinct))
return Expr(f.bit_or(arg.expr, distinct=distinct))


def bit_xor(*args: Expr, distinct: bool = False) -> Expr:
def bit_xor(arg: Expr, distinct: bool = False) -> Expr:
"""Computes the bitwise XOR of the argument."""
args = [arg.expr for arg in args]
return Expr(f.bit_xor(*args, distinct=distinct))
return Expr(f.bit_xor(arg.expr, distinct=distinct))


def bool_and(*args: Expr, distinct: bool = False) -> Expr:
def bool_and(arg: Expr, distinct: bool = False) -> Expr:
"""Computes the boolean AND of the arugment."""
args = [arg.expr for arg in args]
return Expr(f.bool_and(*args, distinct=distinct))
return Expr(f.bool_and(arg.expr, distinct=distinct))


def bool_or(*args: Expr, distinct: bool = False) -> Expr:
def bool_or(arg: Expr, distinct: bool = False) -> Expr:
"""Computes the boolean OR of the arguement."""
args = [arg.expr for arg in args]
return Expr(f.bool_or(*args, distinct=distinct))
return Expr(f.bool_or(arg.expr, distinct=distinct))
55 changes: 50 additions & 5 deletions src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,56 @@ pub fn avg(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
}
}

#[pyfunction]
pub fn bit_and(expr_x: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::bit_and(expr_x.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn bit_or(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::bit_or(expression.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn bit_xor(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::bit_xor(expression.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn bool_and(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::bool_and(expression.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn bool_or(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::bool_or(expression.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn mean(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
// alias for avg
Expand Down Expand Up @@ -937,11 +987,6 @@ array_fn!(range, start stop step);
aggregate_function!(array_agg, ArrayAgg);
aggregate_function!(max, Max);
aggregate_function!(min, Min);
aggregate_function!(bit_and, BitAnd);
aggregate_function!(bit_or, BitOr);
aggregate_function!(bit_xor, BitXor);
aggregate_function!(bool_and, BoolAnd);
aggregate_function!(bool_or, BoolOr);

pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(abs))?;
Expand Down

0 comments on commit bbe02ff

Please sign in to comment.