use vortex_error::VortexResult;
use crate::ArrayRef;
use crate::IntoArray;
use crate::LEGACY_SESSION;
use crate::VortexSessionExecute;
use crate::arrays::scalar_fn::array::ScalarFnArray;
use crate::arrays::scalar_fn::vtable::ArrayExpr;
use crate::arrays::scalar_fn::vtable::FakeEq;
use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
use crate::expr::Expression;
use crate::expr::lit;
use crate::scalar_fn::ScalarFn;
use crate::scalar_fn::VecExecutionArgs;
use crate::scalar_fn::fns::literal::Literal;
use crate::scalar_fn::fns::root::Root;
use crate::validity::Validity;
use crate::vtable::ValidityVTable;
fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult<ArrayRef> {
let mut ctx = LEGACY_SESSION.create_execution_ctx();
if expr.is::<Root>() {
vortex_error::vortex_bail!("Root expression cannot be executed in validity context");
}
if expr.is::<Literal>() {
let scalar = expr.as_::<Literal>();
return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array());
}
let inputs: Vec<ArrayRef> = expr
.children()
.iter()
.map(|child| execute_expr(child, row_count))
.collect::<VortexResult<_>>()?;
let args = VecExecutionArgs::new(inputs, row_count);
Ok(expr.scalar_fn().execute(&args, &mut ctx)?.into_array())
}
impl ValidityVTable<ScalarFnVTable> for ScalarFnVTable {
fn validity(array: &ScalarFnArray) -> VortexResult<Validity> {
let inputs: Vec<_> = array
.children
.iter()
.map(|child| {
if let Some(scalar) = child.as_constant() {
return Ok(lit(scalar));
}
Expression::try_new(ScalarFn::new(ArrayExpr, FakeEq(child.clone())).erased(), [])
})
.collect::<VortexResult<_>>()?;
let expr = Expression::try_new(array.scalar_fn.clone(), inputs)?;
let validity_expr = array.scalar_fn().validity(&expr)?;
Ok(Validity::Array(execute_expr(&validity_expr, array.len())?))
}
}