From d92028871f305e00619610dbfae786056519cf80 Mon Sep 17 00:00:00 2001 From: Ajay Padwal Date: Sat, 4 Jul 2026 04:05:50 +0530 Subject: [PATCH 1/5] Wire S2 covering cell ID parameters --- c/sedona-s2geography/s2geography | 2 +- c/sedona-s2geography/src/kernels.rs | 89 +++++++++++++++++-- docs/reference/sql/s2_coveringcellids.qmd | 18 ++++ .../sedonadb/tests/geography/test_geog_s2.py | 11 +++ 4 files changed, 114 insertions(+), 6 deletions(-) diff --git a/c/sedona-s2geography/s2geography b/c/sedona-s2geography/s2geography index 3452959580..45a0ecf20b 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 06e95b95cc..257b876a16 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 499839e7b6..005d9ab902 100644 --- a/docs/reference/sql/s2_coveringcellids.qmd +++ b/docs/reference/sql/s2_coveringcellids.qmd @@ -21,6 +21,12 @@ description: Returns an array of S2 cell IDs that cover the given geography. kernels: - returns: array args: [geography] + - returns: array + args: [geography, int64] + - returns: array + args: [geography, int64, int64] + - returns: array + args: [geography, int64, int64, int64] --- ## Description @@ -29,8 +35,20 @@ Returns an array of 64-bit S2 cell IDs that form a covering of the given geograp This is useful for spatial indexing and efficient spatial queries. The returned cells may be at different S2 levels depending on the size and shape of the geography. +Optional positional arguments can control the covering: + +| Argument | Description | Default | +| --- | --- | --- | +| `min_level` | Minimum S2 cell level to return. Must be between 0 and 30. | 0 | +| `max_level` | Maximum S2 cell level to return. Must be between 0 and 30. | 30 | +| `max_cells` | Maximum number of cells to return. Must be greater than 0. | 8 | + ## Examples ```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 5b1d170ee5..da359cf539 100644 --- a/python/sedonadb/tests/geography/test_geog_s2.py +++ b/python/sedonadb/tests/geography/test_geog_s2.py @@ -92,3 +92,14 @@ 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() + result = eng.execute_and_collect( + "SELECT S2_CoveringCellIds(" + "ST_GeogFromText('LINESTRING (0 0, 100 50)'), 0, 30, 2)" + ) + df = eng.result_to_pandas(result) + assert len(df.iloc[0, 0]) <= 2 From ac471813614a2658a26eb3bae3d20617d83b0998 Mon Sep 17 00:00:00 2001 From: Ajay Padwal Date: Sat, 4 Jul 2026 04:17:54 +0530 Subject: [PATCH 2/5] Point s2geography submodule at fork pending upstream PR CI can't fetch commit 45a0ecf from paleolimbot/s2geography since it only exists on a fork branch. Temporarily point at the fork (paleolimbot/s2geography#123) until it merges upstream. --- .gitmodules | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index 2923807c67..175fa05a65 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 From 0153f0a1ac659c993e501f062dcb1b8e8a6ee65c Mon Sep 17 00:00:00 2001 From: Ajay Padwal Date: Thu, 9 Jul 2026 05:21:38 +0530 Subject: [PATCH 3/5] Address S2 covering review feedback --- docs/reference/sql/s2_coveringcellids.qmd | 35 +++++++++++++------ .../sedonadb/tests/geography/test_geog_s2.py | 29 +++++++++++---- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/docs/reference/sql/s2_coveringcellids.qmd b/docs/reference/sql/s2_coveringcellids.qmd index 005d9ab902..ffe78a646b 100644 --- a/docs/reference/sql/s2_coveringcellids.qmd +++ b/docs/reference/sql/s2_coveringcellids.qmd @@ -22,11 +22,32 @@ kernels: - returns: array args: [geography] - returns: array - args: [geography, int64] + 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, int64, int64] + 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, int64, int64, int64] + 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 @@ -35,14 +56,6 @@ Returns an array of 64-bit S2 cell IDs that form a covering of the given geograp This is useful for spatial indexing and efficient spatial queries. The returned cells may be at different S2 levels depending on the size and shape of the geography. -Optional positional arguments can control the covering: - -| Argument | Description | Default | -| --- | --- | --- | -| `min_level` | Minimum S2 cell level to return. Must be between 0 and 30. | 0 | -| `max_level` | Maximum S2 cell level to return. Must be between 0 and 30. | 30 | -| `max_cells` | Maximum number of cells to return. Must be greater than 0. | 8 | - ## Examples ```sql diff --git a/python/sedonadb/tests/geography/test_geog_s2.py b/python/sedonadb/tests/geography/test_geog_s2.py index da359cf539..aad9b7d3f8 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( @@ -97,9 +103,20 @@ def test_s2_coveringcellids(eng, geog, expected): @pytest.mark.parametrize("eng", [SedonaDB]) def test_s2_coveringcellids_parameters(eng): eng = eng.create_or_skip() - result = eng.execute_and_collect( - "SELECT S2_CoveringCellIds(" - "ST_GeogFromText('LINESTRING (0 0, 100 50)'), 0, 30, 2)" - ) - df = eng.result_to_pandas(result) - assert len(df.iloc[0, 0]) <= 2 + geog = "ST_GeogFromText('LINESTRING (0 0, 100 50)')" + cases = [ + (f"S2_CoveringCellIds({geog})", 0, 30, 8), + (f"S2_CoveringCellIds({geog}, 4)", 4, 30, 8), + (f"S2_CoveringCellIds({geog}, 4, 6)", 4, 6, 8), + (f"S2_CoveringCellIds({geog}, 4, 6, 2)", 4, 6, 2), + ] + + for expr, min_level, max_level, 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 + assert len(cells) <= max_cells + assert all(min_level <= level <= max_level for level in levels) From 0b08706d3a9f89fd76ed54bc1b0e03a1146910bc Mon Sep 17 00:00:00 2001 From: Ajay Padwal Date: Thu, 9 Jul 2026 05:35:41 +0530 Subject: [PATCH 4/5] Fix S2 covering overload test expectation --- python/sedonadb/tests/geography/test_geog_s2.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/python/sedonadb/tests/geography/test_geog_s2.py b/python/sedonadb/tests/geography/test_geog_s2.py index aad9b7d3f8..081975ca5f 100644 --- a/python/sedonadb/tests/geography/test_geog_s2.py +++ b/python/sedonadb/tests/geography/test_geog_s2.py @@ -105,18 +105,19 @@ 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, 8), - (f"S2_CoveringCellIds({geog}, 4)", 4, 30, 8), - (f"S2_CoveringCellIds({geog}, 4, 6)", 4, 6, 8), + (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}, 4, 6, 2)", 4, 6, 2), ] - for expr, min_level, max_level, max_cells in cases: + 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 - assert len(cells) <= max_cells + if expected_max_cells is not None: + assert len(cells) <= expected_max_cells assert all(min_level <= level <= max_level for level in levels) From 0e8ec6c07a0e2d06a346f76541409da5f3464045 Mon Sep 17 00:00:00 2001 From: Ajay Padwal Date: Fri, 10 Jul 2026 01:16:18 +0530 Subject: [PATCH 5/5] Use permissive levels for S2 max cells test --- python/sedonadb/tests/geography/test_geog_s2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sedonadb/tests/geography/test_geog_s2.py b/python/sedonadb/tests/geography/test_geog_s2.py index 081975ca5f..b5c806846b 100644 --- a/python/sedonadb/tests/geography/test_geog_s2.py +++ b/python/sedonadb/tests/geography/test_geog_s2.py @@ -108,7 +108,7 @@ def test_s2_coveringcellids_parameters(eng): (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}, 4, 6, 2)", 4, 6, 2), + (f"S2_CoveringCellIds({geog}, 0, 30, 2)", 0, 30, 2), ] for expr, min_level, max_level, expected_max_cells in cases: