🐛 Bug: CountEmbedder raises ColumnNotFoundError even when feature column exists

[regislon]

Describe the bug

Using CountEmbedder.transform(...) raises a ColumnNotFoundError on the expected output feature (e.g. "accident_occurred") despite it being present in the features_gdf and correctly indexed by feature_id.

Code :

joint_gdf = IntersectionJoiner().transform(gdf_h3, gdf_accidents)

count_embeddings = CountEmbedder(
    count_subcategories=False,
    expected_output_features=["accident_occurred"]
).transform(gdf_h3, gdf_accidents, joint_gdf)

Error message:

---> FAILED HERE RESOLVING 'sink' <---
DF ["accident_occurred"]; PROJECT */1 COLUMNS

It seems there may be an issue in srai/embedders/count_embedder.py:163, or perhaps I didn't define the columns or index as expected. Any input would be appreciated.

Indexes and columns names

print(gdf_accidents.index.name)       
print(gdf_accidents.columns)
print(gdf_h3.index.name)
print(gdf_h3.columns)

returns

feature_id
Index(['geometry', 'accident_occurred'], dtype='object')
region_id
Index(['geometry', 'type_of_surface', 'roundabout', 'traffic_significance',
       'object_type', 'engineering_structure'],
      dtype='object')

Environment:

• Python: 3.12
• srai version: 0.9.7
• Polars: 1.31.0
• pyarrow : 15.0.0
• OS: macOS

Full error desciption :

---------------------------------------------------------------------------
ColumnNotFoundError                       Traceback (most recent call last)
Cell In[11], [line 19](vscode-notebook-cell:?execution_count=11&line=19)
     10 print(gdf_h3.columns)
     14 joint_gdf = IntersectionJoiner().transform(gdf_h3, gdf_accidents)
     16 count_embeddings = CountEmbedder(
     17     count_subcategories=False,
     18     expected_output_features=["accident_occurred"]
---> [19](vscode-notebook-cell:?execution_count=11&line=19) ).transform(gdf_h3, gdf_accidents, joint_gdf)

File ~/Documents/Projects/EPFL_apply_data_sciences/capstone-project-adsml-ibex-c5-s2-8710-4645/.venv/lib/python3.12/site-packages/srai/embedders/count_embedder.py:163, in CountEmbedder.transform(self, regions_gdf, features_gdf, joint_gdf)
    147 region_embeddings = joint_with_encodings.drop(FEATURES_INDEX).group_by(REGIONS_INDEX).sum()
    148 region_embeddings, feature_columns = self._maybe_filter_to_expected_features(
    149     region_embeddings, feature_columns
    150 )
    152 region_embeddings_df = (
    153     (
    154         regions_df.join(region_embeddings, on=REGIONS_INDEX, how="left")
    155         .fill_null(0)
    156         .with_columns(
    157             [
    158                 pl.col(REGIONS_INDEX),
    159                 *(pl.col(col).cast(pl.Int32) for col in feature_columns),
    160             ]
    161         )
...

Resolved plan until failure:

	---> FAILED HERE RESOLVING 'sink' <---
DF ["accident_occurred"]; PROJECT */1 COLUMNS`

[RaczeQ]

Hello @regislon, sorry for late response.

I have tried to recreate the code, but everything worked on my end: https://colab.research.google.com/drive/1vtTcTVPsCZqFo5a7jEtDEF0rsgf2hrZc?usp=sharing

import geopandas as gpd

from srai.joiners import IntersectionJoiner
from srai.embedders import CountEmbedder
from shapely import box, Point

# Example region gdf
regions_gdf = gpd.GeoDataFrame(
    data={
        "region_id": ["a", "b", "c", "d"],
        "type_of_surface": "asphalt",
        "roundabout": [True, True, False, False],
        "traffic_significance": 0.75,
        "object_type": None,
        "engineering_structure": "tower",
    },
    geometry=[
        box(0, 0, 1, 1),
        box(0, 1, 1, 2),
        box(1, 0, 2, 1),
        box(1, 1, 2, 2),
    ],
    crs=4326,
).set_index("region_id")

# Example feature gdf
features_gdf = gpd.GeoDataFrame(
    data={
        "feature_id": ["0", "1", "2", "3", "5", "6", "7", "8"],
        "accident_occurred": [False, False, True, True, False, True, False, True],
    },
    geometry=[
        Point(0.5, 0.5),
        Point(0.5, 1.5),
        Point(1.5, 0.5),
        Point(1.5, 1.5),
        Point(0.5, 1),
        Point(1, 1.5),
        Point(1.5, 1),
        Point(1.25, 1.5),
    ],
    crs=4326,
).set_index("feature_id")

joint_gdf = IntersectionJoiner().transform(regions_gdf, features_gdf)

count_embeddings = CountEmbedder(
    count_subcategories=False,
    expected_output_features=["accident_occurred"]
).transform(regions_gdf, features_gdf, joint_gdf)

count_embeddings

Could you share the data (snippet / anonymised) for regions_gdf and features_gdf for me to test it with your data?