diff --git a/.gitmodules b/.gitmodules index 2923807c6..175fa05a6 100644 --- a/.gitmodules +++ b/.gitmodules @@ -19,7 +19,9 @@ url = https://github.com/geoarrow/geoarrow-data.git [submodule "c/sedona-s2geography/s2geography"] path = c/sedona-s2geography/s2geography - url = https://github.com/paleolimbot/s2geography.git + # TEMPORARY: points at a fork branch pending paleolimbot/s2geography#123. + # Revert to https://github.com/paleolimbot/s2geography.git once merged. + url = https://github.com/ajaypadwal73/s2geography.git [submodule "c/sedona-s2geography/s2geometry"] path = c/sedona-s2geography/s2geometry url = https://github.com/google/s2geometry.git diff --git a/c/sedona-s2geography/s2geography b/c/sedona-s2geography/s2geography index 345295958..45a0ecf20 160000 --- a/c/sedona-s2geography/s2geography +++ b/c/sedona-s2geography/s2geography @@ -1 +1 @@ -Subproject commit 34529595806b0a2420a389efc03c32e841ec9dda +Subproject commit 45a0ecf20b9c046428158b6459afe58fa77991bd diff --git a/c/sedona-s2geography/src/kernels.rs b/c/sedona-s2geography/src/kernels.rs index 06e95b95c..257b876a1 100644 --- a/c/sedona-s2geography/src/kernels.rs +++ b/c/sedona-s2geography/src/kernels.rs @@ -31,6 +31,14 @@ use sedona_schema::{ use crate::{s2geog_check, s2geography_c_bindgen::*}; +fn covering_cell_ids_type() -> SedonaType { + SedonaType::Arrow(DataType::List(Arc::new(arrow_schema::Field::new( + "item", + DataType::Int64, + true, + )))) +} + pub fn s2_scalar_kernels() -> Result> { let mut ffi_scalar_kernels = Vec::::new(); ffi_scalar_kernels.resize_with(unsafe { S2GeogNumKernels() }, Default::default); @@ -235,14 +243,63 @@ pub fn s2_scalar_kernels() -> Result> { "s2_coveringcellids".to_string(), Arc::new(NullKernelHelper::new(ArgMatcher::new( vec![ArgMatcher::is_null()], - SedonaType::Arrow(DataType::List(Arc::new(arrow_schema::Field::new( - "item", - DataType::Int64, - true, - )))), + covering_cell_ids_type(), ))), )); + // s2_coveringcellids(NULL, ...), s2_coveringcellids(geography, NULL, ...) + for matchers in [ + vec![ArgMatcher::is_null(), ArgMatcher::is_integer()], + vec![ArgMatcher::is_geography(), ArgMatcher::is_null()], + vec![ + ArgMatcher::is_null(), + ArgMatcher::is_integer(), + ArgMatcher::is_integer(), + ], + vec![ + ArgMatcher::is_geography(), + ArgMatcher::is_null(), + ArgMatcher::is_integer(), + ], + vec![ + ArgMatcher::is_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_null(), + ], + vec![ + ArgMatcher::is_null(), + ArgMatcher::is_integer(), + ArgMatcher::is_integer(), + ArgMatcher::is_integer(), + ], + vec![ + ArgMatcher::is_geography(), + ArgMatcher::is_null(), + ArgMatcher::is_integer(), + ArgMatcher::is_integer(), + ], + vec![ + ArgMatcher::is_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_null(), + ArgMatcher::is_integer(), + ], + vec![ + ArgMatcher::is_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_integer(), + ArgMatcher::is_null(), + ], + ] { + kernels.push(( + "s2_coveringcellids".to_string(), + Arc::new(NullKernelHelper::new(ArgMatcher::new( + matchers, + covering_cell_ids_type(), + ))), + )); + } + Ok(kernels) } @@ -305,6 +362,28 @@ mod test { assert!(!kernels.is_empty()); } + #[test] + fn covering_cell_ids_return_type() { + let udf = s2_udf("s2_coveringcellids"); + let int64 = SedonaType::Arrow(DataType::Int64); + let null = SedonaType::Arrow(DataType::Null); + + for arg_types in [ + vec![WKB_GEOGRAPHY], + vec![WKB_GEOGRAPHY, int64.clone()], + vec![WKB_GEOGRAPHY, int64.clone(), int64.clone()], + vec![WKB_GEOGRAPHY, int64.clone(), int64.clone(), int64.clone()], + vec![null.clone()], + vec![null.clone(), int64.clone()], + vec![WKB_GEOGRAPHY, null.clone(), int64.clone()], + vec![WKB_GEOGRAPHY, int64.clone(), null.clone(), int64.clone()], + vec![WKB_GEOGRAPHY, int64.clone(), int64.clone(), null.clone()], + ] { + let tester = ScalarUdfTester::new(udf.clone().into(), arg_types); + assert_eq!(tester.return_type().unwrap(), covering_cell_ids_type()); + } + } + fn s2_udf(name: &str) -> SedonaScalarUDF { let mut udf: Option = None; for (kernel_name, kernel) in s2_scalar_kernels().unwrap() { diff --git a/docs/reference/sql/s2_coveringcellids.qmd b/docs/reference/sql/s2_coveringcellids.qmd index 499839e7b..ffe78a646 100644 --- a/docs/reference/sql/s2_coveringcellids.qmd +++ b/docs/reference/sql/s2_coveringcellids.qmd @@ -21,6 +21,33 @@ description: Returns an array of S2 cell IDs that cover the given geography. kernels: - returns: array args: [geography] + - returns: array + args: + - geography + - name: min_level + type: int64 + description: Minimum S2 cell level to return. Must be between 0 and 30. Defaults to 0. + - returns: array + args: + - geography + - name: min_level + type: int64 + description: Minimum S2 cell level to return. Must be between 0 and 30. Defaults to 0. + - name: max_level + type: int64 + description: Maximum S2 cell level to return. Must be between 0 and 30. Defaults to 30. + - returns: array + args: + - geography + - name: min_level + type: int64 + description: Minimum S2 cell level to return. Must be between 0 and 30. Defaults to 0. + - name: max_level + type: int64 + description: Maximum S2 cell level to return. Must be between 0 and 30. Defaults to 30. + - name: max_cells + type: int64 + description: Maximum number of cells to return. Must be greater than 0. Defaults to 8. --- ## Description @@ -34,3 +61,7 @@ This is useful for spatial indexing and efficient spatial queries. The returned ```sql SELECT S2_CoveringCellIds(ST_GeogFromWKT('LINESTRING (0 0, 1 1)')); ``` + +```sql +SELECT S2_CoveringCellIds(ST_GeogFromWKT('LINESTRING (0 0, 1 1)'), 10, 12, 16); +``` diff --git a/python/sedonadb/tests/geography/test_geog_s2.py b/python/sedonadb/tests/geography/test_geog_s2.py index 5b1d170ee..b5c806846 100644 --- a/python/sedonadb/tests/geography/test_geog_s2.py +++ b/python/sedonadb/tests/geography/test_geog_s2.py @@ -23,6 +23,12 @@ pytest.skip("Python package built without s2geography", allow_module_level=True) +def _s2_cell_level(cell_id): + # The least significant set bit encodes the S2 cell level. + marker_bit_index = (int(cell_id) & -int(cell_id)).bit_length() - 1 + return 30 - (marker_bit_index // 2) + + # S2_CellIdFromPoint tests - returns the S2 cell ID containing a point @pytest.mark.parametrize("eng", [SedonaDB, BigQuery]) @pytest.mark.parametrize( @@ -92,3 +98,26 @@ def test_s2_coveringcellids(eng, geog, expected): ) df = eng.result_to_pandas(result) assert len(df.iloc[0, 0]) == len(expected) + + +@pytest.mark.parametrize("eng", [SedonaDB]) +def test_s2_coveringcellids_parameters(eng): + eng = eng.create_or_skip() + geog = "ST_GeogFromText('LINESTRING (0 0, 100 50)')" + cases = [ + (f"S2_CoveringCellIds({geog})", 0, 30, None), + (f"S2_CoveringCellIds({geog}, 4)", 4, 30, None), + (f"S2_CoveringCellIds({geog}, 4, 6)", 4, 6, None), + (f"S2_CoveringCellIds({geog}, 0, 30, 2)", 0, 30, 2), + ] + + for expr, min_level, max_level, expected_max_cells in cases: + result = eng.execute_and_collect(f"SELECT {expr}") + df = eng.result_to_pandas(result) + cells = df.iloc[0, 0] + levels = [_s2_cell_level(cell_id) for cell_id in cells] + + assert len(cells) > 0 + if expected_max_cells is not None: + assert len(cells) <= expected_max_cells + assert all(min_level <= level <= max_level for level in levels)