diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index dd98029c739a8..2d6a297a44d87 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -1765,6 +1765,92 @@ mod tests { Ok(()) } + #[test] + fn case_with_expr_fixed_size_binary() -> Result<()> { + let schema = + Schema::new(vec![Field::new("a", DataType::FixedSizeBinary(2), true)]); + let mut builder = FixedSizeBinaryBuilder::with_capacity(4, 2); + builder.append_value(b"aa")?; + builder.append_value(b"bb")?; + builder.append_null(); + builder.append_value(b"cc")?; + let values = builder.finish(); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)])?; + + let schema = batch.schema(); + + // CASE a WHEN b"aa" THEN 1 WHEN b"cc" THEN 2 ELSE 3 END + let when1 = lit(ScalarValue::FixedSizeBinary(2, Some(b"aa".to_vec()))); + let then1 = lit(1i32); + let when2 = lit(ScalarValue::FixedSizeBinary(2, Some(b"cc".to_vec()))); + let then2 = lit(2i32); + let else_expr = lit(3i32); + + let expr = case( + Some(col("a", &schema)?), + vec![(when1, then1), (when2, then2)], + Some(else_expr), + )?; + let result = expr + .evaluate(&batch)? + .into_array(batch.num_rows()) + .expect("Failed to convert to array"); + let result = as_int32_array(&result)?; + + let expected = &Int32Array::from(vec![Some(1), Some(3), Some(3), Some(2)]); + + assert_eq!(expected, result); + + Ok(()) + } + + #[test] + fn case_with_expr_fixed_size_binary_dictionary() -> Result<()> { + let schema = Schema::new(vec![Field::new( + "a", + DataType::Dictionary( + Box::new(DataType::UInt8), + Box::new(DataType::FixedSizeBinary(2)), + ), + true, + )]); + let keys = UInt8Array::from(vec![0u8, 1u8, 2u8, 3u8]); + let mut builder = FixedSizeBinaryBuilder::with_capacity(4, 2); + builder.append_value(b"aa")?; + builder.append_value(b"bb")?; + builder.append_null(); + builder.append_value(b"cc")?; + let values = builder.finish(); + let dictionary = DictionaryArray::new(keys, Arc::new(values)); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(dictionary)])?; + + let schema = batch.schema(); + + // CASE a WHEN b"aa" THEN 1 WHEN b"cc" THEN 2 ELSE 3 END + let when1 = lit(ScalarValue::FixedSizeBinary(2, Some(b"aa".to_vec()))); + let then1 = lit(1i32); + let when2 = lit(ScalarValue::FixedSizeBinary(2, Some(b"cc".to_vec()))); + let then2 = lit(2i32); + let else_expr = lit(3i32); + + let expr = case( + Some(col("a", &schema)?), + vec![(when1, then1), (when2, then2)], + Some(else_expr), + )?; + let result = expr + .evaluate(&batch)? + .into_array(batch.num_rows()) + .expect("Failed to convert to array"); + let result = as_int32_array(&result)?; + + let expected = &Int32Array::from(vec![Some(1), Some(3), Some(3), Some(2)]); + + assert_eq!(expected, result); + + Ok(()) + } + // Make sure we are not failing when got literal in case when but input is dictionary encoded #[test] fn case_with_expr_boolean_dictionary() -> Result<()> { diff --git a/datafusion/physical-expr/src/expressions/case/literal_lookup_table/bytes_like_lookup_table.rs b/datafusion/physical-expr/src/expressions/case/literal_lookup_table/bytes_like_lookup_table.rs index e5cf3f84fd919..d5f27bb089ce5 100644 --- a/datafusion/physical-expr/src/expressions/case/literal_lookup_table/bytes_like_lookup_table.rs +++ b/datafusion/physical-expr/src/expressions/case/literal_lookup_table/bytes_like_lookup_table.rs @@ -106,7 +106,9 @@ fn try_get_bytes_iterator( DataType::LargeBinary => Box::new(array.as_binary::().into_iter()), - DataType::FixedSizeBinary(_) => Box::new(array.as_binary::().into_iter()), + DataType::FixedSizeBinary(_) => { + Box::new(array.as_fixed_size_binary().into_iter()) + } DataType::Utf8View => Box::new( array diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 3f9058bce865b..83eaf61b44768 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -836,6 +836,63 @@ NULL 0 0 0 0 0 63 0 0 0 0 0 64 3 3 3 3 3 +# CASE with FixedSizeBinary WHEN literals (regression test for #26012) +query ?II +select + letter_binary as letter + ,CASE letter_fsb + WHEN arrow_cast(X'62', 'FixedSizeBinary(1)') THEN 1 + WHEN arrow_cast(X'61', 'FixedSizeBinary(1)') THEN 2 + WHEN arrow_cast(X'64', 'FixedSizeBinary(1)') THEN 3 + ELSE 0 + END as fsb_literals + ,CASE arrow_cast(letter_fsb, 'Dictionary(Int32, FixedSizeBinary(1))') + WHEN arrow_cast(X'62', 'FixedSizeBinary(1)') THEN 1 + WHEN arrow_cast(X'61', 'FixedSizeBinary(1)') THEN 2 + WHEN arrow_cast(X'64', 'FixedSizeBinary(1)') THEN 3 + ELSE 0 + END as fsb_dict_literals +FROM letters_binary; +---- +61 2 2 +62 1 1 +NULL 0 0 +63 0 0 +61 2 2 +63 0 0 +64 3 3 + +# Minimal reproducer for issue #26012 +query I +SELECT CASE x + WHEN arrow_cast(X'01', 'FixedSizeBinary(1)') THEN 1 + WHEN arrow_cast(X'02', 'FixedSizeBinary(1)') THEN 2 + ELSE 3 + END AS r +FROM (SELECT arrow_cast(X'03', 'FixedSizeBinary(1)') AS x); +---- +3 + +query I +SELECT CASE x + WHEN arrow_cast(X'01', 'FixedSizeBinary(1)') THEN 1 + WHEN arrow_cast(X'02', 'FixedSizeBinary(1)') THEN 2 + ELSE 3 + END AS r +FROM (SELECT arrow_cast(X'01', 'FixedSizeBinary(1)') AS x); +---- +1 + +query I +SELECT CASE x + WHEN arrow_cast(X'01', 'FixedSizeBinary(1)') THEN 1 + WHEN arrow_cast(X'02', 'FixedSizeBinary(1)') THEN 2 + ELSE 3 + END AS r +FROM (SELECT arrow_cast(X'02', 'FixedSizeBinary(1)') AS x); +---- +2 + statement ok drop table source;