Skip to content

Migrate GeoDistanceFeatures to narwhals, add polars support #993

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
solegalli merged 4 commits into from
Aug 24, 2026
Merged

Migrate GeoDistanceFeatures to narwhals, add polars support #993

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1541fc3
Migrate GeoDistanceFeatures to narwhals, add polars support
solegalli Aug 24, 2026
0c087fa
Apply suggestion from @solegalli
solegalli Aug 24, 2026
3fd7459
Fix stale example output throughout GeoDistanceFeatures user guide
solegalli Aug 24, 2026
b07187e
Merge remote-tracking branch 'origin/narwhals-geo-features'
solegalli Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 67 additions & 18 deletions docs/user_guide/creation/GeoDistanceFeatures.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Expand All @@ -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:

Expand All @@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Expand All @@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Expand Down Expand Up @@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Expand Down Expand Up @@ -232,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
--------------------
Expand Down
160 changes: 111 additions & 49 deletions feature_engine/creation/geo_features.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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__(
Expand Down Expand Up @@ -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
Expand All @@ -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)

# Coordinate variables
variables: List[Union[str, int]] = [
Expand All @@ -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."
Expand All @@ -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.
"""

Expand All @@ -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

Expand Down
Loading