jiayuasu commented on code in PR #1163: URL: https://github.com/apache/sedona-db/pull/1163#discussion_r3869629051
########## rust/sedona-functions/src/st_geohash.rs: ########## @@ -0,0 +1,1628 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{collections::HashMap, iter::zip, sync::Arc}; + +use crate::executor::WkbExecutor; +use arrow_array::{builder::StringBuilder, Array}; +use arrow_schema::DataType; +use datafusion_common::{ + cast::{as_int64_array, as_string_view_array, as_struct_array}, + config::ConfigOptions, + error::{DataFusionError, Result}, + exec_err, plan_err, ScalarValue, +}; +use datafusion_expr::{ColumnarValue, Volatility}; +use geo_traits::{GeometryTrait, GeometryType}; +use sedona_common::{option::SedonaOptions, sedona_internal_datafusion_err, sedona_internal_err}; +use sedona_expr::{ + item_crs::parse_item_crs_arg_type, + scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, +}; +use sedona_geometry::{ + bounds::{WkbBounder2D, WkbBounder2DFactory}, + interval::{Interval, IntervalTrait, WraparoundInterval}, + types::Edges, +}; +use sedona_schema::{ + crs::{deserialize_crs, lnglat}, + datatypes::SedonaType, + matchers::ArgMatcher, +}; +use wkb::reader::Wkb; + +/// The base32 alphabet used by geohash encoding (Gustavo Niemeyer's specification) +const BASE32: &[u8; 32] = b"0123456789bcdefghjkmnpqrstuvwxyz"; + +/// The maximum number of geohash characters (matches Apache Sedona's +/// PointGeoHashEncoder, which caps precision at 20) +const MAX_PRECISION: i64 = 20; + +/// ST_GeoHash() scalar UDF +/// +/// Native implementation to compute the geohash of a geometry or geography. +/// The two-argument form hashes at the requested precision (number of base32 +/// characters); the one-argument form hashes a point at [MAX_PRECISION]. +pub fn st_geohash_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_geohash", + // ST_GeoHash cannot use ItemCrsKernel::wrap_impl(): that wrapper strips + // the CRS before calling the inner kernel, which is what makes it free + // for functions like ST_Area() whose answer does not depend on the CRS. + // This one rejects a non-WGS84 CRS, so it needs to see it. The sibling + // kernels below are the same shape ST_SRID() and ST_CRS() use, one + // matching the plain types and one matching item_crs. + vec![ + Arc::new(STGeoHashItemCrs { + matcher: ArgMatcher::new( + vec![ArgMatcher::is_item_crs()], + SedonaType::Arrow(DataType::Utf8), + ), + }) as _, + Arc::new(STGeoHashItemCrs { + matcher: ArgMatcher::new( + vec![ArgMatcher::is_item_crs(), ArgMatcher::is_integer()], + SedonaType::Arrow(DataType::Utf8), + ), + }) as _, + Arc::new(STGeoHash { + matcher: ArgMatcher::new( + vec![ArgMatcher::is_geometry_or_geography()], + SedonaType::Arrow(DataType::Utf8), + ), + }) as _, + Arc::new(STGeoHash { + matcher: ArgMatcher::new( + vec![ + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ], + SedonaType::Arrow(DataType::Utf8), + ), + }) as _, + ], + Volatility::Immutable, + ) +} + +#[derive(Debug)] +struct STGeoHash { + matcher: ArgMatcher, +} + +impl SedonaScalarKernel for STGeoHash { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + let Some(out_type) = self.matcher.match_args(args)? else { + return Ok(None); + }; + + // Checked only after the argument shapes match, so that a call with a + // different arity falls through to the next kernel rather than erroring + // out of kernel resolution here. + ensure_wgs84_crs(&args[0])?; + + Ok(Some(out_type)) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Utf8), 0, None) + } + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + config_options: Option<&ConfigOptions>, + ) -> Result<ColumnarValue> { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = StringBuilder::with_capacity( + executor.num_iterations(), + MAX_PRECISION as usize * executor.num_iterations(), + ); + + // A bounder is a resettable accumulator, so the batch shares one + // instance and clear()s it per row rather than allocating per row. + let mut bounder = bounder_for_arg_type(&arg_types[0], config_options)?; + let mut raw_bounder = raw_bounder_for_arg_type(&arg_types[0]); + + // The CRS is a property of the type, so this is decided once for the + // batch rather than per row. + let wrap = longitude_wrap_for_arg_type(&arg_types[0]); + let mut next_wrap = move || Ok(wrap); + + if args.len() > 1 { + append_geohash_with_precision( + &executor, + args, + bounder.as_mut(), + &mut raw_bounder, + &mut next_wrap, + &mut builder, + )?; + } else { + append_point_geohash( + &executor, + bounder.as_mut(), + &mut raw_bounder, + &mut next_wrap, + &mut builder, + )?; + } + + executor.finish(Arc::new(builder.finish())) + } +} + +/// Resolve the bounder to use for an argument's edge type +/// +/// Planar (geometry) arguments always resolve, falling back to the default +/// Cartesian bounder. Spherical (geography) arguments resolve only when a +/// spherical bounder has been registered on the session runtime, which +/// requires the s2geography-backed bounder; there is no planar fallback, +/// because planar bounds of a geography would silently be wrong. +fn bounder_for_arg_type( + arg_type: &SedonaType, + config_options: Option<&ConfigOptions>, +) -> Result<Box<dyn WkbBounder2D>> { + let edges = match arg_type { + SedonaType::Wkb(edges, _) | SedonaType::WkbView(edges, _) => *edges, + // A literal NULL argument (e.g. ST_GeoHash(NULL, 10)) keeps its Null + // type: every row is null, so the bounder is never used and the choice + // of edge type doesn't matter. + SedonaType::Arrow(DataType::Null) => Edges::Planar, + _ => { + return sedona_internal_err!( + "Expected geometry or geography argument but got {arg_type:?}" + ) + } + }; + + let maybe_bounder = match config_options.and_then(|o| o.extensions.get::<SedonaOptions>()) { + Some(options) => options + .runtime + .bounder_factory() + .bounder_for_edge_type(edges), + None => WkbBounder2DFactory::default().bounder_for_edge_type(edges), + }; + + maybe_bounder.ok_or_else(|| { + DataFusionError::Execution( + "ST_GeoHash() on a geography requires the s2geography-backed spherical bounder, \ + which is not registered in this session" + .to_string(), + ) + }) +} + +/// A planar bounder for range-checking the coordinates a spherical bounder hides +/// +/// s2geography bounds a geography with an `S2LatLngRect`, whose latitudes are +/// constrained to [-90, 90] by construction, so an out-of-range latitude is +/// clamped inside S2 before [`WkbBounder2D::finish`] returns it: a geography at +/// latitude 100, 91 or 270 all come back as 90 and would otherwise encode as +/// the pole. Bounding the same bytes in the plane recovers what the input +/// actually said, so the domain check sees the original value. +/// +/// Only the latitude check uses this. The centre still comes from the spherical +/// bounds, which is the point of bounding a geography on the sphere. Longitude +/// needs no such care: S2 normalizes it into [-180, 180], which is the same +/// wrapping this function applies to a geometry. +/// +/// Returns `None` for a planar argument, whose bounder does no clamping and can +/// be range-checked directly. +fn raw_bounder_for_arg_type(arg_type: &SedonaType) -> Option<Box<dyn WkbBounder2D>> { + let edges = match arg_type { + SedonaType::Wkb(edges, _) | SedonaType::WkbView(edges, _) => *edges, + _ => return None, + }; + + match edges { + Edges::Planar => None, + _ => WkbBounder2DFactory::default().bounder_for_edge_type(Edges::Planar), + } +} + +/// Per-batch cache of the WGS84 verdict for each distinct CRS string +/// +/// Deserializing a CRS is expensive and an item_crs column carries only a +/// handful of distinct values across many rows, so each is resolved once. +/// Mirrors `CachedCrsToSRIDMapping`, which exists for the same reason. +#[derive(Default)] +struct CachedWgs84Check { + /// `None` marks a CRS that is not WGS84, i.e. one this function rejects + cache: HashMap<String, Option<LongitudeWrap>>, +} + +impl CachedWgs84Check { + fn with_capacity(capacity: usize) -> Self { + Self { + cache: HashMap::with_capacity(capacity), + } + } + + /// Resolve one row's CRS into a wrap setting, or reject the row + /// + /// Applies exactly the rule [`ensure_wgs84_crs`] applies to a type-level + /// CRS, so the same geometry is treated the same way whether its CRS is + /// attached to the type or to the row. + fn wrap_for(&mut self, maybe_crs: Option<&str>) -> Result<LongitudeWrap> { + // No CRS on this row: accepted and assumed WGS84, but not wrapped, + // matching an absent type-level CRS. + let Some(crs_str) = maybe_crs else { + return Ok(LongitudeWrap::Disabled); + }; + + if let Some(cached) = self.cache.get(crs_str) { + return cached.map_or_else(|| non_wgs84_err(crs_str), Ok); + } + + // `deserialize_crs` yields None for the "no CRS" sentinels ("0", ""). + let crs = deserialize_crs(crs_str)?; + let verdict = if crs.is_none() { + Some(LongitudeWrap::Disabled) + } else if crs == lnglat() { + Some(LongitudeWrap::Enabled) + } else { + None + }; + + self.cache.insert(crs_str.to_string(), verdict); + verdict.map_or_else(|| non_wgs84_err(crs_str), Ok) + } +} + +fn non_wgs84_err<T>(name: &str) -> Result<T> { + exec_err!( + "ST_GeoHash() requires WGS84 longitude/latitude coordinates but a row has CRS \ + '{name}'. Use ST_Transform() to reproject to EPSG:4326 first." + ) +} + +/// ST_GeoHash() over an item_crs argument, where each row carries its own CRS +/// +/// The rule is the one [`ensure_wgs84_crs`] applies at the type level; only the +/// moment of detection differs. A type-level CRS is known while the query is +/// planned, so it is rejected before any row is read; a row-level CRS is not +/// known until the row is in hand, so it is rejected then. Erroring rather than +/// nulling keeps the two consistent, and matches how the crate already treats a +/// bad row-level CRS elsewhere (`ensure_crs_string_arrays_equal2` raises on the +/// first mismatched row). +/// +/// Seeing the per-row CRS also means a row that declares WGS84 gets longitude +/// wrapping, so item-level and type-level WGS84 agree on the same input rather +/// than one nulling where the other wraps. +#[derive(Debug)] +struct STGeoHashItemCrs { + matcher: ArgMatcher, +} + +impl SedonaScalarKernel for STGeoHashItemCrs { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + self.matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Utf8), 0, None) + } + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + config_options: Option<&ConfigOptions>, + ) -> Result<ColumnarValue> { + let struct_array = match &args[0] { + ColumnarValue::Array(array) => as_struct_array(array)?, + ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => struct_array.as_ref(), + ColumnarValue::Scalar(ScalarValue::Null) => { + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))); + } + _ => return sedona_internal_err!("Unexpected input to ST_GeoHash()"), + }; + + let (item_type, _) = parse_item_crs_arg_type(&arg_types[0])?; + let crs_array = as_string_view_array(struct_array.column(1))?; + + // Rebuild the arguments against the unwrapped item so the geometry path + // below is the same one the type-level kernel takes. A scalar argument + // stays scalar, so that `WkbExecutor` still sizes the batch from + // whichever argument is an array. + let item_arg = match &args[0] { + ColumnarValue::Array(_) => ColumnarValue::Array(struct_array.column(0).clone()), + _ => ColumnarValue::Scalar(ScalarValue::try_from_array(struct_array.column(0), 0)?), + }; + let mut item_arg_types = vec![item_type.clone()]; + let mut item_args = vec![item_arg]; + for (arg_type, arg) in zip(&arg_types[1..], &args[1..]) { + item_arg_types.push(arg_type.clone()); + item_args.push(arg.clone()); + } + + let executor = WkbExecutor::new(&item_arg_types, &item_args); + let mut builder = StringBuilder::with_capacity( + executor.num_iterations(), + MAX_PRECISION as usize * executor.num_iterations(), + ); + let mut bounder = bounder_for_arg_type(&item_type, config_options)?; + let mut raw_bounder = raw_bounder_for_arg_type(&item_type); + let mut checker = CachedWgs84Check::with_capacity(crs_array.len()); + + // A scalar struct carries a single CRS that applies to every iteration; + // an array carries one per row. + let scalar_wrap = match &args[0] { + ColumnarValue::Array(_) => None, + _ => Some(checker.wrap_for(crs_array.iter().next().flatten())?), + }; + let mut crs_iter = crs_array.iter(); + let mut next_wrap = move || match scalar_wrap { + Some(wrap) => Ok(wrap), + // Every row is validated, including one whose geometry is null, so + // that a column fails the same way regardless of where its nulls + // happen to fall. + None => checker.wrap_for(crs_iter.next().flatten()), + }; + + if item_args.len() > 1 { + append_geohash_with_precision( + &executor, + &item_args, + bounder.as_mut(), + &mut raw_bounder, + &mut next_wrap, + &mut builder, + )?; + } else { + append_point_geohash( + &executor, + bounder.as_mut(), + &mut raw_bounder, + &mut next_wrap, + &mut builder, + )?; + } + + executor.finish(Arc::new(builder.finish())) + } +} + +/// Whether an out-of-range longitude may be wrapped back into [-180, 180] +/// +/// [`ensure_wgs84_crs`] has already rejected any CRS that is not WGS84, so this +/// only separates coordinates *declared* to be WGS84 from coordinates carrying +/// no CRS at all. +#[derive(Debug, Clone, Copy, PartialEq)] +enum LongitudeWrap { + /// The argument declares a WGS84 CRS, or is a geography + Enabled, + /// The argument carries no CRS, so its units are assumed but not known + Disabled, +} + +/// Decide whether an argument's longitudes may be wrapped +/// +/// Wrapping requires a *declared* WGS84 CRS (or a geography, which carries one +/// by construction). An absent CRS is accepted by [`ensure_wgs84_crs`] but does +/// not earn wrapping: accepting it is passive -- we cannot know, so we do not +/// break the query -- whereas wrapping it is active, and would invent a location +/// for coordinates whose provenance is unknown. +/// +/// This is also what keeps faith with Apache Sedona. Spark has no CRS concept, +/// so every geometry migrated from it arrives here undeclared; wrapping those +/// would change the out-of-range result from Spark's null to a hash, which is +/// precisely the parity this function is built to preserve. +/// +/// Item-level CRS resolves to `None` here for the reasons given on +/// [`ensure_wgs84_crs`], so it does not wrap either. +fn longitude_wrap_for_arg_type(arg_type: &SedonaType) -> LongitudeWrap { + let edges = match arg_type { + SedonaType::Wkb(edges, _) | SedonaType::WkbView(edges, _) => *edges, + // A literal NULL argument: every row is null, so this is never consulted. + _ => return LongitudeWrap::Disabled, + }; + + if edges == Edges::Spherical || arg_type.crs() == &lnglat() { + LongitudeWrap::Enabled + } else { + LongitudeWrap::Disabled + } +} + +/// Ensure a geohash argument's coordinates are WGS84 longitude/latitude +/// +/// A geohash is defined against the WGS84 datum, so this is the only CRS whose +/// coordinates it can encode correctly. Anything else is rejected rather than +/// hashed: a projected coordinate is in metres, not degrees, and one that +/// happens to land inside [-180, 180] x [-90, 90] would otherwise produce a +/// confident geohash for an unrelated place. EPSG:3857 POINT (10 20) is ten +/// metres east and twenty metres north of the origin, in the Gulf of Guinea, +/// but reads as 10 degrees east and 20 degrees north -- in Chad. +/// +/// `EPSG:4326` and `OGC:CRS84` are the same CRS here (see `authority_codes_equal`), +/// so both are accepted. Other geographic CRSes are *not*, even though their units +/// are also degrees: NAD83 sits one to two metres from WGS84 and NAD27 up to a +/// hundred, which is many cells wide at high precision. Reproject them with +/// `ST_Transform()`. +/// +/// An absent CRS is accepted and assumed to be WGS84. This is what the function +/// already assumes of undeclared coordinates -- the [-180, 180] x [-90, 90] +/// domain check is only meaningful under that reading -- and rejecting it would +/// break every `ST_GeomFromText()` call, which carries no CRS. +/// +/// Item-level CRS is *not* checked: [`ItemCrsKernel`] resolves the per-row CRS +/// outside this kernel and hands the inner kernel an item type whose CRS has +/// been stripped to `None` (`parse_item_crs_arg_type_strip_crs`), which is +/// indistinguishable here from a genuinely absent one. Validating it would mean +/// changing that kernel's contract for every function built on it. +fn ensure_wgs84_crs(arg_type: &SedonaType) -> Result<()> { + let crs = arg_type.crs(); + + if crs.is_none() || crs == &lnglat() { + return Ok(()); + } + + let name = crs + .as_ref() + .map(|crs| crs.to_crs_string()) + .unwrap_or_else(|| "unknown".to_string()); + + plan_err!( + "ST_GeoHash() requires WGS84 longitude/latitude coordinates but the argument has CRS \ + '{name}'. Use ST_Transform() to reproject to EPSG:4326 first." + ) +} + +/// Append the geohash of each geometry at the precision given by the second argument +fn append_geohash_with_precision( + executor: &WkbExecutor<'_, '_>, + args: &[ColumnarValue], + bounder: &mut dyn WkbBounder2D, + raw_bounder: &mut Option<Box<dyn WkbBounder2D>>, + next_wrap: &mut dyn FnMut() -> Result<LongitudeWrap>, + builder: &mut StringBuilder, +) -> Result<()> { + let precision_value = args[1] + .cast_to(&DataType::Int64, None)? Review Comment: Could we handle `UInt64` before this cast? I may be missing another coercion path, but this reproduces for me: ```sql SELECT ST_GeoHash( ST_Point(1, 2), arrow_cast('18446744073709551615', 'UInt64') ); ``` I get `Can't cast value 18446744073709551615 to type Int64`. Since the signature accepts `UInt64` and precision above 20 is capped, I expected the same 20-character hash as precision 20. Maybe values could be clamped before narrowing, or `UInt64` could be excluded from the matcher if it is not meant to be accepted. A boundary test around `i64::MAX` would probably cover it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
