From 30bada9c9efc089fc3e6199fb520e29c63e5c242 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Sat, 11 Sep 2021 07:58:25 +0000
Subject: [PATCH] Fixed error in writing dictionary with extensions

---
 src/io/ipc/convert.rs         | 67 ++++++++++++++++++++---------------
 tests/it/io/ipc/write/file.rs |  6 ++++
 2 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/src/io/ipc/convert.rs b/src/io/ipc/convert.rs
index b0765bf5fa8..c42b72e041f 100644
--- a/src/io/ipc/convert.rs
+++ b/src/io/ipc/convert.rs
@@ -333,6 +333,29 @@ fn write_metadata<'a>(
     }
 }
 
+fn write_extension<'a>(
+    fbb: &mut FlatBufferBuilder<'a>,
+    name: &str,
+    metadata: &Option<String>,
+    kv_vec: &mut Vec<WIPOffset<ipc::KeyValue<'a>>>,
+) {
+    // metadata
+    if let Some(metadata) = metadata {
+        let kv_args = ipc::KeyValueArgs {
+            key: Some(fbb.create_string("ARROW:extension:metadata")),
+            value: Some(fbb.create_string(metadata.as_str())),
+        };
+        kv_vec.push(ipc::KeyValue::create(fbb, &kv_args));
+    }
+
+    // name
+    let kv_args = ipc::KeyValueArgs {
+        key: Some(fbb.create_string("ARROW:extension:name")),
+        value: Some(fbb.create_string(name)),
+    };
+    kv_vec.push(ipc::KeyValue::create(fbb, &kv_args));
+}
+
 /// Create an IPC Field from an Arrow Field
 pub(crate) fn build_field<'a>(
     fbb: &mut FlatBufferBuilder<'a>,
@@ -341,39 +364,16 @@ pub(crate) fn build_field<'a>(
     // custom metadata.
     let mut kv_vec = vec![];
     if let DataType::Extension(name, _, metadata) = field.data_type() {
-        // append extension information.
-
-        // metadata
-        if let Some(metadata) = metadata {
-            let kv_args = ipc::KeyValueArgs {
-                key: Some(fbb.create_string("ARROW:extension:metadata")),
-                value: Some(fbb.create_string(metadata.as_str())),
-            };
-            kv_vec.push(ipc::KeyValue::create(fbb, &kv_args));
-        }
-
-        // name
-        let kv_args = ipc::KeyValueArgs {
-            key: Some(fbb.create_string("ARROW:extension:name")),
-            value: Some(fbb.create_string(name.as_str())),
-        };
-        kv_vec.push(ipc::KeyValue::create(fbb, &kv_args));
+        write_extension(fbb, name, metadata, &mut kv_vec);
     }
-    if let Some(metadata) = field.metadata() {
-        if !metadata.is_empty() {
-            write_metadata(fbb, metadata, &mut kv_vec);
-        }
-    };
-    let fb_metadata = if !kv_vec.is_empty() {
-        Some(fbb.create_vector(&kv_vec))
-    } else {
-        None
-    };
 
     let fb_field_name = fbb.create_string(field.name().as_str());
     let field_type = get_fb_field_type(field.data_type(), field.is_nullable(), fbb);
 
-    let fb_dictionary = if let Dictionary(index_type, _) = field.data_type() {
+    let fb_dictionary = if let Dictionary(index_type, inner) = field.data_type() {
+        if let DataType::Extension(name, _, metadata) = inner.as_ref() {
+            write_extension(fbb, name, metadata, &mut kv_vec);
+        }
         Some(get_fb_dictionary(
             index_type,
             field
@@ -388,6 +388,17 @@ pub(crate) fn build_field<'a>(
         None
     };
 
+    if let Some(metadata) = field.metadata() {
+        if !metadata.is_empty() {
+            write_metadata(fbb, metadata, &mut kv_vec);
+        }
+    };
+    let fb_metadata = if !kv_vec.is_empty() {
+        Some(fbb.create_vector(&kv_vec))
+    } else {
+        None
+    };
+
     let mut field_builder = ipc::FieldBuilder::new(fbb);
     field_builder.add_name(fb_field_name);
     if let Some(dictionary) = fb_dictionary {
diff --git a/tests/it/io/ipc/write/file.rs b/tests/it/io/ipc/write/file.rs
index cc7d292b223..449a057eeca 100644
--- a/tests/it/io/ipc/write/file.rs
+++ b/tests/it/io/ipc/write/file.rs
@@ -162,6 +162,12 @@ fn write_100_decimal() -> Result<()> {
     test_file("1.0.0-bigendian", "generated_decimal")
 }
 
+#[test]
+fn write_100_extension() -> Result<()> {
+    test_file("1.0.0-littleendian", "generated_extension")?;
+    test_file("1.0.0-bigendian", "generated_extension")
+}
+
 #[test]
 fn write_100_union() -> Result<()> {
     test_file("1.0.0-littleendian", "generated_union")?;