This is an automated email from the ASF dual-hosted git repository.
kontinuation pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 20808454 refactor(rust/sedona-schema): centralize SedonaType CRS
helpers (#705)
20808454 is described below
commit 20808454b5b3987658621240a00686de44142e27
Author: Kristin Cowalcijk <[email protected]>
AuthorDate: Wed Mar 11 11:15:09 2026 +0800
refactor(rust/sedona-schema): centralize SedonaType CRS helpers (#705)
## Summary
- add `SedonaType::crs()` and `SedonaType::is_item_crs()` in
`sedona-schema` so CRS extraction and `item_crs` detection live on the type
itself
- replace duplicated matching logic in raster functions, expression
handling, testing helpers, and the Python schema wrapper with the new
`SedonaType` helpers
---
python/sedonadb/src/schema.rs | 31 ++++++++++++----------------
rust/sedona-expr/src/item_crs.rs | 7 ++-----
rust/sedona-functions/src/executor.rs | 4 +---
rust/sedona-raster-functions/src/executor.rs | 21 +++----------------
rust/sedona-schema/src/datatypes.rs | 31 ++++++++++++++++++++++++++++
rust/sedona-schema/src/matchers.rs | 16 ++------------
rust/sedona-testing/src/create.rs | 8 ++-----
rust/sedona-testing/src/testers.rs | 9 ++------
8 files changed, 56 insertions(+), 71 deletions(-)
diff --git a/python/sedonadb/src/schema.rs b/python/sedonadb/src/schema.rs
index d9ead708..9b02cc61 100644
--- a/python/sedonadb/src/schema.rs
+++ b/python/sedonadb/src/schema.rs
@@ -205,24 +205,19 @@ impl PySedonaType {
impl PySedonaType {
#[getter]
fn crs<'py>(&self, py: Python<'py>) -> Result<Option<PyObject>,
PySedonaError> {
- match &self.inner {
- SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) => {
- if let Some(crs) = crs {
- let json = py.import("json")?;
- let geoarrow_types_crs = py.import("geoarrow.types.crs")?;
-
- // Use geoarrow.types.crs.create() so that we don't have
to do that here
- // Parse the JSON into a Python object first, because this
is what create()
- // expects
- let crs_string = crs.to_json();
- let crs_py = json.getattr("loads")?.call1((crs_string,))?;
- let crs_py_obj =
geoarrow_types_crs.getattr("create")?.call1((crs_py,))?;
- Ok(Some(crs_py_obj.into()))
- } else {
- Ok(None)
- }
- }
- _ => Ok(None),
+ if let Some(crs) = self.inner.crs() {
+ let json = py.import("json")?;
+ let geoarrow_types_crs = py.import("geoarrow.types.crs")?;
+
+ // Use geoarrow.types.crs.create() so that we don't have to do
that here
+ // Parse the JSON into a Python object first, because this is what
create()
+ // expects
+ let crs_string = crs.to_json();
+ let crs_py = json.getattr("loads")?.call1((crs_string,))?;
+ let crs_py_obj =
geoarrow_types_crs.getattr("create")?.call1((crs_py,))?;
+ Ok(Some(crs_py_obj.into()))
+ } else {
+ Ok(None)
}
}
diff --git a/rust/sedona-expr/src/item_crs.rs b/rust/sedona-expr/src/item_crs.rs
index 8abf99c7..d941853c 100644
--- a/rust/sedona-expr/src/item_crs.rs
+++ b/rust/sedona-expr/src/item_crs.rs
@@ -555,8 +555,7 @@ pub fn parse_item_crs_arg_type(
sedona_type: &SedonaType,
) -> Result<(SedonaType, Option<SedonaType>)> {
if let SedonaType::Arrow(DataType::Struct(fields)) = sedona_type {
- let field_names = fields.iter().map(|f| f.name()).collect::<Vec<_>>();
- if field_names != ["item", "crs"] {
+ if !sedona_type.is_item_crs() {
return Ok((sedona_type.clone(), None));
}
@@ -578,9 +577,7 @@ pub fn parse_item_crs_arg_type_strip_crs(
match sedona_type {
SedonaType::Wkb(edges, _) => Ok((SedonaType::Wkb(*edges, None), None)),
SedonaType::WkbView(edges, _) => Ok((SedonaType::WkbView(*edges,
None), None)),
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.iter().map(|f| f.name()).collect::<Vec<_>>() ==
vec!["item", "crs"] =>
- {
+ SedonaType::Arrow(DataType::Struct(fields)) if
sedona_type.is_item_crs() => {
let item = SedonaType::from_storage_field(&fields[0])?;
let crs = SedonaType::from_storage_field(&fields[1])?;
Ok((item, Some(crs)))
diff --git a/rust/sedona-functions/src/executor.rs
b/rust/sedona-functions/src/executor.rs
index 8809ac0e..47820486 100644
--- a/rust/sedona-functions/src/executor.rs
+++ b/rust/sedona-functions/src/executor.rs
@@ -357,9 +357,7 @@ impl IterGeo for ArrayRef {
}
SedonaType::Wkb(_, _) => iter_wkb_binary(as_binary_array(self)?,
func),
SedonaType::WkbView(_, _) =>
iter_wkb_binary(as_binary_view_array(self)?, func),
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.len() == 2 && fields[0].name() == "item" &&
fields[1].name() == "crs" =>
- {
+ SedonaType::Arrow(DataType::Struct(fields)) if
sedona_type.is_item_crs() => {
let struct_array = as_struct_array(self)?;
let item_type = SedonaType::from_storage_field(&fields[0])?;
struct_array
diff --git a/rust/sedona-raster-functions/src/executor.rs
b/rust/sedona-raster-functions/src/executor.rs
index aa71f711..75652721 100644
--- a/rust/sedona-raster-functions/src/executor.rs
+++ b/rust/sedona-raster-functions/src/executor.rs
@@ -195,21 +195,6 @@ fn resolve_item_crs(item_crs_str: Option<&str>,
static_crs: &Crs) -> Result<Crs>
}
}
-fn crs_from_sedona_type(sedona_type: &SedonaType) -> Crs {
- match sedona_type {
- SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) => crs.clone(),
- _ => None,
- }
-}
-
-fn is_item_crs_type(sedona_type: &SedonaType) -> bool {
- matches!(
- sedona_type,
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.len() == 2 && fields[0].name() == "item" &&
fields[1].name() == "crs"
- )
-}
-
impl<'a, 'b> RasterExecutor<'a, 'b> {
/// Create a new [RasterExecutor]
pub fn new(arg_types: &'a [SedonaType], args: &'b [ColumnarValue]) -> Self
{
@@ -553,14 +538,14 @@ impl<'a, 'b> RasterExecutor<'a, 'b> {
.get(arg_index)
.ok_or_else(|| sedona_internal_datafusion_err!("Missing
argument"))?;
- if is_item_crs_type(sedona_type) {
+ if sedona_type.is_item_crs() {
let item_type = match sedona_type {
SedonaType::Arrow(DataType::Struct(fields)) => {
SedonaType::from_storage_field(&fields[0])?
}
_ => return sedona_internal_err!("Unexpected item_crs type"),
};
- let item_static_crs = crs_from_sedona_type(&item_type);
+ let item_static_crs = item_type.crs().clone();
match arg {
ColumnarValue::Array(array) => {
@@ -635,7 +620,7 @@ impl<'a, 'b> RasterExecutor<'a, 'b> {
}
}
} else {
- let static_crs = crs_from_sedona_type(sedona_type);
+ let static_crs = sedona_type.crs().clone();
match arg {
ColumnarValue::Array(array) => match sedona_type {
SedonaType::Wkb(_, _) |
SedonaType::Arrow(DataType::Binary) => {
diff --git a/rust/sedona-schema/src/datatypes.rs
b/rust/sedona-schema/src/datatypes.rs
index 8284c248..3bc32350 100644
--- a/rust/sedona-schema/src/datatypes.rs
+++ b/rust/sedona-schema/src/datatypes.rs
@@ -279,6 +279,23 @@ impl SedonaType {
_ => false,
}
}
+
+ /// Return the CRS associated with a geometry/geography type.
+ pub fn crs(&self) -> &Crs {
+ match self {
+ SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) => crs,
+ _ => &Crs::None,
+ }
+ }
+
+ /// Return true if this is an item-level CRS wrapper type.
+ pub fn is_item_crs(&self) -> bool {
+ matches!(
+ self,
+ SedonaType::Arrow(DataType::Struct(fields))
+ if fields.len() == 2 && fields[0].name() == "item" &&
fields[1].name() == "crs"
+ )
+ }
}
// Implementation details for type serialization and display
@@ -563,6 +580,20 @@ mod tests {
);
}
+ #[test]
+ fn sedona_type_crs_and_item_crs_helpers() {
+ let geometry = SedonaType::Wkb(Edges::Planar, lnglat());
+ assert_eq!(geometry.crs(), &lnglat());
+
+ let non_geo = SedonaType::Arrow(DataType::Int32);
+ assert_eq!(non_geo.crs(), &Crs::None);
+
+ let item_crs = SedonaType::new_item_crs(&WKB_GEOMETRY).unwrap();
+ assert!(item_crs.is_item_crs());
+ assert!(!geometry.is_item_crs());
+ assert!(!non_geo.is_item_crs());
+ }
+
#[test]
fn geoarrow_serialize() {
assert_eq!(serialize_edges_and_crs(&Edges::Planar, &Crs::None), "{}");
diff --git a/rust/sedona-schema/src/matchers.rs
b/rust/sedona-schema/src/matchers.rs
index d1a466f5..1910a6c0 100644
--- a/rust/sedona-schema/src/matchers.rs
+++ b/rust/sedona-schema/src/matchers.rs
@@ -49,10 +49,7 @@ impl ArgMatcher {
let geometry_arg_crses = args
.iter()
.filter(|arg_type| IsGeometryOrGeography {}.match_type(arg_type))
- .map(|arg_type| match arg_type {
- SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) =>
crs.clone(),
- _ => None,
- })
+ .map(|arg_type| arg_type.crs().clone())
.collect::<Vec<_>>();
if geometry_arg_crses.is_empty() {
@@ -360,16 +357,7 @@ struct IsItemCrs {}
impl TypeMatcher for IsItemCrs {
fn match_type(&self, arg: &SedonaType) -> bool {
- if let SedonaType::Arrow(DataType::Struct(fields)) = arg {
- let field_names = fields.iter().map(|f|
f.name()).collect::<Vec<_>>();
- if field_names != ["item", "crs"] {
- return false;
- }
-
- return true;
- }
-
- false
+ arg.is_item_crs()
}
}
diff --git a/rust/sedona-testing/src/create.rs
b/rust/sedona-testing/src/create.rs
index 07724a63..caf5faac 100644
--- a/rust/sedona-testing/src/create.rs
+++ b/rust/sedona-testing/src/create.rs
@@ -60,9 +60,7 @@ pub fn create_array_storage(wkt_values: &[Option<&str>],
data_type: &SedonaType)
match data_type {
SedonaType::Wkb(_, _) =>
Arc::new(make_wkb_array::<BinaryArray>(wkt_values)),
SedonaType::WkbView(_, _) =>
Arc::new(make_wkb_array::<BinaryViewArray>(wkt_values)),
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.iter().map(|f| f.name()).collect::<Vec<_>>() ==
vec!["item", "crs"] =>
- {
+ SedonaType::Arrow(DataType::Struct(fields)) if data_type.is_item_crs()
=> {
let item_type =
SedonaType::from_storage_field(&fields[0]).unwrap();
create_array_item_crs(wkt_values, (0..wkt_values.len()).map(|_|
None), &item_type)
}
@@ -102,9 +100,7 @@ pub fn create_scalar_storage(wkt_value: Option<&str>,
data_type: &SedonaType) ->
match data_type {
SedonaType::Wkb(_, _) => ScalarValue::Binary(wkt_value.map(make_wkb)),
SedonaType::WkbView(_, _) =>
ScalarValue::BinaryView(wkt_value.map(make_wkb)),
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.iter().map(|f| f.name()).collect::<Vec<_>>() ==
vec!["item", "crs"] =>
- {
+ SedonaType::Arrow(DataType::Struct(fields)) if data_type.is_item_crs()
=> {
let item_type =
SedonaType::from_storage_field(&fields[0]).unwrap();
create_scalar_item_crs(wkt_value, None, &item_type)
}
diff --git a/rust/sedona-testing/src/testers.rs
b/rust/sedona-testing/src/testers.rs
index 33476d49..f7d50369 100644
--- a/rust/sedona-testing/src/testers.rs
+++ b/rust/sedona-testing/src/testers.rs
@@ -17,7 +17,7 @@
use std::{iter::zip, sync::Arc};
use arrow_array::{ArrayRef, BooleanArray, RecordBatch};
-use arrow_schema::{DataType, FieldRef, Schema};
+use arrow_schema::{FieldRef, Schema};
use datafusion_common::{
arrow::compute::kernels::concat::concat, config::ConfigOptions, Result,
ScalarValue,
};
@@ -665,12 +665,7 @@ impl ScalarUdfTester {
if let Expr::Literal(scalar, _) = arg.lit() {
let is_geometry_or_geography = match sedona_type {
SedonaType::Wkb(_, _) | SedonaType::WkbView(_, _) => true,
- SedonaType::Arrow(DataType::Struct(fields))
- if fields.iter().map(|f| f.name()).collect::<Vec<_>>()
- == vec!["item", "crs"] =>
- {
- true
- }
+ SedonaType::Arrow(_) if sedona_type.is_item_crs() => true,
_ => false,
};
