From a7af50d5e22160ace7f3380dc35d24f266f64b77 Mon Sep 17 00:00:00 2001 From: wirybeaver Date: Sun, 4 Oct 2026 12:53:30 -0700 Subject: [PATCH 1/2] feat: add geospatial read support --- Cargo.lock | 47 ++++ Cargo.toml | 1 + crates/iceberg/Cargo.toml | 2 + crates/iceberg/public-api.txt | 2 + crates/iceberg/src/arrow/schema.rs | 339 +++++++++++++++++++++--- crates/iceberg/src/arrow/value.rs | 46 ++-- crates/iceberg/src/avro/schema.rs | 18 +- crates/iceberg/src/spec/values/datum.rs | 12 +- crates/iceberg/src/spec/values/tests.rs | 16 +- 9 files changed, 401 insertions(+), 82 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9e92c31b3b..3c9bc1f4e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2278,6 +2278,26 @@ dependencies = [ "version_check", ] +[[package]] +name = "geo-traits" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e7c353d12a704ccfab1ba8bfb1a7fe6cb18b665bf89d37f4f7890edcd260206" +dependencies = [ + "geo-types", +] + +[[package]] +name = "geo-types" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5" +dependencies = [ + "approx", + "num-traits", + "serde", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -2739,6 +2759,7 @@ dependencies = [ "once_cell", "ordered-float 4.6.0", "parquet", + "parquet-geospatial", "pretty_assertions", "rand 0.9.5", "regex", @@ -4281,6 +4302,7 @@ dependencies = [ "num-bigint 0.5.1", "num-integer", "num-traits", + "parquet-geospatial", "ring", "seq-macro", "simdutf8", @@ -4290,6 +4312,19 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-geospatial" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "634827365851180ca26bf54c76e970dc1c14c075684c0580ee97d8d0c953faf1" +dependencies = [ + "arrow-schema", + "geo-traits", + "serde", + "serde_json", + "wkb", +] + [[package]] name = "paste" version = "1.0.15" @@ -7439,6 +7474,18 @@ version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" +[[package]] +name = "wkb" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a120b336c7ad17749026d50427c23d838ecb50cd64aaea6254b5030152f890a9" +dependencies = [ + "byteorder", + "geo-traits", + "num_enum", + "thiserror 1.0.69", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/Cargo.toml b/Cargo.toml index bb9770275a..82632a6d36 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,6 +113,7 @@ once_cell = "1.20" opendal = "0.58" ordered-float = "4" parquet = "59.2" +parquet-geospatial = "59.2" pilota = "0.11.10" pretty_assertions = "1.4" proc-macro2 = "1" diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index 99ec5f433f..994377fd06 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -69,7 +69,9 @@ parquet = { workspace = true, features = [ "async", "encryption", "arrow_canonical_extension_types", + "geospatial", ] } +parquet-geospatial = { workspace = true } rand = { workspace = true } reqwest = { workspace = true } roaring = { workspace = true } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index b86e48c778..98ab472aad 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1502,6 +1502,8 @@ pub fn iceberg::spec::geospatial::EdgeInterpolationAlgorithm::clone(&self) -> ic impl core::cmp::Eq for iceberg::spec::geospatial::EdgeInterpolationAlgorithm impl core::cmp::PartialEq for iceberg::spec::geospatial::EdgeInterpolationAlgorithm pub fn iceberg::spec::geospatial::EdgeInterpolationAlgorithm::eq(&self, other: &iceberg::spec::geospatial::EdgeInterpolationAlgorithm) -> bool +impl core::convert::From for iceberg::spec::geospatial::EdgeInterpolationAlgorithm +pub fn iceberg::spec::geospatial::EdgeInterpolationAlgorithm::from(edges: parquet_geospatial::types::Edges) -> Self impl core::default::Default for iceberg::spec::geospatial::EdgeInterpolationAlgorithm pub fn iceberg::spec::geospatial::EdgeInterpolationAlgorithm::default() -> iceberg::spec::geospatial::EdgeInterpolationAlgorithm impl core::fmt::Debug for iceberg::spec::geospatial::EdgeInterpolationAlgorithm diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index 34eb597d47..91b8d687ad 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -32,14 +32,16 @@ use arrow_schema::{ }; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::file::statistics::Statistics; +use parquet_geospatial::{WkbEdges, WkbMetadata, WkbType, WkbTypeHint}; use uuid::Uuid; use crate::error::{Result, invalid_data}; use crate::spec::decimal_utils::i128_from_be_bytes; +use crate::spec::geospatial::EdgeInterpolationAlgorithm; use crate::spec::{ - Datum, FIRST_FIELD_ID, LIST_FIELD_NAME, ListType, MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME, - MapType, NestedField, NestedFieldRef, PrimitiveLiteral, PrimitiveType, Schema, SchemaVisitor, - StructType, Type, VariantType, + Datum, FIRST_FIELD_ID, GeographyType, GeometryType, LIST_FIELD_NAME, ListType, + MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME, MapType, NestedField, NestedFieldRef, + PrimitiveLiteral, PrimitiveType, Schema, SchemaVisitor, StructType, Type, VariantType, }; use crate::{Error, ErrorKind}; @@ -47,6 +49,7 @@ use crate::{Error, ErrorKind}; pub const DEFAULT_MAP_FIELD_NAME: &str = "key_value"; /// UTC time zone for Arrow timestamp type. pub const UTC_TIME_ZONE: &str = "+00:00"; +const UNSET_GEOSPATIAL_CRS: &str = "srid:0"; /// The canonical Arrow [`arrow.parquet.variant`] extension type. /// @@ -101,6 +104,59 @@ impl ExtensionType for VariantExtensionType { } } +fn to_arrow_wkb_edges(algorithm: EdgeInterpolationAlgorithm) -> WkbEdges { + match algorithm { + EdgeInterpolationAlgorithm::Spherical => WkbEdges::Spherical, + EdgeInterpolationAlgorithm::Vincenty => WkbEdges::Vincenty, + EdgeInterpolationAlgorithm::Thomas => WkbEdges::Thomas, + EdgeInterpolationAlgorithm::Andoyer => WkbEdges::Andoyer, + EdgeInterpolationAlgorithm::Karney => WkbEdges::Karney, + } +} + +impl From for EdgeInterpolationAlgorithm { + fn from(edges: WkbEdges) -> Self { + match edges { + WkbEdges::Spherical => Self::Spherical, + WkbEdges::Vincenty => Self::Vincenty, + WkbEdges::Thomas => Self::Thomas, + WkbEdges::Andoyer => Self::Andoyer, + WkbEdges::Karney => Self::Karney, + } + } +} + +fn iceberg_crs_from_wkb_metadata(metadata: &WkbMetadata) -> Result { + match metadata.crs.as_ref() { + None => Ok(UNSET_GEOSPATIAL_CRS.to_string()), + Some(serde_json::Value::String(crs)) => Ok(crs.clone()), + Some(serde_json::Value::Object(crs)) => { + let id = crs.get("id"); + let authority = id + .and_then(|id| id.get("authority")) + .and_then(serde_json::Value::as_str) + .filter(|authority| !authority.is_empty()); + let code = match id.and_then(|id| id.get("code")) { + Some(serde_json::Value::String(code)) => Some(code.clone()), + Some(serde_json::Value::Number(code)) => Some(code.to_string()), + _ => None, + }; + + match (authority, code) { + (Some(authority), Some(code)) => Ok(format!("{authority}:{code}")), + _ => Err(Error::new( + ErrorKind::DataInvalid, + "Cannot determine Iceberg geospatial type from PROJJSON CRS without an embedded authority/code", + )), + } + } + Some(_) => Err(Error::new( + ErrorKind::DataInvalid, + "Geospatial CRS metadata must be a string or PROJJSON object", + )), + } +} + /// A post order arrow schema visitor. /// /// For order of methods called, please refer to the internal `visit_schema` function. @@ -413,7 +469,7 @@ impl ArrowSchemaConverter { let mut results = Vec::with_capacity(fields.len()); for i in 0..fields.len() { let field = &fields[i]; - let field_type = &field_results[i]; + let field_type = self.apply_arrow_extension(field, field_results[i].clone())?; let id = self.get_field_id(field)?; let doc = get_field_doc(field); let nested_field = NestedField { @@ -421,7 +477,7 @@ impl ArrowSchemaConverter { doc, name: field.name().clone(), required: !field.is_nullable(), - field_type: Box::new(field_type.clone()), + field_type: Box::new(field_type), initial_default: None, write_default: None, }; @@ -429,6 +485,51 @@ impl ArrowSchemaConverter { } Ok(results) } + + fn apply_arrow_extension(&self, arrow_field: &FieldRef, iceberg_type: Type) -> Result { + if arrow_field.extension_type_name() != Some(WkbType::NAME) { + return Ok(iceberg_type); + } + + if !matches!(&iceberg_type, Type::Primitive(PrimitiveType::Binary)) { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "WKB extension type on field {} requires binary storage, got {iceberg_type}", + arrow_field.name() + ), + )); + } + + let wkb_type = arrow_field.try_extension_type::().map_err(|err| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Invalid geospatial Arrow extension metadata for field {}", + arrow_field.name() + ), + ) + .with_source(err) + })?; + + let crs = iceberg_crs_from_wkb_metadata(wkb_type.metadata())?; + + match wkb_type.metadata().type_hint() { + WkbTypeHint::Geometry => Ok(Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some(crs))?, + ))), + WkbTypeHint::Geography => Ok(Type::Primitive(PrimitiveType::Geography( + GeographyType::new( + Some(crs), + wkb_type + .metadata() + .algorithm + .unwrap_or(WkbEdges::Spherical) + .into(), + )?, + ))), + } + } } impl ArrowSchemaVisitor for ArrowSchemaConverter { @@ -459,6 +560,7 @@ impl ArrowSchemaVisitor for ArrowSchemaConverter { } }; + let value = self.apply_arrow_extension(element_field, value)?; let id = self.get_field_id(element_field)?; let doc = get_field_doc(element_field); let mut element_field = @@ -480,6 +582,8 @@ impl ArrowSchemaVisitor for ArrowSchemaConverter { let key_field = &fields[0]; let value_field = &fields[1]; + let key_value = self.apply_arrow_extension(key_field, key_value)?; + let value = self.apply_arrow_extension(value_field, value)?; let key_id = self.get_field_id(key_field)?; let key_doc = get_field_doc(key_field); @@ -641,22 +745,32 @@ impl SchemaVisitor for ToArrowSchemaConverter { } else { HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), field.id.to_string())]) }; - let arrow_field = + let mut arrow_field = Field::new(field.name.clone(), ty, !field.required).with_metadata(metadata); - // A variant column's storage is a struct; tag the field with the canonical - // `arrow.parquet.variant` extension type so consumers read it as a Variant, not a struct. - let arrow_field = if field.field_type.is_variant() { - arrow_field.with_extension_type(VariantExtensionType) - } else if matches!( - field.field_type.as_ref(), - Type::Primitive(PrimitiveType::Uuid) - ) { - // A uuid column's storage is FixedSizeBinary(16); tag the field with the canonical - // `arrow.uuid` extension type so consumers read it as a UUID, not opaque bytes. - arrow_field.with_extension_type(UuidExtensionType) - } else { - arrow_field - }; + + match field.field_type.as_ref() { + Type::Variant(_) => { + // A variant column's storage is a struct; tag the field with the canonical + // `arrow.parquet.variant` extension type so consumers read it as a Variant, not a struct. + arrow_field.try_with_extension_type(VariantExtensionType)?; + } + Type::Primitive(PrimitiveType::Uuid) => { + // Preserve the canonical UUID extension on its FixedSizeBinary(16) storage. + arrow_field.try_with_extension_type(UuidExtensionType)?; + } + Type::Primitive(PrimitiveType::Geometry(geometry)) => { + let metadata = WkbMetadata::new(Some(geometry.crs()), None); + arrow_field.try_with_extension_type(WkbType::new(Some(metadata)))?; + } + Type::Primitive(PrimitiveType::Geography(geography)) => { + let metadata = WkbMetadata::new( + Some(geography.crs()), + Some(to_arrow_wkb_edges(geography.algorithm())), + ); + arrow_field.try_with_extension_type(WkbType::new(Some(metadata)))?; + } + _ => {} + } Ok(ArrowSchemaOrFieldOrType::Field(arrow_field)) } @@ -776,11 +890,9 @@ impl SchemaVisitor for ToArrowSchemaConverter { .map(DataType::FixedSizeBinary) .unwrap_or(DataType::LargeBinary), )), - PrimitiveType::Binary => Ok(ArrowSchemaOrFieldOrType::Type(DataType::LargeBinary)), - PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => Err(Error::new( - ErrorKind::FeatureUnsupported, - format!("Converting {p} to an Arrow schema is not supported yet"), - )), + PrimitiveType::Binary | PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { + Ok(ArrowSchemaOrFieldOrType::Type(DataType::LargeBinary)) + } } } @@ -1439,10 +1551,15 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use arrow_schema::extension::ExtensionType; use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use parquet_geospatial::WkbEdges; use super::*; use crate::spec::decimal_utils::decimal_new; + use crate::spec::geospatial::{ + DEFAULT_GEOSPATIAL_CRS, EdgeInterpolationAlgorithm as IcebergEdgeInterpolationAlgorithm, + }; use crate::spec::{Literal, Schema}; /// Create a simple field with metadata. @@ -2379,21 +2496,183 @@ mod tests { } #[test] - fn test_geospatial_arrow_schema_conversion_is_unsupported() { + fn test_geospatial_arrow_schema_roundtrip() { let schema = Schema::builder() + .with_schema_id(1) .with_fields(vec![ - NestedField::optional( + NestedField::required( 1, "geom", - Type::Primitive(PrimitiveType::Geometry(Default::default())), + Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some("EPSG:3857".to_string())).unwrap(), + )), + ) + .into(), + NestedField::optional( + 2, + "geog", + Type::Primitive(PrimitiveType::Geography( + GeographyType::new( + Some("OGC:CRS27".to_string()), + IcebergEdgeInterpolationAlgorithm::Karney, + ) + .unwrap(), + )), + ) + .into(), + NestedField::optional( + 3, + "geom_list", + Type::List(ListType::new( + NestedField::list_element( + 4, + Type::Primitive(PrimitiveType::Geometry(GeometryType::default())), + true, + ) + .into(), + )), ) .into(), ]) .build() .unwrap(); - let err = schema_to_arrow_schema(&schema).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + let geom = arrow_schema.field(0); + assert_eq!(geom.data_type(), &DataType::LargeBinary); + + let geog = arrow_schema.field(1); + let geog_wkb = geog.try_extension_type::().unwrap(); + assert_eq!(geog_wkb.metadata().algorithm, Some(WkbEdges::Karney)); + let geog_metadata: serde_json::Value = + serde_json::from_str(&geog_wkb.serialize_metadata().unwrap()).unwrap(); + assert_eq!(geog_metadata.get("edges").unwrap(), "karney"); + assert!(geog_metadata.get("algorithm").is_none()); + + let list = arrow_schema.field(2); + let DataType::List(element) = list.data_type() else { + panic!("Expected list field"); + }; + assert!( + matches!( + element + .try_extension_type::() + .unwrap() + .metadata() + .type_hint(), + WkbTypeHint::Geometry + ), + "Expected list element to retain WKB extension metadata" + ); + assert_eq!( + element + .try_extension_type::() + .unwrap() + .metadata() + .crs + .as_ref() + .and_then(serde_json::Value::as_str), + Some(DEFAULT_GEOSPATIAL_CRS) + ); + + let converted = arrow_schema_to_schema(&arrow_schema).unwrap(); + assert_eq!(converted.as_struct().fields(), schema.as_struct().fields()); + } + + #[test] + fn test_geospatial_arrow_crs_import() { + let mut geom = simple_field("geom", DataType::LargeBinary, true, "1"); + geom.try_with_extension_type(WkbType::new(Some(WkbMetadata::new( + Some(r#"{"id":{"authority":"EPSG","code":3857}}"#), + None, + )))) + .unwrap(); + let mut unset_geom = simple_field("unset_geom", DataType::LargeBinary, true, "2"); + unset_geom + .try_with_extension_type(WkbType::new(Some(WkbMetadata::new(None, None)))) + .unwrap(); + + let arrow_schema = ArrowSchema::new(vec![geom, unset_geom]); + let schema = arrow_schema_to_schema(&arrow_schema).unwrap(); + + assert_eq!( + schema.field_by_id(1).unwrap().field_type.as_ref(), + &Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some("EPSG:3857".to_string())).unwrap() + )) + ); + assert_eq!( + schema.field_by_id(2).unwrap().field_type.as_ref(), + &Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some(UNSET_GEOSPATIAL_CRS.to_string())).unwrap() + )) + ); + } + + #[test] + fn test_geospatial_arrow_invalid_crs_is_rejected() { + for crs in [r#"{}"#, "42"] { + let mut field = simple_field("geom", DataType::LargeBinary, true, "1"); + field + .try_with_extension_type(WkbType::new(Some(WkbMetadata::new(Some(crs), None)))) + .unwrap(); + + let err = arrow_schema_to_schema(&ArrowSchema::new(vec![field])).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid, "{err}"); + } + } + + #[test] + fn test_geospatial_arrow_extension_requires_binary_storage() { + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; + + let metadata = WkbType::new(Some(WkbMetadata::new(None, None))) + .serialize_metadata() + .unwrap(); + let field = + simple_field("geom", DataType::Int32, true, "1").with_metadata(HashMap::from([ + (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()), + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + WkbType::NAME.to_string(), + ), + (EXTENSION_TYPE_METADATA_KEY.to_string(), metadata), + ])); + + let err = arrow_schema_to_schema(&ArrowSchema::new(vec![field])).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid, "{err}"); + assert!( + err.message().contains("requires binary storage"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_geospatial_parquet_statistics_are_ignored() { + use parquet::data_type::ByteArray; + use parquet::file::statistics::ValueStatistics; + + let stats = Statistics::ByteArray(ValueStatistics::new( + Some(ByteArray::from(vec![1, 2, 3])), + Some(ByteArray::from(vec![4, 5, 6])), + None, + None, + false, + )); + + for primitive_type in [ + PrimitiveType::Geometry(GeometryType::default()), + PrimitiveType::Geography(GeographyType::default()), + ] { + assert_eq!( + get_parquet_stat_min_as_datum(&primitive_type, &stats).unwrap(), + None + ); + assert_eq!( + get_parquet_stat_max_as_datum(&primitive_type, &stats).unwrap(), + None + ); + } } #[test] diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index b3bfb3f8e3..547cff7a06 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -331,7 +331,7 @@ impl SchemaWithPartnerVisitor for ArrowArrayToIcebergStructConverter { .map(|v| v.map(|v| Literal::fixed(v.iter().cloned()))) .collect()) } - PrimitiveType::Binary => { + PrimitiveType::Binary | PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { if let Some(array) = partner.as_any().downcast_ref::() { Ok(array .iter() @@ -346,10 +346,6 @@ impl SchemaWithPartnerVisitor for ArrowArrayToIcebergStructConverter { Err(invalid_data!("The partner is not a binary array")) } } - PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => Err(Error::new( - ErrorKind::FeatureUnsupported, - format!("Converting {p} Arrow array to an Iceberg literal is not supported yet"), - )), } } @@ -575,6 +571,9 @@ pub(crate) fn create_primitive_array_single_element( (DataType::Binary, Some(PrimitiveLiteral::Binary(v))) => { Ok(Arc::new(BinaryArray::from_vec(vec![v.as_slice()]))) } + (DataType::LargeBinary, Some(PrimitiveLiteral::Binary(v))) => { + Ok(Arc::new(LargeBinaryArray::from_vec(vec![v.as_slice()]))) + } (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::Int128(v))) => { let array = Decimal128Array::from(vec![{ *v }]) .with_precision_and_scale(*precision, *scale) @@ -1111,28 +1110,23 @@ mod test { } #[test] - fn test_arrow_geospatial_to_literal_is_unsupported() { - let values = Arc::new(LargeBinaryArray::from_vec(vec![b"wkb".as_slice()])) as ArrayRef; - let struct_array = Arc::new(StructArray::from(vec![( - Arc::new( - Field::new("geom", DataType::LargeBinary, false).with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - "1".to_string(), - )])), - ), - values, - )])) as ArrayRef; - let ty = StructType::new(vec![ - NestedField::required( - 1, - "geom", - Type::Primitive(PrimitiveType::Geometry(Default::default())), - ) - .into(), - ]); + fn test_arrow_geospatial_to_literal() { + let values = Arc::new(LargeBinaryArray::from(vec![ + Some(b"wkb-1".as_slice()), + None, + Some(b"wkb-2".as_slice()), + ])) as ArrayRef; + let mut converter = ArrowArrayToIcebergStructConverter; - let err = arrow_struct_to_literal(&struct_array, &ty).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); + let result = converter + .primitive(&PrimitiveType::Geometry(Default::default()), &values) + .unwrap(); + + assert_eq!(result, vec![ + Some(Literal::binary(b"wkb-1".to_vec())), + None, + Some(Literal::binary(b"wkb-2".to_vec())), + ]); } #[test] diff --git a/crates/iceberg/src/avro/schema.rs b/crates/iceberg/src/avro/schema.rs index 8bdd67c505..c396bbe575 100644 --- a/crates/iceberg/src/avro/schema.rs +++ b/crates/iceberg/src/avro/schema.rs @@ -275,12 +275,8 @@ impl SchemaVisitor for SchemaToAvroSchema { // already accepts both forms. PrimitiveType::Uuid => AvroSchema::Uuid(UuidSchema::String), PrimitiveType::Fixed(len) => self.define_once(avro_fixed_schema((*len) as usize)?), - PrimitiveType::Binary => AvroSchema::Bytes, - PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - format!("Converting {p} to an Avro schema is not supported yet"), - )); + PrimitiveType::Binary | PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { + AvroSchema::Bytes } PrimitiveType::Decimal { precision, scale } => { self.define_once(avro_decimal_schema(*precision as usize, *scale as usize)?) @@ -1412,10 +1408,10 @@ mod tests { } #[test] - fn test_geospatial_to_avro_schema_is_unsupported() { + fn test_geospatial_to_avro_bytes_schema() { let schema = Schema::builder() .with_fields(vec![ - NestedField::optional( + NestedField::required( 1, "geom", Type::Primitive(PrimitiveType::Geometry(Default::default())), @@ -1425,7 +1421,9 @@ mod tests { .build() .unwrap(); - let err = schema_to_avro_schema("t", &schema).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); + let AvroSchema::Record(record) = schema_to_avro_schema("t", &schema).unwrap() else { + panic!("expected Avro record schema"); + }; + assert_eq!(record.fields[0].schema, AvroSchema::Bytes); } } diff --git a/crates/iceberg/src/spec/values/datum.rs b/crates/iceberg/src/spec/values/datum.rs index 09c1e88313..7b38d848ef 100644 --- a/crates/iceberg/src/spec/values/datum.rs +++ b/crates/iceberg/src/spec/values/datum.rs @@ -384,12 +384,6 @@ impl Datum { "Cannot create datum for unknown type from bytes", )); } - PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - "Cannot create geospatial datum from bytes yet", - )); - } PrimitiveType::Boolean => { if bytes.len() == 1 && bytes[0] == 0u8 { PrimitiveLiteral::Boolean(false) @@ -439,8 +433,10 @@ impl Datum { PrimitiveType::Uuid => { PrimitiveLiteral::UInt128(u128::from_be_bytes(bytes.try_into()?)) } - PrimitiveType::Fixed(_) => PrimitiveLiteral::Binary(Vec::from(bytes)), - PrimitiveType::Binary => PrimitiveLiteral::Binary(Vec::from(bytes)), + PrimitiveType::Fixed(_) + | PrimitiveType::Binary + | PrimitiveType::Geometry(_) + | PrimitiveType::Geography(_) => PrimitiveLiteral::Binary(Vec::from(bytes)), PrimitiveType::Decimal { .. } => PrimitiveLiteral::Int128( i128_from_be_bytes(bytes) .ok_or_else(|| invalid_data!("Can't convert bytes to i128: {bytes:?}"))?, diff --git a/crates/iceberg/src/spec/values/tests.rs b/crates/iceberg/src/spec/values/tests.rs index a6945015bc..609904728c 100644 --- a/crates/iceberg/src/spec/values/tests.rs +++ b/crates/iceberg/src/spec/values/tests.rs @@ -549,14 +549,14 @@ fn avro_bytes_string() { } #[test] -fn geospatial_binary_deserialization_is_not_yet_supported() { - for primitive_type in [ - PrimitiveType::Geometry(Default::default()), - PrimitiveType::Geography(Default::default()), - ] { - let err = Datum::try_from_bytes(&[1, 2, 3], primitive_type).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); - } +fn avro_bytes_geometry() { + let bytes = vec![1u8, 2u8, 3u8, 4u8]; + let ty = PrimitiveType::Geometry(Default::default()); + check_avro_bytes_serde( + bytes.clone(), + Datum::new(ty.clone(), PrimitiveLiteral::Binary(bytes)), + &ty, + ); } #[test] From 40d3b73f562066f5a8af50d2008f4867819a4ae2 Mon Sep 17 00:00:00 2001 From: wirybeaver Date: Sun, 4 Oct 2026 12:53:59 -0700 Subject: [PATCH 2/2] feat: add geospatial parquet write support --- .../src/writer/file_writer/parquet_writer.rs | 158 +++++++++++++++++- 1 file changed, 154 insertions(+), 4 deletions(-) diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index a471aea5c6..ef289eef19 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -371,6 +371,12 @@ impl MinMaxColAggregator { )); }; + if matches!(ty, PrimitiveType::Geometry(_) | PrimitiveType::Geography(_)) { + // Statistics are optional, not setting them is fine. + // TODO: Implement bounds for geospatial types. + return Ok(()); + } + if value.min_is_exact() { let Some(min_datum) = get_parquet_stat_min_as_datum(&ty, value)? else { return Err(Error::new( @@ -759,9 +765,12 @@ mod tests { use futures::TryStreamExt; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; - use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; + use parquet::basic::{ + BrotliLevel, Compression, EdgeInterpolationAlgorithm, GzipLevel, LogicalType, ZstdLevel, + }; use parquet::file::statistics::ValueStatistics; use parquet::schema::types::ColumnPath; + use parquet_geospatial::testing::wkb_point_xy; use tempfile::TempDir; use uuid::Uuid; @@ -771,6 +780,7 @@ mod tests { use crate::io::FileIO; use crate::scan::{FileScanTask, FileScanTaskStream}; use crate::spec::decimal_utils::{decimal_mantissa, decimal_new, decimal_scale}; + use crate::spec::geospatial::EdgeInterpolationAlgorithm as IcebergEdgeInterpolationAlgorithm; use crate::spec::{PrimitiveLiteral, Struct, *}; use crate::test_utils::make_encryption_manager; use crate::writer::file_writer::location_generator::{ @@ -1604,7 +1614,7 @@ mod tests { let output_file = file_io.new_output( location_gen.generate_location(None, &file_name_gen.generate_file_name()), )?; - let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone()) + let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema) .build(output_file) .await?; let col0 = Arc::new( @@ -1714,7 +1724,7 @@ mod tests { let output_file = file_io.new_output( location_gen.generate_location(None, &file_name_gen.generate_file_name()), )?; - let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema) + let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone()) .build(output_file) .await?; let col0 = Arc::new( @@ -2565,11 +2575,151 @@ mod tests { ( parquet::basic::Type::FIXED_LEN_BYTE_ARRAY, 16, - Some(&parquet::basic::LogicalType::Uuid) + Some(&LogicalType::Uuid) ) ); } + #[tokio::test] + async fn test_parquet_writer_geospatial_logical_types() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let location_gen = DefaultLocationGenerator::with_data_location( + temp_dir.path().to_str().unwrap().to_string(), + ); + let file_name_gen = + DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet); + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required( + 0, + "geom", + Type::Primitive(PrimitiveType::Geometry(GeometryType::default())), + ) + .into(), + NestedField::optional( + 1, + "unset_geom", + Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some("srid:0".to_string())).unwrap(), + )), + ) + .into(), + NestedField::optional( + 2, + "geog", + Type::Primitive(PrimitiveType::Geography( + GeographyType::new(None, IcebergEdgeInterpolationAlgorithm::Karney) + .unwrap(), + )), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let geom_wkb = wkb_point_xy(1.0, 2.0); + let geog_wkb = wkb_point_xy(3.0, 4.0); + let geom = Arc::new(arrow_array::LargeBinaryArray::from_vec(vec![ + geom_wkb.as_slice(), + ])) as ArrayRef; + let geog = Arc::new(arrow_array::LargeBinaryArray::from_vec(vec![ + geog_wkb.as_slice(), + ])) as ArrayRef; + let unset_geom = geom.clone(); + let to_write = + RecordBatch::try_new(arrow_schema.clone(), vec![geom, unset_geom, geog]).unwrap(); + + let output_file = file_io.new_output( + location_gen.generate_location(None, &file_name_gen.generate_file_name()), + )?; + let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone()) + .build(output_file) + .await?; + + pw.write(&to_write).await?; + let res = pw.close().await?; + assert_eq!(res.len(), 1); + let data_file = res + .into_iter() + .next() + .unwrap() + .content(DataContentType::Data) + .partition(Struct::empty()) + .partition_spec_id(0) + .build() + .unwrap(); + + assert_eq!(data_file.record_count(), 1); + // Geospatial bounds are intentionally omitted until zonemap statistics are implemented. + assert!( + data_file.lower_bounds().is_empty(), + "geospatial lower bounds should be omitted" + ); + assert!( + data_file.upper_bounds().is_empty(), + "geospatial upper bounds should be omitted" + ); + + let input_file = file_io.new_input(data_file.file_path())?; + let file_metadata = input_file.metadata().await?; + let reader = input_file.reader().await?; + let mut parquet_reader = ArrowFileReader::new(file_metadata, reader); + let parquet_metadata = parquet_reader.get_metadata(None).await?; + let schema_descr = parquet_metadata.file_metadata().schema_descr(); + + // Parquet omits CRS when default `OGC:CRS84` is used. + assert_eq!( + schema_descr.column(0).logical_type_ref(), + Some(&LogicalType::geometry(None)) + ); + assert_eq!( + schema_descr.column(1).logical_type_ref(), + Some(&LogicalType::geometry(Some("srid:0".to_string()))) + ); + assert_eq!( + schema_descr.column(2).logical_type_ref(), + Some(&LogicalType::geography( + None, + Some(EdgeInterpolationAlgorithm::KARNEY), + )) + ); + + let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + let task = FileScanTask::builder() + .with_file_size_in_bytes(data_file.file_size_in_bytes()) + .with_start(0) + .with_length(0) + .with_data_file_path(data_file.file_path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![0, 1, 2]) + .with_case_sensitive(false) + .build()?; + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let batches: Vec = reader.read(tasks)?.stream().try_collect().await?; + let read_back = concat_batches(batches[0].schema_ref(), &batches)?; + + for (column, expected) in [ + (0, geom_wkb.as_slice()), + (1, geom_wkb.as_slice()), + (2, geog_wkb.as_slice()), + ] { + let values = read_back + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), expected); + } + + Ok(()) + } + #[test] fn test_min_max_aggregator() { let schema = Arc::new(