Skip to content

Commit

Permalink
Fix InListExpr to return the correct number of rows (#8601)
Browse files Browse the repository at this point in the history
* Fix InListExpr  to return the correct number of rows

* Reduce repetition
  • Loading branch information
alamb authored Dec 22, 2023
1 parent 0ff5305 commit 55121d8
Showing 1 changed file with 53 additions and 4 deletions.
57 changes: 53 additions & 4 deletions datafusion/physical-expr/src/expressions/in_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -349,17 +349,18 @@ impl PhysicalExpr for InListExpr {
}

fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
let num_rows = batch.num_rows();
let value = self.expr.evaluate(batch)?;
let r = match &self.static_filter {
Some(f) => f.contains(value.into_array(1)?.as_ref(), self.negated)?,
Some(f) => f.contains(value.into_array(num_rows)?.as_ref(), self.negated)?,
None => {
let value = value.into_array(batch.num_rows())?;
let value = value.into_array(num_rows)?;
let found = self.list.iter().map(|expr| expr.evaluate(batch)).try_fold(
BooleanArray::new(BooleanBuffer::new_unset(batch.num_rows()), None),
BooleanArray::new(BooleanBuffer::new_unset(num_rows), None),
|result, expr| -> Result<BooleanArray> {
Ok(or_kleene(
&result,
&eq(&value, &expr?.into_array(batch.num_rows())?)?,
&eq(&value, &expr?.into_array(num_rows)?)?,
)?)
},
)?;
Expand Down Expand Up @@ -1267,4 +1268,52 @@ mod tests {

Ok(())
}

#[test]
fn in_list_no_cols() -> Result<()> {
// test logic when the in_list expression doesn't have any columns
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let a = Int32Array::from(vec![Some(1), Some(2), None]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

let list = vec![lit(ScalarValue::from(1i32)), lit(ScalarValue::from(6i32))];

// 1 IN (1, 6)
let expr = lit(ScalarValue::Int32(Some(1)));
in_list!(
batch,
list.clone(),
&false,
// should have three outputs, as the input batch has three rows
vec![Some(true), Some(true), Some(true)],
expr,
&schema
);

// 2 IN (1, 6)
let expr = lit(ScalarValue::Int32(Some(2)));
in_list!(
batch,
list.clone(),
&false,
// should have three outputs, as the input batch has three rows
vec![Some(false), Some(false), Some(false)],
expr,
&schema
);

// NULL IN (1, 6)
let expr = lit(ScalarValue::Int32(None));
in_list!(
batch,
list.clone(),
&false,
// should have three outputs, as the input batch has three rows
vec![None, None, None],
expr,
&schema
);

Ok(())
}
}

0 comments on commit 55121d8

Please sign in to comment.