paleolimbot commented on code in PR #2933:
URL: https://github.com/apache/iceberg-rust/pull/2933#discussion_r3937533244
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -396,22 +418,71 @@ 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_field_extension_type(field,
&field_results[i])?;
let id = self.get_field_id(field)?;
let doc = get_field_doc(field);
let nested_field = NestedField {
id,
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,
};
results.push(Arc::new(nested_field));
}
Ok(results)
}
+
+ fn apply_field_extension_type(&self, field: &FieldRef, field_type: &Type)
-> Result<Type> {
+ if field.extension_type_name() != Some(WkbType::NAME) {
+ return Ok(field_type.clone());
+ }
+
+ let wkb_type = field.try_extension_type::<WkbType>().map_err(|err| {
Review Comment:
I think this can be resolved now 👍
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -100,6 +102,57 @@ impl ExtensionType for VariantExtensionType {
}
}
+fn edge_interpolation_algorithm_to_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,
+ }
+}
+
+fn wkb_edges_to_edge_interpolation_algorithm(edges: WkbEdges) ->
EdgeInterpolationAlgorithm {
+ match edges {
+ WkbEdges::Spherical => EdgeInterpolationAlgorithm::Spherical,
+ WkbEdges::Vincenty => EdgeInterpolationAlgorithm::Vincenty,
+ WkbEdges::Thomas => EdgeInterpolationAlgorithm::Thomas,
+ WkbEdges::Andoyer => EdgeInterpolationAlgorithm::Andoyer,
+ WkbEdges::Karney => EdgeInterpolationAlgorithm::Karney,
+ }
+}
+
+fn iceberg_crs_from_wkb_metadata(crs: Option<&serde_json::Value>) ->
Result<Option<String>> {
+ match crs {
+ None => Ok(None),
Review Comment:
A `None` CRS from WKB metadata maps to `srid:0`
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -2506,6 +2516,96 @@ mod tests {
assert_eq!(std::fs::read_dir(temp_dir.path()).unwrap().count(), 0);
}
+ #[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,
+ "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 to_write = RecordBatch::try_new(arrow_schema.clone(), vec![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)
+ .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);
+ assert!(data_file.lower_bounds().is_empty());
+ assert!(data_file.upper_bounds().is_empty());
+
+ 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();
+
+ assert_eq!(
+ schema_descr.column(0).logical_type_ref(),
+ Some(&LogicalType::geometry(Some("srid:0".to_string())))
+ );
+ assert_eq!(
+ schema_descr.column(1).logical_type_ref(),
+ Some(&LogicalType::geography(
+ Some("srid:0".to_string()),
+ Some(EdgeInterpolationAlgorithm::KARNEY),
+ ))
+ );
Review Comment:
I agree this looks like a bug...I think my other comment here should fix it.
Parameterizing this test or including the `srid:0` cases explicitly would be a
good idea since this checks Iceberg -> Arrow -> Parquet instead of just Iceberg
-> Arrow.
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -100,6 +102,57 @@ impl ExtensionType for VariantExtensionType {
}
}
+fn edge_interpolation_algorithm_to_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,
+ }
+}
+
+fn wkb_edges_to_edge_interpolation_algorithm(edges: WkbEdges) ->
EdgeInterpolationAlgorithm {
+ match edges {
+ WkbEdges::Spherical => EdgeInterpolationAlgorithm::Spherical,
+ WkbEdges::Vincenty => EdgeInterpolationAlgorithm::Vincenty,
+ WkbEdges::Thomas => EdgeInterpolationAlgorithm::Thomas,
+ WkbEdges::Andoyer => EdgeInterpolationAlgorithm::Andoyer,
+ WkbEdges::Karney => EdgeInterpolationAlgorithm::Karney,
+ }
+}
+
+fn iceberg_crs_from_wkb_metadata(crs: Option<&serde_json::Value>) ->
Result<Option<String>> {
+ match crs {
+ None => Ok(None),
+ Some(serde_json::Value::String(crs)) => Ok(Some(crs.clone())),
Review Comment:
A sanity check would be worthwhile here (I believe an earlier version of
this PR or maybe the Go version did a byte length check in the event of a very
long arbitrary string, which is allowed by the GeoArrow spec). Parquet permits
arbitrary strings as CRS values; Iceberg explicitly advises against PROJJSON
values mostly because they are long and clog up the metadata files.
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -628,15 +717,31 @@ 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 {
- 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 =
arrow_field.with_extension_type(VariantExtensionType);
+ }
+ Type::Primitive(PrimitiveType::Geometry(geometry)) => {
+ let metadata = WkbMetadata::new(geometry.crs(), None);
+
arrow_field.try_with_extension_type(WkbType::new(Some(metadata)))?;
+ }
+ Type::Primitive(PrimitiveType::Geography(geography)) => {
+ let metadata = WkbMetadata::new(
+ geography.crs(),
+ Some(edge_interpolation_algorithm_to_wkb_edges(
+ geography.algorithm(),
+ )),
+ );
Review Comment:
This needs to handle the difference between a `None` CRS in GeoArrow (==
`srid:0`) and a `None` CRS in Parquet/Iceberg (== `OGC:CRS84`). I believe and
earlier draft of this PR did so. You can use the list of test cases here to
parameterize a test (although your list may be slightly different because
Iceberg always normalizes PROJJSON to authority/code):
https://github.com/apache/arrow-rs/blob/7e3b403ac6481493b169c6f83144070f6352026b/parquet/tests/geospatial.rs#L535-L610
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
