core-lens 0.1.dev106__tar.gz → 0.1.dev109__tar.gz

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: core-lens
-Version: 0.1.dev106
+Version: 0.1.dev109
 Summary: Query, analyse, and visualise CoreStack's microwatershed and Earth science data through a clean, composable Python API.
 Project-URL: Homepage, https://github.com/ApoorvaKashyap/core-lens
 Project-URL: Issues, https://github.com/ApoorvaKashyap/core-lens/issues
@@ -23,6 +23,7 @@ Requires-Dist: pydantic<3,>=2.13.3; extra == 'core'
 Provides-Extra: full
 Requires-Dist: bokeh<4,>=3.6.0; extra == 'full'
 Requires-Dist: duckdb<2,>=1.5.3; extra == 'full'
+Requires-Dist: geoarrow-rust-core>=0.3.0; extra == 'full'
 Requires-Dist: geopandas<2,>=1.1.3; extra == 'full'
 Requires-Dist: lonboard<1,>=0.16.0; extra == 'full'
 Requires-Dist: matplotlib<4,>=3.10.9; extra == 'full'
@@ -38,6 +39,7 @@ Requires-Dist: bokeh<4,>=3.6.0; extra == 'full-gpu'
 Requires-Dist: cudf-cu13==26.6.*; extra == 'full-gpu'
 Requires-Dist: cudf-polars-cu13==26.6.*; extra == 'full-gpu'
 Requires-Dist: duckdb<2,>=1.5.3; extra == 'full-gpu'
+Requires-Dist: geoarrow-rust-core>=0.3.0; extra == 'full-gpu'
 Requires-Dist: geopandas<2,>=1.1.3; extra == 'full-gpu'
 Requires-Dist: lonboard<1,>=0.16.0; extra == 'full-gpu'
 Requires-Dist: matplotlib<4,>=3.10.9; extra == 'full-gpu'
@@ -70,6 +72,7 @@ Requires-Dist: scipy<2,>=1.17.1; extra == 'stats'
 Requires-Dist: statsmodels<1,>=0.14.6; extra == 'stats'
 Provides-Extra: viz
 Requires-Dist: bokeh<4,>=3.6.0; extra == 'viz'
+Requires-Dist: geoarrow-rust-core>=0.3.0; extra == 'viz'
 Requires-Dist: lonboard<1,>=0.16.0; extra == 'viz'
 Requires-Dist: matplotlib<4,>=3.10.9; extra == 'viz'
 Requires-Dist: polars<2,>=1.39.0; extra == 'viz'

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/pyproject.toml

