From dd026bd3963f249670a9752eec53573c319f3cc0 Mon Sep 17 00:00:00 2001 From: Jia Yu Date: Wed, 15 Jul 2026 00:17:03 -0700 Subject: [PATCH] feat(python/sedonadb-geopandas): add GeoPandas-compatible API package --- python/sedonadb-geopandas/README.md | 53 ++++ python/sedonadb-geopandas/pyproject.toml | 55 ++++ .../python/sedonadb_geopandas/__init__.py | 54 ++++ .../python/sedonadb_geopandas/_context.py | 34 +++ .../python/sedonadb_geopandas/_frame.py | 174 +++++++++++++ .../python/sedonadb_geopandas/_series.py | 139 ++++++++++ python/sedonadb-geopandas/tests/__init__.py | 16 ++ .../tests/test_geopandas_compat.py | 244 ++++++++++++++++++ 8 files changed, 769 insertions(+) create mode 100644 python/sedonadb-geopandas/README.md create mode 100644 python/sedonadb-geopandas/pyproject.toml create mode 100644 python/sedonadb-geopandas/python/sedonadb_geopandas/__init__.py create mode 100644 python/sedonadb-geopandas/python/sedonadb_geopandas/_context.py create mode 100644 python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py create mode 100644 python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py create mode 100644 python/sedonadb-geopandas/tests/__init__.py create mode 100644 python/sedonadb-geopandas/tests/test_geopandas_compat.py diff --git a/python/sedonadb-geopandas/README.md b/python/sedonadb-geopandas/README.md new file mode 100644 index 0000000000..d30d2ef057 --- /dev/null +++ b/python/sedonadb-geopandas/README.md @@ -0,0 +1,53 @@ + + +# sedonadb-geopandas + +A GeoPandas-compatible API on top of [SedonaDB](https://sedona.apache.org/sedonadb/). + +The goal is to let existing GeoPandas code run against SedonaDB's relational +engine with minimal changes, by providing `GeoDataFrame` / `GeoSeries` wrappers +whose methods mirror GeoPandas but delegate to SedonaDB expressions. + +```python +import geopandas +import sedonadb_geopandas as sgpd + +gdf = sgpd.from_geopandas(geopandas.read_file("cities.geojson")) +big = gdf[gdf["pop"] > 1_000_000] # boolean-mask filter +buffered = gdf.geometry.buffer(0.5) # element-wise .geo operation +web = gdf.to_crs("EPSG:3857") # reproject (CRS tracked through) +result = web.to_geopandas() # back to a real GeoDataFrame +``` + +## Intentional differences from GeoPandas + +This is a compatibility layer over a lazy, relational engine, so it is +deliberately *not* identical to GeoPandas: + +- **Lazy, not eager**: operations build a query; data materializes on + `to_geopandas()` / `to_pandas()` / display. +- **No row index / alignment**: there is no pandas `Index`; joins and filters + are positional/relational, not index-aligned. +- **Immutable under the hood**: "in-place" style operations return a new frame. +- **Plotting and arbitrary `apply`**: use the `to_geopandas()` escape hatch and + operate on the materialized result. + +See the SedonaDB "Migrating from GeoPandas" guide for the relational model that +underlies each method. diff --git a/python/sedonadb-geopandas/pyproject.toml b/python/sedonadb-geopandas/pyproject.toml new file mode 100644 index 0000000000..a1d8b059ef --- /dev/null +++ b/python/sedonadb-geopandas/pyproject.toml @@ -0,0 +1,55 @@ +# 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. + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "sedonadb-geopandas" +version = "0.4.0" +description = "GeoPandas-compatible API on top of SedonaDB" +readme = "README.md" +requires-python = ">=3.9" +dependencies = [ + "sedonadb", + "sedonadb-expr", + "pyarrow", +] +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", +] + +[project.optional-dependencies] +# geopandas is only needed for the from_geopandas / to_geopandas interop +# helpers, so it is an opt-in extra rather than a hard dependency. +geopandas = ["geopandas"] +test = [ + "pytest", + "pandas", + "geopandas", +] + +[tool.hatch.build.targets.wheel] +packages = ["python/sedonadb_geopandas"] diff --git a/python/sedonadb-geopandas/python/sedonadb_geopandas/__init__.py b/python/sedonadb-geopandas/python/sedonadb_geopandas/__init__.py new file mode 100644 index 0000000000..a9d86ee397 --- /dev/null +++ b/python/sedonadb-geopandas/python/sedonadb_geopandas/__init__.py @@ -0,0 +1,54 @@ +# 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. +"""GeoPandas-compatible API on top of SedonaDB. + +**EXPERIMENTAL.** This package provides `GeoDataFrame` / `GeoSeries` wrappers +whose methods mirror GeoPandas but delegate to a lazy SedonaDB engine. It is a +compatibility layer, not a drop-in replacement: see the package README for the +intentional differences (laziness, no row index, immutability). The API may +change without notice. +""" + +from sedonadb_geopandas._context import default_context +from sedonadb_geopandas._frame import GeoDataFrame +from sedonadb_geopandas._series import GeoSeries, Series + +__all__ = ["GeoDataFrame", "GeoSeries", "Series", "from_geopandas"] + + +def from_geopandas(data, *, context=None, geometry=None): + """Load a `geopandas.GeoDataFrame` into a SedonaDB-backed `GeoDataFrame`. + + **EXPERIMENTAL.** + + Args: + data: A `geopandas.GeoDataFrame` (or any object accepted by + `SedonaContext.create_data_frame`, which surfaces an error for + anything that cannot be turned into a DataFrame). + context: An optional SedonaDB context. Defaults to a shared, + lazily-created one. + geometry: The active geometry column name. Defaults to SedonaDB's + primary-geometry heuristic (the same one `to_geopandas` uses). + + Returns: + A `GeoDataFrame`. + """ + ctx = context or default_context() + df = ctx.create_data_frame(data) + if geometry is None: + return GeoDataFrame(df) + return GeoDataFrame(df, geometry=geometry) diff --git a/python/sedonadb-geopandas/python/sedonadb_geopandas/_context.py b/python/sedonadb-geopandas/python/sedonadb_geopandas/_context.py new file mode 100644 index 0000000000..0e05d3297f --- /dev/null +++ b/python/sedonadb-geopandas/python/sedonadb_geopandas/_context.py @@ -0,0 +1,34 @@ +# 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. +"""Default SedonaDB context. + +GeoPandas has no notion of a connection, so this package lazily creates a single +shared SedonaDB context the first time it is needed. Callers that want an +explicit context can always pass one through the public constructors. +""" + +_DEFAULT = None + + +def default_context(): + """Return the process-wide default SedonaDB context, creating it if needed.""" + global _DEFAULT + if _DEFAULT is None: + import sedonadb + + _DEFAULT = sedonadb.connect() + return _DEFAULT diff --git a/python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py b/python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py new file mode 100644 index 0000000000..63375166ba --- /dev/null +++ b/python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py @@ -0,0 +1,174 @@ +# 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. +"""GeoPandas-style GeoDataFrame backed by a lazy SedonaDB frame.""" + +from sedonadb_geopandas._series import GeoSeries, Series + +# Rows to collect for the Jupyter rich-text (`_repr_html_`) preview. +_REPR_HTML_ROWS = 10 + +# Default for the `geometry` argument, distinguishing "not specified, apply the +# heuristic" from an explicit `None` meaning "this frame has no active geometry". +_DERIVE = object() + + +def _geometry_column_names(df): + names = df.schema.names + return {names[i] for i in df.schema.geometry_column_indices} + + +class GeoDataFrame: + """A lazy SedonaDB frame in the shape of a `geopandas.GeoDataFrame`. + + **EXPERIMENTAL.** Wraps a SedonaDB `DataFrame` and tracks the active + geometry column. Row selection, column access, and geometry operations + mirror GeoPandas but build a query rather than computing eagerly; call + `to_geopandas()` to materialize. + """ + + def __init__(self, df, geometry=_DERIVE): + self._df = df + if geometry is _DERIVE: + # Fall back to SedonaDB's primary-geometry heuristic (same one + # `to_geopandas` uses); `None` when the frame has no geometry. + geometry = df._impl.primary_geometry_column() + elif geometry is not None and geometry not in _geometry_column_names(df): + if geometry not in df.schema.names: + raise KeyError( + f"Geometry column {geometry!r} not found; columns: " + f"{df.schema.names}" + ) + raise ValueError(f"Column {geometry!r} is not a geometry column") + self._geometry_name = geometry + + @property + def geometry(self): + """The active geometry column as a `GeoSeries`.""" + if self._geometry_name is None: + raise AttributeError("This GeoDataFrame has no active geometry column") + return GeoSeries(self._df, self._df[self._geometry_name], self._geometry_name) + + @property + def crs(self): + """The CRS of the active geometry column, or `None` if there is none.""" + if self._geometry_name is None: + return None + return self._df.schema.field(self._geometry_name).type.crs + + @property + def columns(self): + """Column names, mirroring `GeoDataFrame.columns`.""" + return list(self._df.schema.names) + + def __getitem__(self, key): + # Boolean mask -> row filter (gdf[gdf["pop"] > 1000]). + if isinstance(key, Series): + return GeoDataFrame(self._df.filter(key._expr), self._geometry_name) + + # Column subset -> GeoDataFrame. Matching GeoPandas, the active geometry + # column is persisted when it survives the subset (rather than being + # re-derived, which could silently pick a different geometry column) and + # is dropped when it does not. GeoPandas returns a plain DataFrame in + # that case; here the result keeps its type but has no active geometry, + # so `.geometry` raises just as it does there. + if isinstance(key, list): + geometry = self._geometry_name if self._geometry_name in key else None + return GeoDataFrame(self._df.select(*key), geometry) + + # Single column -> (Geo)Series. + if isinstance(key, str): + expr = self._df[key] + if key == self._geometry_name: + return GeoSeries(self._df, expr, key) + return Series(self._df, expr, key) + + if isinstance(key, slice): + raise TypeError( + "Positional row slicing isn't supported: this frame has no row " + "index, and row order isn't guaranteed. Use head(n) for a " + "bounded number of rows, or filter on a column." + ) + + if isinstance(key, int): + # Matches GeoPandas/pandas, where an integer key is a column label. + raise KeyError( + f"Column {key!r} not found (an integer key is a column label, " + f"not a row position). Columns: {self.columns}" + ) + + raise TypeError( + f"GeoDataFrame indices must be a column name, list of names, or " + f"boolean mask, not {type(key).__name__}" + ) + + def head(self, n=5): + """Return a `GeoDataFrame` of at most `n` rows. + + Note that this applies a limit without an ordering, so *which* rows come + back isn't guaranteed — the frame has no inherent row order. + """ + return GeoDataFrame(self._df.limit(n), self._geometry_name) + + def to_crs(self, crs): + """Reproject the geometry column to `crs` (`ST_Transform`).""" + if self._geometry_name is None: + raise ValueError("to_crs() requires an active geometry column") + from sedonadb.expr import lit + + transformed = self._df[self._geometry_name].geo.transform(lit(crs)) + new_df = self._df.mutate(**{self._geometry_name: transformed}) + return GeoDataFrame(new_df, self._geometry_name) + + def to_geopandas(self): + """Execute and return a `geopandas.GeoDataFrame` (or plain DataFrame). + + The active geometry column is carried over, so a frame whose geometry + column is not the one SedonaDB's own heuristic would pick (for example a + column named `geom` alongside one named `geometry`) still comes back with + the expected column active. + """ + result = self._df.to_pandas() + if self._geometry_name is not None and hasattr(result, "set_geometry"): + try: + active = result.geometry.name + except Exception: + active = None + if active != self._geometry_name: + result = result.set_geometry(self._geometry_name) + return result + + # Alias: results carry geometry, so this returns a GeoDataFrame too. + to_pandas = to_geopandas + + def __len__(self): + return self._df.count() + + def __repr__(self): + # Cheap: no execution. IDEs/consoles call repr frequently. + return f"GeoDataFrame(columns={self.columns}, geometry={self._geometry_name!r})" + + def _repr_html_(self): + # Rich Jupyter display: collect only a small preview. + try: + preview = self._df.limit(_REPR_HTML_ROWS).to_pandas() + table = preview._repr_html_() + except Exception: + return None # fall back to __repr__ + return ( + f"
GeoDataFrame (preview of up to " + f"{_REPR_HTML_ROWS} rows)
{table}" + ) diff --git a/python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py b/python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py new file mode 100644 index 0000000000..8fd27a6ac0 --- /dev/null +++ b/python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py @@ -0,0 +1,139 @@ +# 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. +"""pandas/GeoPandas-style Series backed by a SedonaDB expression.""" + + +def _operand(df, other): + """Coerce the right-hand side of an operator into something usable. + + A `Series` is unwrapped to its expression, but only if it came from the same + source frame as `df`: combining columns from two different frames has no + defined meaning here (there is no row alignment) and would otherwise build a + plan that silently returns wrong rows. A raw SedonaDB `Expr` or `Literal` is + passed through as-is (`lit()` is a useful escape hatch for specifying a + literal that carries a CRS); other scalars pass through unchanged. A + pandas/numpy array-like is rejected with a clear message, since it would + otherwise fail obscurely as a multi-element literal. + """ + from sedonadb.expr import Expr, Literal + + if isinstance(other, Series): + if other._df is not df: + raise ValueError( + "Cannot combine Series that come from different DataFrames: " + "there is no row alignment, so the result would be silently " + "wrong. Reference columns of a single frame, or join the two " + "frames first." + ) + return other._expr + if isinstance(other, (Expr, Literal)): + return other + if hasattr(other, "__array__"): + raise TypeError( + "Operating against a pandas/numpy array-like isn't supported " + "(there is no row alignment). Operate within this frame, or collect " + "with to_pandas() first." + ) + return other + + +class Series: + """A single column of a lazy SedonaDB frame, in the shape of a pandas Series. + + **EXPERIMENTAL.** A `Series` pairs a source SedonaDB `DataFrame` with an + expression over its columns. Comparisons produce a boolean `Series` usable + as a filter mask (`gdf[gdf["pop"] > 1000]`). Nothing is computed until + `to_pandas()`. + """ + + def __init__(self, df, expr, name): + self._df = df + self._expr = expr + self._name = name + + # -- element-wise comparisons -> boolean mask -------------------------- + def __gt__(self, other): + return Series(self._df, self._expr > _operand(self._df, other), self._name) + + def __ge__(self, other): + return Series(self._df, self._expr >= _operand(self._df, other), self._name) + + def __lt__(self, other): + return Series(self._df, self._expr < _operand(self._df, other), self._name) + + def __le__(self, other): + return Series(self._df, self._expr <= _operand(self._df, other), self._name) + + def __eq__(self, other): + return Series(self._df, self._expr == _operand(self._df, other), self._name) + + def __ne__(self, other): + return Series(self._df, self._expr != _operand(self._df, other), self._name) + + # -- boolean composition of masks -------------------------------------- + def __and__(self, other): + return Series(self._df, self._expr & _operand(self._df, other), self._name) + + def __or__(self, other): + return Series(self._df, self._expr | _operand(self._df, other), self._name) + + def __invert__(self): + return Series(self._df, ~self._expr, self._name) + + __hash__ = None + + # -- materialization --------------------------------------------------- + def to_pandas(self): + """Execute and return this column as a pandas (or GeoPandas) Series.""" + return self._df.select(self._expr.alias(self._name)).to_pandas()[self._name] + + def __repr__(self): + # Cheap: show the underlying expression rather than executing. + return f"<{type(self).__name__} {self._expr!r} (lazy; call .to_pandas())>" + + +class GeoSeries(Series): + """A geometry column, in the shape of a `geopandas.GeoSeries`. + + **EXPERIMENTAL.** Element-wise geometry operations (`buffer`, `centroid`, …) + return a new `GeoSeries`; measures (`area`, `length`) return a numeric + `Series`. Each delegates to the corresponding `ST_*` function via SedonaDB's + `.geo` accessor. + """ + + def buffer(self, distance): + """Buffer each geometry by `distance` (`ST_Buffer`).""" + return GeoSeries(self._df, self._expr.geo.buffer(distance), self._name) + + @property + def centroid(self): + """The centroid of each geometry (`ST_Centroid`).""" + return GeoSeries(self._df, self._expr.geo.centroid(), self._name) + + @property + def area(self): + """The area of each geometry (`ST_Area`) as a numeric `Series`.""" + return Series(self._df, self._expr.geo.area(), "area") + + @property + def length(self): + """The length/perimeter of each geometry (`ST_Length`).""" + return Series(self._df, self._expr.geo.length(), "length") + + def to_geopandas(self): + """Execute and return this column as a `geopandas.GeoSeries`.""" + return self.to_pandas() diff --git a/python/sedonadb-geopandas/tests/__init__.py b/python/sedonadb-geopandas/tests/__init__.py new file mode 100644 index 0000000000..13a83393a9 --- /dev/null +++ b/python/sedonadb-geopandas/tests/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/python/sedonadb-geopandas/tests/test_geopandas_compat.py b/python/sedonadb-geopandas/tests/test_geopandas_compat.py new file mode 100644 index 0000000000..4b2f6a4c96 --- /dev/null +++ b/python/sedonadb-geopandas/tests/test_geopandas_compat.py @@ -0,0 +1,244 @@ +# 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. + +import geopandas as gpd +import pytest + +import sedonadb_geopandas as sgpd +from sedonadb_geopandas import GeoDataFrame, GeoSeries, Series + + +@pytest.fixture +def cities(): + return gpd.GeoDataFrame( + {"name": ["A", "B", "C"], "pop": [100, 200, 300]}, + geometry=gpd.GeoSeries.from_wkt(["POINT (0 0)", "POINT (1 1)", "POINT (5 5)"]), + crs="EPSG:4326", + ) + + +@pytest.fixture +def con_free_geom_frame(): + """A frame whose active geometry column is *not* the heuristic's pick. + + Two geometry columns, `geom` and `geometry`, with `geom` marked active — so + anything that re-derives the geometry column instead of persisting it would + wrongly land on `geometry`. + """ + df = sgpd.default_context().sql( + "SELECT ST_SetSRID(ST_Point(0.0, 0.0), 3857) AS geom, " + "ST_SetSRID(ST_Point(9.0, 9.0), 3857) AS geometry, 1 AS a" + ) + return GeoDataFrame(df, geometry="geom") + + +def assert_geopandas_expr_equal(gdf, op, *, sort_by): + """Assert an operation gives the same result on GeoPandas and on the wrapper. + + Applies `op` to `gdf` (GeoPandas) and to `sgpd.from_geopandas(gdf)`, then + compares the materialized results. Rows are sorted by `sort_by` and the + index reset, since the relational engine preserves neither row order nor a + row index. `check_crs` is left on, so CRS propagation is asserted too. + Useful for throwing a corpus of GeoPandas ops at the wrapper. + """ + from geopandas.testing import assert_geodataframe_equal + + expected = op(gdf).sort_values(sort_by).reset_index(drop=True) + got = ( + op(sgpd.from_geopandas(gdf)) + .to_geopandas() + .sort_values(sort_by) + .reset_index(drop=True) + ) + assert_geodataframe_equal(got, expected, check_like=True, check_crs=True) + + +def test_from_geopandas_returns_geodataframe(cities): + gdf = sgpd.from_geopandas(cities) + assert isinstance(gdf, GeoDataFrame) + assert gdf.columns == ["name", "pop", "geometry"] + assert len(gdf) == 3 + + +def test_roundtrip_preserves_data(cities): + out = sgpd.from_geopandas(cities).to_geopandas().sort_values("name") + assert list(out["name"]) == ["A", "B", "C"] + assert list(out["pop"]) == [100, 200, 300] + assert out.geometry.to_wkt().tolist() == cities.geometry.to_wkt().tolist() + + +def test_filter_boolean_mask(cities): + gdf = sgpd.from_geopandas(cities) + out = gdf[gdf["pop"] > 150].to_geopandas().sort_values("name") + # Same result as GeoPandas boolean-mask indexing. + expected = cities[cities["pop"] > 150].sort_values("name") + assert list(out["name"]) == list(expected["name"]) + + +def test_filter_boolean_composition(cities): + gdf = sgpd.from_geopandas(cities) + out = gdf[(gdf["pop"] > 150) & (gdf["pop"] < 300)].to_geopandas() + assert list(out["name"]) == ["B"] + + +def test_getitem_column_types(cities): + gdf = sgpd.from_geopandas(cities) + assert isinstance(gdf["geometry"], GeoSeries) + assert isinstance(gdf["pop"], Series) + # A non-geometry Series materializes to a plain pandas Series. + assert sorted(gdf["pop"].to_pandas().tolist()) == [100, 200, 300] + + +def test_geoseries_centroid_of_points_is_identity(cities): + gdf = sgpd.from_geopandas(cities) + got = gdf.geometry.centroid.to_geopandas().to_wkt().tolist() + assert got == cities.geometry.to_wkt().tolist() + + +def test_geoseries_buffer_area(cities): + gdf = sgpd.from_geopandas(cities) + areas = gdf.geometry.buffer(0.5).area.to_pandas().tolist() + # A radius-0.5 buffer has area ~= pi/4; segmentation differs slightly from + # GeoPandas, so compare approximately. + assert areas == pytest.approx([0.785, 0.785, 0.785], abs=0.01) + + +def test_to_crs(cities): + gdf = sgpd.from_geopandas(cities) + web = gdf.to_crs("EPSG:3857") + assert isinstance(web, GeoDataFrame) + # `.crs` is read cheaply from the schema (SedonaDB's CRS representation). + assert "3857" in str(web.crs) + + +def test_column_subset_keeps_geometry(cities): + gdf = sgpd.from_geopandas(cities) + sub = gdf[["name", "geometry"]] + assert isinstance(sub, GeoDataFrame) + assert sub.columns == ["name", "geometry"] + # Geometry is still usable after subsetting. + assert isinstance(sub.geometry, GeoSeries) + + +def test_filter_matches_geopandas(cities): + # Same op applied to GeoPandas and to the wrapper yields the same result. + assert_geopandas_expr_equal(cities, lambda df: df[df["pop"] > 150], sort_by="name") + + +def test_to_crs_matches_geopandas(cities): + # Exercises CRS propagation through an operation (asserted via check_crs). + assert_geopandas_expr_equal( + cities, lambda df: df.to_crs("EPSG:3857"), sort_by="name" + ) + + +@pytest.mark.parametrize("source_crs", ["EPSG:4326", "EPSG:32633"]) +def test_crs_propagates_from_projected_and_geographic(source_crs): + # CRS survives the round trip from either a geographic or projected source. + gdf = gpd.GeoDataFrame( + {"name": ["A", "B"]}, + geometry=gpd.GeoSeries.from_wkt(["POINT (0 0)", "POINT (1 1)"]), + crs=source_crs, + ) + assert_geopandas_expr_equal(gdf, lambda df: df, sort_by="name") + + +def test_operand_accepts_literal(cities): + # lit() is a supported escape hatch (and the way to carry a CRS). + from sedonadb.expr import lit + + gdf = sgpd.from_geopandas(cities) + out = gdf[gdf["pop"] > lit(150)].to_geopandas().sort_values("name") + assert list(out["name"]) == ["B", "C"] + + +def test_repr_is_lazy(cities): + # repr() must not execute the query (IDEs/consoles call it constantly). + gdf = sgpd.from_geopandas(cities) + assert ( + repr(gdf) + == "GeoDataFrame(columns=['name', 'pop', 'geometry'], geometry='geometry')" + ) + assert "GeoSeries" in repr(gdf.geometry) + assert "Series" in repr(gdf["pop"]) + + +def test_operand_rejects_arraylike(cities): + gdf = sgpd.from_geopandas(cities) + with pytest.raises(TypeError, match="array-like"): + gdf["pop"] > cities["pop"] # a real pandas Series + + +def test_invalid_geometry_column_raises(cities): + df = sgpd.default_context().create_data_frame(cities) + # A non-geometry column named as the geometry is rejected. + with pytest.raises(ValueError, match="not a geometry column"): + GeoDataFrame(df, geometry="pop") + with pytest.raises(KeyError, match="not found"): + GeoDataFrame(df, geometry="nope") + + +def test_no_geometry_frame(cities): + # Dropping the geometry column yields a frame with no active geometry. + # GeoPandas returns a plain DataFrame here, whose .geometry also raises. + gdf = sgpd.from_geopandas(cities) + plain = gdf[["name", "pop"]] + assert plain.crs is None + with pytest.raises(AttributeError, match="no active geometry"): + plain.geometry + assert not hasattr(cities[["name", "pop"]], "geometry") + + +def test_column_subset_persists_custom_geometry_name(con_free_geom_frame): + # The active geometry column is persisted through a subset, not re-derived + # (re-deriving would pick "geometry" over the active "geom"). + gdf = con_free_geom_frame + assert gdf["geom"] is not None + sub = gdf[["a", "geom"]] + assert sub._geometry_name == "geom" + assert sub.crs is not None + + +def test_to_geopandas_persists_custom_geometry_name(con_free_geom_frame): + # A custom active geometry column survives the trip back to GeoPandas, even + # when another column would win SedonaDB's primary-geometry heuristic. + out = con_free_geom_frame.to_geopandas() + assert out.geometry.name == "geom" + + +def test_cross_frame_series_raises(cities): + # Combining Series from two different frames has no row alignment, so it + # must error rather than silently produce wrong rows. + a = sgpd.from_geopandas(cities) + b = sgpd.from_geopandas(cities) + with pytest.raises(ValueError, match="different DataFrames"): + a["pop"] > b["pop"] + + +def test_slice_and_integer_keys(cities): + gdf = sgpd.from_geopandas(cities) + with pytest.raises(TypeError, match="Positional row slicing"): + gdf[0:2] + with pytest.raises(KeyError, match="column label"): + gdf[0] + + +def test_head(cities): + gdf = sgpd.from_geopandas(cities) + assert len(gdf.head(2)) == 2 + # head() keeps the active geometry column. + assert isinstance(gdf.head(2).geometry, GeoSeries)