diff --git a/Cargo.lock b/Cargo.lock index c05aae6923e..3541870d8a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9606,6 +9606,7 @@ dependencies = [ "insta", "inventory", "itertools 0.14.0", + "itoa", "jiff", "memchr", "mimalloc", @@ -9623,6 +9624,7 @@ dependencies = [ "rstest", "rstest_reuse", "rustc-hash", + "ryu", "serde", "serde_json", "serde_test", diff --git a/Cargo.toml b/Cargo.toml index 36aa5b2ac9e..d44899c4e92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -176,6 +176,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" @@ -238,6 +239,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 746ce079a6a..5f1c7b01e64 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..8207b6d8049 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use num_traits::One; use num_traits::Zero; use vortex_buffer::BufferMut; @@ -13,7 +15,9 @@ use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; +use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; use crate::match_each_native_ptype; use crate::scalar_fn::fns::cast::CastKernel; @@ -53,6 +57,40 @@ 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 values = array.to_bit_buffer(); + let true_view = BinaryView::new_inlined(b"true"); + let false_view = BinaryView::new_inlined(b"false"); + let true_count = values.true_count(); + + let views = if true_count <= len - true_count { + let mut views = BufferMut::full(false_view, len); + values.for_each_set_index(|index| views[index] = true_view); + views + } else { + let mut views = BufferMut::full(true_view, len); + (!&values).for_each_set_index(|index| views[index] = false_view); + views + }; + + // SAFETY: every view is one of two known-valid inlined UTF-8 strings, no view + // references an external buffer, and cast_nullability returns matching validity. + return Ok(Some(unsafe { + VarBinViewArray::new_unchecked( + views.freeze(), + Arc::from([]), + dtype.clone(), + new_validity, + ) + .into_array() + })); + } + let DType::Primitive(new_ptype, new_nullability) = dtype else { return Ok(None); }; @@ -79,12 +117,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 +158,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..8d424cf6b4b 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; @@ -65,4 +66,27 @@ mod tests { Some(DecimalValue::I128(4200)) ); } + + #[rstest] + #[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) + )] + fn test_cast_bool_and_primitive_constants_to_utf8( + #[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 8f5a0f95fcd..6765fe4e058 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::CheckedMul; use num_traits::NumCast; @@ -24,6 +26,8 @@ use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -34,6 +38,7 @@ use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::ToI256; +use crate::dtype::half::f16; use crate::dtype::i256; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -85,10 +90,11 @@ impl CastKernel for Primitive { if let DType::Decimal(decimal_dtype, nullability) = dtype { return cast_to_decimal(array, *decimal_dtype, *nullability, ctx).map(Some); } - 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 @@ -616,6 +622,119 @@ 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_each_native_ptype!( + array.ptype(), + integral: |T| { + append_integer_values_to_utf8::(&mut builder, array.as_slice::(), &mask) + }, + floating: |T| { + append_float_values_to_utf8::(&mut builder, array.as_slice::(), &mask) + } + ); + + Ok(builder.finish_into_varbinview().into_array()) +} + +fn append_integer_values_to_utf8(builder: &mut VarBinViewBuilder, values: &[T], mask: &Mask) +where + T: NativePType + itoa::Integer, +{ + let mut formatter = itoa::Buffer::new(); + append_values_to_utf8(builder, values, mask, |builder, value| { + builder.append_value(formatter.format(value)); + }); +} + +trait FloatToUtf8: NativePType { + type Formatter; + + fn formatter() -> Self::Formatter; + + fn format(self, formatter: &mut Self::Formatter) -> &str; +} + +impl FloatToUtf8 for f16 { + type Formatter = String; + + fn formatter() -> Self::Formatter { + String::with_capacity(16) + } + + fn format(self, formatter: &mut Self::Formatter) -> &str { + formatter.clear(); + // Writing to a String is infallible. + let _ = write!(formatter, "{self}"); + formatter.as_str() + } +} + +macro_rules! impl_float_to_utf8 { + ($($ty:ty),+ $(,)?) => { + $( + impl FloatToUtf8 for $ty { + type Formatter = ryu::Buffer; + + fn formatter() -> Self::Formatter { + ryu::Buffer::new() + } + + fn format(self, formatter: &mut Self::Formatter) -> &str { + formatter.format(self) + } + } + )+ + }; +} + +impl_float_to_utf8!(f32, f64); + +fn append_float_values_to_utf8(builder: &mut VarBinViewBuilder, values: &[T], mask: &Mask) +where + T: FloatToUtf8, +{ + let mut formatter = T::formatter(); + append_values_to_utf8(builder, values, mask, |builder, value| { + builder.append_value(value.format(&mut formatter)); + }); +} + +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 rstest::rstest; @@ -623,14 +742,17 @@ mod test { 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::DecimalArray; 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; @@ -639,8 +761,10 @@ mod test { use crate::dtype::DecimalType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::half::f16; use crate::dtype::i256; use crate::expr::stats::Stat; + use crate::match_each_native_ptype; use crate::validity::Validity; #[test] @@ -1155,4 +1279,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..e9c40e770e7 100644 --- a/vortex-array/src/scalar/typed_view/bool.rs +++ b/vortex-array/src/scalar/typed_view/bool.rs @@ -83,15 +83,16 @@ 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::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 +205,21 @@ mod test { } #[test] - fn test_bool_cast_to_non_bool_fails() { - use crate::dtype::PType; - - let bool_scalar = Scalar::bool(true, NonNullable); - let bool = bool_scalar.as_bool(); + fn test_bool_cast_to_utf8() -> VortexResult<()> { + 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::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 1b4b1126ede..8b0567ef802 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`,