dannycjones commented on PR #2933:
URL: https://github.com/apache/iceberg-rust/pull/2933#issuecomment-5721720174
Here's an example of a test in parquet_writer.rs I drafted for testing some
of this.
```rust
/// Write geometry and geography columns through the Iceberg writer and
read them back through
/// the Iceberg reader, checking that both the WKB payloads and the
geospatial type parameters
/// (CRS, edge algorithm) survive the Iceberg -> Arrow -> Parquet ->
Arrow -> Iceberg trip.
#[tokio::test]
async fn test_parquet_writer_geospatial_data_roundtrip() -> 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 geom_type = PrimitiveType::Geometry(GeometryType::default());
let geog_type = PrimitiveType::Geography(
GeographyType::new(
Some("OGC:CRS84".to_string()),
IcebergEdgeInterpolationAlgorithm::Karney,
)
.unwrap(),
);
let schema: SchemaRef = Arc::new(
Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(1, "geom",
Type::Primitive(geom_type.clone())).into(),
NestedField::optional(2, "geog",
Type::Primitive(geog_type.clone())).into(),
])
.build()
.unwrap(),
);
let arrow_schema: ArrowSchemaRef =
Arc::new(schema_to_arrow_schema(&schema).unwrap());
let geom_wkb: Vec<Vec<u8>> = vec![wkb_point_xy(1.0, 2.0),
wkb_point_xy(3.0, 4.0)];
let geog_wkb: Vec<Option<Vec<u8>>> = vec![Some(wkb_point_xy(5.0,
6.0)), None];
let geom_array =
Arc::new(LargeBinaryArray::from_iter_values(geom_wkb.iter())) as ArrayRef;
let geog_array =
Arc::new(LargeBinaryArray::from_iter(geog_wkb.iter())) as ArrayRef;
let to_write =
RecordBatch::try_new(arrow_schema.clone(), vec![geom_array,
geog_array]).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 data_file = pw
.close()
.await?
.into_iter()
.next()
.unwrap()
.content(DataContentType::Data)
.partition(Struct::empty())
.partition_spec_id(0)
.build()
.unwrap();
assert_eq!(data_file.record_count(), 2);
// The Iceberg schema recovered from the file's Arrow schema must
carry the geospatial
// parameters, not just fall back to plain binary.
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 file_arrow_schema = parquet_to_arrow_schema(
parquet_metadata.file_metadata().schema_descr(),
parquet_metadata.file_metadata().key_value_metadata(),
)
.unwrap();
let recovered = arrow_schema_to_schema(&file_arrow_schema).unwrap();
assert_eq!(
recovered.field_by_id(1).unwrap().field_type.as_ref(),
&Type::Primitive(geom_type)
);
assert_eq!(
recovered.field_by_id(2).unwrap().field_type.as_ref(),
&Type::Primitive(geog_type)
);
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![1, 2])
.with_case_sensitive(false)
.build()
.unwrap();
let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as
FileScanTaskStream;
let batches: Vec<RecordBatch> = reader
.read(tasks)
.unwrap()
.stream()
.try_collect()
.await
.unwrap();
let read_back = concat_batches(batches[0].schema_ref(),
&batches).unwrap();
assert_eq!(read_back.num_rows(), 2);
let read_geom = read_back
.column(0)
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
assert_eq!(
read_geom.iter().collect::<Vec<_>>(),
geom_wkb
.iter()
.map(|wkb| Some(wkb.as_slice()))
.collect::<Vec<_>>()
);
let read_geog = read_back
.column(1)
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
assert_eq!(
read_geog.iter().collect::<Vec<_>>(),
geog_wkb
.iter()
.map(|wkb| wkb.as_deref())
.collect::<Vec<_>>()
);
Ok(())
}
```
--
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]
