vortex-geo 0.79.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! `ST_Intersects`: OGC intersection test between two native geometries.

use geo::Intersects;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::IntoArray;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::Constant;
use vortex_array::arrays::ConstantArray;
use vortex_array::arrays::ScalarFnArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::scalar::Scalar;
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_error::VortexResult;
use vortex_error::vortex_ensure_eq;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::extension::geometries;
use crate::extension::single_geometry;
use crate::extension::validate_geometry_operands;

/// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry
/// operands, each a column or a constant literal.
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct GeoIntersects;

impl GeoIntersects {
    /// A lazy `ScalarFnArray` computing per-row whether operands `a` and `b` intersect; either may
    /// be constant. The output length is taken from `a`.
    pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult<ScalarFnArray> {
        ScalarFnArray::try_new(
            TypedScalarFnInstance::new(GeoIntersects, EmptyOptions).erased(),
            vec![a, b],
        )
    }
}

impl ScalarFnVTable for GeoIntersects {
    type Options = EmptyOptions;

    fn id(&self) -> ScalarFnId {
        static ID: CachedId = CachedId::new("vortex.geo.intersects");
        *ID
    }

    fn serialize(&self, _: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
        Ok(Some(vec![]))
    }

    fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult<Self::Options> {
        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!("intersects has exactly two children"),
        }
    }

    fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
        validate_geometry_operands(dtypes)?;
        Ok(DType::Bool(Nullability::NonNullable))
    }

    fn execute(
        &self,
        _: &Self::Options,
        args: &dyn ExecutionArgs,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<ArrayRef> {
        let a = args.get(0)?;
        let b = args.get(1)?;
        match (a.as_opt::<Constant>(), b.as_opt::<Constant>()) {
            (Some(qa), Some(qb)) => {
                let ga = single_geometry(qa.scalar(), ctx)?;
                let gb = single_geometry(qb.scalar(), ctx)?;
                Ok(ConstantArray::new(
                    Scalar::bool(ga.intersects(&gb), Nullability::NonNullable),
                    a.len(),
                )
                .into_array())
            }
            (Some(query), None) => intersects_constant(&b, query.scalar(), ctx),
            (None, Some(query)) => intersects_constant(&a, query.scalar(), ctx),
            (None, None) => {
                vortex_ensure_eq!(
                    a.len(),
                    b.len(),
                    "geo intersects: operand length mismatch {} vs {}",
                    a.len(),
                    b.len()
                );
                let ag = geometries(&a, ctx)?;
                let bg = geometries(&b, ctx)?;
                let hits = ag.iter().zip(&bg).map(|(x, y)| x.intersects(y));
                Ok(BoolArray::from_iter(hits).into_array())
            }
        }
    }
}

/// Whether each row of `operand` intersects the constant `query` geometry. Intersection is
/// symmetric, so this serves a constant on either side.
fn intersects_constant(
    operand: &ArrayRef,
    query: &Scalar,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    let query = single_geometry(query, ctx)?;
    let geoms = geometries(operand, ctx)?;
    let hits = geoms.iter().map(|g| g.intersects(&query));
    Ok(BoolArray::from_iter(hits).into_array())
}

#[cfg(test)]
mod tests {
    use geo_types::Coord;
    use geo_types::Geometry;
    use geo_types::LineString;
    use geo_types::MultiPolygon;
    use geo_types::Polygon;
    use rstest::rstest;
    use vortex_array::ArrayRef;
    use vortex_array::Canonical;
    use vortex_array::ExecutionCtx;
    use vortex_array::IntoArray;
    use vortex_array::VortexSessionExecute;
    use vortex_array::arrays::BoolArray;
    use vortex_array::arrays::ConstantArray;
    use vortex_array::assert_arrays_eq;
    use vortex_array::dtype::DType;
    use vortex_array::dtype::Nullability;
    use vortex_array::dtype::PType;
    use vortex_array::scalar_fn::EmptyOptions;
    use vortex_array::scalar_fn::ScalarFnVTable;
    use vortex_error::VortexResult;
    use vortex_error::vortex_err;
    use wkb::writer::WriteOptions;

    use super::GeoIntersects;
    use crate::test_harness::point_column;

