vortex-geo 0.77.0

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

//! `ST_Distance`: planar (Euclidean) distance between two native geometries via the `geo` crate.

use geo::Distance;
use geo::Euclidean;
use vortex_array::ArrayRef;
use vortex_array::ExecutionCtx;
use vortex_array::IntoArray;
use vortex_array::arrays::Constant;
use vortex_array::arrays::ConstantArray;
use vortex_array::arrays::ExtensionArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::ScalarFnArray;
use vortex_array::arrays::StructArray;
use vortex_array::arrays::extension::ExtensionArrayExt;
use vortex_array::arrays::struct_::StructArrayExt;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
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;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::extension::Point;
use crate::extension::coordinate::coordinate_from_struct;
use crate::extension::geometries;
use crate::extension::single_geometry;

/// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry operands.
/// Each is a column or a constant literal; `geo` computes the distance between each pair.
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct GeoDistance;

impl GeoDistance {
    /// A lazy `ScalarFnArray` computing the per-row distance between operands `a` and `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(GeoDistance, EmptyOptions).erased(),
            vec![a, b],
        )
    }
}

impl ScalarFnVTable for GeoDistance {
    type Options = EmptyOptions;

    fn id(&self) -> ScalarFnId {
        static ID: CachedId = CachedId::new("vortex.geo.distance");
        *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!("distance has exactly two children"),
        }
    }

    fn return_dtype(&self, _: &Self::Options, _: &[DType]) -> VortexResult<DType> {
        Ok(DType::Primitive(PType::F64, 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)?;
                let distance = Euclidean.distance(&ga, &gb);
                Ok(ConstantArray::new(
                    Scalar::primitive(distance, Nullability::NonNullable),
                    a.len(),
                )
                .into_array())
            }
            (Some(query), None) => distances_to_constant(&b, query.scalar(), ctx),
            (None, Some(query)) => distances_to_constant(&a, query.scalar(), ctx),
            (None, None) => {
                vortex_ensure!(
                    a.len() == b.len(),
                    "geo distance: operand length mismatch {} vs {}",
                    a.len(),
                    b.len()
                );
                // Fast path: two Point columns, distance straight over their `x`/`y` f64 buffers.
                if is_nonnull_point(a.dtype()) && is_nonnull_point(b.dtype()) {
                    let (xa, ya) = point_xy(&a, ctx)?;
                    let (xb, yb) = point_xy(&b, ctx)?;
                    return Ok(point_distances(
                        xa.as_slice::<f64>().iter().copied(),
                        ya.as_slice::<f64>().iter().copied(),
                        xb.as_slice::<f64>().iter().copied(),
                        yb.as_slice::<f64>().iter().copied(),
                    ));
                }
                let ag = geometries(&a, ctx)?;
                let bg = geometries(&b, ctx)?;
                let distances = ag.iter().zip(&bg).map(|(x, y)| Euclidean.distance(x, y));
                Ok(PrimitiveArray::from_iter(distances).into_array())
            }
        }
    }
}

/// Distance from each row of `operand` to a constant `query` geometry, decoded once and broadcast.
/// Distance is symmetric, so this serves a constant on either side.
fn distances_to_constant(
    operand: &ArrayRef,
    query: &Scalar,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    // Fast path: Point column vs constant Point, `x`/`y` f64 buffers, broadcasting the constant.
    if is_nonnull_point(operand.dtype()) && is_point(query.dtype()) {
        let q = coordinate_from_struct(&query.as_extension().to_storage_scalar())?;
        let (xs, ys) = point_xy(operand, ctx)?;
        return Ok(point_distances(
            xs.as_slice::<f64>().iter().copied(),
            ys.as_slice::<f64>().iter().copied(),
            std::iter::repeat(q.x),
            std::iter::repeat(q.y),
        ));
    }

    let query = single_geometry(query, ctx)?;
    let geoms = geometries(operand, ctx)?;
    let distances = geoms.iter().map(|g| Euclidean.distance(g, &query));
    Ok(PrimitiveArray::from_iter(distances).into_array())
}

