jiayuasu commented on code in PR #1163: URL: https://github.com/apache/sedona-db/pull/1163#discussion_r3859863167
########## rust/sedona-functions/src/st_geohash.rs: ########## @@ -0,0 +1,1102 @@ +// 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::sync::Arc; + +use crate::executor::WkbExecutor; +use arrow_array::builder::StringBuilder; +use arrow_schema::DataType; +use datafusion_common::{ + cast::as_int64_array, + config::ConfigOptions, + error::{DataFusionError, Result}, + exec_err, +}; +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::ItemCrsKernel, + scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, +}; +use sedona_geometry::{ + bounds::{WkbBounder2D, WkbBounder2DFactory}, + interval::{Interval, IntervalTrait, WraparoundInterval}, + types::Edges, +}; +use sedona_schema::{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", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STGeoHash { + matcher: ArgMatcher::new( + vec![ArgMatcher::is_geometry_or_geography()], + SedonaType::Arrow(DataType::Utf8), + ), + }), + Arc::new(STGeoHash { + matcher: ArgMatcher::new( + vec![ + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ], + SedonaType::Arrow(DataType::Utf8), + ), + }), + ]), + Volatility::Immutable, + ) +} + +#[derive(Debug)] +struct STGeoHash { + matcher: ArgMatcher, +} + +impl SedonaScalarKernel for STGeoHash { + 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 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)?; + + // The CRS is a property of the type, so whether longitudes may be + // wrapped is decided once for the batch rather than per row. + let wrap = longitude_wrap_for_arg_type(&arg_types[0])?; + + if args.len() > 1 { + append_geohash_with_precision(&executor, args, bounder.as_mut(), wrap, &mut builder)?; + } else { + append_point_geohash(&executor, bounder.as_mut(), 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(), + ) + }) +} + +/// Whether an out-of-range longitude may be wrapped back into [-180, 180] +/// +/// Wrapping is only meaningful when the coordinates are known to be longitude +/// and latitude in degrees: 181 is then an unambiguous spelling of -179, and +/// hashing it is better than dropping the row. Applied to a projected CRS the +/// same arithmetic would turn a coordinate in metres into a plausible-looking +/// geohash for somewhere it has nothing to do with, so it stays off unless the +/// units are known. +#[derive(Debug, Clone, Copy, PartialEq)] +enum LongitudeWrap { + /// The argument is known to be in longitude/latitude degrees + Enabled, + /// The units are projected, or simply not known, so an out-of-range + /// coordinate is not necessarily a wrapped longitude + Disabled, +} + +/// Decide whether an argument's coordinates are known to be lon/lat degrees +/// +/// A geography is lon/lat by definition. A geometry qualifies only when its +/// type-level CRS resolves to geographic parameters (EPSG:4326, OGC:CRS84, +/// EPSG:4269, ...), which is the same test `st_setsrid` uses to decide whether +/// a CRS is usable as a geography. +/// +/// An absent CRS does *not* qualify. Undeclared coordinates are the common case +/// for `ST_GeomFromText()` and could be anything, so they keep the pre-existing +/// null-on-out-of-range behavior rather than being silently reinterpreted. +/// Item-level CRS does not qualify either: [`ItemCrsKernel`] resolves the +/// per-row CRS outside this kernel and hands the inner kernel an item type with +/// no CRS attached, so there is nothing to inspect here. +fn longitude_wrap_for_arg_type(arg_type: &SedonaType) -> Result<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 Ok(LongitudeWrap::Disabled), + }; + + if edges == Edges::Spherical { + return Ok(LongitudeWrap::Enabled); + } + + match arg_type.crs() { + Some(crs) if crs.geographic_params()?.is_some() => Ok(LongitudeWrap::Enabled), + _ => Ok(LongitudeWrap::Disabled), + } +} + +/// 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, + wrap: LongitudeWrap, + builder: &mut StringBuilder, +) -> Result<()> { + let precision_value = args[1] + .cast_to(&DataType::Int64, None)? + .to_array(executor.num_iterations())?; + let precision_array = as_int64_array(&precision_value)?; + let mut precision_iter = precision_array.iter(); + + executor.execute_wkb_void(|maybe_wkb| { + match (maybe_wkb, precision_iter.next().unwrap()) { + (Some(wkb), Some(precision)) => match invoke_scalar(wkb, precision, bounder, wrap)? { + Some(geohash) => builder.append_value(geohash), + // Geometry was empty or outside the lon/lat bounds + None => builder.append_null(), + }, + _ => builder.append_null(), + } + Ok(()) + }) +} + +/// Append the geohash of each point at [MAX_PRECISION] +/// +/// This is the one-argument overload. PostGIS' one-argument ST_GeoHash() +/// derives a precision from the extent of the geometry (the smallest cell that +/// contains it, or a level-20 cell for a point); only the point case is +/// implemented here, where the answer is unambiguous. Anything else errors so +/// that a precision must be stated rather than guessed. +fn append_point_geohash( + executor: &WkbExecutor<'_, '_>, + bounder: &mut dyn WkbBounder2D, + wrap: LongitudeWrap, + builder: &mut StringBuilder, +) -> Result<()> { + executor.execute_wkb_void(|maybe_wkb| { + match maybe_wkb { + Some(wkb) => { + // Only a single POINT counts: a MULTIPOINT, even one holding + // exactly one point, takes the non-point path. PostGIS' rule is + // really about a zero-area bounding box, but "POINT only" is + // simpler to state and to predict. + if !matches!(wkb.as_type(), GeometryType::Point(_)) { + return exec_err!( + "ST_GeoHash(geometry) is only defined for POINT; pass a precision to \ + hash the bounding box center of a non-point geometry" + ); + } + + match invoke_scalar(wkb, MAX_PRECISION, bounder, wrap)? { + Some(geohash) => builder.append_value(geohash), + // Point was empty or outside the lon/lat bounds + None => builder.append_null(), + } + } + None => builder.append_null(), + } + Ok(()) + }) +} + +/// Compute the geohash of a geometry +/// +/// Follows Apache Sedona's GeometryGeoHashEncoder.calculate(): the point that +/// is hashed is the center of the geometry's bounding box. Unlike Sedona's Java +/// implementation (where an empty geometry yields JTS' "null envelope" and thus +/// an accidental hash of (-0.5, -0.5)), empty geometries return null here. +/// +/// Out-of-range coordinates yield null rather than an error, matching Apache +/// Sedona (GeometryGeoHashEncoder.calculate returns null) rather than PostGIS' +/// `geometry` overload (which raises "Geohash requires inputs in decimal +/// degrees"). The divergence is deliberate: a Spark query that returns nulls +/// for out-of-range input should keep returning nulls here rather than start +/// failing partway through a large scan. +/// +/// When `wrap` is [`LongitudeWrap::Enabled`], an out-of-range *longitude* is +/// first wrapped back into [-180, 180] instead of nulling the row, so a +/// longitude of 181 hashes as -179. Latitude is never wrapped; see +/// [`normalize_longitude`] and [`in_latitude_range`] for why the two axes are +/// treated differently. +/// +/// The bounding box comes from `bounder`, which the caller resolved from the +/// argument's edge type, so a geography is bounded on the sphere rather than +/// in the plane. +fn invoke_scalar( + geom: &Wkb, + precision: i64, + bounder: &mut dyn WkbBounder2D, + wrap: LongitudeWrap, +) -> Result<Option<String>> { + bounder.clear(); + bounder Review Comment: Could we validate the raw geography latitude before invoking the spherical bounder? I may be missing another check, but this reproduces for me with the S2-backed CLI: ```sql SELECT ST_GeoHash(ST_GeogFromText('POINT (50 100)'), 12) AS invalid_lat, ST_GeoHash(ST_GeogFromText('POINT (50 90)'), 12) AS pole; -- vpgxczbzuryp | vpgxczbzuryp ``` It looks like S2 clamps 100° to 90° before the check below, so the original out-of-range latitude is no longer visible. Based on the comment below, I expected `invalid_lat` to be `NULL`. A test with the real S2 bounder would probably catch this; the current helper substitutes the planar bounder. ########## rust/sedona-functions/src/st_geohash.rs: ########## @@ -0,0 +1,1102 @@ +// 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::sync::Arc; + +use crate::executor::WkbExecutor; +use arrow_array::builder::StringBuilder; +use arrow_schema::DataType; +use datafusion_common::{ + cast::as_int64_array, + config::ConfigOptions, + error::{DataFusionError, Result}, + exec_err, +}; +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::ItemCrsKernel, + scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, +}; +use sedona_geometry::{ + bounds::{WkbBounder2D, WkbBounder2DFactory}, + interval::{Interval, IntervalTrait, WraparoundInterval}, + types::Edges, +}; +use sedona_schema::{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", + ItemCrsKernel::wrap_impl(vec![ Review Comment: Could this preserve the per-row CRS when deciding whether to wrap longitude? This gives different results for the same EPSG:4326 value depending on how the SRID is represented: ```sql WITH t(srid) AS (VALUES (CAST(4326 AS BIGINT))) SELECT ST_SRID(ST_SetSRID(ST_GeomFromText('POINT (190 50)'), srid)) AS srid, ST_GeoHash(ST_SetSRID(ST_GeomFromText('POINT (190 50)'), srid), 12) AS item_crs, ST_GeoHash(ST_SetSRID(ST_GeomFromText('POINT (190 50)'), 4326), 12) AS type_crs FROM t; -- 4326 | NULL | b0zh7w1z0gs3 ``` My guess is `ItemCrsKernel` strips the row CRS before `longitude_wrap_for_arg_type()` runs. Could we either inspect the CRS per row here or keep the wrapping policy consistent between the two representations? -- 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]