@@ -55,6 +55,7 @@ spatial = [
 viz = [
     "core-lens[core]",
     "lonboard>=0.16.0,<1",
+    "geoarrow-rust-core>=0.3.0",
     "matplotlib>=3.10.9,<4",
     "bokeh>=3.6.0,<4",
 ]
@@ -107,3 +108,6 @@ explicit = true
 [tool.uv.sources]
 cudf-cu13 = { index = "nvidia" }
 cudf-polars-cu13 = { index = "nvidia" }
+
+[tool.coverage.run]
+source = ["src/core_lens"]

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.dev106'
-__version_tuple__ = version_tuple = (0, 1, 'dev106')
+__version__ = version = '0.1.dev109'
+__version_tuple__ = version_tuple = (0, 1, 'dev109')
 
 __commit_id__ = commit_id = None

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/aoi.py

@@ -247,6 +247,7 @@ class AoI:
 
         layers = []
 
+        # Base layer: single AoI boundary (one Shapely geom — GeoPandas is fine here).
         aoi_gdf = gpd.GeoDataFrame(geometry=[self.geometry], crs="EPSG:4326")
         base_layer = lonboard.PolygonLayer.from_geopandas(
             aoi_gdf,
@@ -261,33 +262,19 @@ class AoI:
             if hasattr(overlay, "has_geometry") and not overlay.has_geometry:
                 overlay = overlay.with_geometry()
 
-            if hasattr(overlay, "gdf"):
-                gdf = overlay.gdf()
-
-                # Lonboard PolygonLayer only supports Polygon and MultiPolygon.
-                # Extract polygon parts from GeometryCollections rather than dropping them.
-                def _extract_polygons(geom: Any) -> Any:
-                    if geom is None:
-                        return None
-                    if geom.geom_type in ("Polygon", "MultiPolygon"):
-                        return geom
-                    if geom.geom_type == "GeometryCollection":
-                        from shapely.geometry import MultiPolygon
-
-                        polys = [
-                            g
-                            for g in getattr(geom, "geoms", [])
-                            if g.geom_type in ("Polygon", "MultiPolygon")
-                        ]
-                        if polys:
-                            return MultiPolygon(polys)
-                    return None
-
-                gdf["geometry"] = gdf.geometry.apply(_extract_polygons)
-                gdf = gdf[gdf.geometry.notna()].copy()
-
-                overlay_layer = lonboard.PolygonLayer.from_geopandas(
-                    gdf,
+            if hasattr(overlay, "entity") and hasattr(overlay, "data"):
+                # --- Direct Polars → GeoArrow → Lonboard (no Shapely, no GeoPandas) ---
+                from core_lens.base.namespaces.plot import _wkb_to_arrow_table
+
+                geom_col = overlay.entity.geometry_col
+                key_cols = (
+                    list(overlay.key_cols) if hasattr(overlay, "key_cols") else []
+                )
+                arrow_table = _wkb_to_arrow_table(
+                    overlay.data, geom_col, extra_cols=key_cols
+                )
+                overlay_layer = lonboard.PolygonLayer(
+                    arrow_table,
                     get_fill_color=[255, 0, 0, 100],
                     get_line_color=[255, 0, 0, 200],
                     line_width_min_pixels=1,
@@ -332,7 +319,6 @@ class AoI:
             ValueError: If the filters match zero rows.
         """
         import shapely.ops as sops
-        import shapely.wkb as swkb
 
         # Find the registered entity whose key_col or known attribute column
         # matches one of the filter keys.
@@ -374,7 +360,7 @@ class AoI:
                 or col in candidate.schema_profile.extra_static_cols
             ):
                 if isinstance(val, list):
-                    filter_expr = filter_expr & pl.col(col).is_in(val)
+                    lf = lf.join(pl.LazyFrame({col: val}), on=col, how="inner")
                 else:
                     filter_expr = filter_expr & (pl.col(col) == val)
 
@@ -386,7 +372,9 @@ class AoI:
                 f"in {candidate.static_path!r}."
             )
 
-        geoms = [swkb.loads(row) for row in df[geom_col].to_list()]
+        import shapely
+
+        geoms = shapely.from_wkb(df[geom_col].to_numpy())
         return sops.unary_union(geoms) if len(geoms) > 1 else geoms[0]
 
     @classmethod

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/base/namespaces/plot.py

@@ -45,6 +45,110 @@ _PALETTE = [
 ]
 
 
+def _wkb_to_arrow_table(
+    df: Any,
+    geom_col: str,
+    extra_cols: list[str] | None = None,
+) -> Any:
+    """Convert a Polars DataFrame with a WKB geometry column to a GeoArrow PyArrow Table.
+
+    Uses ``geoarrow.rust.core.from_wkb`` — a zero-copy C extension bundled
+    transitively by lonboard via ``geoarrow-rust-core``.  No Shapely objects
+    are ever instantiated for the common case (Polygon / MultiPolygon rows).
+
+    GeometryCollection rows are handled by extracting their polygon parts via
+    Shapely, matching the behaviour of the previous GeoPandas code path.  Because
+    GeometryCollections are rare in practice the Shapely overhead is bounded.
+
+    WKB type byte encoding (after the 1-byte endian flag):
+        1=Point, 2=LineString, 3=Polygon, 4=MultiPoint, 5=MultiLineString,
+        6=MultiPolygon, 7=GeometryCollection.
+
+    Args:
+        df: A Polars DataFrame containing the WKB binary column.
+        geom_col: Name of the WKB binary column in *df*.\
+        extra_cols: Additional columns to include in the output table.
+
+    Returns:
+        An ``arro3.core.Table`` whose geometry column is GeoArrow-typed.
+    """
+    import geoarrow.rust.core as ga
+    import polars as pl
+
+    keep = [geom_col] + (extra_cols or [])
+    df = df.select([c for c in keep if c in df.columns])
+
+    def _wkb_geom_type(b: bytes) -> int:
+        """Return the WKB geometry type integer (1-7), or 0 on invalid input."""
+        if not b or len(b) < 5:
+            return 0
+        endian = b[0]  # 1 = little-endian, 0 = big-endian
+        raw = int.from_bytes(b[1:5], "little" if endian == 1 else "big")
+        return raw & 0xFFFF  # strip ISO Z/M/SRID flags
+
+    geom_type_series = df[geom_col].map_elements(_wkb_geom_type, return_dtype=pl.Int32)
+
+    # --- Fast path: pure Polygon / MultiPolygon rows (no Shapely) ---
+    poly_mask = geom_type_series.is_in([3, 6])
+    df_poly = df.filter(poly_mask)
+
+    # --- Slow path: GeometryCollection rows — extract polygon parts via Shapely ---
+    gc_mask = geom_type_series == 7
+    df_gc = df.filter(gc_mask)
+
+    if df_gc.height > 0:
+        import shapely.wkb as swkb
+        from shapely.geometry import MultiPolygon
+
+        repaired_rows: list[dict[str, Any]] = []
+        other_cols = [c for c in df_gc.columns if c != geom_col]
+
+        for row in df_gc.iter_rows(named=True):
+            raw_wkb = row[geom_col]
+            geom = swkb.loads(raw_wkb)
+            polys = [
+                g
+                for g in getattr(geom, "geoms", [])
+                if g.geom_type in ("Polygon", "MultiPolygon")
+            ]
+            if not polys:
+                continue
+            merged = MultiPolygon(polys) if len(polys) > 1 else polys[0]
+            new_row = {geom_col: swkb.dumps(merged)}
+            for col in other_cols:
+                new_row[col] = row[col]
+            repaired_rows.append(new_row)
+
+        if repaired_rows:
+            schema = {geom_col: pl.Binary}
+            for col in other_cols:
+                schema[col] = df_gc[col].dtype
+            df_gc_repaired = pl.DataFrame(repaired_rows, schema=schema)
+            df_poly = pl.concat([df_poly, df_gc_repaired], how="diagonal")
+
+    import arro3.core as ac
+
+    if df_poly.height == 0:
+        # Return an empty table with the correct schema rather than crashing.
+        arrow_table = df_poly.to_arrow()
+        return ac.Table.from_arrow(arrow_table)
+
+    arrow_table = df_poly.to_arrow()
+    wkb_col = arrow_table.column(geom_col).combine_chunks()
+    geo_col = ga.from_wkb(wkb_col)
+
+    idx = arrow_table.schema.get_field_index(geom_col)
+    schema = arrow_table.schema.remove(idx)
+
+    tbl_cols = [ac.ChunkedArray(arrow_table.column(c)) for c in schema.names]
+    tbl_cols.insert(idx, ac.ChunkedArray([ac.Array(geo_col)]))
+    names = list(schema.names)
+    names.insert(idx, geom_col)
+
+    tbl = ac.Table.from_arrays(tbl_cols, names=names)
+    return tbl
+
+
 def _color_for(index: int) -> str:
     """Get a color from the palette by index.
 
@@ -137,64 +241,52 @@ class PlotNamespace:
         subplot_col = subplot_on.value if subplot_on is not None else None
 
         res = self.result if self.result.has_geometry else self.result.with_geometry()
-        gdf = res.gdf()
 
-        if column not in gdf.columns:
+        if column not in res.columns:
             raise ValueError(f"Column {column!r} not found in Result.")
 
-        if subplot_col is not None and subplot_col not in gdf.columns:
+        if subplot_col is not None and subplot_col not in res.columns:
             raise ValueError(
                 f"PlotNamespace.choropleth: subplot_on column {subplot_col!r} not found "
                 "in Result. Ensure fortnightly data is materialised and temporal columns "
                 "are present (they are added automatically by the materialisation layer)."
             )
 
-        # Lonboard PolygonLayer only supports Polygon and MultiPolygon.
-        # Extract polygon parts from GeometryCollections rather than dropping them.
-        def _extract_polygons(geom: Any) -> Any:
-            if geom is None:
-                return None
-            if geom.geom_type in ("Polygon", "MultiPolygon"):
-                return geom
-            if geom.geom_type == "GeometryCollection":
-                from shapely.geometry import MultiPolygon
-
-                polys = [
-                    g
-                    for g in getattr(geom, "geoms", [])
-                    if g.geom_type in ("Polygon", "MultiPolygon")
-                ]
-                if polys:
-                    return MultiPolygon(polys)
-            return None
-
-        gdf["geometry"] = gdf.geometry.apply(_extract_polygons)
-        gdf = gdf[gdf.geometry.notna()].copy()
-
         import matplotlib as mpl
         import numpy as np
+        import polars as pl
+
+        geom_col = res.entity.geometry_col
+        extra = [column] + (res.key_cols or [])
+        if subplot_col is not None:
+            extra.append(subplot_col)
+
+        df = res.data
 
         if subplot_col is not None:
-            # Render the most-recent/first unique value of subplot_on so that the
-            # map is still useful.  True multi-panel subplots are not supported by
-            # Lonboard's single-Map API.
-            unique_vals = sorted(gdf[subplot_col].to_list())
+            # Render most-recent unique value of subplot_on only.
+            # True multi-panel subplots not supported by Lonboard's single-Map API.
+            unique_vals = sorted(df[subplot_col].drop_nulls().unique().to_list())
             if unique_vals:
-                gdf = gdf[gdf[subplot_col] == unique_vals[-1]].copy()
+                df = df.filter(pl.col(subplot_col) == unique_vals[-1])
 
-        values = np.asarray(gdf[column])
+        # --- Direct Polars → GeoArrow → Lonboard path (zero Shapely, zero GeoPandas) ---
+        arrow_table = _wkb_to_arrow_table(df, geom_col, extra_cols=extra)
+
+        # Extract colormap values from the already-filtered arrow_table.
+        # arrow_table has the same row order as df post-filter, so indices align.
+        values = np.asarray(arrow_table.column(column).to_pylist())
         v_min, v_max = np.nanmin(values), np.nanmax(values)
-        if v_max > v_min:
-            norm_values = (values - v_min) / (v_max - v_min)
-        else:
-            norm_values = np.zeros_like(values)
+        norm_values = (
+            (values - v_min) / (v_max - v_min)
+            if v_max > v_min
+            else np.zeros_like(values)
+        )
 
         cmap = mpl.colormaps["plasma"]
 
-        import geopandas as gpd
-
-        layer = lonboard.PolygonLayer.from_geopandas(
-            gpd.GeoDataFrame(gdf),
+        layer = lonboard.PolygonLayer(
+            arrow_table,
             get_fill_color=apply_continuous_cmap(norm_values, cmap),
         )
         return lonboard.Map(layers=[layer])
@@ -368,7 +460,9 @@ class PlotNamespace:
         unique_entities = entity_df[key_col].unique()
         if len(unique_entities) > top_n:
             entities = unique_entities.limit(top_n).to_list()
-            entity_df = entity_df.filter(pl.col(key_col).is_in(entities))
+            entity_df = entity_df.join(
+                pl.DataFrame({key_col: entities}), on=key_col, how="inner"
+            )
         entity_pdf = entity_df.to_pandas()
 
         fig_entity = _make_fig(entity_pdf, f"Timeseries — {y_col} (Per Entity)")
@@ -420,7 +514,7 @@ class PlotNamespace:
         unique_entities = df[key_col].unique()
         if len(unique_entities) > top_n:
             entities = unique_entities.limit(top_n).to_list()
-            df = df.filter(pl.col(key_col).is_in(entities))
+            df = df.join(pl.DataFrame({key_col: entities}), on=key_col, how="inner")
 
         df_pd = df.to_pandas()
 
@@ -491,7 +585,7 @@ class PlotNamespace:
         unique_entities = df[key_col].unique()
         if len(unique_entities) > top_n:
             entities = unique_entities.limit(top_n).to_list()
-            df = df.filter(pl.col(key_col).is_in(entities))
+            df = df.join(pl.DataFrame({key_col: entities}), on=key_col, how="inner")
 
         df_pd = df.to_pandas()
 

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/base/result.py

@@ -148,11 +148,7 @@ class Result:
         geo_df = collect_lf(
             pl.scan_parquet(static_path)
             .select(key_cols + [geom_col])
-            .filter(
-                pl.col(key_cols[0]).is_in(self.data[key_cols[0]].to_list())
-                if len(key_cols) == 1
-                else pl.lit(True)
-            )
+            .join(self.data.select(key_cols).unique().lazy(), on=key_cols, how="inner")
         )
 
         joined = self.data.join(geo_df, on=key_cols, how="left")

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/export/formats.py

@@ -4,6 +4,18 @@ from typing import TYPE_CHECKING, Any
 if TYPE_CHECKING:
     from core_lens.base.result import Result
 
+_DUCKDB_CONN = None
+
+
+def _get_duckdb_conn() -> Any:
+    global _DUCKDB_CONN
+    if _DUCKDB_CONN is None:
+        import duckdb
+
+        _DUCKDB_CONN = duckdb.connect()
+        _DUCKDB_CONN.execute("INSTALL spatial; LOAD spatial;")
+    return _DUCKDB_CONN
+
 
 def parquet(result: "Result", path: str | pathlib.Path, **kwargs: Any) -> None:
     """Export the result data to a standard Parquet file.
@@ -105,8 +117,6 @@ def geoparquet(result: "Result", path: str | pathlib.Path, **kwargs: Any) -> Non
             "Call .with_geometry() first to join the static geometry column before exporting to geoparquet."
         )
 
-    import duckdb
-
     df = result.df()
     geom_col = result.entity.geometry_col
 
@@ -119,8 +129,7 @@ def geoparquet(result: "Result", path: str | pathlib.Path, **kwargs: Any) -> Non
 
     select_clause = ", ".join(cols)
 
-    conn = duckdb.connect()
-    conn.execute("INSTALL spatial; LOAD spatial;")
+    conn = _get_duckdb_conn()
 
     # We must escape single quotes in path if any exist.
     path_str = str(path).replace("'", "''")
@@ -173,8 +182,6 @@ def geojson(result: "Result", path: str | pathlib.Path, **kwargs: Any) -> None:
             "Call .with_geometry() first to join the static geometry column before exporting to geojson."
         )
 
-    import duckdb
-
     df = result.df()
     geom_col = result.entity.geometry_col
 
@@ -187,8 +194,7 @@ def geojson(result: "Result", path: str | pathlib.Path, **kwargs: Any) -> None:
 
     select_clause = ", ".join(cols)
 
-    conn = duckdb.connect()
-    conn.execute("INSTALL spatial; LOAD spatial;")
+    conn = _get_duckdb_conn()
 
     path_str = str(path).replace("'", "''")
 

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/utils/polars_utils.py

@@ -102,16 +102,10 @@ def scan_with_key_filter(
     """
     lf = pl.scan_parquet(path)
 
-    if len(key_cols) == 1:
-        key = key_cols[0]
-        values = key_values[key].to_list()
-        lf = lf.filter(pl.col(key).is_in(values))
-    else:
-        # Composite key: filter each column independently. A small over-selection
-        # is acceptable because the subsequent join at collect time is exact.
-        for key in key_cols:
-            values = key_values[key].to_list()
-            lf = lf.filter(pl.col(key).is_in(values))
+    # Use an inner join to filter the parquet file down to the exact requested keys.
+    # This avoids building a massive literal expression tree (which consumes gigabytes
+    # of RAM) that occurs when passing large lists to pl.col().is_in().
+    lf = lf.join(key_values.lazy(), on=key_cols, how="inner")
 
     if time_expr is not None:
         lf = lf.filter(time_expr)

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/src/core_lens/utils/spatial.py

@@ -6,11 +6,10 @@ import pathlib
 from typing import TYPE_CHECKING, Any
 
 import polars as pl
-import shapely.wkb as swkb
-import shapely.wkt as swkt
+import shapely
 
 if TYPE_CHECKING:
-    import shapely
+    pass
 
 
 def resolve_path(path: str) -> str:
@@ -81,19 +80,27 @@ def build_bbox_index(
     df = pl.read_parquet(static_path, columns=cols_to_read)
     geom_series = df[geometry_col]
 
+    geom_array = geom_series.to_numpy()
+
     if geometry_type == "wkb":
-        geoms = [swkb.loads(v) for v in geom_series.to_list()]
+        geoms = shapely.from_wkb(geom_array)
     else:
-        geoms = [swkt.loads(v) for v in geom_series.to_list()]
-
-    bounds = [g.bounds for g in geoms]
-    minx, miny, maxx, maxy = zip(*bounds) if bounds else ([], [], [], [])
+        geoms = shapely.from_wkt(geom_array)
+
+    if len(geoms) > 0:
+        bnds = shapely.bounds(geoms)
+        return df.select(key_cols).with_columns(
+            pl.Series("minx", bnds[:, 0], dtype=pl.Float64),
+            pl.Series("miny", bnds[:, 1], dtype=pl.Float64),
+            pl.Series("maxx", bnds[:, 2], dtype=pl.Float64),
+            pl.Series("maxy", bnds[:, 3], dtype=pl.Float64),
+        )
 
     return df.select(key_cols).with_columns(
-        pl.Series("minx", list(minx), dtype=pl.Float64),
-        pl.Series("miny", list(miny), dtype=pl.Float64),
-        pl.Series("maxx", list(maxx), dtype=pl.Float64),
-        pl.Series("maxy", list(maxy), dtype=pl.Float64),
+        pl.Series("minx", [], dtype=pl.Float64),
+        pl.Series("miny", [], dtype=pl.Float64),
+        pl.Series("maxx", [], dtype=pl.Float64),
+        pl.Series("maxy", [], dtype=pl.Float64),
     )
 
 
@@ -175,23 +182,19 @@ def exact_spatial_filter(
     if candidates.is_empty():
         return candidates.select(key_cols)
 
-    # Build a filter expression to read only the candidate rows.
-    full_df = pl.read_parquet(static_path, columns=key_cols + [geometry_col])
-
-    # Filter to candidate keys.
-    if len(key_cols) == 1:
-        key = key_cols[0]
-        candidate_keys = candidates[key].to_list()
-        full_df = full_df.filter(pl.col(key).is_in(candidate_keys))
-    else:
-        # Composite key: join on all key columns.
-        full_df = full_df.join(candidates.select(key_cols), on=key_cols, how="inner")
+    # Build a lazy scan and push down an inner join to load only required geometries.
+    full_df = (
+        pl.scan_parquet(static_path)
+        .select(key_cols + [geometry_col])
+        .join(candidates.select(key_cols).lazy(), on=key_cols, how="inner")
+        .collect()
+    )
 
-    geom_series = full_df[geometry_col].to_list()
+    geom_array = full_df[geometry_col].to_numpy()
     if geometry_type == "wkb":
-        geoms = [swkb.loads(v) for v in geom_series]
+        geoms = shapely.from_wkb(geom_array)
     else:
-        geoms = [swkt.loads(v) for v in geom_series]
+        geoms = shapely.from_wkt(geom_array)
 
     if relationship == "centroid":
         # Centroid mode: entity centroid must lie within the AoI geometry.
@@ -209,10 +212,10 @@ def exact_spatial_filter(
         hit_indices = []
         for i in intersect_idx:
             entity_geom = geoms[i]
-            entity_area = entity_geom.area
+            entity_area = entity_geom.area  # pyright: ignore[reportAttributeAccessIssue]
             if entity_area == 0.0:
                 continue
-            inter_area = entity_geom.intersection(aoi_geometry).area
+            inter_area = entity_geom.intersection(aoi_geometry).area  # pyright: ignore[reportAttributeAccessIssue]
             if inter_area / entity_area >= threshold:
                 hit_indices.append(i)
 
@@ -267,15 +270,24 @@ def execute_spatial_join(
     other_df = pl.read_parquet(other_static, columns=other_cols)
 
     # Decode other geometries.
-    other_geom_series = other_df[other_geom_col].to_list()
+    other_geom_array = other_df[other_geom_col].to_numpy()
     if other_geom_type == "wkb":
-        other_geoms = [swkb.loads(v) for v in other_geom_series]
+        other_geoms = shapely.from_wkb(other_geom_array)
     else:
-        other_geoms = [swkt.loads(v) for v in other_geom_series]
+        other_geoms = shapely.from_wkt(other_geom_array)
 
     # Build STRtree from other entity geometries.
     other_tree = shapely.STRtree(other_geoms)
-    other_df_np = other_df.to_pandas()  # for row-level indexing during agg
+
+    # Only convert the columns we actually need to aggregate to Pandas, avoiding
+    # a massive memory copy of the WKB geometry strings and unused keys.
+    agg_cols_only = [
+        c for c in agg.keys() if c not in ("count", "area") and c in other_df.columns
+    ]
+    if agg_cols_only:
+        other_df_np = other_df.select(agg_cols_only).to_pandas()
+    else:
+        other_df_np = None
 
     # Decode primary geometries.
     if primary_geom_col not in primary_df.columns:
@@ -285,18 +297,28 @@ def execute_spatial_join(
             "with_geometry() first."
         )
 
-    primary_geom_series = primary_df[primary_geom_col].to_list()
+    primary_geom_array = primary_df[primary_geom_col].to_numpy()
     if primary_geom_type == "wkb":
-        primary_geoms = [swkb.loads(v) for v in primary_geom_series]
+        primary_geoms = shapely.from_wkb(primary_geom_array)
     else:
-        primary_geoms = [swkt.loads(v) for v in primary_geom_series]
+        primary_geoms = shapely.from_wkt(primary_geom_array)
 
     # For each primary entity, find overlapping other entities and aggregate.
     result_rows: list[dict[str, Any]] = []
     primary_keys = primary_df.select(primary_key_cols).to_dicts()
 
+    # Vectorized STRtree query over all primary geometries at once
+    # Returns a 2D array of shape (2, N): [[primary_indices], [other_indices]]
+    idx_pairs = other_tree.query(primary_geoms, predicate="intersects")
+
+    from collections import defaultdict
+
+    hits_by_primary = defaultdict(list)
+    for p_idx, o_idx in zip(idx_pairs[0], idx_pairs[1]):
+        hits_by_primary[p_idx].append(o_idx)
+
     for i, (pgeom, pkey) in enumerate(zip(primary_geoms, primary_keys)):
-        hit_idx = other_tree.query(pgeom, predicate="intersects").tolist()
+        hit_idx = hits_by_primary.get(i, [])
         row: dict[str, Any] = dict(pkey)
 
         if not hit_idx:
@@ -306,7 +328,10 @@ def execute_spatial_join(
             result_rows.append(row)
             continue
 
-        matched_other = other_df_np.iloc[hit_idx]
+        if other_df_np is not None:
+            matched_other = other_df_np.iloc[hit_idx]
+        else:
+            matched_other = None
 
         for col, fn in agg.items():
             out_col = f"{other_entity_name}_{col}"
@@ -314,10 +339,10 @@ def execute_spatial_join(
                 row[out_col] = len(hit_idx)
             elif fn == "area":
                 # Sum of intersection areas (geographic, approx).
-                total = sum(pgeom.intersection(other_geoms[j]).area for j in hit_idx)
+                total = sum(pgeom.intersection(other_geoms[j]).area for j in hit_idx)  # pyright: ignore[reportAttributeAccessIssue]
                 row[out_col] = total
-            elif col in matched_other.columns:
-                vals = matched_other[col].dropna().values
+            elif matched_other is not None and col in matched_other.columns:
+                vals = matched_other[col].dropna()
                 if fn == "mean":
                     row[out_col] = float(vals.mean()) if len(vals) else None
                 elif fn == "sum":

{core_lens-0.1.dev106 → core_lens-0.1.dev109}/uv.lock

@@ -554,6 +554,7 @@ core = [
 full = [
     { name = "bokeh" },
     { name = "duckdb" },
+    { name = "geoarrow-rust-core" },
     { name = "geopandas" },
     { name = "lonboard" },
     { name = "matplotlib" },
@@ -570,6 +571,7 @@ full-gpu = [
     { name = "cudf-cu13" },
     { name = "cudf-polars-cu13" },
     { name = "duckdb" },
+    { name = "geoarrow-rust-core" },
     { name = "geopandas" },
     { name = "lonboard" },
     { name = "matplotlib" },
@@ -606,6 +608,7 @@ stats = [
 ]
 viz = [
     { name = "bokeh" },
+    { name = "geoarrow-rust-core" },
     { name = "lonboard" },
     { name = "matplotlib" },
     { name = "polars" },
@@ -653,6 +656,9 @@ requires-dist = [
     { name = "duckdb", marker = "extra == 'full'", specifier = ">=1.5.3,<2" },
     { name = "duckdb", marker = "extra == 'full-gpu'", specifier = ">=1.5.3,<2" },
     { name = "duckdb", marker = "extra == 'spatial'", specifier = ">=1.5.3,<2" },
+    { name = "geoarrow-rust-core", marker = "extra == 'full'", specifier = ">=0.3.0" },
+    { name = "geoarrow-rust-core", marker = "extra == 'full-gpu'", specifier = ">=0.3.0" },
+    { name = "geoarrow-rust-core", marker = "extra == 'viz'", specifier = ">=0.3.0" },
     { name = "geopandas", marker = "extra == 'full'", specifier = ">=1.1.3,<2" },
     { name = "geopandas", marker = "extra == 'full-gpu'", specifier = ">=1.1.3,<2" },
     { name = "geopandas", marker = "extra == 'spatial'", specifier = ">=1.1.3,<2" },

core_lens-0.1.dev106/plot_wb.py

@@ -1,29 +0,0 @@
-import polars as pl
-from core_lens import AoI
-from core_lens.base.entity import BaseEntity
-
-from core_lens.entities.waterbody import WaterbodyEntity
-
-def main():
-    # Register the entity
-    AoI.register(WaterbodyEntity)
-    
-    # Init AoI with a bounding box for India (or you can adjust bbox)
-    DATA_ROOT = "data/"
-    aoi = AoI(DATA_ROOT, bbox=(68.0, 6.0, 98.0, 38.0))
-    
-    # Materialise the static waterbodies data
-    print("Loading waterbodies...")
-    wb_result = aoi.waterbody.static
-    
-    # Plot using the area column for choropleth colors
-    print("Plotting waterbodies...")
-    choropleth_map = wb_result.plot.choropleth("area_ha")
-    
-    # Save the map to HTML
-    output_file = "wb_map.html"
-    choropleth_map.to_html(output_file)
-    print(f"Saved waterbodies map to {output_file}")
-
-if __name__ == "__main__":
-    main()

core_lens-0.1.dev106/wb_annual.txt

@@ -1,11 +0,0 @@
-tehsil                         String
-district                       String
-state                          String
-area_in_ha                     Float64
-wb_id                          String
-ba_name                        String
-year                           Int32
-area_ha                        Float64
-kharif                         Float64
-karif_rabi                     Float64
-kharif_rabi_zaid               Float64

core_lens-0.1.dev106/wb_static.txt

@@ -1,11 +0,0 @@
-tehsil                         String
-district                       String
-state                          String
-area_in_ha                     Float64
-wb_id                          String
-ba_name                        String
-year                           Int32
-area_ha                        Float64
-kharif                         Float64
-karif_rabi                     Float64
-kharif_rabi_zaid               Float64