This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3342-200ee6593dee84654a3b98aed724c7d094aa18cf in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit 2079e5851768bc4ef24df707d7faa664406375bf Author: Xuanyi Li <[email protected]> AuthorDate: Wed Nov 4 04:17:28 2026 +0000 feat: add geospatial parquet write support [3/3] (#3342) * feat: add geospatial parquet write support Aggregate Parquet Geometry statistics into Iceberg file bounds using the spec-defined coordinate encoding. Carry the field type with the bounds state, keep incomplete bounds unknown, and omit Geography bounds until geodesic bounds are supported. Export an unset CRS without Arrow CRS metadata. Keep spatial bounds out of scalar metrics comparisons. Cover Parquet payload and metadata round-trips, manifest bounds, and scalar filtering separately, with focused tests for coordinate dimensions, wrapping Geometry boxes, and missing statistics. * fix: preserve supplied geography bounds Consume existing Parquet bounding boxes for Geography as well as Geometry. Bounds generation and edge interpolation remain the Parquet writer responsibility. Preserve wraparound and serialize full longitude coverage using the canonical range. --- crates/iceberg/src/arrow/schema.rs | 49 +- .../expr/visitors/inclusive_metrics_evaluator.rs | 5 +- crates/iceberg/src/expr/visitors/mod.rs | 12 +- .../src/expr/visitors/strict_metrics_evaluator.rs | 5 +- .../src/writer/file_writer/parquet_writer.rs | 558 ++++++++++++++++++++- 5 files changed, 602 insertions(+), 27 deletions(-) diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index 24fc93f03..6e8c1f99e 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -157,6 +157,10 @@ fn iceberg_crs_from_wkb_metadata(metadata: &WkbMetadata) -> Result<String> { } } +fn arrow_crs_from_iceberg(crs: &str) -> Option<&str> { + (!crs.eq_ignore_ascii_case(UNSET_GEOSPATIAL_CRS)).then_some(crs) +} + /// A post order arrow schema visitor. /// /// For order of methods called, please refer to the internal `visit_schema` function. @@ -759,12 +763,12 @@ impl SchemaVisitor for ToArrowSchemaConverter { arrow_field.try_with_extension_type(UuidExtensionType)?; } Type::Primitive(PrimitiveType::Geometry(geometry)) => { - let metadata = WkbMetadata::new(Some(geometry.crs()), None); + let metadata = WkbMetadata::new(arrow_crs_from_iceberg(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()), + arrow_crs_from_iceberg(geography.crs()), Some(to_arrow_wkb_edges(geography.algorithm())), ); arrow_field.try_with_extension_type(WkbType::new(Some(metadata)))?; @@ -2610,6 +2614,47 @@ mod tests { ); } + #[test] + fn test_geospatial_arrow_unset_crs_export() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "geom", + Type::Primitive(PrimitiveType::Geometry( + GeometryType::new(Some(UNSET_GEOSPATIAL_CRS.to_string())).unwrap(), + )), + ) + .into(), + NestedField::optional( + 2, + "geog", + Type::Primitive(PrimitiveType::Geography( + GeographyType::new( + Some(UNSET_GEOSPATIAL_CRS.to_string()), + IcebergEdgeInterpolationAlgorithm::Spherical, + ) + .unwrap(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + for field in arrow_schema.fields() { + assert!( + field + .try_extension_type::<WkbType>() + .unwrap() + .metadata() + .crs + .is_none() + ); + } + assert_eq!(arrow_schema_to_schema(&arrow_schema).unwrap(), schema); + } + #[test] fn test_geospatial_arrow_invalid_crs_is_rejected() { for crs in [r#"{}"#, "42"] { diff --git a/crates/iceberg/src/expr/visitors/inclusive_metrics_evaluator.rs b/crates/iceberg/src/expr/visitors/inclusive_metrics_evaluator.rs index c1b61cd6a..37546fa2c 100644 --- a/crates/iceberg/src/expr/visitors/inclusive_metrics_evaluator.rs +++ b/crates/iceberg/src/expr/visitors/inclusive_metrics_evaluator.rs @@ -19,6 +19,7 @@ use std::borrow::Cow; use fnv::FnvHashSet; +use super::scalar_bound; use crate::expr::visitors::bound_predicate_visitor::{BoundPredicateVisitor, visit}; use crate::expr::{BoundPredicate, BoundReference}; use crate::spec::{DataFile, Datum, PrimitiveLiteral}; @@ -67,11 +68,11 @@ impl<'a> InclusiveMetricsEvaluator<'a> { } fn lower_bound(&self, field_id: i32) -> Option<&Datum> { - self.data_file.lower_bounds.get(&field_id) + scalar_bound(self.data_file.lower_bounds.get(&field_id)) } fn upper_bound(&self, field_id: i32) -> Option<&Datum> { - self.data_file.upper_bounds.get(&field_id) + scalar_bound(self.data_file.upper_bounds.get(&field_id)) } fn contains_nans_only(&self, field_id: i32) -> bool { diff --git a/crates/iceberg/src/expr/visitors/mod.rs b/crates/iceberg/src/expr/visitors/mod.rs index 2fe92340b..e29c959a1 100644 --- a/crates/iceberg/src/expr/visitors/mod.rs +++ b/crates/iceberg/src/expr/visitors/mod.rs @@ -17,7 +17,7 @@ use fnv::FnvHashSet; -use crate::spec::Datum; +use crate::spec::{Datum, PrimitiveType}; pub(crate) mod bloom_filter_evaluator; pub(crate) mod bound_predicate_visitor; @@ -32,6 +32,16 @@ pub(crate) mod row_group_metrics_evaluator; pub(crate) mod strict_metrics_evaluator; pub(crate) mod strict_projection; +fn scalar_bound(bound: Option<&Datum>) -> Option<&Datum> { + // Geospatial bounds encode bounding-box coordinates, not ordered column values. + bound.filter(|datum| { + !matches!( + datum.data_type(), + PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) + ) + }) +} + /// Returns true if every literal is provably outside the inclusive `[lower, upper]` /// range, so rows cannot match. Missing bounds are treated as unbounded on that side. /// diff --git a/crates/iceberg/src/expr/visitors/strict_metrics_evaluator.rs b/crates/iceberg/src/expr/visitors/strict_metrics_evaluator.rs index 6cb9f01c5..1c8a4a674 100644 --- a/crates/iceberg/src/expr/visitors/strict_metrics_evaluator.rs +++ b/crates/iceberg/src/expr/visitors/strict_metrics_evaluator.rs @@ -17,6 +17,7 @@ use fnv::FnvHashSet; +use super::scalar_bound; use crate::Result; use crate::error::invalid_data; use crate::expr::visitors::bound_predicate_visitor::{BoundPredicateVisitor, visit}; @@ -74,11 +75,11 @@ impl<'a> StrictMetricsEvaluator<'a> { } fn lower_bound(&self, field_id: i32) -> Option<&Datum> { - self.data_file.lower_bounds.get(&field_id) + scalar_bound(self.data_file.lower_bounds.get(&field_id)) } fn upper_bound(&self, field_id: i32) -> Option<&Datum> { - self.data_file.upper_bounds.get(&field_id) + scalar_bound(self.data_file.upper_bounds.get(&field_id)) } fn contains_nans_only(&self, field_id: i32) -> bool { diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index 485a94200..d1a021d8e 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -29,9 +29,11 @@ use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter; use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use parquet::encryption::encrypt::FileEncryptionProperties; -use parquet::file::metadata::ParquetMetaData; +use parquet::file::metadata::{ColumnChunkMetaData, ParquetMetaData}; use parquet::file::properties::{CdcOptions, WriterProperties}; use parquet::file::statistics::Statistics; +use parquet::geospatial::bounding_box::BoundingBox; +use parquet_geospatial::interval::{Interval, IntervalTrait, WraparoundInterval}; use super::{FileWriter, FileWriterBuilder}; use crate::arrow::{ @@ -44,8 +46,9 @@ use crate::error::invalid_data; use crate::io::{FileIO, FileWrite, OutputFile}; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, MapType, - NestedFieldRef, PartitionSpec, PrimitiveType, Schema, SchemaRef, SchemaVisitor, Struct, - StructType, TableMetadata, TableProperties, Type, VariantType, visit_schema, + NestedFieldRef, PartitionSpec, PrimitiveLiteral, PrimitiveType, Schema, SchemaRef, + SchemaVisitor, Struct, StructType, TableMetadata, TableProperties, Type, VariantType, + visit_schema, }; use crate::transform::create_transform_function; use crate::writer::{CurrentFileStatus, DataFile}; @@ -318,15 +321,91 @@ pub struct ParquetWriter { struct MinMaxColAggregator { lower_bounds: HashMap<i32, Datum>, upper_bounds: HashMap<i32, Datum>, + // None means at least one non-null row group has no bounding box: file bounds are unknown. + geo_bounds: HashMap<i32, (PrimitiveType, Option<GeospatialBounds>)>, schema: SchemaRef, } +struct GeospatialBounds { + x: WraparoundInterval, + y: Interval, + z: Interval, + m: Interval, +} + +impl From<&BoundingBox> for GeospatialBounds { + fn from(bbox: &BoundingBox) -> Self { + Self { + x: WraparoundInterval::new(bbox.get_xmin(), bbox.get_xmax()), + y: Interval::new(bbox.get_ymin(), bbox.get_ymax()), + z: bbox + .get_zmin() + .zip(bbox.get_zmax()) + .map(|(lo, hi)| Interval::new(lo, hi)) + .unwrap_or_else(Interval::empty), + m: bbox + .get_mmin() + .zip(bbox.get_mmax()) + .map(|(lo, hi)| Interval::new(lo, hi)) + .unwrap_or_else(Interval::empty), + } + } +} + +impl GeospatialBounds { + fn merge(&mut self, bbox: &BoundingBox) { + let other = Self::from(bbox); + self.x = self.x.merge_interval(&other.x); + self.y.update_interval(&other.y); + self.z.update_interval(&other.z); + self.m.update_interval(&other.m); + } + + fn into_datums(self, ty: PrimitiveType) -> (Datum, Datum) { + let (xmin, xmax) = if matches!(ty, PrimitiveType::Geography(_)) { + // A full interval union uses infinities; Geography requires canonical longitudes. + (self.x.lo().max(-180.0), self.x.hi().min(180.0)) + } else if self.x.is_wraparound() { + // Parquet permits wrapping geometry boxes; Iceberg only permits wrapping geography. + (f64::NEG_INFINITY, f64::INFINITY) + } else { + (self.x.lo(), self.x.hi()) + }; + let dimensions = if !self.m.is_empty() { + 4 + } else if !self.z.is_empty() { + 3 + } else { + 2 + }; + let (zmin, zmax) = if self.z.is_empty() { + (f64::NAN, f64::NAN) + } else { + (self.z.lo(), self.z.hi()) + }; + // Iceberg bound serialization is XY[Z[M]] doubles, not WKB (spec Appendix D). + let encode = |point: [f64; 4]| { + PrimitiveLiteral::Binary( + point[..dimensions] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(), + ) + }; + ( + Datum::new(ty.clone(), encode([xmin, self.y.lo(), zmin, self.m.lo()])), + Datum::new(ty, encode([xmax, self.y.hi(), zmax, self.m.hi()])), + ) + } +} + impl MinMaxColAggregator { /// Creates new and empty `MinMaxColAggregator` fn new(schema: SchemaRef) -> Self { Self { lower_bounds: HashMap::new(), upper_bounds: HashMap::new(), + geo_bounds: HashMap::new(), schema, } } @@ -353,8 +432,8 @@ impl MinMaxColAggregator { .or_insert(datum); } - /// Update statistics - fn update(&mut self, field_id: i32, value: &Statistics) -> Result<()> { + /// Update bounds from a column chunk. + fn update(&mut self, field_id: i32, column: &ColumnChunkMetaData) -> Result<()> { let Some(ty) = self .schema .field_by_id(field_id) @@ -364,13 +443,47 @@ impl MinMaxColAggregator { // Ignore the field if it is not in schema. return Ok(()); }; - let Type::Primitive(ty) = ty.clone() else { + let Type::Primitive(ty) = ty else { return Err(Error::new( ErrorKind::Unexpected, format!("Composed type {ty} is not supported for min max aggregation."), )); }; + match ty { + PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => { + // Null-only row groups do not contribute coordinates. + if column.num_values() == 0 + || column + .statistics() + .and_then(Statistics::null_count_opt) + .is_some_and(|nulls| nulls == column.num_values() as u64) + { + return Ok(()); + } + // Parquet bounds already account for the logical type's edge interpolation. + let bbox = column + .geo_statistics() + .and_then(|stats| stats.bounding_box()); + self.geo_bounds + .entry(field_id) + .and_modify(|(_, state)| { + if let (Some(bounds), Some(bbox)) = (state.as_mut(), bbox) { + bounds.merge(bbox); + } else { + *state = None; + } + }) + .or_insert_with(|| (ty.clone(), bbox.map(GeospatialBounds::from))); + return Ok(()); + } + _ => {} + } + let Some(value) = column.statistics() else { + return Ok(()); + }; + let ty = ty.clone(); + if value.min_is_exact() { let Some(min_datum) = get_parquet_stat_min_as_datum(&ty, value)? else { return Err(Error::new( @@ -402,7 +515,14 @@ impl MinMaxColAggregator { } /// Returns lower and upper bounds - fn produce(self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) { + fn produce(mut self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) { + for (field_id, (ty, bounds)) in self.geo_bounds { + if let Some(bounds) = bounds { + let (lower, upper) = bounds.into_datums(ty); + self.lower_bounds.insert(field_id, lower); + self.upper_bounds.insert(field_id, upper); + } + } (self.lower_bounds, self.upper_bounds) } } @@ -480,12 +600,13 @@ impl ParquetWriter { *per_col_val_num.entry(field_id).or_insert(0) += column_chunk_metadata.num_values() as u64; - if let Some(statistics) = column_chunk_metadata.statistics() { - if let Some(null_count) = statistics.null_count_opt() { - *per_col_null_val_num.entry(field_id).or_insert(0) += null_count; - } + min_max_agg.update(field_id, column_chunk_metadata)?; - min_max_agg.update(field_id, statistics)?; + if let Some(null_count) = column_chunk_metadata + .statistics() + .and_then(Statistics::null_count_opt) + { + *per_col_null_val_num.entry(field_id).or_insert(0) += null_count; } } } @@ -764,9 +885,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; @@ -776,6 +900,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::{ @@ -2570,11 +2695,380 @@ mod tests { ( parquet::basic::Type::FIXED_LEN_BYTE_ARRAY, 16, - Some(&parquet::basic::LogicalType::Uuid) + Some(&LogicalType::Uuid) ) ); } + struct GeospatialTestFile { + temp_dir: TempDir, + file_io: FileIO, + schema: SchemaRef, + batch: RecordBatch, + data_file: DataFile, + } + + async fn write_geospatial_test_file() -> Result<GeospatialTestFile> { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let geom_type = PrimitiveType::Geometry(GeometryType::new(Some("srid:0".to_string()))?); + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "geom", Type::Primitive(geom_type)).into(), + NestedField::optional( + 2, + "geog", + Type::Primitive(PrimitiveType::Geography(GeographyType::new( + None, + IcebergEdgeInterpolationAlgorithm::Karney, + )?)), + ) + .into(), + ]) + .build()?, + ); + let geometry = [ + parquet_geospatial::testing::wkb_point_xyzm(1.0, 20.0, f64::NAN, 40.0), + parquet_geospatial::testing::wkb_point_xyzm(-3.0, 4.0, f64::NAN, 70.0), + ]; + let geography = wkb_point_xy(3.0, 4.0); + let batch = RecordBatch::try_new(Arc::new(schema_to_arrow_schema(&schema)?), vec![ + Arc::new(arrow_array::LargeBinaryArray::from_iter_values(&geometry)), + Arc::new(arrow_array::LargeBinaryArray::from(vec![ + Some(geography.as_slice()), + None, + ])), + ])?; + let output = file_io.new_output(temp_dir.path().join("geo.parquet").to_str().unwrap())?; + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .build(); + let mut writer = ParquetWriterBuilder::new(props, schema.clone()) + .build(output) + .await?; + writer.write(&batch).await?; + let data_file = writer + .close() + .await? + .into_iter() + .next() + .unwrap() + .content(DataContentType::Data) + .partition(Struct::empty()) + .partition_spec_id(0) + .build() + .unwrap(); + + Ok(GeospatialTestFile { + temp_dir, + file_io, + schema, + batch, + data_file, + }) + } + + #[tokio::test] + async fn test_geospatial_parquet_roundtrip() -> Result<()> { + let file = write_geospatial_test_file().await?; + let input = file.file_io.new_input(file.data_file.file_path())?; + let mut parquet_reader = + ArrowFileReader::new(input.metadata().await?, input.reader().await?); + let metadata = parquet_reader.get_metadata(None).await?; + let schema_descr = metadata.file_metadata().schema_descr(); + assert_eq!(metadata.num_row_groups(), 2); + let logical_types: Vec<_> = schema_descr + .columns() + .iter() + .map(|column| column.logical_type_ref()) + .collect(); + let expected = [ + LogicalType::geometry(Some("srid:0".to_string())), + LogicalType::geography(None, Some(EdgeInterpolationAlgorithm::KARNEY)), + ]; + assert_eq!(logical_types, expected.iter().map(Some).collect::<Vec<_>>()); + + let arrow_schema = parquet::arrow::parquet_to_arrow_schema( + schema_descr, + metadata.file_metadata().key_value_metadata(), + )?; + let unset_metadata: serde_json::Value = + serde_json::from_str(arrow_schema.field(0).extension_type_metadata().unwrap())?; + assert!(unset_metadata.get("crs").is_none()); + + let task = FileScanTask::builder() + .with_file_size_in_bytes(file.data_file.file_size_in_bytes()) + .with_start(0) + .with_length(0) + .with_data_file_path(file.data_file.file_path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(file.schema) + .with_project_field_ids(vec![1, 2]) + .with_case_sensitive(false) + .build()?; + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let reader = ArrowReaderBuilder::new(file.file_io, Runtime::current()).build(); + let batches: Vec<RecordBatch> = reader.read(tasks)?.stream().try_collect().await?; + let read_back = concat_batches(batches[0].schema_ref(), &batches)?; + assert_eq!(read_back.columns(), file.batch.columns()); + Ok(()) + } + + fn coordinate_bytes(coordinates: &[f64]) -> Vec<u8> { + coordinates + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + #[tokio::test] + async fn test_geometry_bounds_manifest_roundtrip() -> Result<()> { + let file = write_geospatial_test_file().await?; + let path = file.temp_dir.path().join("manifest.avro"); + let output = file.file_io.new_output(path.to_str().unwrap())?; + let mut writer = ManifestWriterBuilder::new( + output, + Some(1), + file.schema.clone(), + PartitionSpec::builder(file.schema).build()?, + ) + .build_v3_data(); + writer.add_file(file.data_file, 1)?; + writer.write_manifest_file().await?; + let manifest = Manifest::parse_avro(&std::fs::read(path).unwrap())?; + let data_file = manifest.entries()[0].data_file(); + let ty = PrimitiveType::Geometry(GeometryType::new(Some("srid:0".to_string()))?); + let expected_lower = HashMap::from([( + 1, + Datum::new( + ty.clone(), + PrimitiveLiteral::Binary(coordinate_bytes(&[-3.0, 4.0, f64::NAN, 40.0])), + ), + )]); + let expected_upper = HashMap::from([( + 1, + Datum::new( + ty, + PrimitiveLiteral::Binary(coordinate_bytes(&[1.0, 20.0, f64::NAN, 70.0])), + ), + )]); + assert_eq!( + (data_file.lower_bounds(), data_file.upper_bounds()), + (&expected_lower, &expected_upper), + ); + Ok(()) + } + + #[tokio::test] + async fn test_geometry_bounds_do_not_drive_scalar_metrics() -> Result<()> { + use crate::expr::visitors::inclusive_metrics_evaluator::InclusiveMetricsEvaluator; + use crate::expr::visitors::strict_metrics_evaluator::StrictMetricsEvaluator; + use crate::expr::{Bind, Reference}; + + let file = write_geospatial_test_file().await?; + let values = file + .batch + .column(0) + .as_any() + .downcast_ref::<arrow_array::LargeBinaryArray>() + .unwrap(); + let literal = Datum::try_from_bytes( + values.value(0), + file.schema + .field_by_id(1) + .unwrap() + .field_type + .as_primitive_type() + .unwrap() + .clone(), + )?; + let inequality = Reference::new("geom") + .less_than(literal.clone()) + .bind(file.schema.clone(), true)?; + let exclusion = Reference::new("geom") + .not_equal_to(literal) + .bind(file.schema, true)?; + assert_eq!( + ( + InclusiveMetricsEvaluator::eval(&inequality, &file.data_file, true)?, + StrictMetricsEvaluator::eval(&exclusion, &file.data_file)?, + ), + (true, false), + ); + Ok(()) + } + + fn geometry_bound_bytes(bbox: BoundingBox) -> (Vec<u8>, Vec<u8>) { + let ty = PrimitiveType::Geometry(GeometryType::default()); + let (lower, upper) = GeospatialBounds::from(&bbox).into_datums(ty); + ( + lower.to_bytes().unwrap().into_vec(), + upper.to_bytes().unwrap().into_vec(), + ) + } + + #[test] + fn test_geometry_xy_bounds() { + assert_eq!( + geometry_bound_bytes(BoundingBox::new(-3.0, 1.0, 4.0, 20.0)), + ( + coordinate_bytes(&[-3.0, 4.0]), + coordinate_bytes(&[1.0, 20.0]) + ), + ); + } + + #[test] + fn test_geometry_xyz_bounds() { + assert_eq!( + geometry_bound_bytes(BoundingBox::new(-3.0, 1.0, 4.0, 20.0).with_zrange(2.0, 9.0)), + ( + coordinate_bytes(&[-3.0, 4.0, 2.0]), + coordinate_bytes(&[1.0, 20.0, 9.0]) + ), + ); + } + + #[test] + fn test_wrapping_geometry_bounds_cover_all_x() { + assert_eq!( + geometry_bound_bytes(BoundingBox::new(170.0, -170.0, -10.0, 10.0)), + ( + coordinate_bytes(&[f64::NEG_INFINITY, -10.0]), + coordinate_bytes(&[f64::INFINITY, 10.0]), + ), + ); + } + + fn geospatial_stats_fixture(ty: PrimitiveType) -> (SchemaRef, ColumnChunkMetaData) { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "geo", Type::Primitive(ty)).into(), + ]) + .build() + .unwrap(), + ); + let parquet_schema = parquet::arrow::ArrowSchemaConverter::new() + .convert(&schema_to_arrow_schema(&schema).unwrap()) + .unwrap(); + let column = ColumnChunkMetaData::builder(parquet_schema.column(0)) + .set_num_values(1) + .build() + .unwrap(); + (schema, column) + } + + fn with_geo_box(column: &ColumnChunkMetaData, bbox: BoundingBox) -> ColumnChunkMetaData { + use parquet::geospatial::statistics::GeospatialStatistics; + column + .clone() + .into_builder() + .set_geo_statistics(Box::new(GeospatialStatistics::new(Some(bbox), None))) + .build() + .unwrap() + } + + #[test] + fn test_geometry_bounds_merge_across_null_row_group() -> Result<()> { + let (schema, column) = + geospatial_stats_fixture(PrimitiveType::Geometry(GeometryType::default())); + let first = with_geo_box( + &column, + BoundingBox::new(-3.0, 1.0, 4.0, 20.0).with_zrange(4.0, 9.0), + ); + let second = with_geo_box( + &column, + BoundingBox::new(0.0, 5.0, -20.0, 5.0).with_mrange(1.0, 3.0), + ); + let nulls = column + .into_builder() + .set_statistics(Statistics::ByteArray(ValueStatistics::new( + None, + None, + None, + Some(1), + false, + ))) + .build()?; + let mut bounds = MinMaxColAggregator::new(schema); + for column in [&first, &nulls, &second] { + bounds.update(1, column)?; + } + let (lower, upper) = bounds.produce(); + assert_eq!( + ( + lower[&1].to_bytes()?.into_vec(), + upper[&1].to_bytes()?.into_vec() + ), + ( + coordinate_bytes(&[-3.0, -20.0, 4.0, 1.0]), + coordinate_bytes(&[5.0, 20.0, 9.0, 3.0]) + ), + ); + Ok(()) + } + + #[test] + fn test_geometry_bounds_require_all_non_null_row_groups() -> Result<()> { + let (schema, unknown) = + geospatial_stats_fixture(PrimitiveType::Geometry(GeometryType::default())); + let known = with_geo_box(&unknown, BoundingBox::new(1.0, 2.0, 3.0, 4.0)); + let mut bounds = MinMaxColAggregator::new(schema); + for column in [&known, &unknown, &known] { + bounds.update(1, column)?; + } + assert_eq!(bounds.produce(), (HashMap::new(), HashMap::new())); + Ok(()) + } + + #[test] + fn test_geography_bounds_merge_from_parquet_statistics() -> Result<()> { + let (schema, column) = + geospatial_stats_fixture(PrimitiveType::Geography(GeographyType::default())); + let first = with_geo_box(&column, BoundingBox::new(170.0, -170.0, -20.0, 20.0)); + let second = with_geo_box(&column, BoundingBox::new(175.0, 178.0, -10.0, 30.0)); + let mut bounds = MinMaxColAggregator::new(schema); + bounds.update(1, &first)?; + bounds.update(1, &second)?; + let (lower, upper) = bounds.produce(); + assert_eq!( + ( + lower[&1].to_bytes()?.into_vec(), + upper[&1].to_bytes()?.into_vec() + ), + ( + coordinate_bytes(&[170.0, -20.0]), + coordinate_bytes(&[-170.0, 30.0]) + ), + ); + Ok(()) + } + + #[test] + fn test_geography_bounds_full_longitude_range() -> Result<()> { + let (schema, column) = + geospatial_stats_fixture(PrimitiveType::Geography(GeographyType::default())); + let first = with_geo_box(&column, BoundingBox::new(170.0, -170.0, -20.0, 20.0)); + let second = with_geo_box(&column, BoundingBox::new(-175.0, 175.0, -10.0, 30.0)); + let mut bounds = MinMaxColAggregator::new(schema); + bounds.update(1, &first)?; + bounds.update(1, &second)?; + let (lower, upper) = bounds.produce(); + assert_eq!( + ( + lower[&1].to_bytes()?.into_vec(), + upper[&1].to_bytes()?.into_vec() + ), + ( + coordinate_bytes(&[-180.0, -20.0]), + coordinate_bytes(&[180.0, 30.0]) + ), + ); + Ok(()) + } + #[test] fn test_min_max_aggregator() { let schema = Arc::new( @@ -2588,9 +3082,18 @@ mod tests { .build() .expect("Failed to create schema"), ); + let parquet_schema = parquet::arrow::ArrowSchemaConverter::new() + .convert(&schema_to_arrow_schema(&schema).unwrap()) + .unwrap(); let mut min_max_agg = MinMaxColAggregator::new(schema); - let create_statistics = - |min, max| Statistics::Int32(ValueStatistics::new(min, max, None, None, false)); + let create_statistics = |min, max| { + ColumnChunkMetaData::builder(parquet_schema.column(0)) + .set_statistics(Statistics::Int32(ValueStatistics::new( + min, max, None, None, false, + ))) + .build() + .unwrap() + }; min_max_agg .update(0, &create_statistics(None, Some(42))) .unwrap(); @@ -2622,11 +3125,26 @@ mod tests { .build() .expect("Failed to create schema"), ); + let parquet_schema = parquet::arrow::ArrowSchemaConverter::new() + .convert(&schema_to_arrow_schema(&schema).unwrap()) + .unwrap(); let mut min_max_agg = MinMaxColAggregator::new(schema); - let float_stats = - |min, max| Statistics::Float(ValueStatistics::new(min, max, None, None, false)); - let double_stats = - |min, max| Statistics::Double(ValueStatistics::new(min, max, None, None, false)); + let float_stats = |min, max| { + ColumnChunkMetaData::builder(parquet_schema.column(0)) + .set_statistics(Statistics::Float(ValueStatistics::new( + min, max, None, None, false, + ))) + .build() + .unwrap() + }; + let double_stats = |min, max| { + ColumnChunkMetaData::builder(parquet_schema.column(1)) + .set_statistics(Statistics::Double(ValueStatistics::new( + min, max, None, None, false, + ))) + .build() + .unwrap() + }; min_max_agg .update(0, &float_stats(Some(-f32::NAN), Some(f32::NAN)))
