This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 698c1b28 fix: unwrap dictionaries in CreateNamedStruct (#754)
698c1b28 is described below
commit 698c1b282be98bb87563e511f5c5739b51de5068
Author: Andy Grove <[email protected]>
AuthorDate: Fri Aug 2 11:47:58 2024 -0600
fix: unwrap dictionaries in CreateNamedStruct (#754)
---
.../datafusion/expressions/create_named_struct.rs | 78 ++++++++++++++++++++--
.../org/apache/comet/CometExpressionSuite.scala | 16 +++++
2 files changed, 87 insertions(+), 7 deletions(-)
diff --git
a/native/core/src/execution/datafusion/expressions/create_named_struct.rs
b/native/core/src/execution/datafusion/expressions/create_named_struct.rs
index 17d763cc..c06aacf1 100644
--- a/native/core/src/execution/datafusion/expressions/create_named_struct.rs
+++ b/native/core/src/execution/datafusion/expressions/create_named_struct.rs
@@ -15,6 +15,14 @@
// specific language governing permissions and limitations
// under the License.
+use arrow::compute::take;
+use arrow::record_batch::RecordBatch;
+use arrow_array::types::Int32Type;
+use arrow_array::{Array, DictionaryArray, StructArray};
+use arrow_schema::{DataType, Schema};
+use datafusion::logical_expr::ColumnarValue;
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_physical_expr::PhysicalExpr;
use std::{
any::Any,
fmt::{Display, Formatter},
@@ -22,13 +30,6 @@ use std::{
sync::Arc,
};
-use arrow::record_batch::RecordBatch;
-use arrow_array::StructArray;
-use arrow_schema::{DataType, Schema};
-use datafusion::logical_expr::ColumnarValue;
-use datafusion_common::{DataFusionError, Result as DataFusionResult};
-use datafusion_physical_expr::PhysicalExpr;
-
use crate::execution::datafusion::expressions::utils::down_cast_any_ref;
#[derive(Debug, Hash)]
@@ -63,6 +64,21 @@ impl PhysicalExpr for CreateNamedStruct {
.map(|expr| expr.evaluate(batch))
.collect::<datafusion_common::Result<Vec<_>>>()?;
let arrays = ColumnarValue::values_to_arrays(&values)?;
+ // TODO it would be more efficient if we could preserve dictionaries
within the
+ // struct array but for now we unwrap them to avoid runtime errors
+ // https://github.com/apache/datafusion-comet/issues/755
+ let arrays = arrays
+ .iter()
+ .map(|array| {
+ if let Some(dict_array) =
+ array.as_any().downcast_ref::<DictionaryArray<Int32Type>>()
+ {
+ take(dict_array.values().as_ref(), dict_array.keys(), None)
+ } else {
+ Ok(Arc::clone(array))
+ }
+ })
+ .collect::<Result<Vec<_>, _>>()?;
let fields = match &self.data_type {
DataType::Struct(fields) => fields,
_ => {
@@ -125,3 +141,51 @@ impl PartialEq<dyn Any> for CreateNamedStruct {
.unwrap_or(false)
}
}
+
+#[cfg(test)]
+mod test {
+ use super::CreateNamedStruct;
+ use arrow_array::{Array, DictionaryArray, Int32Array, RecordBatch,
StringArray};
+ use arrow_schema::{DataType, Field, Fields, Schema};
+ use datafusion_common::Result;
+ use datafusion_expr::ColumnarValue;
+ use datafusion_physical_expr_common::expressions::column::Column;
+ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+ use std::sync::Arc;
+
+ #[test]
+ fn test_create_struct_from_dict_encoded_i32() -> Result<()> {
+ let keys = Int32Array::from(vec![0, 1, 2]);
+ let values = Int32Array::from(vec![0, 111, 233]);
+ let dict = DictionaryArray::try_new(keys, Arc::new(values))?;
+ let data_type = DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Int32));
+ let schema = Schema::new(vec![Field::new("a", data_type, false)]);
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(dict)])?;
+ let data_type =
+ DataType::Struct(Fields::from(vec![Field::new("a",
DataType::Int32, false)]));
+ let x = CreateNamedStruct::new(vec![Arc::new(Column::new("a", 0))],
data_type);
+ let ColumnarValue::Array(x) = x.evaluate(&batch)? else {
+ unreachable!()
+ };
+ assert_eq!(3, x.len());
+ Ok(())
+ }
+
+ #[test]
+ fn test_create_struct_from_dict_encoded_string() -> Result<()> {
+ let keys = Int32Array::from(vec![0, 1, 2]);
+ let values = StringArray::from(vec!["a".to_string(), "b".to_string(),
"c".to_string()]);
+ let dict = DictionaryArray::try_new(keys, Arc::new(values))?;
+ let data_type = DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8));
+ let schema = Schema::new(vec![Field::new("a", data_type, false)]);
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(dict)])?;
+ let data_type =
+ DataType::Struct(Fields::from(vec![Field::new("a", DataType::Utf8,
false)]));
+ let x = CreateNamedStruct::new(vec![Arc::new(Column::new("a", 0))],
data_type);
+ let ColumnarValue::Array(x) = x.evaluate(&batch)? else {
+ unreachable!()
+ };
+ assert_eq!(3, x.len());
+ Ok(())
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index dd02be4f..f169538a 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -1887,4 +1887,20 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
}
+
+ test("struct and named_struct with dictionary") {
+ Seq(true, false).foreach { dictionaryEnabled =>
+ withParquetTable(
+ (0 until 100).map(i =>
+ (
+ i,
+ if (i % 2 == 0) { "even" }
+ else { "odd" })),
+ "tbl",
+ withDictionary = dictionaryEnabled) {
+ checkSparkAnswerAndOperator("SELECT struct(_1, _2) FROM tbl")
+ checkSparkAnswerAndOperator("SELECT named_struct('a', _1, 'b', _2)
FROM tbl")
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
