Skip to content

refactor(vortex-spatial): execute geo kernels with RowFn #9623

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
connortsui20 merged 2 commits into from
Aug 26, 2026
Merged

refactor(vortex-spatial): execute geo kernels with RowFn #9623

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3adca63
refactor(vortex-spatial): execute geo kernels with RowFn
HarukiMoriarty Aug 25, 2026
f73fe6e
perf(spatial): keep length on native storage
HarukiMoriarty Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 47 additions & 21 deletions vortex-spatial/src/extension/polygon.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,28 +11,30 @@ use arrow_array::ArrayRef as ArrowArrayRef;
use arrow_schema::DataType;
use arrow_schema::Field;
use arrow_schema::extension::ExtensionType;
use geo::HasDimensions;
use geo_traits::to_geo::ToGeoGeometry;
use geo_types::Geometry;
use geoarrow::array::GeoArrowArray;
use geoarrow::array::GeoArrowArrayAccessor;
use geoarrow::array::IntoArrow;
use geoarrow::array::PolygonArray;
use geoarrow::array::PolygonBuilder;
use geoarrow::datatypes::CoordType;
use geoarrow::datatypes::PolygonType;
use prost::Message;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::IntoArray;
use vortex_array::arrays::ExtensionArray;
use vortex_array::arrays::ListArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::StructArray;
use vortex_array::arrays::extension::ExtensionArrayExt;
use vortex_array::builtins::ArrayBuiltins;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::extension::ExtDType;
use vortex_array::dtype::extension::ExtId;
use vortex_array::dtype::extension::ExtVTable;
use vortex_array::scalar::ScalarValue;
use vortex_array::validity::Validity;
use vortex_arrow::ArrowExport;
use vortex_arrow::ArrowExportVTable;
use vortex_arrow::ArrowImport;
Expand Down Expand Up @@ -115,25 +117,49 @@ fn polygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> Pol
PolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata))
}

/// Build a native 2-D [`Polygon`] array from row-oriented `geo_types` polygons.
pub(crate) fn build_polygon_array(
polygons: &[Option<geo_types::Polygon<f64>>],
metadata: SpatialMetadata,
nullability: Nullability,
session: &ArrowSession,
/// Build canonical non-nullable 2-D polygon storage from row-oriented `geo_types` polygons.
pub(crate) fn build_polygon_storage(
polygons: &[geo_types::Polygon<f64>],
) -> VortexResult<ArrayRef> {
let polygons =
PolygonBuilder::from_nullable_polygons(polygons, polygon_type(&metadata, Dimension::Xy))
.finish();
let storage_dtype = polygon_storage_dtype(Dimension::Xy, nullability);
let storage = session
.from_arrow_array(
polygons.to_array_ref(),
nullability == Nullability::Nullable,
)?
.cast(storage_dtype.clone())?;
let ext_dtype = ExtDType::<Polygon>::try_new(metadata, storage_dtype)?;
Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array())
let mut xs = Vec::new();
let mut ys = Vec::new();
let mut ring_offsets = vec![0_u64];
let mut polygon_offsets = vec![0_u64];

for polygon in polygons {
let exterior = (!polygon.exterior().is_empty()).then_some(polygon.exterior());
for ring in exterior.into_iter().chain(polygon.interiors()) {
xs.extend(ring.0.iter().map(|coord| coord.x));
ys.extend(ring.0.iter().map(|coord| coord.y));
ring_offsets.push(
u64::try_from(xs.len())
.map_err(|_| vortex_err!("spatial: polygon coordinate count exceeds u64"))?,
);
}
polygon_offsets.push(
u64::try_from(ring_offsets.len() - 1)
.map_err(|_| vortex_err!("spatial: polygon ring count exceeds u64"))?,
);
}

let coordinates = StructArray::from_fields(&[
("x", PrimitiveArray::from_iter(xs).into_array()),
("y", PrimitiveArray::from_iter(ys).into_array()),
])?
.into_array();
let rings = ListArray::try_new(
coordinates,
PrimitiveArray::from_iter(ring_offsets).into_array(),
Validity::NonNullable,
)?
.into_array();
let storage = ListArray::try_new(
rings,
PrimitiveArray::from_iter(polygon_offsets).into_array(),
Validity::NonNullable,
)?
.into_array();
Ok(storage)
}

/// Decode `Polygon` storage (`List<List<coordinate>>`) to `geo_types` polygons, for the spatial scalar
Expand Down
79 changes: 13 additions & 66 deletions vortex-spatial/src/scalar_fn/area.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,41 +5,18 @@

use geo::Area;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::arrays::ScalarFnArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::PType;
use vortex_array::expr::Expression;
use vortex_array::expr::union_child_validities;
use vortex_array::scalar_fn::Arity;
use vortex_array::scalar_fn::ChildName;
use vortex_array::scalar_fn::EmptyOptions;
use vortex_array::scalar_fn::ExecutionArgs;
use vortex_array::scalar_fn::ScalarFnId;
use vortex_array::scalar_fn::ScalarFnVTable;
use vortex_array::scalar_fn::TypedScalarFnInstance;
use vortex_array::scalar_fn::unstable::row::RowFn;
use vortex_array::scalar_fn::unstable::row::RowVisitor;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::extension::is_native_geometry;
use crate::scalar_fn::execute::execute_unary_geo_types;

/// Validate the native geometry operand accepted by `ST_Area`.
fn validate_area_operand(dtypes: &[DType]) -> VortexResult<()> {
vortex_ensure!(
dtypes.len() == 1,
"spatial: area requires exactly one geometry operand, got {}",
dtypes.len()
);
vortex_ensure!(
is_native_geometry(&dtypes[0]),
"spatial: area operand {} is not a native geometry",
dtypes[0]
);
Ok(())
}
use crate::scalar_fn::row::GeometryRow;

