Set arrow.json extension type on geozero JSON type (#819)

Closes #730
geoarrow · Oct 9, 2024 · 21892fb · 21892fb
1 parent 87992cc
commit 21892fb
Show file tree

Hide file tree

Showing 5 changed files with 114 additions and 149 deletions.
diff --git a/src/io/flatgeobuf/reader/async.rs b/src/io/flatgeobuf/reader/async.rs
@@ -53,7 +53,7 @@ pub async fn read_flatgeobuf_async(
         options.coord_type,
         true,
         options.batch_size,
-        Some(Arc::new(schema.finish())),
+        Some(schema),
         features_count,
         Default::default(),
     );

diff --git a/src/io/flatgeobuf/reader/common.rs b/src/io/flatgeobuf/reader/common.rs
@@ -1,4 +1,7 @@
-use arrow_schema::{DataType, Field, SchemaBuilder, TimeUnit};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use arrow_schema::{DataType, Field, SchemaBuilder, SchemaRef, TimeUnit};
 use flatgeobuf::{ColumnType, Header};
 
 use crate::array::CoordType;
@@ -28,7 +31,7 @@ impl Default for FlatGeobufReaderOptions {
     }
 }
 
-pub(super) fn infer_schema(header: Header<'_>) -> SchemaBuilder {
+pub(super) fn infer_schema(header: Header<'_>) -> SchemaRef {
     let columns = header.columns().unwrap();
     let mut schema = SchemaBuilder::with_capacity(columns.len());
 
@@ -46,7 +49,11 @@ pub(super) fn infer_schema(header: Header<'_>) -> SchemaBuilder {
             ColumnType::Float => Field::new(col.name(), DataType::Float32, col.nullable()),
             ColumnType::Double => Field::new(col.name(), DataType::Float64, col.nullable()),
             ColumnType::String => Field::new(col.name(), DataType::Utf8, col.nullable()),
-            ColumnType::Json => Field::new(col.name(), DataType::Utf8, col.nullable()),
+            ColumnType::Json => {
+                let mut metadata = HashMap::with_capacity(1);
+                metadata.insert("ARROW:extension:name".to_string(), "arrow.json".to_string());
+                Field::new(col.name(), DataType::Utf8, col.nullable()).with_metadata(metadata)
+            }
             ColumnType::DateTime => Field::new(
                 col.name(),
                 DataType::Timestamp(TimeUnit::Microsecond, None),
@@ -60,5 +67,5 @@ pub(super) fn infer_schema(header: Header<'_>) -> SchemaBuilder {
         schema.push(field);
     }
 
-    schema
+    Arc::new(schema.finish())
 }
diff --git a/src/io/flatgeobuf/reader/sync.rs b/src/io/flatgeobuf/reader/sync.rs
@@ -28,7 +28,6 @@ use crate::io::geozero::table::{GeoTableBuilder, GeoTableBuilderOptions};
 use crate::table::Table;
 use flatgeobuf::{FgbReader, GeometryType};
 use std::io::{Read, Seek};
-use std::sync::Arc;
 
 /// Read a FlatGeobuf file to a Table
 pub fn read_flatgeobuf<R: Read + Seek>(
@@ -60,7 +59,7 @@ pub fn read_flatgeobuf<R: Read + Seek>(
         options.coord_type,
         true,
         options.batch_size,
-        Some(Arc::new(schema.finish())),
+        Some(schema),
         features_count,
         Default::default(),
     );