Fix handling JSON type when reading FlatGeobuf (#901)

Closes #821
geoarrow · Dec 4, 2024 · 901c8ab · 901c8ab
1 parent 29d3491
commit 901c8ab
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 22 deletions.
diff --git a/fixtures/flatgeobuf/alldatatypes.fgb b/fixtures/flatgeobuf/alldatatypes.fgb
diff --git a/rust/geoarrow/src/io/flatgeobuf/reader/sync.rs b/rust/geoarrow/src/io/flatgeobuf/reader/sync.rs
@@ -189,29 +189,32 @@ mod test {
         ));
     }
 
-    #[ignore = "fails on JSON columns"]
     #[test]
     fn test_all_datatypes() {
         let mut filein =
             BufReader::new(File::open("fixtures/flatgeobuf/alldatatypes.fgb").unwrap());
         let table = read_flatgeobuf(&mut filein, Default::default()).unwrap();
 
-        let _geom_col = table.geometry_column(None).unwrap();
-        // assert!(matches!(geom_col.data_type(), NativeType::Polygon(_, _)));
-
-        // let (batches, schema) = table.into_inner();
-        // assert_eq!(batches[0].num_rows(), 10);
-        // assert!(matches!(
-        //     schema.field_with_name("AREA").unwrap().data_type(),
-        //     DataType::Float64
-        // ));
-        // assert!(matches!(
-        //     schema.field_with_name("EAS_ID").unwrap().data_type(),
-        //     DataType::Int64
-        // ));
-        // assert!(matches!(
-        //     schema.field_with_name("PRFEDEA").unwrap().data_type(),
-        //     DataType::Utf8
-        // ));
+        let geom_col = table.geometry_column(None).unwrap();
+        assert!(matches!(geom_col.data_type(), NativeType::Point(_, _)));
+
+        let (batches, schema) = table.into_inner();
+        assert_eq!(batches[0].num_rows(), 1);
+        assert!(matches!(
+            schema.field_with_name("byte").unwrap().data_type(),
+            DataType::Int8
+        ));
+        assert!(matches!(
+            schema.field_with_name("float").unwrap().data_type(),
+            DataType::Float32
+        ));
+        assert!(matches!(
+            schema.field_with_name("json").unwrap().data_type(),
+            DataType::Utf8
+        ));
+        assert!(matches!(
+            schema.field_with_name("binary").unwrap().data_type(),
+            DataType::Binary
+        ));
     }
 }
diff --git a/rust/geoarrow/src/io/geozero/table/builder/anyvalue.rs b/rust/geoarrow/src/io/geozero/table/builder/anyvalue.rs
@@ -125,9 +125,17 @@ impl AnyBuilder {
         }
     }
 
-    pub fn from_data_type_with_capacity(data_type: &DataType, capacity: usize) -> Self {
+    pub fn from_field_with_capacity(field: &Field, capacity: usize) -> Self {
         use AnyBuilder::*;
-        match data_type {
+
+        // Short circuit check for JSON type
+        if let Some(ext_val) = field.metadata().get("ARROW:extension:name") {
+            if ext_val.as_str() == "arrow.json" {
+                return Json(StringBuilder::with_capacity(capacity, 0));
+            }
+        }
+
+        match field.data_type() {
             DataType::Boolean => Bool(BooleanBuilder::with_capacity(capacity)),
             DataType::Int8 => Int8(Int8Builder::with_capacity(capacity)),
             DataType::UInt8 => Uint8(UInt8Builder::with_capacity(capacity)),
@@ -146,7 +154,7 @@ impl AnyBuilder {
                 TimestampMicrosecondBuilder::with_capacity(capacity),
                 tz.clone(),
             )),
-            _ => todo!("Unsupported type {data_type}"),
+            _ => todo!("Unsupported type {}", field.data_type()),
         }
     }
 

diff --git a/rust/geoarrow/src/io/geozero/table/builder/properties.rs b/rust/geoarrow/src/io/geozero/table/builder/properties.rs
@@ -98,7 +98,7 @@ impl PropertiesBatchBuilder {
         for field in schema.fields().iter() {
             columns.insert(
                 field.name().clone(),
-                AnyBuilder::from_data_type_with_capacity(field.data_type(), capacity),
+                AnyBuilder::from_field_with_capacity(field, capacity),
             );
         }