revert, nested

apache · Sep 8, 2024 · 95195e0 · 95195e0
1 parent 6222276
commit 95195e0
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 56 deletions.
diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs
@@ -2375,6 +2375,7 @@ async fn boolean_dictionary_as_filter() {
         keys_array,
         Arc::new(values_array) as Arc<dyn Array>,
     );
+    let array = Arc::new(array);
 
     let field = Field::new(
         "my_dict",
@@ -2383,7 +2384,7 @@ async fn boolean_dictionary_as_filter() {
     );
     let schema = Arc::new(Schema::new(vec![field]));
 
-    let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
+    let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap();
 
     let ctx = SessionContext::new();
 
@@ -2407,7 +2408,6 @@ async fn boolean_dictionary_as_filter() {
     ];
     assert_batches_eq!(expected, &df.clone().collect().await.unwrap());
 
-    // filter where is null
     let result_df = df.clone().filter(col("my_dict")).unwrap();
     let expected = [
         "+---------+",
@@ -2418,5 +2418,51 @@ async fn boolean_dictionary_as_filter() {
         "| true    |",
         "+---------+",
     ];
-    assert_batches_eq!(expected, &result_df.collect().await.unwrap());    
+    assert_batches_eq!(expected, &result_df.collect().await.unwrap());
+
+    // test nested dictionary
+    let keys = vec![0, 2]; // 0 -> true, 2 -> false
+    let keys_array = Int8Array::from(keys);
+    let nested_array = DictionaryArray::new(
+        keys_array,
+        array,
+    );
+
+    let field = Field::new(
+        "my_nested_dict",
+        DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Boolean)))),
+        true,
+    );
+
+    let schema = Arc::new(Schema::new(vec![field]));
+
+    let batch = RecordBatch::try_new(schema, vec![Arc::new(nested_array)]).unwrap();
+
+    ctx.register_batch("nested_dict_batch", batch).unwrap();
+
+    let df = ctx.table("nested_dict_batch").await.unwrap();
+
+    // view_all
+    let expected = [
+        "+----------------+",
+        "| my_nested_dict |",
+        "+----------------+",
+        "| true           |",
+        "| false          |",
+        "+----------------+",
+    ];
+
+    assert_batches_eq!(expected, &df.clone().collect().await.unwrap());
+
+    let result_df = df.clone().filter(col("my_nested_dict")).unwrap();
+    let expected = [
+        "+----------------+",
+        "| my_nested_dict |",
+        "+----------------+",
+        "| true           |",
+        "+----------------+",
+    ];
+
+    assert_batches_eq!(expected, &result_df.collect().await.unwrap());
+
 }
diff --git a/datafusion/physical-expr/src/expressions/is_not_null.rs b/datafusion/physical-expr/src/expressions/is_not_null.rs
@@ -117,7 +117,7 @@ mod tests {
     use super::*;
     use crate::expressions::col;
     use arrow::{
-        array::{BooleanArray, StringArray, StringDictionaryBuilder},
+        array::{BooleanArray, StringArray},
         datatypes::*,
     };
     use arrow_array::{Array, Float64Array, Int32Array, UnionArray};
@@ -189,38 +189,4 @@ mod tests {
 
         assert_eq!(expected, actual);
     }
-
-    #[test]
-    fn dictionary_is_not_null() {
-        let mut builder = StringDictionaryBuilder::<Int8Type>::new();
-        builder.append("a").unwrap();
-        builder.append("").unwrap();
-        builder.append_null();
-        builder.append("a").unwrap();
-        let array = builder.finish();
-
-        let field = Field::new(
-            "my_dict",
-            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
-            true,
-        );
-
-        let schema = Schema::new(vec![field]);
-        let expr = is_not_null(col("my_dict", &schema).unwrap()).unwrap();
-        let batch =
-            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
-
-        // expression: "a is not null"
-        let actual = expr
-            .evaluate(&batch)
-            .unwrap()
-            .into_array(batch.num_rows())
-            .expect("Failed to convert to array");
-
-        let actual = as_boolean_array(&actual).unwrap();
-
-        let expected = &BooleanArray::from(vec![true, true, false, true]);
-
-        assert_eq!(expected, actual);
-    }
 }
diff --git a/datafusion/physical-expr/src/expressions/is_null.rs b/datafusion/physical-expr/src/expressions/is_null.rs
@@ -181,7 +181,7 @@ mod tests {
     use super::*;
     use crate::expressions::col;
     use arrow::{
-        array::{BooleanArray, StringArray, StringDictionaryBuilder},
+        array::{BooleanArray, StringArray},
         datatypes::*,
     };
     use arrow_array::{Float64Array, Int32Array};
@@ -278,21 +278,4 @@ mod tests {
         let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
         assert_eq!(expected, &result);
     }
-
-    #[test]
-    fn dictionary_is_null() {
-        let mut builder = StringDictionaryBuilder::<Int8Type>::new();
-        builder.append("a").unwrap();
-        builder.append("").unwrap();
-        builder.append_null();
-        builder.append("a").unwrap();
-        let array = builder.finish();
-
-        let array_ref = Arc::new(array) as ArrayRef;
-        let result = compute_is_null(array_ref).unwrap();
-
-        let expected = &BooleanArray::from(vec![false, false, true, false]);
-
-        assert_eq!(expected, &result);
-    }
 }