paleolimbot commented on code in PR #1116:
URL: https://github.com/apache/sedona-db/pull/1116#discussion_r3738737677
##########
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:
I added the deferred error mechanism to make sure this gets caught
##########
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:
I added this to the test!
--
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]
