rluvaton commented on code in PR #27160:
URL: https://github.com/apache/datafusion/pull/27160#discussion_r4503055305
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -1644,6 +1684,118 @@ mod tests {
}
}
+ #[test]
+ fn literal_when_preserves_row_indices_and_short_circuits() -> Result<()> {
+ let batch = RecordBatch::try_from_iter(vec![
+ (
+ "key",
+ Arc::new(Int32Array::from(vec![
+ Some(2),
+ None,
+ Some(1),
+ Some(2),
+ Some(3),
+ Some(4),
+ Some(0),
+ ])) as ArrayRef,
+ ),
+ (
+ "value",
+ Arc::new(Int32Array::from(vec![20, 99, 10, 21, 30, 40, 0])) as
ArrayRef,
+ ),
+ ])?;
+ let schema = batch.schema();
+ let invalid = binary(lit(1i32), Operator::Divide, lit(0i32), &schema)?;
+ for with_else in [false, true] {
+ let expr = case(
+ Some(col("key", &schema)?),
+ vec![
+ (lit(99i32), Arc::clone(&invalid)),
+ (lit(ScalarValue::Int32(None)), Arc::clone(&invalid)),
+ (lit(2i32), col("value", &schema)?),
+ (lit(1i32), col("value", &schema)?),
+ (lit(2i32), Arc::clone(&invalid)),
+ (lit(3i32), lit(300i32)),
+ ],
+ with_else.then(|| col("value", &schema)).transpose()?,
+ )?;
+ let expected = if with_else {
+ vec![
+ Some(20),
+ Some(99),
+ Some(10),
+ Some(21),
+ Some(300),
+ Some(40),
+ Some(0),
+ ]
+ } else {
+ vec![Some(20), None, Some(10), Some(21), Some(300), None, None]
+ };
+ // Four rows also exercise the branch that matches every remaining
row.
+ for count in [4, 7] {
+ let result = expr.evaluate(&batch.slice(0,
count))?.into_array(count)?;
+ assert_eq!(
+ as_int32_array(&result)?,
+ &Int32Array::from(expected[..count].to_vec())
+ );
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn taking_case_rows_preserves_zero_column_batch_length() -> Result<()> {
+ let batch = RecordBatch::try_new_with_options(
+ Arc::new(Schema::empty()),
+ vec![],
+ &RecordBatchOptions::new().with_row_count(Some(5)),
+ )?;
+ let selected = take_case_rows(&batch, &UInt32Array::from(vec![1, 3]))?;
+ assert_eq!(selected.num_rows(), 2);
+ Ok(())
+ }
+
+ #[test]
+ fn literal_when_preserves_nested_field_names() -> Result<()> {
+ let mut payload =
ListBuilder::new(Int32Builder::new()).with_field(Field::new(
+ "value",
+ DataType::Int32,
+ true,
+ ));
+ for value in [10, 20] {
+ payload.values().append_value(value);
+ payload.append(true);
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int32, false),
+ Field::new(
+ "payload",
+ DataType::List(Arc::new(Field::new("item", DataType::Int32,
true))),
+ false,
+ ),
+ ]));
+ let batch = RecordBatch::try_new_with_options(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])),
+ Arc::new(payload.finish()),
+ ],
+ &RecordBatchOptions::new().with_match_field_names(false),
+ )?;
Review Comment:
You should not build this batch with bad schema
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
