james-willis commented on code in PR #1116:
URL: https://github.com/apache/sedona-db/pull/1116#discussion_r3738271197
##########
rust/sedona-geoparquet/src/format.rs:
##########
@@ -391,7 +391,34 @@ impl FileFormat for GeoParquetFormat {
)
.unwrap();
source.options = self.options.clone();
- Arc::new(source)
+
+ // DataFusion 52 has an issue where field metadata (like
ARROW:extension:name)
+ // is stripped when evaluating embedded projections in ParquetOpener.
This is
+ // because the batch schema comes from the parquet reader (which
doesn't have
+ // extension metadata), and Column::return_field() looks up fields
from that schema.
+ // This isn't a bug in DataFusion because we're the ones that
advertised the table
+ // schema as having metadata'd expressions in the first place.
+ //
+ // We fix this by wrapping Column expressions with
MetadataPreservingColumn,
+ // which stores the correct field from the file schema and returns it
from
+ // return_field() regardless of the input schema.
+ let initial_projection =
source.projection().cloned().unwrap_or_else(|| {
+ let indices: Vec<usize> =
+
(0..source.table_schema().table_schema().fields().len()).collect();
+ ProjectionExprs::from_indices(&indices,
source.table_schema().table_schema())
+ });
+ let projection_with_types = wrap_columns_with_metadata_preserving(
+ initial_projection,
+ source.table_schema().table_schema(),
+ )
+ .map(|projection_with_types|
source.try_pushdown_projection(&projection_with_types));
+
+ // These are both failable but we can't fail here, so fall back to the
original
+ // source.
+ match projection_with_types {
+ Ok(Ok(Some(modified))) => modified,
+ _ => Arc::new(source),
Review Comment:
Is it ok to silently not wrap? When do we expect this case to hit?
##########
rust/sedona-geoparquet/src/format.rs:
##########
@@ -691,8 +702,15 @@ fn wrap_expr_columns(
expr.transform_down(|node| {
if let Some(column) = node.as_any().downcast_ref::<Column>() {
let index = column.index();
+
+ if index >= file_schema.fields().len() {
+ return sedona_internal_err!(
+ "Unexpected projection expression in GeoParquet source:
index {index} out of bounds"
+ );
Review Comment:
This would get swallowed by the match at line 418.
Really wouldn't expect this to even happen since schema and expr come from
the same place
##########
rust/sedona-geoparquet/src/format.rs:
##########
@@ -391,7 +391,34 @@ impl FileFormat for GeoParquetFormat {
)
.unwrap();
source.options = self.options.clone();
- Arc::new(source)
+
+ // DataFusion 52 has an issue where field metadata (like
ARROW:extension:name)
+ // is stripped when evaluating embedded projections in ParquetOpener.
This is
+ // because the batch schema comes from the parquet reader (which
doesn't have
+ // extension metadata), and Column::return_field() looks up fields
from that schema.
+ // This isn't a bug in DataFusion because we're the ones that
advertised the table
+ // schema as having metadata'd expressions in the first place.
+ //
+ // We fix this by wrapping Column expressions with
MetadataPreservingColumn,
+ // which stores the correct field from the file schema and returns it
from
Review Comment:
field from table schema right? thats kind of the idea here that the table
might have type metadata that isnt correctly reflected in the file?
##########
rust/sedona-geoparquet/src/format.rs:
##########
@@ -1081,4 +1099,96 @@ mod test {
let geo_source_with_predicate = geo_source.with_predicate(predicate);
assert!(geo_source_with_predicate.inner.filter().is_some());
}
+
+ /// Test that columns with extension metadata are correctly wrapped
+ #[test]
+ fn test_wrap_expr_columns_wraps_geometry_column() {
+ let mut metadata = HashMap::new();
+ metadata.insert(
+ "ARROW:extension:name".to_string(),
+ "geoarrow.wkb".to_string(),
+ );
+ let file_schema = Schema::new(vec![
+ Field::new("geometry", DataType::Binary,
true).with_metadata(metadata)
+ ]);
+
+ // Column expression for the geometry column (index 0 in file schema)
+ let geometry_column: Arc<dyn PhysicalExpr> =
Arc::new(Column::new("geometry", 0));
+
+ let result = wrap_expr_columns(geometry_column, &file_schema).unwrap();
+
+ // The result should be wrapped in MetadataPreservingColumn
+ assert!(result
+ .as_any()
+ .downcast_ref::<MetadataPreservingColumn>()
+ .is_some());
+ }
+
+ /// Test that columns without extension metadata are not wrapped
+ #[test]
+ fn test_wrap_expr_columns_skips_non_geometry_column() {
+ let mut metadata = HashMap::new();
+ metadata.insert(
+ "ARROW:extension:name".to_string(),
+ "geoarrow.wkb".to_string(),
+ );
+ let file_schema = Schema::new(vec![
+ Field::new("name", DataType::Utf8, true),
+ Field::new("geometry", DataType::Binary,
true).with_metadata(metadata),
+ ]);
+
+ // Column expression for a non-geometry column (no extension metadata)
+ let name_column: Arc<dyn PhysicalExpr> = Arc::new(Column::new("name",
0));
+
+ let result = wrap_expr_columns(name_column, &file_schema).unwrap();
+
+ // The result should NOT be wrapped (still a Column)
+ assert!(result.as_any().downcast_ref::<Column>().is_some());
+ }
+
+ /// Test that column index out of file schema bounds returns an error
+ #[test]
+ fn test_wrap_expr_columns_errors_on_out_of_bounds_index() {
+ let file_schema = Schema::new(vec![
+ Field::new("name", DataType::Utf8, true),
+ Field::new("age", DataType::Int32, true),
+ ]);
+
+ // Column expression with index 5, but schema only has 2 fields
(indices 0 and 1)
+ let out_of_bounds_column: Arc<dyn PhysicalExpr> =
Arc::new(Column::new("phantom", 5));
+
+ let result = wrap_expr_columns(out_of_bounds_column, &file_schema);
+
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("index 5 out of bounds"),
+ "Error message should mention out of bounds index, got: {err_msg}"
+ );
+ }
+
+ /// Integration test for projection with multiple derived columns plus
geometry accessor
+ /// Regression test for https://github.com/apache/sedona-db/issues/1115
+ #[tokio::test]
+ async fn projection_with_derived_columns_and_geometry_accessor() {
+ let ctx = setup_context();
+ let example = test_geoparquet("example", "geometry").unwrap();
+
+ // Query that adds multiple literal columns alongside the scanned
geometry
+ // This mimics: SELECT 'a' AS c1, 'b' AS c2, geometry FROM ...
+ let df = ctx
+ .sql(&format!(
+ "SELECT 'a' AS c1, 'b' AS c2, geometry FROM '{}' LIMIT 1",
+ example
+ ))
+ .await
+ .unwrap();
+
+ // This should not panic - the issue was "index out of bounds" when
+ // projection pushdown tried to wrap columns with indices beyond the
+ // file schema's bounds
+ let batches = df.collect().await.unwrap();
+ assert!(!batches.is_empty());
Review Comment:
should we assert the metadata survives as well?
##########
rust/sedona-geoparquet/src/format.rs:
##########
@@ -391,7 +391,34 @@ impl FileFormat for GeoParquetFormat {
)
.unwrap();
source.options = self.options.clone();
- Arc::new(source)
+
+ // DataFusion 52 has an issue where field metadata (like
ARROW:extension:name)
+ // is stripped when evaluating embedded projections in ParquetOpener.
This is
+ // because the batch schema comes from the parquet reader (which
doesn't have
+ // extension metadata), and Column::return_field() looks up fields
from that schema.
+ // This isn't a bug in DataFusion because we're the ones that
advertised the table
+ // schema as having metadata'd expressions in the first place.
+ //
+ // We fix this by wrapping Column expressions with
MetadataPreservingColumn,
+ // which stores the correct field from the file schema and returns it
from
+ // return_field() regardless of the input schema.
+ let initial_projection =
source.projection().cloned().unwrap_or_else(|| {
+ let indices: Vec<usize> =
+
(0..source.table_schema().table_schema().fields().len()).collect();
+ ProjectionExprs::from_indices(&indices,
source.table_schema().table_schema())
Review Comment:
Isn't this just a copy of what is already in `ParquetSource.new`? When would
we hit this case?
https://github.com/apache/datafusion/blob/branch-52/datafusion/datasource-parquet/src/source.rs#L308
--
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]
