fix returning of values vs array

Cargo.toml

@@ -15,6 +15,7 @@ datafusion = "41"
 jiter = "0.5"
 paste = "1"
 log = "0.4"
+arrow-cast = "52.0.0"
 
 [dev-dependencies]
 codspeed-criterion-compat = "2.3"

src/common.rs

@@ -3,6 +3,8 @@ use std::str::Utf8Error;
 use datafusion::arrow::array::{
     Array, ArrayRef, AsArray, Int64Array, LargeStringArray, StringArray, StringViewArray, UInt64Array,
 };
+use datafusion::arrow::error::ArrowError;
+use arrow_cast::cast;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::common::{exec_err, plan_err, Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::ColumnarValue;
@@ -46,6 +48,16 @@ fn is_int(d: &DataType) -> bool {
     matches!(d, DataType::UInt64 | DataType::Int64)
 }
 
+/// Convert a dict array to a non-dict array.
+fn unpack_dict_array(array: ArrayRef) -> Result<ArrayRef, ArrowError> {
+    match array.data_type() {
+        DataType::Dictionary(_, value_type) => {
+            cast(array.as_ref(), value_type)
+        }
+        _ => Ok(array)
+    }
+}
+
 fn undict(d: &DataType) -> &DataType {
     if let DataType::Dictionary(_, value) = d {
         value.as_ref()
@@ -129,7 +141,8 @@ fn invoke_array<C: FromIterator<Option<I>> + 'static, I>(
     jiter_find: impl Fn(Option<&str>, &[JsonPath]) -> Result<I, GetError>,
 ) -> DataFusionResult<ArrayRef> {
     if let Some(d) = needle_array.as_any_dictionary_opt() {
-        invoke_array(json_array, d.values(), to_array, jiter_find)
+        let values = invoke_array(json_array, d.values(), to_array, jiter_find)?;
+        return unpack_dict_array(d.with_values(values)).map_err(Into::into);
     } else if let Some(str_path_array) = needle_array.as_any().downcast_ref::<StringArray>() {
         let paths = str_path_array.iter().map(|opt_key| opt_key.map(JsonPath::Key));
         zip_apply(json_array, paths, to_array, jiter_find, true)
@@ -160,7 +173,8 @@ fn zip_apply<'a, P: Iterator<Item = Option<JsonPath<'a>>>, C: FromIterator<Optio
     if let Some(d) = json_array.as_any_dictionary_opt() {
         // NOTE we do NOT map back to an dictionary as that doesn't work for `is null` or filtering
         // see https://github.com/apache/datafusion/issues/12380
-        return zip_apply(d.values(), path_array, to_array, jiter_find, object_lookup);
+        let values = zip_apply(d.values(), path_array, to_array, jiter_find, object_lookup)?;
+        return unpack_dict_array(d.with_values(values)).map_err(Into::into);
     }
 
     let c = if let Some(string_array) = json_array.as_any().downcast_ref::<StringArray>() {
@@ -226,7 +240,8 @@ fn scalar_apply<C: FromIterator<Option<I>>, I>(
 ) -> DataFusionResult<ArrayRef> {
     if let Some(d) = json_array.as_any_dictionary_opt() {
         // as above, don't return a dict
-        return scalar_apply(d.values(), path, to_array, jiter_find);
+        let values = scalar_apply(d.values(), path, to_array, jiter_find)?;
+        return unpack_dict_array(d.with_values(values)).map_err(Into::into);
     }
 
     let c = if let Some(string_array) = json_array.as_any().downcast_ref::<StringArray>() {

tests/main.rs

@@ -1,10 +1,14 @@
-use datafusion::arrow::datatypes::DataType;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, RecordBatch};
+use datafusion::arrow::datatypes::{Field, Int8Type, Schema};
+use datafusion::arrow::{array::StringDictionaryBuilder, datatypes::DataType};
 use datafusion::assert_batches_eq;
 use datafusion::common::ScalarValue;
 use datafusion::logical_expr::ColumnarValue;
 
 use datafusion_functions_json::udfs::json_get_str_udf;
-use utils::{display_val, logical_plan, run_query, run_query_large, run_query_params};
+use utils::{create_context, display_val, logical_plan, run_query, run_query_large, run_query_params};
 
 mod utils;
 
@@ -1261,3 +1265,40 @@ async fn test_dict_get_int() {
     let batches = run_query(sql).await.unwrap();
     assert_batches_eq!(expected, &batches);
 }
+
+
+#[tokio::test]
+async fn test_dict_filter() {
+    let mut builder = StringDictionaryBuilder::<Int8Type>::new();
+    builder.append(r#"{"foo": "bar"}"#).unwrap();
+    builder.append(r#"{"baz": "fizz"}"#).unwrap();
+    builder.append("nah").unwrap();
+    builder.append_null();
+
+    let dict = builder.finish();
+    let array = Arc::new(dict) as ArrayRef;
+
+    let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), true)])); 
+
+    let data = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
+
+    let ctx = create_context().await.unwrap();
+    ctx.register_batch("data", data).unwrap();
+
+    let sql = "select json_get(x, 'baz') v from data";
+    let expected = [
+        "+------------+",
+        "| v          |",
+        "+------------+",
+        "| {null=}    |",
+        "| {str=fizz} |",
+        "| {null=}    |",
+        "| {null=}    |",
+        "+------------+",
+    ];
+
+    let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+
+    assert_batches_eq!(expected, &batches);
+
+}

tests/utils/mod.rs

@@ -13,10 +13,16 @@ use datafusion::execution::context::SessionContext;
 use datafusion::prelude::SessionConfig;
 use datafusion_functions_json::register_all;
 
-async fn create_test_table(large_utf8: bool) -> Result<SessionContext> {
+
+pub async fn create_context() -> Result<SessionContext> {
     let config = SessionConfig::new().set_str("datafusion.sql_parser.dialect", "postgres");
     let mut ctx = SessionContext::new_with_config(config);
     register_all(&mut ctx)?;
+    Ok(ctx)
+}
+
+async fn create_test_table(large_utf8: bool) -> Result<SessionContext> {
+    let ctx = create_context().await?;
 
     let test_data = [
         ("object_foo", r#" {"foo": "abc"} "#),