fixed dict bug

Author: savs02

arrow-pg/src/encoder.rs

@@ -294,27 +294,21 @@ pub fn encode_value<T: Encoder>(
     type_: &Type,
     format: FieldFormat,
 ) -> PgWireResult<()> {
-    println!("[DATAFUSION POSTGRES ENCODER] encode_value: idx={}, type={:?}", idx, arr.data_type());
     match arr.data_type() {
         DataType::Null => encoder.encode_field_with_type_and_format(&None::<i8>, type_, format)?,
         DataType::Boolean => {
-            println!("BOOLEAN BRANCH");
             encoder.encode_field_with_type_and_format(&get_bool_value(arr, idx), type_, format)?
         }
         DataType::Int8 => {
-            println!("BOOLEAN BRANCH");
             encoder.encode_field_with_type_and_format(&get_i8_value(arr, idx), type_, format)?
         }
         DataType::Int16 => {
-            println!("BOOLEAN BRANCH");
             encoder.encode_field_with_type_and_format(&get_i16_value(arr, idx), type_, format)?
         }
         DataType::Int32 => {
-            println!("BOOLEAN BRANCH");
             encoder.encode_field_with_type_and_format(&get_i32_value(arr, idx), type_, format)?
         }
         DataType::Int64 => {
-            println!("BOOLEAN BRANCH");
             encoder.encode_field_with_type_and_format(&get_i64_value(arr, idx), type_, format)?
         }
         DataType::UInt8 => encoder.encode_field_with_type_and_format(
@@ -496,7 +490,6 @@ pub fn encode_value<T: Encoder>(
             }
         },
         DataType::List(_) | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => {
-            println!("LIST BRANCH");
             if arr.is_null(idx) {
                 return encoder.encode_field_with_type_and_format(&None::<&[i8]>, type_, format);
             }
@@ -505,7 +498,6 @@ pub fn encode_value<T: Encoder>(
             encoder.encode_field_with_type_and_format(&value, type_, format)?
         }
         DataType::Struct(_) => {
-            println!("STRUCT BRANCH");
             let fields = match type_.kind() {
                 postgres_types::Kind::Composite(fields) => fields,
                 _ => {
@@ -518,66 +510,34 @@ pub fn encode_value<T: Encoder>(
             encoder.encode_field_with_type_and_format(&value, type_, format)?
         }
         DataType::Dictionary(_, value_type) => {
-            println!("DICT BRANCH");
             if arr.is_null(idx) {
                 return encoder.encode_field_with_type_and_format(&None::<i8>, type_, format);
             }
-
-            macro_rules! get_value_index {
-                ($key_ty:ty) => {
-                    arr.as_any()
-                        .downcast_ref::<DictionaryArray<$key_ty>>()
-                        .map(|dict| dict.keys().value(idx) as usize)
-                };
-            }
-
-            let value_idx = get_value_index!(Int8Type)
-                .or_else(|| get_value_index!(Int16Type))
-                .or_else(|| get_value_index!(Int32Type))
-                .or_else(|| get_value_index!(Int64Type))
-                .or_else(|| get_value_index!(UInt8Type))
-                .or_else(|| get_value_index!(UInt16Type))
-                .or_else(|| get_value_index!(UInt32Type))
-                .or_else(|| get_value_index!(UInt64Type))
-                .ok_or_else(|| {
-                    ToSqlError::from(format!(
-                        "Unsupported dictionary key type"
-                    ))
-                })?;
-
-
-            // Get the dictionary values, ignoring keys
-            // We'll use Int32Type as a common key type, but we're only interested in values
-            macro_rules! get_dict_values {
+            // Get the dictionary values and the mapped row index
+            macro_rules! get_dict_values_and_index {
                 ($key_type:ty) => {
                     arr.as_any()
                         .downcast_ref::<DictionaryArray<$key_type>>()
-                        .map(|dict| dict.values().clone())
+                        .map(|dict| (dict.values(), dict.keys().value(idx) as usize))
                 };
             }
 
             // Try to extract values using different key types
-            let values = get_dict_values!(Int8Type)
-                .or_else(|| get_dict_values!(Int16Type))
-                .or_else(|| get_dict_values!(Int32Type))
-                .or_else(|| get_dict_values!(Int64Type))
-                .or_else(|| get_dict_values!(UInt8Type))
-                .or_else(|| get_dict_values!(UInt16Type))
-                .or_else(|| get_dict_values!(UInt32Type))
-                .or_else(|| get_dict_values!(UInt64Type))
+            let (values, idx) = get_dict_values_and_index!(Int8Type)
+                .or_else(|| get_dict_values_and_index!(Int16Type))
+                .or_else(|| get_dict_values_and_index!(Int32Type))
+                .or_else(|| get_dict_values_and_index!(Int64Type))
+                .or_else(|| get_dict_values_and_index!(UInt8Type))
+                .or_else(|| get_dict_values_and_index!(UInt16Type))
+                .or_else(|| get_dict_values_and_index!(UInt32Type))
+                .or_else(|| get_dict_values_and_index!(UInt64Type))
                 .ok_or_else(|| {
                     ToSqlError::from(format!(
-                        "Unsupported dictionary key type"
+                        "Unsupported dictionary key type for value type {value_type}"
                     ))
                 })?;
 
-            // If the dictionary has only one value, treat it as a primitive
-            if values.len() == 1 {
-                encode_value(encoder, &values, 0, type_, format)?
-            } else {
-                // Otherwise, use value directly indexed by values array
-                encode_value(encoder, &values, idx, type_, format)?
-            }
+            encode_value(encoder, values, idx, type_, format)?
         }
         _ => {
             return Err(PgWireError::ApiError(ToSqlError::from(format!(
@@ -590,3 +550,56 @@ pub fn encode_value<T: Encoder>(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests{
+    use super::*;
+
+    #[test]
+    fn encodes_dictionary_array() {
+
+        #[derive(Default)]
+        struct MockEncoder {
+            encoded_value: String,
+        }
+
+        impl Encoder for MockEncoder {
+            fn encode_field_with_type_and_format<T>(
+                &mut self,
+                value: &T,
+                data_type: &Type,
+                _format: FieldFormat,
+            ) -> PgWireResult<()>
+            where
+                T: ToSql + ToSqlText + Sized {
+                let mut bytes = BytesMut::new();
+                let _sql_text = value.to_sql_text(data_type, &mut bytes);
+                let string = String::from_utf8((&bytes).to_vec());
+                self.encoded_value = string.unwrap();
+                Ok(())
+            }
+        }
+
+        let val = "~!@&$[]()@@!!";
+        let value = StringArray::from_iter_values([val]);
+        let keys = Int8Array::from_iter_values([0, 0, 0, 0]);
+        let dict_arr: Arc<dyn Array> = Arc::new(DictionaryArray::<Int8Type>::try_new(keys, Arc::new(value)).unwrap());
+
+
+        let mut encoder = MockEncoder::default();
+
+        let result = encode_value(
+            &mut encoder,
+            &dict_arr,
+            2,
+            &Type::TEXT,
+            FieldFormat::Text,
+        );
+
+        assert!(result.is_ok());
+
+        assert!(encoder.encoded_value == val);
+
+    }
+}
+