Repository: apache/arrow-rs
Pull request #6650 (2024-10-30 11:44:42)

fix: include chrono-tz in flight sql cli
# Which issue does this PR close?
\-

# Rationale for this change
Otherwise you get an error like:

```text
Error: format results

Caused by:
    Parser error: Invalid timezone \"Europe/Berlin\": only offset based timezones supported without chrono-tz feature
```

# What changes are included in this PR?
- enable required feature
- add test

# Are there any user-facing changes?
Printing results w/ time zones now works.


Patch:
diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml
index 5473c4816..702cb1360 100644
--- a/arrow-flight/Cargo.toml
+++ b/arrow-flight/Cargo.toml
@@ -65,7 +65,7 @@ flight-sql-experimental = ["arrow-arith", "arrow-data", "arrow-ord", "arrow-row"
 tls = ["tonic/tls"]
 
 # Enable CLI tools
-cli = ["anyhow", "arrow-cast/prettyprint", "clap", "tracing-log", "tracing-subscriber", "tonic/tls-webpki-roots"]
+cli = ["anyhow", "arrow-array/chrono-tz", "arrow-cast/prettyprint", "clap", "tracing-log", "tracing-subscriber", "tonic/tls-webpki-roots"]
 
 [dev-dependencies]
 arrow-cast = { workspace = true, features = ["prettyprint"] }
diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs
index c05f5e81c..acfbd9b53 100644
--- a/arrow-flight/src/encode.rs
+++ b/arrow-flight/src/encode.rs
@@ -506,6 +506,15 @@ fn prepare_field_for_flight(
                 .with_metadata(field.metadata().clone())
             }
         }
+        DataType::Map(inner, sorted) => Field::new(
+            field.name(),
+            DataType::Map(
+                prepare_field_for_flight(inner, dictionary_tracker, send_dictionaries).into(),
+                *sorted,
+            ),
+            field.is_nullable(),
+        )
+        .with_metadata(field.metadata().clone()),
         _ => field.as_ref().clone(),
     }
 }
