Skip to content

Commit d7f07de

Browse files
authored
vortex-array: Attempt casting inner array to target for extensions (#7469)
## Summary Getting ``` No CastKernel to cast canonical array vortex.ext from vortex.timestamp[ns, tz=UTC](i64) to i64 ``` When running a datafusion query like: ``` SELECT cast(timestamp as bigint) as timestamp from table; ``` Where the timestamp array is a utc timestamp type column. Since the inner array of a timestamp array is int64 already, all we really need to do is try to attempt to cast the inner array of the extension type to the target type. That's exactly what this PR does. ## Testing Added a unit test. Signed-off-by: Frederic Branczyk <fbranczyk@gmail.com>
1 parent e4e7667 commit d7f07de

1 file changed

Lines changed: 26 additions & 1 deletion

File tree

  • vortex-array/src/arrays/extension/compute

vortex-array/src/arrays/extension/compute/cast.rs

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ impl CastReduce for Extension {
1717
dtype: &DType,
1818
) -> vortex_error::VortexResult<Option<ArrayRef>> {
1919
if !array.dtype().eq_ignore_nullability(dtype) {
20-
return Ok(None);
20+
// Target is not the same extension type.
21+
// Delegate to the storage array's cast.
22+
return Ok(Some(array.storage_array().cast(dtype.clone())?));
2123
}
2224

2325
let DType::Extension(ext_dtype) = dtype else {
@@ -51,9 +53,12 @@ mod tests {
5153
use super::*;
5254
use crate::IntoArray;
5355
use crate::arrays::PrimitiveArray;
56+
use crate::assert_arrays_eq;
5457
use crate::builtins::ArrayBuiltins;
5558
use crate::compute::conformance::cast::test_cast_conformance;
59+
use crate::dtype::DType;
5660
use crate::dtype::Nullability;
61+
use crate::dtype::PType;
5762
use crate::extension::datetime::TimeUnit;
5863
use crate::extension::datetime::Timestamp;
5964

@@ -108,6 +113,26 @@ mod tests {
108113
);
109114
}
110115

116+
#[test]
117+
fn cast_timestamp_to_i64() -> vortex_error::VortexResult<()> {
118+
let ext_dtype = Timestamp::new_with_tz(
119+
TimeUnit::Nanoseconds,
120+
Some("UTC".into()),
121+
Nullability::NonNullable,
122+
)
123+
.erased();
124+
let storage = buffer![1i64, 2, 3].into_array();
125+
let arr = ExtensionArray::new(ext_dtype, storage).into_array();
126+
127+
let result = arr.cast(DType::Primitive(PType::I64, Nullability::NonNullable))?;
128+
assert_eq!(
129+
result.dtype(),
130+
&DType::Primitive(PType::I64, Nullability::NonNullable)
131+
);
132+
assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array());
133+
Ok(())
134+
}
135+
111136
#[rstest]
112137
#[case(create_timestamp_array(TimeUnit::Milliseconds, false))]
113138
#[case(create_timestamp_array(TimeUnit::Microseconds, true))]

0 commit comments

Comments
 (0)