    /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes.
    fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon {
        Polygon::new(
            LineString::from(vec![(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]),
            vec![],
        )
    }

    /// The query region shared by the point tests: a `10x10` square with a `4..6` square hole.
    fn donut() -> Geometry {
        Geometry::Polygon(Polygon::new(
            LineString::from(vec![(0.0, 0.0), (10.0, 0.0), (10.0, 10.0), (0.0, 10.0)]),
            vec![LineString::from(vec![
                (4.0, 4.0),
                (6.0, 4.0),
                (6.0, 6.0),
                (4.0, 6.0),
            ])],
        ))
    }

    /// Probes for [`donut`], one per verdict class: interior, exterior, on the outer boundary,
    /// and inside the hole.
    fn donut_probes() -> VortexResult<ArrayRef> {
        point_column(vec![2.0, 20.0, 0.0, 5.0], vec![2.0, 20.0, 5.0, 5.0])
    }

    /// [`donut_probes`]'s verdicts: interior and boundary contact intersect; exterior and
    /// in-hole do not.
    const DONUT_EXPECTED: [bool; 4] = [true, false, true, false];

    /// A constant column of length `len`, every row the native form of `geometry`.
    fn geometry_constant(geometry: &Geometry, len: usize) -> VortexResult<ArrayRef> {
        let mut buf = Vec::new();
        wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default())
            .map_err(|e| vortex_err!("writing WKB failed: {e}"))?;
        let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)?
            .ok_or_else(|| vortex_err!("unsupported geometry type"))?;
        Ok(ConstantArray::new(scalar, len).into_array())
    }

    /// Materialize `array` so it is no longer a `Constant`, forcing the non-constant kernel paths.
    fn materialize(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
        Ok(array.execute::<Canonical>(ctx)?.into_array())
    }

    /// Execute `GeoIntersects(a, b)` and assert the per-row verdicts equal `expected`.
    fn assert_intersects(
        a: ArrayRef,
        b: ArrayRef,
        expected: impl IntoIterator<Item = bool>,
    ) -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();
        let intersects = GeoIntersects::try_new_array(a, b)?.into_array();
        assert_arrays_eq!(intersects, BoolArray::from_iter(expected), &mut ctx);
        Ok(())
    }

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

    /// Constant vs constant: overlapping and touching (edge or corner) polygons intersect,
    /// disjoint ones do not; every output row carries the same verdict.
    #[rstest]
    #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)]
    #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)]
    #[case::touching_corner(rect_polygon(4.0, 4.0, 8.0, 8.0), true)]
    #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)]
    fn constant_vs_constant_polygons(
        #[case] other: Polygon,
        #[case] expected: bool,
    ) -> VortexResult<()> {
        let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?;
        let b = geometry_constant(&Geometry::Polygon(other), 3)?;
        assert_intersects(a, b, [expected; 3])
    }

    /// Point column vs constant polygon, with the constant on either side since intersection
    /// is symmetric: the [`DONUT_EXPECTED`] verdicts.
    #[test]
    fn point_column_vs_constant_polygon() -> VortexResult<()> {
        assert_intersects(
            donut_probes()?,
            geometry_constant(&donut(), 4)?,
            DONUT_EXPECTED,
        )?;
        assert_intersects(
            geometry_constant(&donut(), 4)?,
            donut_probes()?,
            DONUT_EXPECTED,
        )
    }

    /// Point column vs constant multipolygon: a point in either part intersects, a point in
    /// neither does not.
    #[test]
    fn point_column_vs_constant_multipolygon() -> VortexResult<()> {
        let query = Geometry::MultiPolygon(MultiPolygon::new(vec![
            rect_polygon(0.0, 0.0, 2.0, 2.0),
            rect_polygon(10.0, 10.0, 12.0, 12.0),
        ]));
        let points = point_column(vec![1.0, 11.0, 5.0], vec![1.0, 11.0, 5.0])?;
        assert_intersects(points, geometry_constant(&query, 3)?, [true, true, false])
    }

    /// Polygon column vs constant: the same pairs and verdicts as
    /// [`constant_vs_constant_polygons`].
    #[rstest]
    #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)]
    #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)]
    #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)]
    fn polygon_column_vs_constant(
        #[case] other: Polygon,
        #[case] expected: bool,
    ) -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();
        let column = materialize(geometry_constant(&Geometry::Polygon(other), 2)?, &mut ctx)?;
        let query = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?;
        assert_intersects(column, query, [expected; 2])
    }

    /// Point column vs point column: rows intersect exactly at equal coordinates.
    #[test]
    fn point_column_vs_point_column() -> VortexResult<()> {
        let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?;
        let b = point_column(vec![0.0, 2.0], vec![0.0, 1.0])?;
        assert_intersects(a, b, [true, false])
    }

    /// Point column vs polygon column pairwise: agrees with point column vs constant on the
    /// same data.
    #[test]
    fn columns_agree_with_constant() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let constant = geometry_constant(&donut(), 4)?;
        let column = materialize(constant.clone(), &mut ctx)?;

        let against_constant =
            GeoIntersects::try_new_array(donut_probes()?, constant)?.into_array();
        let pairwise = GeoIntersects::try_new_array(donut_probes()?, column)?.into_array();

        assert_arrays_eq!(against_constant, pairwise, &mut ctx);
        Ok(())
    }

    /// An empty constant geometry intersects nothing.
    #[test]
    fn empty_constant_intersects_nothing() -> VortexResult<()> {
        let empty = Geometry::LineString(LineString::new(Vec::<Coord>::new()));
        let points = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?;
        assert_intersects(points, geometry_constant(&empty, 2)?, [false, false])
    }

    /// Geometry arrays are never nullable, so a nullable operand dtype is rejected.
    #[test]
    fn nullable_operand_is_rejected() -> VortexResult<()> {
        let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone();
        let result = GeoIntersects.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]);
        assert!(result.is_err());
        Ok(())
    }

    /// A non-geometry operand dtype is rejected up front, before execution.
    #[test]
    fn non_geometry_operand_is_rejected() -> VortexResult<()> {
        let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone();
        let numeric = DType::Primitive(PType::I32, Nullability::NonNullable);
        let result = GeoIntersects.return_dtype(&EmptyOptions, &[geo, numeric]);
        assert!(result.is_err());
        Ok(())
    }
}