Skip to content

Commit

Permalink
Fix InListExpr to return the correct number of rows
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Dec 20, 2023
1 parent 8d72196 commit a77d6a2
Showing 1 changed file with 50 additions and 1 deletion.
51 changes: 50 additions & 1 deletion datafusion/physical-expr/src/expressions/in_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -349,9 +349,10 @@ 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 found = self.list.iter().map(|expr| expr.evaluate(batch)).try_fold(
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 a77d6a2

Please sign in to comment.