/// Unsigned planar `ST_Area` of native geometries.
///
Expand All @@ -58,9 +35,12 @@ impl SpatialArea {
}
}

impl ScalarFnVTable for SpatialArea {
impl RowFn for SpatialArea {
type Options = EmptyOptions;

const ARG_NAMES: &'static [&'static str] = &["geometry"];
const INFALLIBLE: bool = true;

fn id(&self) -> ScalarFnId {
static ID: CachedId = CachedId::new("vortex.st.area");
*ID
Expand All @@ -74,46 +54,13 @@ impl ScalarFnVTable for SpatialArea {
Ok(EmptyOptions)
}

fn arity(&self, _: &Self::Options) -> Arity {
Arity::Exact(1)
}

fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName {
match child_idx {
0 => ChildName::from("geometry"),
_ => unreachable!("area has exactly one child"),
}
}

fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
validate_area_operand(dtypes)?;
Ok(DType::Primitive(PType::F64, dtypes[0].nullability()))
}

fn execute(
fn dispatch<V: RowVisitor>(
&self,
_: &Self::Options,
args: &dyn ExecutionArgs,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
let array = args.get(0)?;
execute_unary_geo_types(&array, Area::unsigned_area, ctx)
}

fn validity(
&self,
_: &Self::Options,
expression: &Expression,
) -> VortexResult<Option<Expression>> {
union_child_validities(expression)
}

fn is_strict(&self, _: &Self::Options) -> bool {
true
}

fn is_infallible(&self, _: &Self::Options) -> bool {
true
_options: &Self::Options,
_args: &[DType],
visitor: V,
) -> VortexResult<V::VisitResult> {
visitor.visit::<(GeometryRow,), f64>(|(geometry,)| geometry.unsigned_area())
}
}

Expand Down
91 changes: 15 additions & 76 deletions vortex-spatial/src/scalar_fn/contains.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,42 +5,18 @@

use geo::Contains;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::arrays::ScalarFnArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::expr::Expression;
use vortex_array::expr::union_child_validities;
use vortex_array::scalar_fn::Arity;
use vortex_array::scalar_fn::ChildName;
use vortex_array::scalar_fn::EmptyOptions;
use vortex_array::scalar_fn::ExecutionArgs;
use vortex_array::scalar_fn::ScalarFnId;
use vortex_array::scalar_fn::ScalarFnVTable;
use vortex_array::scalar_fn::TypedScalarFnInstance;
use vortex_array::scalar_fn::unstable::row::RowFn;
use vortex_array::scalar_fn::unstable::row::RowVisitor;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::extension::is_native_geometry;
use crate::scalar_fn::execute::execute_binary_geo_types;

/// Validate the two native geometry operands accepted by `ST_Contains`.
fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> {
vortex_ensure!(
dtypes.len() == 2,
"spatial: contains requires exactly two geometry operands, got {}",
dtypes.len()
);
for dtype in dtypes {
vortex_ensure!(
is_native_geometry(dtype),
"spatial: contains operand {dtype} is not a native geometry type"
);
}
Ok(())
}
use crate::scalar_fn::row::visit_binary_geo_predicate;

/// OGC `ST_Contains` between two native geometry operands, each a column or a constant
/// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone
Expand All @@ -59,9 +35,12 @@ impl SpatialContains {
}
}

impl ScalarFnVTable for SpatialContains {
impl RowFn for SpatialContains {
type Options = EmptyOptions;

const ARG_NAMES: &'static [&'static str] = &["a", "b"];
const INFALLIBLE: bool = true;

fn id(&self) -> ScalarFnId {
static ID: CachedId = CachedId::new("vortex.st.contains");
*ID
Expand All @@ -75,59 +54,21 @@ impl ScalarFnVTable for SpatialContains {
Ok(EmptyOptions)
}

fn arity(&self, _: &Self::Options) -> Arity {
Arity::Exact(2)
}

fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName {
match child_idx {
0 => ChildName::from("a"),
1 => ChildName::from("b"),
_ => unreachable!("contains has exactly two children"),
}
}

fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
validate_contains_operands(dtypes)?;
let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable));
Ok(DType::Bool(nullability))
}

fn execute(
fn dispatch<V: RowVisitor>(
&self,
_: &Self::Options,
args: &dyn ExecutionArgs,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
let a = args.get(0)?;
let b = args.get(1)?;
_options: &Self::Options,
_args: &[DType],
visitor: V,
) -> VortexResult<V::VisitResult> {
// Containment is not symmetric: `a` is always the container and `b` the contained. A
// container's rect must cover the contained's rect (`Rect::contains` is the closed
// test), so a contained rect poking outside proves the row false.
execute_binary_geo_types(
&a,
&b,
visit_binary_geo_predicate(
visitor,
|a, b| a.contains(b),
Some(|ra, rb| (!ra.contains(rb)).then_some(false)),
ctx,
|a, b| (!a.contains(b)).then_some(false),
)
}

fn validity(
&self,
_: &Self::Options,
expression: &Expression,
) -> VortexResult<Option<Expression>> {
union_child_validities(expression)
}

fn is_strict(&self, _: &Self::Options) -> bool {
true
}

fn is_infallible(&self, _: &Self::Options) -> bool {
true
}
}

#[cfg(test)]
Expand Down Expand Up @@ -201,8 +142,6 @@ mod tests {
Ok(())
}

// The tests cover each `execute` dispatch arm in match order, then the edge cases.

/// Constant vs constant: a polygon contains a nested polygon but not a partially
/// overlapping or disjoint one; every output row carries the same verdict.
#[rstest]
Expand Down
Loading
Loading