fix(query): preserve Arrow types when decoding dictionary output [Backport release/v1.3] (#9867)

fix(query): preserve Arrow types when decoding dictionary output (#9858)


(cherry picked from commit 67bd279586)

Signed-off-by: luofucong <luofc@foxmail.com>

src/common/recordbatch/src/lib.rs

@@ -223,14 +223,19 @@ pub fn map_dictionary_to_values_schema(schema: SchemaRef) -> (SchemaRef, bool) {
         .arrow_schema()
         .fields()
         .iter()
+        .zip(schema.column_schemas())
         .zip(&columns)
-        .map(|(field, column)| {
-            Arc::new(
-                field
-                    .as_ref()
-                    .clone()
-                    .with_data_type(column.data_type.as_arrow_type()),
-            )
+        .map(|((field, original), column)| {
+            if original.data_type == column.data_type {
+                field.clone()
+            } else {
+                Arc::new(
+                    field
+                        .as_ref()
+                        .clone()
+                        .with_data_type(column.data_type.as_arrow_type()),
+                )
+            }
         })
         .collect::<Vec<_>>();
     let arrow_schema =
@@ -1234,6 +1239,42 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_map_dictionary_to_values_preserves_timestamp_timezone() {
+        let dictionary = DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![Some(0), None]),
+            Arc::new(StringArray::from(vec!["a"])),
+        );
+        let timestamps = datatypes::arrow::array::TimestampSecondArray::from(vec![Some(42), None])
+            .with_timezone("Europe/Paris");
+        let arrow_schema = Arc::new(ArrowSchema::new(vec![
+            Field::new("label", dictionary.data_type().clone(), true),
+            Field::new("instant", timestamps.data_type().clone(), true),
+        ]));
+        let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap());
+        let batch = DfRecordBatch::try_new(
+            arrow_schema.clone(),
+            vec![Arc::new(dictionary), Arc::new(timestamps)],
+        )
+        .unwrap();
+        let batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
+
+        let (mapped_schema, apply_mapper) = map_dictionary_to_values_schema(schema.clone());
+        assert!(apply_mapper);
+        assert_eq!(arrow_schema.field(1), mapped_schema.arrow_schema().field(1));
+        let mapped = map_dictionary_to_values(batch.clone(), &schema, &mapped_schema).unwrap();
+        assert_eq!(&ArrowDataType::Utf8, mapped.column(0).data_type());
+        assert_eq!(
+            &StringArray::from(vec![Some("a"), None]),
+            mapped
+                .column(0)
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap()
+        );
+        assert!(Arc::ptr_eq(batch.column(1), mapped.column(1)));
+    }
+
     #[test]
     fn test_map_dictionary_to_values_schema_with_duplicate_columns() {
         let dictionary_type = ArrowDataType::Dictionary(

tests/cases/standalone/common/types/timestamp/timestamp_tz.result

@@ -30,3 +30,50 @@ SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::TIME;
 | 15:50:09                                                         |
 +------------------------------------------------------------------+
 
+CREATE TABLE dictionary_output (
+    label STRING PRIMARY KEY,
+    epoch_seconds BIGINT,
+    ts TIMESTAMP TIME INDEX
+);
+
+Affected Rows: 0
+
+INSERT INTO dictionary_output VALUES ('a', 42, 0), ('b', NULL, 1000);
+
+Affected Rows: 2
+
+SELECT from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
++---------------------------+
+| instant                   |
++---------------------------+
+| 1970-01-01T01:00:42+01:00 |
+|                           |
++---------------------------+
+
+SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
++-------+---------------------------+
+| label | instant                   |
++-------+---------------------------+
+| a     | 1970-01-01T01:00:42+01:00 |
+| b     |                           |
++-------+---------------------------+
+
+-- SQLNESS PROTOCOL MYSQL
+SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
++-------+---------------------+
+| label | instant             |
++-------+---------------------+
+| a     | 1970-01-01 00:00:42 |
+| b     |                     |
++-------+---------------------+
+
+DROP TABLE dictionary_output;
+
+Affected Rows: 0
+

tests/cases/standalone/common/types/timestamp/timestamp_tz.sql

@@ -5,3 +5,23 @@ select '2021-11-15 02:30:00'::TIMESTAMP::TIMESTAMPTZ;
 SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::DATE;
 
 SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::TIME;
+
+CREATE TABLE dictionary_output (
+    label STRING PRIMARY KEY,
+    epoch_seconds BIGINT,
+    ts TIMESTAMP TIME INDEX
+);
+
+INSERT INTO dictionary_output VALUES ('a', 42, 0), ('b', NULL, 1000);
+
+SELECT from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
+SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
+-- SQLNESS PROTOCOL MYSQL
+SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant
+FROM dictionary_output ORDER BY ts;
+
+DROP TABLE dictionary_output;