chore(query): allow arrow dictionary array into databend (#17521)

feat(query): allow arrow dictionary array into databend

Author: sundy-li

src/query/expression/src/converts/arrow/from.rs

@@ -150,6 +150,11 @@ impl TryFrom<&Field> for TableField {
                         fields_type,
                     }
                 }
+                ArrowDataType::Dictionary(_, b) => {
+                    let inner_f =
+                        Field::new(arrow_f.name(), b.as_ref().clone(), arrow_f.is_nullable());
+                    return Self::try_from(&inner_f);
+                }
                 arrow_type => {
                     return Err(ErrorCode::Internal(format!(
                         "Unsupported Arrow type: {:?}",
@@ -235,7 +240,11 @@ impl Column {
         Field::from(&f)
     }
 
-    pub fn from_arrow_rs(array: ArrayRef, data_type: &DataType) -> Result<Self> {
+    pub fn from_arrow_rs(mut array: ArrayRef, data_type: &DataType) -> Result<Self> {
+        if let ArrowDataType::Dictionary(_, v) = array.data_type() {
+            array = arrow_cast::cast(array.as_ref(), v.as_ref())?;
+        }
+
         let column = match data_type {
             DataType::Null => Column::Null { len: array.len() },
             DataType::EmptyArray => Column::EmptyArray { len: array.len() },

src/query/expression/tests/it/arrow.rs

@@ -0,0 +1,48 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use arrow_array::ArrayRef;
+use arrow_array::DictionaryArray;
+use arrow_array::Int32Array;
+use arrow_array::StringArray;
+use databend_common_expression::types::DataType;
+use databend_common_expression::Column;
+
+#[test]
+fn test_from_arrow_rs() {
+    let array = Arc::new(StringArray::from(vec![
+        Some("foo"),
+        Some("bar"),
+        None,
+        Some("baz"),
+    ])) as ArrayRef;
+
+    let array_k = Int32Array::from(vec![0, 3, 2, 2]);
+    let array_v = Arc::new(StringArray::from(vec![
+        Some("foo"),
+        Some("bar"),
+        None,
+        Some("baz"),
+    ])) as ArrayRef;
+
+    let array_dict = Arc::new(DictionaryArray::new(array_k, array_v)) as _;
+
+    for arr in [array, array_dict] {
+        let column = Column::from_arrow_rs(arr.clone(), &DataType::String.wrap_nullable()).unwrap();
+        let c = column.into_arrow_rs();
+        assert_eq!(c.len(), arr.len());
+    }
+}

src/query/expression/tests/it/main.rs

@@ -24,6 +24,7 @@ use databend_common_expression::DataBlock;
 
 extern crate core;
 
+mod arrow;
 mod block;
 mod common;
 mod decimal;