paleolimbot commented on code in PR #1052: URL: https://github.com/apache/sedona-db/pull/1052#discussion_r3667463271
########## python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py: ########## @@ -0,0 +1,129 @@ +# 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 + + +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=None): + self._df = df + if geometry is None: + # 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 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; re-derive the geometry column since the + # subset may have dropped it. + if isinstance(key, list): + return GeoDataFrame(self._df.select(*key)) Review Comment: It's worth checking what GeoPandas does for the case where the subset drops the geometry and make sure you match it here. It's either that the geometry column is preserved (sticky even if not included in the column subset), that the active geometry column is persisted if it still exists (or a regular pandas DataFrame is returned if it's dropped). I would be surprised if the geometry column is rederived (but you can test and find out). ########## python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py: ########## @@ -0,0 +1,129 @@ +# 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 + + +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=None): + self._df = df + if geometry is None: + # 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 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; re-derive the geometry column since the + # subset may have dropped it. + if isinstance(key, list): + return GeoDataFrame(self._df.select(*key)) + + # 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) + + raise TypeError( + f"GeoDataFrame indices must be a column name, list of names, or " + f"boolean mask, not {type(key).__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).""" + return self._df.to_pandas() Review Comment: Persist the custom geometry column name here? ########## python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py: ########## @@ -0,0 +1,131 @@ +# 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(other): + """Coerce the right-hand side of an operator into something usable. + + A `Series` is unwrapped to its expression; 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 — + there is no row alignment, and it would otherwise fail obscurely as a + multi-element literal. + """ + from sedonadb.expr import Expr, Literal + + if isinstance(other, Series): + 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(other), self._name) + + def __ge__(self, other): + return Series(self._df, self._expr >= _operand(other), self._name) + + def __lt__(self, other): + return Series(self._df, self._expr < _operand(other), self._name) + + def __le__(self, other): + return Series(self._df, self._expr <= _operand(other), self._name) + + def __eq__(self, other): + return Series(self._df, self._expr == _operand(other), self._name) + + def __ne__(self, other): + return Series(self._df, self._expr != _operand(other), self._name) + + # -- boolean composition of masks -------------------------------------- + def __and__(self, other): + return Series(self._df, self._expr & _operand(other), self._name) + + def __or__(self, other): + return Series(self._df, self._expr | _operand(other), self._name) + + def __invert__(self): + return Series(self._df, ~self._expr, self._name) + + __hash__ = None + + # -- materialization --------------------------------------------------- + # -- materialization --------------------------------------------------- Review Comment: ```suggestion # -- materialization --------------------------------------------------- ``` ########## python/sedonadb-geopandas/python/sedonadb_geopandas/_frame.py: ########## @@ -0,0 +1,129 @@ +# 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 + + +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=None): + self._df = df + if geometry is None: + # 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 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; re-derive the geometry column since the + # subset may have dropped it. + if isinstance(key, list): + return GeoDataFrame(self._df.select(*key)) + + # 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) Review Comment: Do integers or slices work here? ########## python/sedonadb-geopandas/tests/test_geopandas_compat.py: ########## @@ -0,0 +1,97 @@ +# 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 + + [email protected] +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", + ) + + Review Comment: I don't see this used in this PR but I may have missed it (no need to to include it here if it's not used yet) ########## python/sedonadb-geopandas/python/sedonadb_geopandas/_series.py: ########## @@ -0,0 +1,131 @@ +# 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(other): + """Coerce the right-hand side of an operator into something usable. + + A `Series` is unwrapped to its expression; 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 — + there is no row alignment, and it would otherwise fail obscurely as a + multi-element literal. + """ + from sedonadb.expr import Expr, Literal + + if isinstance(other, Series): + return other._expr Review Comment: Sorry, I keep thinking of other things that could happen here. Do we want to allow Series that derive from different DataFrames to be mixed with each other? Is it reasonable to error for that case or will that be annoying? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
