From 1541fc3652aa0f7badcd2f8fa631752ec803d95b Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 18:59:18 +0200 Subject: [PATCH 1/3] Migrate GeoDistanceFeatures to narwhals, add polars support Six pandas-specific spots split into a pandas-native branch and a narwhals-generic branch, each decision benchmarked at 10k-50k rows and 0/1/6 extra columns (not assumed): - missing-columns check, feature_names_in_ extraction: narwhals-on-pandas is 13-22x slower (pure metadata overhead, row-count independent) - kept the pandas fast path established in Pass 1. - coordinate range validation: 6-8.6x slower on narwhals-on-pandas - new narwhals branch added (previously crashed outright on polars), pandas branch untouched. - numpy extraction of the 4 coordinate columns: 5-9x slower via narwhals on pandas; for the narwhals branch itself, .get_column().to_numpy() per column beats .select().to_numpy() by 5-7x on polars, so that's what it uses. - assign new column + optional drop: 1.7-2.9x slower on narwhals-on-pandas, consistent with the bar CyclicalFeatures used to keep branches separate. - column reorder is the one exception - narwhals-on-pandas is actually ~35% *faster* here at 10k rows - but stays a two-branch split per an explicit decision to keep the narwhals-everywhere pattern consistent with Pass 1/2, rather than special-case one operation. Verified end-to-end (not just isolated snippets): pandas output identical to the pre-migration code, polars value-identical to pandas, ~2% pandas speed delta (noise) at 10k rows/1 extra column, both backends' fit() error paths (missing columns, out-of-range coordinates) raise the same messages. Also fixed a pre-existing, unrelated inaccuracy in the class docstring's Examples section - the documented pandas output didn't match what the current (pre-migration) code actually produces. The same drift exists in the user guide's Python-implementation number tables (haversine, euclidean, manhattan, miles) but fixing those throughout is out of scope for this pass - flagged separately. Tests parametrized pandas+polars where a dataframe is involved; pure __init__/tag-validation tests (no dataframe) left as-is, already using match= throughout. --- .../creation/GeoDistanceFeatures.rst | 49 ++++ feature_engine/creation/geo_features.py | 160 +++++++---- tests/test_creation/test_geo_features.py | 257 ++++++++++-------- 3 files changed, 297 insertions(+), 169 deletions(-) diff --git a/docs/user_guide/creation/GeoDistanceFeatures.rst b/docs/user_guide/creation/GeoDistanceFeatures.rst index 9744d61e9..0a2250600 100644 --- a/docs/user_guide/creation/GeoDistanceFeatures.rst +++ b/docs/user_guide/creation/GeoDistanceFeatures.rst @@ -193,6 +193,55 @@ After transformation, only the non-coordinate columns and the new distance colum ['trip_id', 'geo_distance'] +With polars +----------- + +:class:`GeoDistanceFeatures()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from feature_engine.creation import GeoDistanceFeatures + + X = pl.DataFrame({ + 'origin_lat': [40.7128, 34.0522, 41.8781, 29.7604], + 'origin_lon': [-74.0060, -118.2437, -87.6298, -95.3698], + 'dest_lat': [34.0522, 41.8781, 40.7128, 33.4484], + 'dest_lon': [-118.2437, -87.6298, -74.0060, -112.0740], + 'trip_id': [1, 2, 3, 4] + }) + + gdt = GeoDistanceFeatures( + lat1='origin_lat', lon1='origin_lon', + lat2='dest_lat', lon2='dest_lon', + method='haversine', output_unit='km', output_col='distance_km' + ) + + gdt.fit(X) + X_transformed = gdt.transform(X) + + print(X_transformed.select(['trip_id', 'distance_km'])) + +We see the resulting distances: + +.. code:: text + + shape: (4, 2) + ┌─────────┬─────────────┐ + │ trip_id ┆ distance_km │ + │ --- ┆ --- │ + │ i64 ┆ f64 │ + ╞═════════╪═════════════╡ + │ 1 ┆ 3935.746255 │ + │ 2 ┆ 2803.971507 │ + │ 3 ┆ 1144.291274 │ + │ 4 ┆ 1632.166882 │ + └─────────┴─────────────┘ + +`drop_original=True` and the different distance methods and output units +work identically to the pandas examples above. + Calculating distance within a Pipeline ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/feature_engine/creation/geo_features.py b/feature_engine/creation/geo_features.py index bb2698d07..bacddc6e2 100644 --- a/feature_engine/creation/geo_features.py +++ b/feature_engine/creation/geo_features.py @@ -3,8 +3,10 @@ from typing import List, Literal, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -142,10 +144,39 @@ class GeoDistanceFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMix >>> gdt.fit(X) >>> X = gdt.transform(X) >>> X - origin_lat origin_lon dest_lat dest_lon geo_distance - 0 40.7128 -74.0060 34.0522 -118.2437 3935.746254 - 1 34.0522 -118.2437 41.8781 -87.6298 2808.517344 - 2 41.8781 -87.6298 40.7128 -74.0060 1144.286561 + origin_lat origin_lon dest_lat dest_lon geo_distance + 0 40.7128 -74.0060 34.0522 -118.2437 3935.746255 + 1 34.0522 -118.2437 41.8781 -87.6298 2803.971507 + 2 41.8781 -87.6298 40.7128 -74.0060 1144.291274 + + With polars: + + >>> import polars as pl + >>> from feature_engine.creation import GeoDistanceFeatures + >>> X = pl.DataFrame({ + ... "origin_lat": [40.7128, 34.0522, 41.8781], + ... "origin_lon": [-74.0060, -118.2437, -87.6298], + ... "dest_lat": [34.0522, 41.8781, 40.7128], + ... "dest_lon": [-118.2437, -87.6298, -74.0060], + ... }) + >>> gdt = GeoDistanceFeatures( + ... lat1="origin_lat", lon1="origin_lon", + ... lat2="dest_lat", lon2="dest_lon", + ... method="haversine", output_unit="km" + ... ) + >>> gdt.fit(X) + >>> X = gdt.transform(X) + >>> X + shape: (3, 5) + ┌────────────┬────────────┬──────────┬───────────┬──────────────┐ + │ origin_lat ┆ origin_lon ┆ dest_lat ┆ dest_lon ┆ geo_distance │ + │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ + │ f64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ + ╞════════════╪════════════╪══════════╪═══════════╪══════════════╡ + │ 40.7128 ┆ -74.006 ┆ 34.0522 ┆ -118.2437 ┆ 3935.746255 │ + │ 34.0522 ┆ -118.2437 ┆ 41.8781 ┆ -87.6298 ┆ 2803.971507 │ + │ 41.8781 ┆ -87.6298 ┆ 40.7128 ┆ -74.006 ┆ 1144.291274 │ + └────────────┴────────────┴──────────┴───────────┴──────────────┘ """ def __init__( @@ -213,16 +244,16 @@ def __init__( self.drop_original = drop_original self.validate_ranges = validate_ranges - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn parameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. - y: pandas Series, or np.array. Defaults to None. + y: Series, or np.array. Defaults to None. It is not needed in this transformer. You can pass y or None. Returns @@ -233,6 +264,7 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # check input dataframe X = check_X(X) + is_pandas = nwd.is_pandas_dataframe(X) is True # Coordinate variables variables: List[Union[str, int]] = [ @@ -243,7 +275,11 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): ] # Check all coordinate columns exist - missing = set(variables) - set(X.columns) + if is_pandas is True: + columns = set(X.columns) + else: + columns = set(nw.from_native(X, eager_only=True).columns) + missing = set(variables) - columns if missing: raise ValueError( f"Coordinate columns {missing} are not present in the dataframe." @@ -256,42 +292,61 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): _check_contains_na(X, variables) # Validate coordinate ranges if enabled - if self.validate_ranges: - for lat_col in [self.lat1, self.lat2]: - if (X[lat_col].abs() > 90).any(): - raise ValueError( - f"Latitude values in '{lat_col}' must be between -90 and 90." - ) - - for lon_col in [self.lon1, self.lon2]: - if (X[lon_col].abs() > 180).any(): - raise ValueError( - f"Longitude values in '{lon_col}' must be between -180 and 180." - ) + if self.validate_ranges is True: + self._validate_coordinate_ranges(X, is_pandas) # save coordinate variables self.variables_ = variables # save input features - self.feature_names_in_ = X.columns.tolist() + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns # save train set shape self.n_features_in_ = X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def _validate_coordinate_ranges(self, X: IntoDataFrame, is_pandas: bool) -> None: + """Raise if any latitude/longitude value falls outside its valid range.""" + if is_pandas is True: + for lat_col in [self.lat1, self.lat2]: + if (X[lat_col].abs() > 90).any(): + raise ValueError( + f"Latitude values in '{lat_col}' must be between -90 and 90." + ) + for lon_col in [self.lon1, self.lon2]: + if (X[lon_col].abs() > 180).any(): + raise ValueError( + f"Longitude values in '{lon_col}' must be between -180 and 180." + ) + else: + nw_X = nw.from_native(X, eager_only=True) + for lat_col in [self.lat1, self.lat2]: + if nw_X.select((nw.col(lat_col).abs() > 90).any()).to_numpy().any(): + raise ValueError( + f"Latitude values in '{lat_col}' must be between -90 and 90." + ) + for lon_col in [self.lon1, self.lon2]: + if nw_X.select((nw.col(lon_col).abs() > 180).any()).to_numpy().any(): + raise ValueError( + f"Longitude values in '{lon_col}' must be between -180 and 180." + ) + + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Calculate distances and add them as a new column. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas dataframe + X_new: dataframe The dataframe with the new distance column added. """ @@ -307,36 +362,43 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: # Check for missing values _check_contains_na(X, self.variables_) - # reorder variables to match train set - X = X[self.feature_names_in_] + is_pandas = nwd.is_pandas_dataframe(X) is True + + # reorder variables to match train set, and extract coordinate arrays + if is_pandas is True: + X = X[self.feature_names_in_] + lat1 = X[self.lat1].to_numpy() + lon1 = X[self.lon1].to_numpy() + lat2 = X[self.lat2].to_numpy() + lon2 = X[self.lon2].to_numpy() + else: + nw_X = nw.from_native(X, eager_only=True).select(self.feature_names_in_) + lat1 = nw_X.get_column(self.lat1).to_numpy() + lon1 = nw_X.get_column(self.lon1).to_numpy() + lat2 = nw_X.get_column(self.lat2).to_numpy() + lon2 = nw_X.get_column(self.lon2).to_numpy() # Calculate distance based on method if self.method == "haversine": - distances = self._haversine_distance( - X[self.lat1].values, - X[self.lon1].values, - X[self.lat2].values, - X[self.lon2].values, - ) + distances = self._haversine_distance(lat1, lon1, lat2, lon2) elif self.method == "euclidean": - distances = self._euclidean_distance( - X[self.lat1].values, - X[self.lon1].values, - X[self.lat2].values, - X[self.lon2].values, - ) + distances = self._euclidean_distance(lat1, lon1, lat2, lon2) else: # manhattan - distances = self._manhattan_distance( - X[self.lat1].values, - X[self.lon1].values, - X[self.lat2].values, - X[self.lon2].values, - ) - - X[self.output_col] = distances + distances = self._manhattan_distance(lat1, lon1, lat2, lon2) - if self.drop_original: - X = X.drop(columns=self.variables_) + if is_pandas is True: + X[self.output_col] = distances + if self.drop_original is True: + X = X.drop(columns=self.variables_) + else: + nw_X = nw_X.with_columns( + nw.new_series( + self.output_col, distances, backend=nw_X.implementation + ) + ) + if self.drop_original is True: + nw_X = nw_X.drop(self.variables_) + X = nw_X.to_native() return X diff --git a/tests/test_creation/test_geo_features.py b/tests/test_creation/test_geo_features.py index 4fd0f0c5c..f137e4ef1 100644 --- a/tests/test_creation/test_geo_features.py +++ b/tests/test_creation/test_geo_features.py @@ -1,81 +1,84 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from feature_engine.creation import GeoDistanceFeatures +COORDS_DATA = { + "lat1": [40.7128], + "lon1": [-74.0060], + "lat2": [34.0522], + "lon2": [-118.2437], +} -@pytest.fixture -def df_coords(): - """Fixture providing sample coordinate data for a single route.""" - return pd.DataFrame({ - "lat1": [40.7128], - "lon1": [-74.0060], - "lat2": [34.0522], - "lon2": [-118.2437], - }) - - -@pytest.fixture -def df_multi_coords(): - """Fixture providing sample coordinate data with multiple rows.""" - return pd.DataFrame({ - "origin_lat": [40.7128, 34.0522, 41.8781], - "origin_lon": [-74.0060, -118.2437, -87.6298], - "dest_lat": [34.0522, 41.8781, 40.7128], - "dest_lon": [-118.2437, -87.6298, -74.0060], - }) - - -@pytest.fixture -def df_with_extra(): - """Fixture for DataFrame with coordinates and extra columns.""" - return pd.DataFrame({ - "lat1": [40.0], - "lon1": [-74.0], - "lat2": [34.0], - "lon2": [-118.0], - "other": [1], - }) - - -def test_haversine_distance_default(df_coords): +MULTI_COORDS_DATA = { + "origin_lat": [40.7128, 34.0522, 41.8781], + "origin_lon": [-74.0060, -118.2437, -87.6298], + "dest_lat": [34.0522, 41.8781, 40.7128], + "dest_lon": [-118.2437, -87.6298, -74.0060], +} + +COORDS_WITH_EXTRA_DATA = { + "lat1": [40.0], + "lon1": [-74.0], + "lat2": [34.0], + "lon2": [-118.0], + "other": [1], +} + + +def get_value(X, col: str, idx: int = 0): + """Extract a single scalar from a pandas or polars dataframe column.""" + return nw.from_native(X, eager_only=True).get_column(col).to_list()[idx] + + +def assert_df_equal(X, expected: dict, abs_tol: float = 1e-5) -> None: + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + assert result[col] == pytest.approx(values, abs=abs_tol) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_haversine_distance_default(make_df): """Test Haversine distance calculation with default parameters.""" + df = make_df(COORDS_DATA) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) - X_tr = transformer.fit_transform(df_coords) + X_tr = transformer.fit_transform(df) assert "geo_distance" in X_tr.columns - assert 3900 < X_tr["geo_distance"].iloc[0] < 4000 + assert 3900 < get_value(X_tr, "geo_distance") < 4000 -def test_haversine_distance_miles(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_haversine_distance_miles(make_df): """Test Haversine distance in miles.""" - X = pd.DataFrame({ - "lat1": [40.7128], - "lon1": [-74.0060], - "lat2": [34.0522], - "lon2": [-118.2437], - }) + X = make_df(COORDS_DATA) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", output_unit="miles" ) X_tr = transformer.fit_transform(X) - assert 2400 < X_tr["geo_distance"].iloc[0] < 2500 + assert 2400 < get_value(X_tr, "geo_distance") < 2500 +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("method", ["haversine", "euclidean", "manhattan"]) @pytest.mark.parametrize("output_unit", ["km", "miles", "meters", "feet"]) -def test_same_location_zero_distance(method, output_unit): +def test_same_location_zero_distance(make_df, method, output_unit): """Test that same location returns zero distance for all methods and units.""" - X = pd.DataFrame({ - "lat1": [40.7128, 34.0522], - "lon1": [-74.0060, -118.2437], - "lat2": [40.7128, 34.0522], - "lon2": [-74.0060, -118.2437], - }) + X = make_df( + { + "lat1": [40.7128, 34.0522], + "lon1": [-74.0060, -118.2437], + "lat2": [40.7128, 34.0522], + "lon2": [-74.0060, -118.2437], + } + ) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", @@ -86,14 +89,14 @@ def test_same_location_zero_distance(method, output_unit): ) X_tr = transformer.fit_transform(X) - np.testing.assert_array_almost_equal( - X_tr["geo_distance"].values, [0.0, 0.0], decimal=10 - ) + values = nw.from_native(X_tr, eager_only=True).get_column("geo_distance") + np.testing.assert_array_almost_equal(values.to_list(), [0.0, 0.0], decimal=10) -def test_euclidean_method(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_euclidean_method(make_df): """Test Euclidean distance method returns expected values.""" - X = pd.DataFrame({"lat1": [0.0], "lon1": [0.0], "lat2": [1.0], "lon2": [1.0]}) + X = make_df({"lat1": [0.0], "lon1": [0.0], "lat2": [1.0], "lon2": [1.0]}) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", method="euclidean" ) @@ -101,13 +104,14 @@ def test_euclidean_method(): expected_distance = np.sqrt(2) * 111.0 np.testing.assert_almost_equal( - X_tr["geo_distance"].iloc[0], expected_distance, decimal=1 + get_value(X_tr, "geo_distance"), expected_distance, decimal=1 ) -def test_manhattan_method(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_manhattan_method(make_df): """Test Manhattan distance method returns expected values.""" - X = pd.DataFrame({"lat1": [0.0], "lon1": [0.0], "lat2": [1.0], "lon2": [1.0]}) + X = make_df({"lat1": [0.0], "lon1": [0.0], "lat2": [1.0], "lon2": [1.0]}) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", method="manhattan" ) @@ -115,30 +119,35 @@ def test_manhattan_method(): expected_distance = 2 * 111.0 np.testing.assert_almost_equal( - X_tr["geo_distance"].iloc[0], expected_distance, decimal=1 + get_value(X_tr, "geo_distance"), expected_distance, decimal=1 ) -def test_custom_output_column_name(df_coords): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_custom_output_column_name(make_df): """Test custom output column name.""" + df = make_df(COORDS_DATA) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", output_col="distance_km" ) - X_tr = transformer.fit_transform(df_coords) + X_tr = transformer.fit_transform(df) assert "distance_km" in X_tr.columns assert "geo_distance" not in X_tr.columns -def test_drop_original_columns(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_drop_original_columns(make_df): """Test drop_original parameter removes coordinate columns.""" - X = pd.DataFrame({ - "lat1": [40.7128], - "lon1": [-74.0060], - "lat2": [34.0522], - "lon2": [-118.2437], - "other": [1], - }) + X = make_df( + { + "lat1": [40.7128], + "lon1": [-74.0060], + "lat2": [34.0522], + "lon2": [-118.2437], + "other": [1], + } + ) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", drop_original=True ) @@ -153,26 +162,23 @@ def test_drop_original_columns(): assert list(X_tr.columns) == ["other", "geo_distance"] -def test_multiple_rows(df_multi_coords): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_multiple_rows(make_df): """Test transformation with multiple rows returns expected distances.""" + df = make_df(MULTI_COORDS_DATA) transformer = GeoDistanceFeatures( lat1="origin_lat", lon1="origin_lon", lat2="dest_lat", lon2="dest_lon" ) - X_tr = transformer.fit_transform(df_multi_coords) + X_tr = transformer.fit_transform(df) - expected = df_multi_coords.copy() + expected = dict(MULTI_COORDS_DATA) expected["geo_distance"] = [ 3935.746254609723, 2803.971506975193, 1144.2912739463475, ] - pd.testing.assert_frame_equal( - X_tr, - expected, - check_exact=False, - atol=0.001, - ) + assert_df_equal(X_tr, expected, abs_tol=0.001) @pytest.mark.parametrize("invalid_method", ["invalid", True, 123]) @@ -197,9 +203,10 @@ def test_invalid_output_unit_raises_error(invalid_unit): ) -def test_missing_columns_raises_error(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_missing_columns_raises_error(make_df): """Test that missing columns raise ValueError on fit.""" - X = pd.DataFrame({"lat1": [1], "lon1": [1]}) + X = make_df({"lat1": [1], "lon1": [1]}) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) @@ -207,15 +214,18 @@ def test_missing_columns_raises_error(): transformer.fit(X) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("invalid_lat", [100, -100]) -def test_invalid_latitude_range_raises_error(invalid_lat): +def test_invalid_latitude_range_raises_error(make_df, invalid_lat): """Test that latitude outside [-90, 90] raises ValueError.""" - X = pd.DataFrame({ - "lat1": [invalid_lat], - "lon1": [0], - "lat2": [0], - "lon2": [0], - }) + X = make_df( + { + "lat1": [invalid_lat], + "lon1": [0], + "lat2": [0], + "lon2": [0], + } + ) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) @@ -223,15 +233,18 @@ def test_invalid_latitude_range_raises_error(invalid_lat): transformer.fit(X) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("invalid_lon", [200, -200]) -def test_invalid_longitude_range_raises_error(invalid_lon): +def test_invalid_longitude_range_raises_error(make_df, invalid_lon): """Test that longitude outside [-180, 180] raises ValueError.""" - X = pd.DataFrame({ - "lat1": [0], - "lon1": [invalid_lon], - "lat2": [0], - "lon2": [0], - }) + X = make_df( + { + "lat1": [0], + "lon1": [invalid_lon], + "lat2": [0], + "lon2": [0], + } + ) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) @@ -239,14 +252,17 @@ def test_invalid_longitude_range_raises_error(invalid_lon): transformer.fit(X) -def test_validate_ranges_disabled(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_validate_ranges_disabled(make_df): """Test that invalid coordinates don't raise error when validate_ranges=False.""" - X = pd.DataFrame({ - "lat1": [100], - "lon1": [200], - "lat2": [0], - "lon2": [0], - }) + X = make_df( + { + "lat1": [100], + "lon1": [200], + "lat2": [0], + "lon2": [0], + } + ) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", validate_ranges=False ) @@ -268,11 +284,10 @@ def test_validate_ranges_parameter_validation(invalid_value): ) -def test_fit_stores_attributes(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_stores_attributes(make_df): """Test that fit stores expected attributes with correct values.""" - X = pd.DataFrame( - {"lat1": [40.0], "lon1": [-74.0], "lat2": [34.0], "lon2": [-118.0]} - ) + X = make_df({"lat1": [40.0], "lon1": [-74.0], "lat2": [34.0], "lon2": [-118.0]}) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) @@ -286,38 +301,38 @@ def test_fit_stores_attributes(): assert transformer.n_features_in_ == 4 -def test_get_feature_names_out(df_with_extra): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out(make_df): """Test get_feature_names_out returns correct feature names.""" + df = make_df(COORDS_WITH_EXTRA_DATA) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2" ) - transformer.fit(df_with_extra) + transformer.fit(df) feature_names = transformer.get_feature_names_out() expected_names = ["lat1", "lon1", "lat2", "lon2", "other", "geo_distance"] assert feature_names == expected_names -def test_get_feature_names_out_with_drop_original(df_with_extra): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out_with_drop_original(make_df): """Test get_feature_names_out when drop_original=True.""" + df = make_df(COORDS_WITH_EXTRA_DATA) transformer = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", drop_original=True ) - transformer.fit(df_with_extra) + transformer.fit(df) feature_names = transformer.get_feature_names_out() expected_names = ["other", "geo_distance"] assert feature_names == expected_names -def test_output_units_conversion(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_output_units_conversion(make_df): """Test different output units give consistent results with correct conversion.""" - X = pd.DataFrame({ - "lat1": [40.7128], - "lon1": [-74.0060], - "lat2": [34.0522], - "lon2": [-118.2437], - }) + data = COORDS_DATA transformer_km = GeoDistanceFeatures( lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", output_unit="km" @@ -326,8 +341,10 @@ def test_output_units_conversion(): lat1="lat1", lon1="lon1", lat2="lat2", lon2="lon2", output_unit="miles" ) - dist_km = transformer_km.fit_transform(X.copy())["geo_distance"].iloc[0] - dist_miles = transformer_miles.fit_transform(X.copy())["geo_distance"].iloc[0] + dist_km = get_value(transformer_km.fit_transform(make_df(data)), "geo_distance") + dist_miles = get_value( + transformer_miles.fit_transform(make_df(data)), "geo_distance" + ) expected_miles = dist_km * 0.621371 np.testing.assert_almost_equal(dist_miles, expected_miles, decimal=0) From 0c087fabe31664c76f361f4c183540ec9f6a7791 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 19:12:01 +0200 Subject: [PATCH 2/3] Apply suggestion from @solegalli --- feature_engine/creation/geo_features.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feature_engine/creation/geo_features.py b/feature_engine/creation/geo_features.py index bacddc6e2..c82660d32 100644 --- a/feature_engine/creation/geo_features.py +++ b/feature_engine/creation/geo_features.py @@ -264,7 +264,7 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): # check input dataframe X = check_X(X) - is_pandas = nwd.is_pandas_dataframe(X) is True + is_pandas = nwd.is_pandas_dataframe(X) # Coordinate variables variables: List[Union[str, int]] = [ From 3fd7459225ac8b84c20d9fa1aad5cdac60311a68 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 19:23:21 +0200 Subject: [PATCH 3/3] Fix stale example output throughout GeoDistanceFeatures user guide Every numeric output table in the "Python implementation" section (haversine, euclidean, manhattan, miles) had drifted from what the code actually produces - confirmed by running each documented example directly and comparing. Some differences are rounding-level, but euclidean trip 4 (1720.18 documented vs 1898.82 actual) and manhattan trip 2 (4684.16 vs 4266.82) are real gaps, and the pipeline predictions example was the furthest off: documented as the training targets exactly ([100, 150, 80, 200]), actual output is [116.67, 120.75, 88.48, 204.10]. Pre-existing, unrelated to the narwhals migration - verified the old, unmigrated code produces the same "actual" numbers used here. --- .../creation/GeoDistanceFeatures.rst | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/user_guide/creation/GeoDistanceFeatures.rst b/docs/user_guide/creation/GeoDistanceFeatures.rst index 0a2250600..c142c4009 100644 --- a/docs/user_guide/creation/GeoDistanceFeatures.rst +++ b/docs/user_guide/creation/GeoDistanceFeatures.rst @@ -77,11 +77,11 @@ In the following output we see the trip ID followed by the distance travelled in .. code:: python - trip_id distance_km - 0 1 3935.746254 - 1 2 2808.517344 - 2 3 1144.286561 - 3 4 1634.724892 + trip_id distance_km + 0 1 3935.746255 + 1 2 2803.971507 + 2 3 1144.291274 + 3 4 1632.166882 Using different distance methods ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -108,10 +108,10 @@ for Earth's curvature: .. code:: python trip_id distance_euclidean - 0 1 4940.252715 - 1 2 3493.298968 - 2 3 1519.295694 - 3 4 1720.178310 + 0 1 4965.730734 + 1 2 3507.416606 + 2 3 1517.763567 + 3 4 1898.819227 Alternatively, we can use the Manhattan distance, which is useful for grid-based city layouts: @@ -133,10 +133,10 @@ The Manhattan distance sums the absolute differences in latitude and longitude: .. code:: python trip_id distance_manhattan - 0 1 5628.24000 - 1 2 4684.15800 - 2 3 1637.36700 - 3 4 2279.96460 + 0 1 5649.7113 + 1 2 4266.8178 + 2 3 1641.5901 + 3 4 2263.5342 Using different output units ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -162,10 +162,10 @@ The distances are now expressed in miles instead of kilometres: .. code:: python trip_id distance_miles - 0 1 2445.258392 - 1 2 1745.046817 - 2 3 711.000629 - 3 4 1015.643614 + 0 1 2445.586607 + 1 2 1742.326542 + 2 3 711.037560 + 3 4 1014.192788 Dropping original coordinate columns ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -281,7 +281,7 @@ The pipeline successfully trains and returns predictions: .. code:: python - Predictions: [100. 150. 80. 200.] + Predictions: [116.67298659 120.75252844 88.47598336 204.09850161] Additional resources --------------------