feat: parameterize S2_CoveringCellIds (min_level, max_level, max_cells)

.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

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<Vec<(String, ScalarKernelRef)>> {
     let mut ffi_scalar_kernels = Vec::<SedonaCScalarKernel>::new();
     ffi_scalar_kernels.resize_with(unsafe { S2GeogNumKernels() }, Default::default);
@@ -235,14 +243,63 @@ pub fn s2_scalar_kernels() -> Result<Vec<(String, ScalarKernelRef)>> {
         "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<SedonaScalarUDF> = None;
         for (kernel_name, kernel) in s2_scalar_kernels().unwrap() {

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);
+```

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}, 4, 6, 2)", 4, 6, 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)