This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new fa0874b1e4 test: add regression test for unnesting dictionary encoded
columns (#14395)
fa0874b1e4 is described below
commit fa0874b1e44aecfeb0f1f3e89a011175fbdc7a60
Author: Duong Cong Toai <[email protected]>
AuthorDate: Sun Feb 2 19:45:07 2025 +0700
test: add regression test for unnesting dictionary encoded columns (#14395)
* chore: add regression test for unnest dict encoded cols
* chore: use dataframe api for testing
* chore: rm unused dep
---
datafusion/core/tests/dataframe/mod.rs | 67 ++++++++++++++++++++++++++++++++++
1 file changed, 67 insertions(+)
diff --git a/datafusion/core/tests/dataframe/mod.rs
b/datafusion/core/tests/dataframe/mod.rs
index c43fcf93f1..e570ec75c6 100644
--- a/datafusion/core/tests/dataframe/mod.rs
+++ b/datafusion/core/tests/dataframe/mod.rs
@@ -39,6 +39,7 @@ use datafusion_functions_aggregate::count::count_udaf;
use datafusion_functions_aggregate::expr_fn::{
array_agg, avg, count, count_distinct, max, median, min, sum,
};
+use datafusion_functions_nested::make_array::make_array_udf;
use datafusion_functions_window::expr_fn::{first_value, row_number};
use object_store::local::LocalFileSystem;
use sqlparser::ast::NullTreatment;
@@ -3358,6 +3359,72 @@ async fn unnest_columns() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn unnest_dict_encoded_columns() -> Result<()> {
+ let strings = vec!["x", "y", "z"];
+ let keys = Int32Array::from_iter(0..strings.len() as i32);
+
+ let utf8_values = StringArray::from(strings.clone());
+ let utf8_dict = DictionaryArray::new(keys.clone(), Arc::new(utf8_values));
+
+ let make_array_udf_expr1 = make_array_udf().call(vec![col("column1")]);
+ let batch =
+ RecordBatch::try_from_iter(vec![("column1", Arc::new(utf8_dict) as
ArrayRef)])?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("test", batch)?;
+ let df = ctx
+ .table("test")
+ .await?
+ .select(vec![
+ make_array_udf_expr1.alias("make_array_expr"),
+ col("column1"),
+ ])?
+ .unnest_columns(&["make_array_expr"])?;
+
+ let results = df.collect().await.unwrap();
+ let expected = [
+ "+-----------------+---------+",
+ "| make_array_expr | column1 |",
+ "+-----------------+---------+",
+ "| x | x |",
+ "| y | y |",
+ "| z | z |",
+ "+-----------------+---------+",
+ ];
+ assert_batches_eq!(expected, &results);
+
+ // make_array(dict_encoded_string,literal string)
+ let make_array_udf_expr2 = make_array_udf().call(vec![
+ col("column1"),
+ lit(ScalarValue::new_utf8("fixed_string")),
+ ]);
+ let df = ctx
+ .table("test")
+ .await?
+ .select(vec![
+ make_array_udf_expr2.alias("make_array_expr"),
+ col("column1"),
+ ])?
+ .unnest_columns(&["make_array_expr"])?;
+
+ let results = df.collect().await.unwrap();
+ let expected = [
+ "+-----------------+---------+",
+ "| make_array_expr | column1 |",
+ "+-----------------+---------+",
+ "| x | x |",
+ "| fixed_string | x |",
+ "| y | y |",
+ "| fixed_string | y |",
+ "| z | z |",
+ "| fixed_string | z |",
+ "+-----------------+---------+",
+ ];
+ assert_batches_eq!(expected, &results);
+ Ok(())
+}
+
#[tokio::test]
async fn unnest_column_nulls() -> Result<()> {
let df = table_with_lists_and_nulls().await?;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
