andygrove commented on code in PR #627:
URL: https://github.com/apache/datafusion-comet/pull/627#discussion_r1670554215
##########
native/core/src/execution/datafusion/expressions/cast.rs:
##########
@@ -501,18 +504,24 @@ impl Cast {
let array = array_with_timezone(array, self.timezone.clone(),
Some(to_type))?;
let from_type = array.data_type().clone();
- // unpack dictionary string arrays first
- // TODO: we are unpacking a dictionary-encoded array and then
performing
- // the cast. We could potentially improve performance here by casting
the
- // dictionary values directly without unpacking the array first,
although this
- // would add more complexity to the code
let array = match &from_type {
DataType::Dictionary(key_type, value_type)
if key_type.as_ref() == &DataType::Int32
&& (value_type.as_ref() == &DataType::Utf8
|| value_type.as_ref() == &DataType::LargeUtf8) =>
{
- cast_with_options(&array, value_type.as_ref(), &CAST_OPTIONS)?
+ let dict_array = array
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .expect("Expected a dictionary array");
+ let values = dict_array.values();
+ let cast_values = self.cast_array(values.clone())?;
+ let cast_array = Arc::new(DictionaryArray::<Int32Type>::new(
+ dict_array.keys().clone(),
+ cast_values,
+ )) as ArrayRef;
+
+ cast_with_options(&cast_array, value_type, &CAST_OPTIONS)?
Review Comment:
This is converting the output dictionary array back into a normal array,
which we do not want. I think you just want to return the `dict_array` instead.
We really should have unit tests so that we can be sure that the code change
is doing what we think it is doing. I think this test does what we need:
```rust
#[test]
fn test_cast_dict_string_to_dict_timestamp() -> DataFusionResult<()> {
// prepare input data
let keys = Int32Array::from(vec![0, 1]);
let values: ArrayRef = Arc::new(StringArray::from(vec![
Some("2020-01-01T12:34:56.123456"),
Some("T2"),
]));
let dict_array = Arc::new(DictionaryArray::new(keys, values));
// prepare cast expression
let timezone = "UTC".to_string();
let expr = Arc::new(Column::new("a", 0)); // this is not used by the
test
let cast = Cast::new(
expr,
DataType::Timestamp(TimeUnit::Microsecond,
Some(timezone.clone().into())),
EvalMode::Legacy,
timezone.clone(),
);
// test casting string dictionary array to timestamp dictionary array
let result = cast.cast_array(dict_array)?;
assert_eq!(
result.data_type(),
&DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(DataType::Timestamp(
TimeUnit::Microsecond,
Some(timezone.into())
))
)
);
assert_eq!(result.len(), 2);
Ok(())
}
```
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
