dannycjones commented on code in PR #2933:
URL: https://github.com/apache/iceberg-rust/pull/2933#discussion_r4040748059
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -396,22 +452,57 @@ 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| {
+ Error::new(
+ ErrorKind::DataInvalid,
+ format!(
+ "Invalid geospatial Arrow extension metadata for field {}",
+ field.name()
+ ),
+ )
+ .with_source(err)
+ })?;
+
+ let crs =
iceberg_crs_from_wkb_metadata(wkb_type.metadata().crs.as_ref())?;
+
+ match wkb_type.metadata().type_hint() {
+ WkbTypeHint::Geometry =>
Ok(Type::Primitive(PrimitiveType::Geometry(
+ GeometryType::new(crs)?,
+ ))),
+ WkbTypeHint::Geography =>
Ok(Type::Primitive(PrimitiveType::Geography(
+ GeographyType::new(
+ crs,
+ wkb_type
+ .metadata()
+ .algorithm
+ .unwrap_or(WkbEdges::Spherical)
+ .into(),
+ )?,
+ ))),
+ }
+ }
Review Comment:
This function is extracting an Iceberg type from an Arrow field based on its
extension. We have already inferred an Iceberg type though, so I think we
should verify some basic rules about what Arrow types are allowed for the given
Iceberg type.
Also, I think we can clean up the function signature a little and give it a
name that better explains its purpose.
```suggestion
fn apply_arrow_extension(&self, arrow_field: &FieldRef, iceberg_type:
Type) -> Result<Type> {
if arrow_field.extension_type_name() != Some(WkbType::NAME) {
return Ok(iceberg_type);
}
if !matches!(iceberg_type,
Type::PrimitiveType(PrimitiveType::Binary)) {
return Error::new(
ErrorKind::DataInvalid,
format!(
"WKB extension type on field {} requires binary storage,
got {field_type}",
arrow_field.name()
)
);
}
let wkb_type =
arrow_field.try_extension_type::<WkbType>().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().crs.as_ref())?;
let iceberg_type = match wkb_type.metadata().type_hint() {
WkbTypeHint::Geometry => Type::Primitive(PrimitiveType::Geometry(
GeometryType::new(crs)?,
)),
WkbTypeHint::Geography =>
Type::Primitive(PrimitiveType::Geography(
GeographyType::new(
crs,
wkb_type
.metadata()
.algorithm
.unwrap_or(WkbEdges::Spherical)
.into(),
)?,
)),
}
Ok(iceberg_type)
}
```
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -100,6 +103,59 @@ impl ExtensionType for VariantExtensionType {
}
}
+fn wkb_edges_from_edge_interpolation_algorithm(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<WkbEdges> 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(crs: Option<&serde_json::Value>) ->
Result<Option<String>> {
+ match crs {
+ None => Ok(Some(UNSET_GEOSPATIAL_CRS.to_string())),
+ Some(serde_json::Value::String(crs)) => Ok(Some(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(Some(format!("{authority}:{code}"))),
+ _ => Err(Error::new(
+ ErrorKind::DataInvalid,
+ "Cannot write PROJJSON CRS without an embedded
authority/code to Iceberg",
+ )),
+ }
+ }
+ Some(_) => Err(Error::new(
+ ErrorKind::DataInvalid,
+ "Geospatial CRS metadata must be a string or PROJJSON object",
+ )),
+ }
+}
Review Comment:
I think we need to rethink the extraction of the CRS from the WKB metadata.
If we find `WkbType` extension but we cannot parse it, should we fail or
ignore? I'm leaning on fail at the moment but I'm not sure that's right.
In any case, I think we'd be better off taking the WKB metadata as the
argument and returning a result. Also, we always return a value, so let's not
bother with `Option`.
```suggestion
fn iceberg_crs_from_wkb_metadata(metadata: &WkbType::Metadata) ->
Result<String> {
let Some(crs) = metadata.crs else {
return Ok(UNSET_GEOSPATIAL_CRS.to_string())
};
match crs {
serde_json::Value::String(crs) => Ok(crs.clone()),
serde_json::Value::Object(crs) => {
// Parse from PROJJSON
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",
)),
}
}
```
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -377,8 +377,12 @@ impl MinMaxColAggregator {
));
};
+ if matches!(ty, PrimitiveType::Geometry(_) |
PrimitiveType::Geography(_)) {
+ return Ok(());
+ }
Review Comment:
```suggestion
if matches!(ty, PrimitiveType::Geometry(_) |
PrimitiveType::Geography(_)) {
// Statistics are optional, not setting them is fine.
// TODO: Implement bounds for geospatial types.
return Ok(());
}
```
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -100,6 +103,59 @@ impl ExtensionType for VariantExtensionType {
}
}
+fn wkb_edges_from_edge_interpolation_algorithm(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,
+ }
+}
Review Comment:
nitpicky, but easier to digest what we're converting.
```suggestion
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,
}
}
```
##########
crates/iceberg/src/writer/file_writer/parquet_writer.rs:
##########
@@ -2506,6 +2516,117 @@ 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,
+ "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)
+ .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();
+
+ 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),
+ ))
+ );
Review Comment:
```suggestion
// 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),
))
);
```
##########
crates/catalog/hms/src/schema.rs:
##########
@@ -117,7 +117,10 @@ impl SchemaVisitor for HiveSchemaBuilder {
PrimitiveType::Time | PrimitiveType::String | PrimitiveType::Uuid
=> {
"string".to_string()
}
- PrimitiveType::Binary | PrimitiveType::Fixed(_) =>
"binary".to_string(),
+ PrimitiveType::Binary
+ | PrimitiveType::Fixed(_)
+ | PrimitiveType::Geometry(_)
+ | PrimitiveType::Geography(_) => "binary".to_string(),
Review Comment:
Why is this the right type?
##########
crates/iceberg/src/arrow/schema.rs:
##########
@@ -2217,6 +2335,120 @@ mod tests {
);
}
+ #[test]
+ fn test_geospatial_arrow_schema_roundtrip() {
+ let schema = Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(
+ 1,
+ "geom",
+ 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 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::<WkbType>().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());
Review Comment:
nit: ditch this, this was just an old artifact. main thing is that `"edges"`
is set.
##########
crates/catalog/hms/src/schema.rs:
##########
@@ -117,7 +117,10 @@ impl SchemaVisitor for HiveSchemaBuilder {
PrimitiveType::Time | PrimitiveType::String | PrimitiveType::Uuid
=> {
"string".to_string()
}
- PrimitiveType::Binary | PrimitiveType::Fixed(_) =>
"binary".to_string(),
+ PrimitiveType::Binary
+ | PrimitiveType::Fixed(_)
+ | PrimitiveType::Geometry(_)
+ | PrimitiveType::Geography(_) => "binary".to_string(),
Review Comment:
If we're not sure, I'd prefer we bail out like the timestamptz types for now.
```suggestion
PrimitiveType::Binary | PrimitiveType::Fixed(_) =>
"binary".to_string(),
PrimitiveType::Geometry(_) | PrimitiveType::Geography(_) => {
return Err(Error::new(
ErrorKind::FeatureUnsupported,
format!("Conversion from {p:?} is not supported"),
));
}
```
Same for AWS Glue.
--
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]
