diff --git a/Cargo.lock b/Cargo.lock index 45d447e3cd4..07d9488d7c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9596,6 +9596,7 @@ dependencies = [ "insta", "inventory", "itertools 0.14.0", + "itoa", "jiff", "memchr", "mimalloc", @@ -9613,6 +9614,7 @@ dependencies = [ "rstest", "rstest_reuse", "rustc-hash", + "ryu", "serde", "serde_json", "serde_test", diff --git a/Cargo.toml b/Cargo.toml index 1484b75de11..f700cddbd1c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -175,6 +175,7 @@ indicatif = "0.18.0" insta = "1.43" inventory = "0.3.20" itertools = "0.14.0" +itoa = "1.0.18" jiff = "0.2.28" jni = { version = "0.22.0" } kanal = "0.1.1" @@ -236,6 +237,7 @@ rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" rustix = { version = "1.1", features = ["fs"] } +ryu = "1.0.23" serde = "1.0.221" serde_json = "1.0.138" serde_test = "1.0.176" diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 22a8768af5d..3bfacf8e3be 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -36,6 +36,7 @@ half = { workspace = true, features = ["num-traits"] } humansize = { workspace = true } inventory = { workspace = true } itertools = { workspace = true } +itoa = { workspace = true } jiff = { workspace = true } memchr = { workspace = true } num-traits = { workspace = true } @@ -53,6 +54,7 @@ regex-syntax = { workspace = true } rstest = { workspace = true, optional = true } rstest_reuse = { workspace = true, optional = true } rustc-hash = { workspace = true } +ryu = { workspace = true } serde = { workspace = true, optional = true, features = ["derive", "rc"] } simdutf8 = { workspace = true } smallvec = { workspace = true } diff --git a/vortex-array/src/arrays/bool/compute/cast.rs b/vortex-array/src/arrays/bool/compute/cast.rs index 36849f55a18..0c49bb75ced 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -5,6 +5,7 @@ use num_traits::One; use num_traits::Zero; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -14,6 +15,8 @@ use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::match_each_native_ptype; use crate::scalar_fn::fns::cast::CastKernel; @@ -53,6 +56,36 @@ impl CastKernel for Bool { )); } + if let DType::Utf8(new_nullability) = dtype { + let len = array.len(); + let new_validity = array + .validity()? + .cast_nullability(*new_nullability, len, ctx)?; + let mask = new_validity.execute_mask(len, ctx)?; + let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), len); + + match &mask { + Mask::AllTrue(_) => { + for value in array.to_bit_buffer().iter() { + builder.append_value(if value { "true" } else { "false" }); + } + } + Mask::AllFalse(_) => builder.append_nulls(len), + Mask::Values(validity) => { + let bits = array.to_bit_buffer(); + for (value, valid) in bits.iter().zip(validity.bit_buffer().iter()) { + if valid { + builder.append_value(if value { "true" } else { "false" }); + } else { + builder.append_null(); + } + } + } + } + + return Ok(Some(builder.finish_into_varbinview().into_array())); + } + let DType::Primitive(new_ptype, new_nullability) = dtype else { return Ok(None); }; @@ -79,12 +112,15 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::BoolArray; + use crate::arrays::VarBinViewArray; + use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; @@ -117,6 +153,81 @@ mod tests { assert!(result.is_err(), "Expected error, got: {result:?}"); } + #[test] + fn cast_bool_to_utf8() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let actual = BoolArray::from_iter([true, false, true]) + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))?; + let expected = VarBinViewArray::from_iter_str(["true", "false", "true"]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn cast_nullable_bool_to_utf8() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let actual = BoolArray::from_iter([Some(true), None, Some(false)]) + .into_array() + .cast(DType::Utf8(Nullability::Nullable))?; + let expected = VarBinViewArray::from_iter_nullable_str([Some("true"), None, Some("false")]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn cast_all_null_bool_to_utf8() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let actual = BoolArray::from_iter([None, None]) + .into_array() + .cast(DType::Utf8(Nullability::Nullable))?; + let expected = VarBinViewArray::from_iter_nullable_str([None::<&str>, None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn cast_nullable_bool_with_null_to_non_nullable_utf8_fails() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let result = BoolArray::from_iter([Some(true), None]) + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))? + .execute::(&mut ctx); + + assert!(result.is_err(), "Expected error, got: {result:?}"); + Ok(()) + } + + #[test] + fn cast_all_valid_nullable_bool_to_non_nullable_utf8() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let actual = BoolArray::from_iter([Some(true), Some(false)]) + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))?; + let expected = VarBinViewArray::from_iter_str(["true", "false"]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn cast_bool_to_binary_is_unsupported() { + let mut ctx = SESSION.create_execution_ctx(); + let result = BoolArray::from_iter([true, false]) + .into_array() + .cast(DType::Binary(Nullability::NonNullable)) + .and_then(|array| { + array + .execute::(&mut ctx) + .map(|canonical| canonical.into_array()) + }); + + assert!(result.is_err(), "Expected error, got: {result:?}"); + } + #[rstest] #[case(BoolArray::from_iter(vec![true, false, true, true, false]))] #[case(BoolArray::from_iter(vec![Some(true), Some(false), None, Some(true), None]))] diff --git a/vortex-array/src/arrays/constant/compute/cast.rs b/vortex-array/src/arrays/constant/compute/cast.rs index 439bf8367b8..188752491eb 100644 --- a/vortex-array/src/arrays/constant/compute/cast.rs +++ b/vortex-array/src/arrays/constant/compute/cast.rs @@ -23,6 +23,7 @@ impl CastReduce for Constant { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; @@ -33,6 +34,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::scalar::DecimalValue; use crate::scalar::Scalar; @@ -65,4 +67,42 @@ mod tests { Some(DecimalValue::I128(4200)) ); } + + #[rstest] + #[case( + Scalar::from(true), + DType::Primitive(PType::I32, Nullability::NonNullable), + Scalar::primitive(1i32, Nullability::NonNullable) + )] + #[case( + Scalar::from(false), + DType::Utf8(Nullability::Nullable), + Scalar::utf8("false", Nullability::Nullable) + )] + #[case( + Scalar::from(-42i64), + DType::Utf8(Nullability::NonNullable), + Scalar::utf8("-42", Nullability::NonNullable) + )] + #[case( + Scalar::from(100.0f64), + DType::Utf8(Nullability::NonNullable), + Scalar::utf8("100.0", Nullability::NonNullable) + )] + #[case( + Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)), + DType::Utf8(Nullability::Nullable), + Scalar::null(DType::Utf8(Nullability::Nullable)) + )] + fn test_cast_bool_and_primitive_constants( + #[case] source: Scalar, + #[case] target: DType, + #[case] expected: Scalar, + ) -> VortexResult<()> { + let casted = ConstantArray::new(source, 5).into_array().cast(target)?; + + assert_eq!(casted.len(), 5); + assert_eq!(casted.as_constant(), Some(expected)); + Ok(()) + } } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 09defa442b8..b8cf2cd08f0 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt::Write; + use num_traits::AsPrimitive; use num_traits::NumCast; use vortex_buffer::Buffer; @@ -21,12 +23,16 @@ use crate::array::ArrayView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::half::f16; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; +use crate::match_each_integer_ptype; use crate::match_each_native_ptype; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; @@ -68,10 +74,11 @@ impl CastKernel for Primitive { dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let DType::Primitive(new_ptype, new_nullability) = dtype else { - return Ok(None); + let (new_ptype, new_nullability) = match dtype { + DType::Primitive(new_ptype, new_nullability) => (*new_ptype, *new_nullability), + DType::Utf8(_) => return Ok(Some(cast_primitive_to_utf8(array, dtype, ctx)?)), + _ => return Ok(None), }; - let (new_ptype, new_nullability) = (*new_ptype, *new_nullability); let src_ptype = array.ptype(); let new_validity = array @@ -262,25 +269,118 @@ fn cached_values_fit_in(array: ArrayView<'_, Primitive>, target_dtype: &DType) - Some(min.cast(target_dtype).is_ok() && max.cast(target_dtype).is_ok()) } +fn cast_primitive_to_utf8( + array: ArrayView<'_, Primitive>, + dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let new_validity = array + .validity()? + .cast_nullability(dtype.nullability(), len, ctx)?; + let mask = new_validity.execute_mask(len, ctx)?; + let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), len); + + // Match Arrow's formatting: ryu for f32/f64, Display for f16, and itoa for integers + // (Arrow's lexical_core produces the same output as itoa/Display for integers). + match array.ptype() { + PType::F16 => { + let mut scratch = String::with_capacity(16); + append_values_to_utf8( + &mut builder, + array.as_slice::(), + &mask, + |builder, value| { + scratch.clear(); + // Writing to a String is infallible. + let _ = write!(scratch, "{value}"); + builder.append_value(scratch.as_str()); + }, + ); + } + PType::F32 => { + let mut formatter = ryu::Buffer::new(); + append_values_to_utf8( + &mut builder, + array.as_slice::(), + &mask, + |builder, value| builder.append_value(formatter.format(value)), + ); + } + PType::F64 => { + let mut formatter = ryu::Buffer::new(); + append_values_to_utf8( + &mut builder, + array.as_slice::(), + &mask, + |builder, value| builder.append_value(formatter.format(value)), + ); + } + ptype => match_each_integer_ptype!(ptype, |T| { + let mut formatter = itoa::Buffer::new(); + append_values_to_utf8( + &mut builder, + array.as_slice::(), + &mask, + |builder, value| builder.append_value(formatter.format(value)), + ); + }), + } + + Ok(builder.finish_into_varbinview().into_array()) +} + +fn append_values_to_utf8( + builder: &mut VarBinViewBuilder, + values: &[T], + mask: &Mask, + mut append: impl FnMut(&mut VarBinViewBuilder, T), +) { + match mask { + Mask::AllTrue(_) => { + for &value in values { + append(builder, value); + } + } + Mask::AllFalse(_) => builder.append_nulls(values.len()), + Mask::Values(validity) => { + for (&value, valid) in values.iter().zip(validity.bit_buffer().iter()) { + if valid { + append(builder, value); + } else { + builder.append_null(); + } + } + } + } +} + #[cfg(test)] mod test { + use num_traits::NumCast; use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexError; + use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::half::f16; + use crate::match_each_native_ptype; use crate::validity::Validity; #[test] @@ -425,7 +525,7 @@ mod test { /// Same-width integer cast where all values fit: should reinterpret the /// buffer without allocation (pointer identity). #[test] - fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { + fn cast_same_width_int_reinterprets_buffer() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -458,7 +558,7 @@ mod test { /// All-null array cast between same-width types should succeed without /// touching the buffer contents. #[test] - fn cast_same_width_all_null() -> vortex_error::VortexResult<()> { + fn cast_same_width_all_null() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new(buffer![0xFFu8, 0xFF], Validity::AllInvalid); let casted = arr @@ -473,7 +573,7 @@ mod test { /// Same-width integer cast with nullable values: out-of-range nulls should /// not prevent the cast from succeeding. #[test] - fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + fn cast_same_width_int_nullable_with_out_of_range_nulls() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. @@ -494,7 +594,7 @@ mod test { } #[test] - fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + fn cast_u32_to_u8_with_out_of_range_nulls() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], @@ -531,4 +631,161 @@ mod test { fn test_cast_primitive_conformance(#[case] array: ArrayRef) { test_cast_conformance(&array, &mut array_session().create_execution_ctx()); } + + #[rstest] + #[case(PType::U8)] + #[case(PType::U16)] + #[case(PType::U32)] + #[case(PType::U64)] + #[case(PType::I8)] + #[case(PType::I16)] + #[case(PType::I32)] + #[case(PType::I64)] + #[case(PType::F16)] + #[case(PType::F32)] + #[case(PType::F64)] + fn cast_each_primitive_type_to_utf8(#[case] ptype: PType) -> VortexResult<()> { + let array = match_each_native_ptype!(ptype, |T| { + let zero = ::from(0u8) + .ok_or_else(|| vortex_err!("Cannot construct zero as {ptype}"))?; + let one = ::from(1u8) + .ok_or_else(|| vortex_err!("Cannot construct one as {ptype}"))?; + let answer = ::from(42u8) + .ok_or_else(|| vortex_err!("Cannot construct 42 as {ptype}"))?; + PrimitiveArray::from_iter([zero, one, answer]).into_array() + }); + let actual = array.cast(DType::Utf8(Nullability::NonNullable))?; + let expected = if matches!(ptype, PType::F32 | PType::F64) { + VarBinViewArray::from_iter_str(["0.0", "1.0", "42.0"]) + } else { + VarBinViewArray::from_iter_str(["0", "1", "42"]) + }; + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_nullable_primitive_to_utf8() -> VortexResult<()> { + let actual = PrimitiveArray::from_option_iter([Some(100i64), None, Some(-42)]) + .into_array() + .cast(DType::Utf8(Nullability::Nullable))?; + let expected = VarBinViewArray::from_iter_nullable_str([Some("100"), None, Some("-42")]); + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_all_null_primitive_to_utf8() -> VortexResult<()> { + let actual = PrimitiveArray::from_option_iter([None::, None]) + .into_array() + .cast(DType::Utf8(Nullability::Nullable))?; + let expected = VarBinViewArray::from_iter_nullable_str([None::<&str>, None]); + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_nullable_primitive_with_null_to_non_nullable_utf8_fails() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let result = PrimitiveArray::from_option_iter([Some(1i64), None]) + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))? + .execute::(&mut ctx); + + assert!(result.is_err(), "Expected error, got: {result:?}"); + Ok(()) + } + + #[test] + fn cast_all_valid_nullable_primitive_to_non_nullable_utf8() -> VortexResult<()> { + let actual = PrimitiveArray::from_option_iter([Some(1i64), Some(-42)]) + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))?; + let expected = VarBinViewArray::from_iter_str(["1", "-42"]); + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_f64_to_utf8_matches_arrow_formatting() -> VortexResult<()> { + let actual = buffer![ + 0.0f64, + -0.0, + 1.5, + 100.0, + 1e20, + 1e-20, + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY + ] + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))?; + let expected = VarBinViewArray::from_iter_str([ + "0.0", "-0.0", "1.5", "100.0", "1e20", "1e-20", "NaN", "inf", "-inf", + ]); + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_f16_to_utf8_matches_arrow_formatting() -> VortexResult<()> { + let actual = buffer![ + f16::from_f32(0.0), + f16::from_f32(-42.5), + f16::NAN, + f16::INFINITY, + f16::NEG_INFINITY + ] + .into_array() + .cast(DType::Utf8(Nullability::NonNullable))?; + let expected = VarBinViewArray::from_iter_str(["0", "-42.5", "NaN", "inf", "-inf"]); + + assert_arrays_eq!( + actual, + expected, + &mut array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn cast_primitive_to_binary_is_unsupported() { + let mut ctx = array_session().create_execution_ctx(); + let result = buffer![1i64, 2, 3] + .into_array() + .cast(DType::Binary(Nullability::NonNullable)) + .and_then(|array| { + array + .execute::(&mut ctx) + .map(|canonical| canonical.into_array()) + }); + + assert!(result.is_err(), "Expected error, got: {result:?}"); + } } diff --git a/vortex-array/src/scalar/tests/nested.rs b/vortex-array/src/scalar/tests/nested.rs index d02bf43c631..ad1c9182344 100644 --- a/vortex-array/src/scalar/tests/nested.rs +++ b/vortex-array/src/scalar/tests/nested.rs @@ -7,6 +7,8 @@ mod tests { use std::sync::Arc; + use vortex_error::VortexResult; + use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -505,7 +507,7 @@ mod tests { } #[test] - fn test_list_cast_incompatible_element_types() { + fn test_list_cast_bool_and_primitive_elements_to_utf8() -> VortexResult<()> { // Create a list of integers. let int_list = Scalar::list( Arc::from(DType::Primitive(PType::I32, Nullability::NonNullable)), @@ -513,12 +515,43 @@ mod tests { Nullability::NonNullable, ); - // Try to cast to list of strings - should fail. - let target = DType::List( - Arc::from(DType::Utf8(Nullability::NonNullable)), + let utf8_dtype = DType::Utf8(Nullability::NonNullable); + let target = DType::List(Arc::from(utf8_dtype.clone()), Nullability::NonNullable); + let casted = int_list.cast(&target)?; + let expected = Scalar::list( + utf8_dtype, + vec![Scalar::utf8("1", Nullability::NonNullable)], + Nullability::NonNullable, + ); + + assert_eq!(casted, expected); + + let bool_list = Scalar::list( + Arc::from(DType::Bool(Nullability::NonNullable)), + vec![ + Scalar::bool(true, Nullability::NonNullable), + Scalar::bool(false, Nullability::NonNullable), + ], + Nullability::NonNullable, + ); + let casted = bool_list.cast(&target)?; + let expected = Scalar::list( + DType::Utf8(Nullability::NonNullable), + vec![ + Scalar::utf8("true", Nullability::NonNullable), + Scalar::utf8("false", Nullability::NonNullable), + ], + Nullability::NonNullable, + ); + + assert_eq!(casted, expected); + + let unsupported_target = DType::List( + Arc::from(DType::Bool(Nullability::NonNullable)), Nullability::NonNullable, ); - assert!(int_list.cast(&target).is_err()); + assert!(int_list.cast(&unsupported_target).is_err()); + Ok(()) } #[test] diff --git a/vortex-array/src/scalar/typed_view/bool.rs b/vortex-array/src/scalar/typed_view/bool.rs index c971d1169c6..e9f7e82a068 100644 --- a/vortex-array/src/scalar/typed_view/bool.rs +++ b/vortex-array/src/scalar/typed_view/bool.rs @@ -7,11 +7,14 @@ use std::cmp::Ordering; use std::fmt::Display; use std::fmt::Formatter; +use num_traits::One; +use num_traits::Zero; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::dtype::DType; +use crate::match_each_native_ptype; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -83,15 +86,19 @@ impl<'a> BoolScalar<'a> { /// Casts this scalar to the given `dtype`. pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { - if !matches!(dtype, DType::Bool(..)) { - vortex_bail!( - "Cannot cast bool to {dtype}: boolean scalars can only be cast to boolean types with different nullability" - ) + let value = self.value.vortex_expect("nullness handled in Scalar::cast"); + + match dtype { + DType::Bool(nullability) => Ok(Scalar::bool(value, *nullability)), + DType::Primitive(ptype, nullability) => Ok(match_each_native_ptype!(*ptype, |T| { + Scalar::primitive(if value { T::one() } else { T::zero() }, *nullability) + })), + DType::Utf8(nullability) => Ok(Scalar::utf8( + if value { "true" } else { "false" }, + *nullability, + )), + _ => vortex_bail!("Cannot cast bool scalar to {dtype}"), } - Ok(Scalar::bool( - self.value.vortex_expect("nullness handled in Scalar::cast"), - dtype.nullability(), - )) } /// Returns a new boolean scalar with the inverted value. @@ -204,14 +211,31 @@ mod test { } #[test] - fn test_bool_cast_to_non_bool_fails() { + fn test_bool_cast_to_primitive_and_utf8() -> VortexResult<()> { use crate::dtype::PType; - let bool_scalar = Scalar::bool(true, NonNullable); - let bool = bool_scalar.as_bool(); + let true_scalar = Scalar::bool(true, NonNullable); + let false_scalar = Scalar::bool(false, NonNullable); - let result = bool.cast(&DType::Primitive(PType::I32, NonNullable)); - assert!(result.is_err()); + assert_eq!( + true_scalar.cast(&DType::Primitive(PType::I32, NonNullable))?, + Scalar::primitive(1i32, NonNullable) + ); + assert_eq!( + false_scalar.cast(&DType::Primitive(PType::F64, Nullable))?, + Scalar::primitive(0.0f64, Nullable) + ); + assert_eq!( + true_scalar.cast(&DType::Utf8(NonNullable))?, + Scalar::utf8("true", NonNullable) + ); + assert_eq!( + false_scalar.cast(&DType::Utf8(Nullable))?, + Scalar::utf8("false", Nullable) + ); + assert!(true_scalar.cast(&DType::Binary(NonNullable)).is_err()); + + Ok(()) } #[test] diff --git a/vortex-array/src/scalar/typed_view/primitive/scalar.rs b/vortex-array/src/scalar/typed_view/primitive/scalar.rs index 3ca4d337a45..4e9cbb4c6ef 100644 --- a/vortex-array/src/scalar/typed_view/primitive/scalar.rs +++ b/vortex-array/src/scalar/typed_view/primitive/scalar.rs @@ -181,6 +181,17 @@ impl<'a> PrimitiveScalar<'a> { *decimal_dtype, *nullability, )), + DType::Utf8(nullability) => { + // Match Arrow's formatting: ryu for f32/f64, Display for f16 and integers. + let value = match self.ptype { + PType::F32 => ryu::Buffer::new().format(pvalue.cast::()?).to_owned(), + PType::F64 => ryu::Buffer::new().format(pvalue.cast::()?).to_owned(), + ptype => { + match_each_native_ptype!(ptype, |T| { pvalue.cast::()?.to_string() }) + } + }; + Ok(Scalar::utf8(value, *nullability)) + } _ => vortex_bail!("Cannot cast primitive scalar to {dtype}"), } } diff --git a/vortex-array/src/scalar/typed_view/primitive/tests.rs b/vortex-array/src/scalar/typed_view/primitive/tests.rs index b0c0063a3df..c61a2ecb5ff 100644 --- a/vortex-array/src/scalar/typed_view/primitive/tests.rs +++ b/vortex-array/src/scalar/typed_view/primitive/tests.rs @@ -6,6 +6,7 @@ use std::cmp::Ordering; use num_traits::CheckedSub; use rstest::rstest; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_utils::aliases::hash_set::HashSet; use super::pvalue::CoercePValue; @@ -18,6 +19,7 @@ use crate::dtype::ToBytes; use crate::dtype::half::f16; use crate::scalar::PValue; use crate::scalar::PrimitiveScalar; +use crate::scalar::Scalar; use crate::scalar::ScalarValue; #[test] @@ -165,6 +167,33 @@ fn test_primitive_cast( } } +#[rstest] +#[case(Scalar::primitive(42u8, Nullability::NonNullable), "42")] +#[case(Scalar::primitive(-42i64, Nullability::NonNullable), "-42")] +#[case(Scalar::primitive(f16::from_f32(42.0), Nullability::NonNullable), "42")] +#[case(Scalar::primitive(100.0f32, Nullability::NonNullable), "100.0")] +#[case(Scalar::primitive(-0.0f64, Nullability::NonNullable), "-0.0")] +#[case(Scalar::primitive(f64::NAN, Nullability::NonNullable), "NaN")] +#[case(Scalar::primitive(f64::INFINITY, Nullability::NonNullable), "inf")] +#[case(Scalar::primitive(f64::NEG_INFINITY, Nullability::NonNullable), "-inf")] +fn test_primitive_cast_to_utf8(#[case] scalar: Scalar, #[case] expected: &str) -> VortexResult<()> { + let actual = scalar.cast(&DType::Utf8(Nullability::Nullable))?; + + assert_eq!(actual, Scalar::utf8(expected, Nullability::Nullable)); + Ok(()) +} + +#[test] +fn test_primitive_cast_to_binary_fails() { + let scalar = Scalar::primitive(42i64, Nullability::NonNullable); + + assert!( + scalar + .cast(&DType::Binary(Nullability::NonNullable)) + .is_err() + ); +} + #[test] fn test_as_conversion_success() { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index f9c4c65a460..b500663835e 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -728,8 +728,13 @@ mod tests { use arrow_schema::Schema; use arrow_schema::TimeUnit as ArrowTimeUnit; use datafusion::arrow::array::AsArray; + use datafusion::arrow::array::BooleanArray; + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::Int32Type; use datafusion_common::ScalarValue; + use datafusion_common::assert_batches_eq; use datafusion_common::config::ConfigOptions; use datafusion_expr::Operator as DFOperator; use datafusion_expr::ScalarUDF; @@ -1195,7 +1200,7 @@ mod tests { .show() .await?; - // This fails as it pushes string cast to the scan + // Exercise the fallback path with projection pushdown disabled. ctx.session .sql(r#"select cast(id as string) from 'example.vortex'"#) .await? @@ -1205,6 +1210,59 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_cast_to_string_with_projection_pushdown() -> anyhow::Result<()> { + let ctx = TestSessionContext::new(true); + let batch = RecordBatch::try_from_iter([ + ( + "bool_col", + Arc::new(BooleanArray::from(vec![Some(true), Some(false), None])) as _, + ), + ( + "int_col", + Arc::new(Int64Array::from(vec![Some(42), Some(-7), None])) as _, + ), + ( + "float_col", + Arc::new(Float64Array::from(vec![Some(1.5), Some(-2.25), None])) as _, + ), + ])?; + ctx.write_arrow_batch("files/cast_to_string.vortex", &batch) + .await?; + let provider = ctx + .table_provider("cast_to_string", "/files/", batch.schema()) + .await?; + ctx.session.register_table("cast_to_string", provider)?; + + let actual = ctx + .session + .sql( + "SELECT \ + CAST(bool_col AS STRING) AS b, \ + CAST(int_col AS STRING) AS i, \ + CAST(float_col AS STRING) AS f \ + FROM cast_to_string", + ) + .await? + .collect() + .await?; + + assert_batches_eq!( + [ + "+-------+----+-------+", + "| b | i | f |", + "+-------+----+-------+", + "| true | 42 | 1.5 |", + "| false | -7 | -2.25 |", + "| | | |", + "+-------+----+-------+", + ], + &actual + ); + + Ok(()) + } + /// A cast whose target is a UUID-tagged `FixedSizeBinary(16)` must resolve /// through the dtype extension registry (UUID is registered on the default /// session) instead of the static, non-plugin-aware `DType::from_arrow`, diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 36f3b5f4cd6..d57c1f64def 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -80,6 +80,16 @@ fn calculate_physical_field_type( DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => { if dtype.is_binary() || dtype.is_utf8() { logical_type.clone() + } else if matches!(logical_type, DataType::Utf8 | DataType::LargeUtf8) + && (dtype.is_int() || dtype.is_float() || dtype.is_boolean()) + { + // Preserve the file's physical type so the expression adapter can insert the + // cast to the unified logical string type. + arrow_session + .to_arrow_field("", dtype) + .map_err(|e| exec_datafusion_err!("Failed to convert dtype to arrow: {e}"))? + .data_type() + .clone() } else { return Err(exec_datafusion_err!( "Failed to convert dtype to arrow: Vortex DType is {dtype} which is not compatible with {logical_type}" @@ -267,6 +277,7 @@ mod tests { use std::sync::Arc; use arrow_schema::Fields; + use rstest::rstest; use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::StructFields; @@ -355,26 +366,83 @@ mod tests { assert_eq!(physical_schema.field(3).data_type(), &DataType::LargeBinary); } - #[test] - fn test_failing_conversion_incompatible_types() { - let logical_schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]); + #[rstest] + #[case( + DType::Primitive(PType::I32, Nullability::NonNullable), + DataType::Utf8, + DataType::Int32 + )] + #[case( + DType::Primitive(PType::F64, Nullability::Nullable), + DataType::Utf8, + DataType::Float64 + )] + #[case( + DType::Bool(Nullability::NonNullable), + DataType::Utf8, + DataType::Boolean + )] + #[case( + DType::Primitive(PType::I64, Nullability::Nullable), + DataType::LargeUtf8, + DataType::Int64 + )] + fn test_bool_and_numeric_file_column_under_utf8_logical_type( + #[case] physical_dtype: DType, + #[case] logical_type: DataType, + #[case] expected_physical_type: DataType, + ) -> DFResult<()> { + let logical_schema = Schema::new(vec![Field::new("col", logical_type, true)]); + let dtype = DType::Struct( + StructFields::from_iter([("col", physical_dtype)]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + assert_eq!( + physical_schema.field(0).data_type(), + &expected_physical_type + ); + Ok(()) + } + + #[rstest] + #[case( + DType::Primitive(PType::I32, Nullability::NonNullable), + DataType::Binary + )] + #[case(DType::Bool(Nullability::NonNullable), DataType::LargeBinary)] + #[case( + DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ), + DataType::Utf8 + )] + fn test_incompatible_file_column_under_string_or_binary_logical_type( + #[case] physical_dtype: DType, + #[case] logical_type: DataType, + ) { + let logical_schema = Schema::new(vec![Field::new("col", logical_type, false)]); let dtype = DType::Struct( - StructFields::from_iter([( - "col", - DType::Primitive(PType::I32, Nullability::NonNullable), - )]), + StructFields::from_iter([("col", physical_dtype)]), Nullability::NonNullable, ); let result = calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default()); + assert!( result .unwrap_err() .to_string() .contains("not compatible with") ); + } + #[test] + fn test_failing_conversion_incompatible_types() { // Test struct vs non-struct mismatch let logical_schema = Schema::new(vec![Field::new( "col", diff --git a/vortex-datafusion/src/tests/schema_evolution.rs b/vortex-datafusion/src/tests/schema_evolution.rs index 093fed6c148..03009530fbe 100644 --- a/vortex-datafusion/src/tests/schema_evolution.rs +++ b/vortex-datafusion/src/tests/schema_evolution.rs @@ -11,8 +11,12 @@ use arrow_schema::Fields; use arrow_schema::Schema; use datafusion::arrow::array::Array; use datafusion::arrow::array::ArrayRef as ArrowArrayRef; +use datafusion::arrow::array::BooleanArray; use datafusion::arrow::array::DictionaryArray; +use datafusion::arrow::array::Float64Array; +use datafusion::arrow::array::Int64Array; use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::StringArray; use datafusion::arrow::array::StructArray; use datafusion::arrow::datatypes::UInt16Type; use datafusion::arrow::datatypes::UInt32Type; @@ -167,6 +171,194 @@ async fn test_filter_schema_evolution_order( Ok(()) } +#[rstest] +#[tokio::test] +async fn test_schema_evolution_type_widening_chain( + #[values(false, true)] projection_pushdown: bool, +) -> anyhow::Result<()> { + let ctx = TestSessionContext::new(projection_pushdown); + + let bool_batch = RecordBatch::try_from_iter([( + "value", + Arc::new(BooleanArray::from(vec![Some(true), Some(false), None])) as ArrowArrayRef, + )])?; + let int_batch = RecordBatch::try_from_iter([( + "value", + Arc::new(Int64Array::from(vec![Some(42), Some(-7), None])) as ArrowArrayRef, + )])?; + let float_batch = RecordBatch::try_from_iter([( + "value", + Arc::new(Float64Array::from(vec![Some(1.5), Some(-2.25), None])) as ArrowArrayRef, + )])?; + let utf8_batch = RecordBatch::try_from_iter([( + "value", + Arc::new(StringArray::from(vec![Some("alpha"), Some("beta"), None])) as ArrowArrayRef, + )])?; + + ctx.write_arrow_batch("widening/int/bool.vortex", &bool_batch) + .await?; + ctx.write_arrow_batch("widening/int/int.vortex", &int_batch) + .await?; + + let int_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + true, + )])); + let int_provider = ctx + .table_provider("widening_int", "/widening/int/", Arc::clone(&int_schema)) + .await?; + let int_table = ctx.session.read_table(int_provider)?; + + assert_eq!(int_table.schema().as_arrow(), int_schema.as_ref()); + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| |", + "| |", + "| -7 |", + "| 0 |", + "| 1 |", + "| 42 |", + "+-------+", + ], + &int_table.clone().collect().await? + ); + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 1 |", + "| 42 |", + "+-------+", + ], + &int_table + .filter(col("value").eq(lit(1_i64)).or(col("value").eq(lit(42_i64))),)? + .collect() + .await? + ); + + ctx.write_arrow_batch("widening/float/bool.vortex", &bool_batch) + .await?; + ctx.write_arrow_batch("widening/float/int.vortex", &int_batch) + .await?; + ctx.write_arrow_batch("widening/float/float.vortex", &float_batch) + .await?; + + let float_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); + let float_provider = ctx + .table_provider( + "widening_float", + "/widening/float/", + Arc::clone(&float_schema), + ) + .await?; + let float_table = ctx.session.read_table(float_provider)?; + + assert_eq!(float_table.schema().as_arrow(), float_schema.as_ref()); + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| |", + "| |", + "| |", + "| -2.25 |", + "| -7.0 |", + "| 0.0 |", + "| 1.0 |", + "| 1.5 |", + "| 42.0 |", + "+-------+", + ], + &float_table.clone().collect().await? + ); + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 1.5 |", + "| 42.0 |", + "+-------+", + ], + &float_table + .filter( + col("value") + .eq(lit(1.5_f64)) + .or(col("value").eq(lit(42.0_f64))), + )? + .collect() + .await? + ); + + ctx.write_arrow_batch("widening/utf8/bool.vortex", &bool_batch) + .await?; + ctx.write_arrow_batch("widening/utf8/int.vortex", &int_batch) + .await?; + ctx.write_arrow_batch("widening/utf8/float.vortex", &float_batch) + .await?; + ctx.write_arrow_batch("widening/utf8/utf8.vortex", &utf8_batch) + .await?; + + let utf8_schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + let utf8_provider = ctx + .table_provider("widening_utf8", "/widening/utf8/", Arc::clone(&utf8_schema)) + .await?; + let utf8_table = ctx.session.read_table(utf8_provider)?; + + assert_eq!(utf8_table.schema().as_arrow(), utf8_schema.as_ref()); + + let full_scan = utf8_table.clone().collect().await?; + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| |", + "| |", + "| |", + "| |", + "| -2.25 |", + "| -7 |", + "| 1.5 |", + "| 42 |", + "| alpha |", + "| beta |", + "| false |", + "| true |", + "+-------+", + ], + &full_scan + ); + + let filtered_scan = utf8_table + .filter(col("value").eq(lit("true")).or(col("value").eq(lit("42"))))? + .collect() + .await?; + assert_batches_sorted_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 42 |", + "| true |", + "+-------+", + ], + &filtered_scan + ); + + Ok(()) +} + /// Test for correct schema evolution behavior in the presence of nested struct fields. /// We use a hypothetical schema of some observability data with "wide records", struct columns /// with nullable payloads that may or may not be present for every file.