@@ -685,6 +694,7 @@ mod tests {
     use arrow_cast::pretty::pretty_format_batches;
     use arrow_ipc::MetadataVersion;
     use arrow_schema::{UnionFields, UnionMode};
+    use builder::{GenericStringBuilder, MapBuilder};
     use std::collections::HashMap;
 
     use super::*;
@@ -1276,6 +1286,157 @@ mod tests {
         verify_flight_round_trip(vec![batch1, batch2, batch3]).await;
     }
 
+    #[tokio::test]
+    async fn test_dictionary_map_hydration() {
+        let mut builder = MapBuilder::new(
+            None,
+            StringDictionaryBuilder::<UInt16Type>::new(),
+            StringDictionaryBuilder::<UInt16Type>::new(),
+        );
+
+        // {"k1":"a","k2":null,"k3":"b"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("a");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("b");
+        builder.append(true).unwrap();
+
+        let arr1 = builder.finish();
+
+        // {"k1":"c","k2":null,"k3":"d"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("c");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("d");
+        builder.append(true).unwrap();
+
+        let arr2 = builder.finish();
+
+        let schema = Arc::new(Schema::new(vec![Field::new_map(
+            "dict_map",
+            "entries",
+            Field::new_dictionary("keys", DataType::UInt16, DataType::Utf8, false),
+            Field::new_dictionary("values", DataType::UInt16, DataType::Utf8, true),
+            false,
+            false,
+        )]));
+
+        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
+        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();
+
+        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);
+
+        let encoder = FlightDataEncoderBuilder::default().build(stream);
+
+        let mut decoder = FlightDataDecoder::new(encoder);
+        let expected_schema = Schema::new(vec![Field::new_map(
+            "dict_map",
+            "entries",
+            Field::new("keys", DataType::Utf8, false),
+            Field::new("values", DataType::Utf8, true),
+            false,
+            false,
+        )]);
+
+        let expected_schema = Arc::new(expected_schema);
+
+        // Builder without dictionary fields
+        let mut builder = MapBuilder::new(
+            None,
+            GenericStringBuilder::<i32>::new(),
+            GenericStringBuilder::<i32>::new(),
+        );
+
+        // {"k1":"a","k2":null,"k3":"b"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("a");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("b");
+        builder.append(true).unwrap();
+
+        let arr1 = builder.finish();
+
+        // {"k1":"c","k2":null,"k3":"d"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("c");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("d");
+        builder.append(true).unwrap();
+
+        let arr2 = builder.finish();
+
+        let mut expected_arrays = vec![arr1, arr2].into_iter();
+
+        while let Some(decoded) = decoder.next().await {
+            let decoded = decoded.unwrap();
+            match decoded.payload {
+                DecodedPayload::None => {}
+                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
+                DecodedPayload::RecordBatch(b) => {
+                    assert_eq!(b.schema(), expected_schema);
+                    let expected_array = expected_arrays.next().unwrap();
+                    let map_array =
+                        downcast_array::<MapArray>(b.column_by_name("dict_map").unwrap());
+
+                    assert_eq!(map_array, expected_array);
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_dictionary_map_resend() {
+        let mut builder = MapBuilder::new(
+            None,
+            StringDictionaryBuilder::<UInt16Type>::new(),
+            StringDictionaryBuilder::<UInt16Type>::new(),
+        );
+
+        // {"k1":"a","k2":null,"k3":"b"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("a");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("b");
+        builder.append(true).unwrap();
+
+        let arr1 = builder.finish();
+
+        // {"k1":"c","k2":null,"k3":"d"}
+        builder.keys().append_value("k1");
+        builder.values().append_value("c");
+        builder.keys().append_value("k2");
+        builder.values().append_null();
+        builder.keys().append_value("k3");
+        builder.values().append_value("d");
+        builder.append(true).unwrap();
+
+        let arr2 = builder.finish();
+
+        let schema = Arc::new(Schema::new(vec![Field::new_map(
+            "dict_map",
+            "entries",
+            Field::new_dictionary("keys", DataType::UInt16, DataType::Utf8, false),
+            Field::new_dictionary("values", DataType::UInt16, DataType::Utf8, true),
+            false,
+            false,
+        )]));
+
+        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
+        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();
+
+        verify_flight_round_trip(vec![batch1, batch2]).await;
+    }
+
     async fn verify_flight_round_trip(mut batches: Vec<RecordBatch>) {
         let expected_schema = batches.first().unwrap().schema();
 


Test patch:
diff --git a/arrow-flight/tests/flight_sql_client_cli.rs b/arrow-flight/tests/flight_sql_client_cli.rs
index 6e1f6142c..c8e9190e2 100644
--- a/arrow-flight/tests/flight_sql_client_cli.rs
+++ b/arrow-flight/tests/flight_sql_client_cli.rs
@@ -20,7 +20,7 @@ mod common;
 use std::{pin::Pin, sync::Arc};
 
 use crate::common::fixture::TestFixture;
-use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray, TimestampNanosecondArray};
 use arrow_flight::{
     decode::FlightRecordBatchStream,
     encode::FlightDataEncoderBuilder,
@@ -37,7 +37,7 @@ use arrow_flight::{
     HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
 };
 use arrow_ipc::writer::IpcWriteOptions;
-use arrow_schema::{ArrowError, DataType, Field, Schema};
+use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit};
 use assert_cmd::Command;
 use bytes::Bytes;
 use futures::{Stream, TryStreamExt};
@@ -77,13 +77,13 @@ async fn test_simple() {
 
     assert_eq!(
         std::str::from_utf8(&stdout).unwrap().trim(),
-        "+--------------+-----------+\
-        \n| field_string | field_int |\
-        \n+--------------+-----------+\
-        \n| Hello        | 42        |\
-        \n| lovely       |           |\
-        \n| FlightSQL!   | 1337      |\
-        \n+--------------+-----------+",
+        "+--------------+-----------+---------------------------+-----------------------------+\
+        \n| field_string | field_int | field_timestamp_nano_notz | field_timestamp_nano_berlin |\
+        \n+--------------+-----------+---------------------------+-----------------------------+\
+        \n| Hello        | 42        |                           |                             |\
+        \n| lovely       |           | 1970-01-01T00:00:00       | 1970-01-01T01:00:00+01:00   |\
+        \n| FlightSQL!   | 1337      | 2024-10-30T11:36:57       | 2024-10-30T12:36:57+01:00   |\
+        \n+--------------+-----------+---------------------------+-----------------------------+",
     );
 }
 
@@ -321,13 +321,13 @@ async fn test_do_put_prepared_statement(test_server: FlightSqlServiceImpl) {
 
     assert_eq!(
         std::str::from_utf8(&stdout).unwrap().trim(),
-        "+--------------+-----------+\
-        \n| field_string | field_int |\
-        \n+--------------+-----------+\
-        \n| Hello        | 42        |\
-        \n| lovely       |           |\
-        \n| FlightSQL!   | 1337      |\
-        \n+--------------+-----------+",
+        "+--------------+-----------+---------------------------+-----------------------------+\
+        \n| field_string | field_int | field_timestamp_nano_notz | field_timestamp_nano_berlin |\
+        \n+--------------+-----------+---------------------------+-----------------------------+\
+        \n| Hello        | 42        |                           |                             |\
+        \n| lovely       |           | 1970-01-01T00:00:00       | 1970-01-01T01:00:00+01:00   |\
+        \n| FlightSQL!   | 1337      | 2024-10-30T11:36:57       | 2024-10-30T12:36:57+01:00   |\
+        \n+--------------+-----------+---------------------------+-----------------------------+",
     );
 }
 
@@ -371,28 +371,46 @@ impl FlightSqlServiceImpl {
         FlightServiceServer::new(self.clone())
     }
 
-    fn fake_result() -> Result<RecordBatch, ArrowError> {
-        let schema = Schema::new(vec![
+    fn schema() -> Schema {
+        Schema::new(vec![
             Field::new("field_string", DataType::Utf8, false),
             Field::new("field_int", DataType::Int64, true),
-        ]);
+            Field::new(
+                "field_timestamp_nano_notz",
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+                true,
+            ),
+            Field::new(
+                "field_timestamp_nano_berlin",
+                DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("Europe/Berlin"))),
+                true,
+            ),
+        ])
+    }
+
+    fn fake_result() -> Result<RecordBatch, ArrowError> {
+        let schema = Self::schema();
 
         let string_array = StringArray::from(vec!["Hello", "lovely", "FlightSQL!"]);
         let int_array = Int64Array::from(vec![Some(42), None, Some(1337)]);
+        let timestamp_array =
+            TimestampNanosecondArray::from(vec![None, Some(0), Some(1730288217000000000)]);
+        let timestamp_ts_array = timestamp_array
+            .clone()
+            .with_timezone(Arc::from("Europe/Berlin"));
 
         let cols = vec![
             Arc::new(string_array) as ArrayRef,
             Arc::new(int_array) as ArrayRef,
+            Arc::new(timestamp_array) as ArrayRef,
+            Arc::new(timestamp_ts_array) as ArrayRef,
         ];
         RecordBatch::try_new(Arc::new(schema), cols)
     }
 
     fn create_fake_prepared_stmt() -> Result<ActionCreatePreparedStatementResult, ArrowError> {
         let handle = PREPARED_STATEMENT_HANDLE.to_string();
-        let schema = Schema::new(vec![
-            Field::new("field_string", DataType::Utf8, false),
-            Field::new("field_int", DataType::Int64, true),
-        ]);
+        let schema = Self::schema();
 
         let parameter_schema = Schema::new(vec![
             Field::new("$1", DataType::Utf8, false),