vortex-geo 0.79.0

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

//! `ST_Contains`: OGC containment test between two native geometries.

use geo::Contains;
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_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
/// does not count). Containment is not symmetric; the operand order is significant.
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct GeoContains;

impl GeoContains {
    /// A lazy `ScalarFnArray` computing per-row whether operand `a` contains operand `b`;
    /// 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(GeoContains, EmptyOptions).erased(),
            vec![a, b],
        )
    }
}

impl ScalarFnVTable for GeoContains {
    type Options = EmptyOptions;

    fn id(&self) -> ScalarFnId {
        static ID: CachedId = CachedId::new("vortex.geo.contains");
        *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!("contains 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)?;
        // Containment is not symmetric: `a` is always the container and `b` the contained.
        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.contains(&gb), Nullability::NonNullable),
                    a.len(),
                )
                .into_array())
            }
            (Some(qa), None) => constant_contains_column(qa.scalar(), &b, ctx),
            (None, Some(qb)) => column_contains_constant(&a, qb.scalar(), ctx),
            (None, None) => {
                vortex_ensure_eq!(
                    a.len(),
                    b.len(),
                    "geo contains: 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.contains(y));
                Ok(BoolArray::from_iter(hits).into_array())
            }
        }
    }
}

/// Whether the constant `container` contains each row of `contained`.
fn constant_contains_column(
    container: &Scalar,
    contained: &ArrayRef,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    let container = single_geometry(container, ctx)?;
    let geoms = geometries(contained, ctx)?;
    let hits = geoms.iter().map(|g| container.contains(g));
    Ok(BoolArray::from_iter(hits).into_array())
}

/// Whether each row of `container` contains the constant `contained`.
fn column_contains_constant(
    container: &ArrayRef,
    contained: &Scalar,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    let contained = single_geometry(contained, ctx)?;
    let geoms = geometries(container, ctx)?;
    let hits = geoms.iter().map(|g| g.contains(&contained));
    Ok(BoolArray::from_iter(hits).into_array())
}

#[cfg(test)]
mod tests {
    use geo_types::Geometry;
    use geo_types::LineString;
    use geo_types::Point;
    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::GeoContains;
    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![],
        )
    }

    /// 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 `GeoContains(a, b)` and assert the per-row verdicts equal `expected`.
    fn assert_contains(
        a: ArrayRef,
        b: ArrayRef,
        expected: impl IntoIterator<Item = bool>,
    ) -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();
        let contains = GeoContains::try_new_array(a, b)?.into_array();
        assert_arrays_eq!(contains, BoolArray::from_iter(expected), &mut ctx);
        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]
    #[case::nested(rect_polygon(1.0, 1.0, 3.0, 3.0), true)]
    #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), false)]
    #[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 container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?;
        let other = geometry_constant(&Geometry::Polygon(other), 3)?;
        assert_contains(container, other, [expected; 3])
    }

    /// Partially overlapping polygons contain each other in neither direction.
    #[test]
    fn overlapping_polygons_contain_neither_way() -> VortexResult<()> {
        let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?;
        let b = geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 2)?;
        assert_contains(a.clone(), b.clone(), [false; 2])?;
        assert_contains(b, a, [false; 2])
    }

    /// Containment is not symmetric: a polygon contains an interior point, but the point does
    /// not contain the polygon.
    #[test]
    fn contains_is_asymmetric() -> VortexResult<()> {
        let polygon = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?;
        let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?;
        assert_contains(polygon.clone(), point.clone(), [true; 2])?;
        assert_contains(point, polygon, [false; 2])
    }

    /// Constant polygon vs point column: a strictly interior point is contained; points outside
    /// or exactly on the boundary are not (OGC contains excludes the boundary).
    #[test]
    fn constant_polygon_vs_point_column() -> VortexResult<()> {
        let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?;
        let points = point_column(vec![2.0, 10.0, 0.0], vec![2.0, 10.0, 2.0])?;
        assert_contains(container, points, [true, false, false])
    }

    /// Polygon column vs constant point: only the polygon around the point contains it.
    #[test]
    fn polygon_column_vs_constant_point() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let around = materialize(
            geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?,
            &mut ctx,
        )?;
        let away = materialize(
            geometry_constant(&Geometry::Polygon(rect_polygon(20.0, 20.0, 24.0, 24.0)), 2)?,
            &mut ctx,
        )?;
        let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?;

        assert_contains(around, point.clone(), [true; 2])?;
        assert_contains(away, point, [false; 2])
    }

    /// Column vs column pairs rows: each polygon row is tested against the point row at the
    /// same position.
    #[test]
    fn polygon_column_vs_point_column() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let polygons = materialize(
            geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?,
            &mut ctx,
        )?;
        let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?;
        assert_contains(polygons, points, [true, 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 = GeoContains.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 = GeoContains.return_dtype(&EmptyOptions, &[geo, numeric]);
        assert!(result.is_err());
        Ok(())
    }
}