diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 0ccb8a4c28e..26ddfd0c1c6 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/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::>(); let arrow_schema = @@ -1234,6 +1239,42 @@ mod tests { ); } + #[test] + fn test_map_dictionary_to_values_preserves_timestamp_timezone() { + let dictionary = DictionaryArray::::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::() + .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( diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_tz.result b/tests/cases/standalone/common/types/timestamp/timestamp_tz.result index 8b4833066b7..73a279a4fec 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_tz.result +++ b/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 + diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql b/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql index cfb757f02f3..3783363fa63 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql +++ b/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;