/// Extract the `x` and `y` `f64` columns from a native `Point` operand, for the columnar fast paths.
fn point_xy(
    operand: &ArrayRef,
    ctx: &mut ExecutionCtx,
) -> VortexResult<(PrimitiveArray, PrimitiveArray)> {
    let storage = operand
        .clone()
        .execute::<ExtensionArray>(ctx)?
        .storage_array()
        .clone()
        .execute::<StructArray>(ctx)?;
    let xs = storage
        .unmasked_field_by_name("x")?
        .clone()
        .execute::<PrimitiveArray>(ctx)?;
    let ys = storage
        .unmasked_field_by_name("y")?
        .clone()
        .execute::<PrimitiveArray>(ctx)?;
    Ok((xs, ys))
}

/// Per-row planar distance `sqrt(dx^2 + dy^2)` over two `(x, y)` f64 streams; a constant side is fed
/// as `repeat(c)`.
fn point_distances(
    xa: impl Iterator<Item = f64>,
    ya: impl Iterator<Item = f64>,
    xb: impl Iterator<Item = f64>,
    yb: impl Iterator<Item = f64>,
) -> ArrayRef {
    let distances = xa.zip(ya).zip(xb.zip(yb)).map(|((xa, ya), (xb, yb))| {
        let (dx, dy) = (xa - xb, ya - yb);
        (dx * dx + dy * dy).sqrt()
    });
    PrimitiveArray::from_iter(distances).into_array()
}

/// Whether `dtype` is the native `Point` extension (eligible for the columnar fast path).
fn is_point(dtype: &DType) -> bool {
    dtype
        .as_extension_opt()
        .is_some_and(|ext| ext.is::<Point>())
}

/// A non-nullable native `Point`, a column operand the fast path can read straight from `x`/`y`.
fn is_nonnull_point(dtype: &DType) -> bool {
    is_point(dtype) && !dtype.is_nullable()
}

#[cfg(test)]
mod tests {
    use vortex_array::ArrayRef;
    use vortex_array::Canonical;
    use vortex_array::ExecutionCtx;
    use vortex_array::IntoArray;
    use vortex_array::VortexSessionExecute;
    use vortex_array::arrays::ConstantArray;
    use vortex_error::VortexResult;

    use super::GeoDistance;
    use crate::test_harness::point_column;

    /// A constant `Point` column of length `len`, every row at `(x, y)`.
    fn point_constant(
        x: f64,
        y: f64,
        len: usize,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<ArrayRef> {
        let single = point_column(vec![x], vec![y])?.execute_scalar(0, ctx)?;
        Ok(ConstantArray::new(single, len).into_array())
    }

    /// Execute a `GeoDistance` array and read back its per-row `f64` distances.
    fn distances(distance: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Vec<f64>> {
        Ok(distance
            .execute::<Canonical>(ctx)?
            .into_primitive()
            .as_slice::<f64>()
            .to_vec())
    }

    /// `GeoDistance` returns the per-row distance between a point column and a constant query point
    /// (3–4–5 triangles), computed via the geo crate.
    #[test]
    fn distance_over_points() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let a = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?;
        let b = point_constant(0.0, 0.0, 4, &mut ctx)?;
        let distance = GeoDistance::try_new_array(a, b)?.into_array();

        assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]);
        Ok(())
    }

    /// Column-to-column distance pairs corresponding rows of the two columns.
    #[test]
    fn distance_between_columns() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?;
        let b = point_column(vec![3.0, 1.0], vec![4.0, 1.0])?;
        let distance = GeoDistance::try_new_array(a, b)?.into_array();

        assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 0.0]);
        Ok(())
    }

    /// The constant query point may be either operand; distance is symmetric.
    #[test]
    fn distance_with_constant_first_operand() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let a = point_constant(0.0, 0.0, 4, &mut ctx)?;
        let b = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?;
        let distance = GeoDistance::try_new_array(a, b)?.into_array();

        assert_eq!(distances(distance, &mut ctx)?, vec![0.0, 3.0, 4.0, 5.0]);
        Ok(())
    }

    /// Two constant operands: every row has the same distance.
    #[test]
    fn distance_between_two_constants() -> VortexResult<()> {
        let session = vortex_array::array_session();
        let mut ctx = session.create_execution_ctx();

        let a = point_constant(0.0, 0.0, 3, &mut ctx)?;
        let b = point_constant(3.0, 4.0, 3, &mut ctx)?;
        let distance = GeoDistance::try_new_array(a, b)?.into_array();

        assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 5.0, 5.0]);
        Ok(())
    }
}