feat: opt-in split_locations for events() — flat x/y(/z) coordinate c…

statsbombpy/helpers.py

@@ -2,6 +2,14 @@
 
 import pandas as pd
 
+LOCATION_COLUMNS = [
+    "location",
+    "pass_end_location",
+    "carry_end_location",
+    "shot_end_location",
+    "goalkeeper_end_location",
+]
+
 PLURALS = {
     "Starting XI": "starting_xis",
     "Half Start": "half_starts",
@@ -117,3 +125,42 @@ def merge_events_and_frames(
             if key in drop_keys:
                 del event[key]
     return events
+
+
+def split_location_cols(df: pd.DataFrame, location_columns=None) -> pd.DataFrame:
+    """Split list-like coordinate columns into flat float columns.
+
+    StatsBomb locations come through as Python lists (e.g. [x, y] or
+    [x, y, z] for shots), which are awkward to work with downstream:
+    they are unhashable (drop_duplicates fails), they block vectorised
+    operations, and they do not round-trip cleanly through Parquet
+    (lists come back as numpy arrays). This helper adds ``<col>_x``,
+    ``<col>_y`` (and ``<col>_z`` where a third coordinate exists)
+    columns next to each list-like location column.
+
+    Missing or malformed values yield NaN. Original columns are kept.
+    """
+    if location_columns is None:
+        location_columns = LOCATION_COLUMNS
+
+    def _coord(value, i):
+        if isinstance(value, (list, tuple)) and len(value) > i:
+            return value[i]
+        # numpy arrays appear when a dataframe has been cached to Parquet
+        if hasattr(value, "__len__") and not isinstance(value, (str, dict)):
+            try:
+                return value[i] if len(value) > i else float("nan")
+            except (TypeError, IndexError):
+                return float("nan")
+        return float("nan")
+
+    for col in location_columns:
+        if col not in df.columns:
+            continue
+        values = df[col]
+        df[f"{col}_x"] = pd.to_numeric(values.map(lambda v: _coord(v, 0)), errors="coerce")
+        df[f"{col}_y"] = pd.to_numeric(values.map(lambda v: _coord(v, 1)), errors="coerce")
+        z = pd.to_numeric(values.map(lambda v: _coord(v, 2)), errors="coerce")
+        if z.notna().any():
+            df[f"{col}_z"] = z
+    return df

statsbombpy/sb.py

@@ -7,7 +7,8 @@
 from statsbombpy import api_client, public
 from statsbombpy.config import DEFAULT_CREDS, MAX_CONCURRENCY
 from statsbombpy.helpers import (filter_and_group_events,
-                                 merge_events_and_frames, reduce_events)
+                                 merge_events_and_frames, reduce_events,
+                                 split_location_cols)
 
 
 def competitions(fmt="dataframe", creds: dict = DEFAULT_CREDS, version: str = None):
@@ -152,6 +153,7 @@ def events(
     creds: dict = DEFAULT_CREDS,
     include_360_metrics=False,
     version: str = None,
+    split_locations: bool = False,
 ) -> Union[pd.DataFrame, dict]:
 
     if not api_client.has_auth(creds) and include_360_metrics:
@@ -169,6 +171,8 @@ def events(
         events = filter_and_group_events(events, filters, fmt, flatten_attrs)
         for ev_type, evs in events.items():
             events[ev_type] = pd.DataFrame(evs)
+            if split_locations:
+                events[ev_type] = split_location_cols(events[ev_type])
         if split is False:
             events = pd.concat([*events.values()], axis=0, ignore_index=True, sort=True)
     return events

tests/test_split_locations.py

@@ -0,0 +1,62 @@
+import numpy as np
+import pandas as pd
+
+from statsbombpy.helpers import split_location_cols
+
+
+def _events_df():
+    return pd.DataFrame(
+        {
+            "type": ["Pass", "Shot", "Half Start", "Pass"],
+            "location": [[10.0, 20.0], [100.0, 40.0], None, [50.0, 60.0]],
+            "pass_end_location": [[30.0, 25.0], None, None, [55.0, 62.0]],
+            "shot_end_location": [None, [120.0, 38.0, 2.4], None, None],
+        }
+    )
+
+
+def test_splits_xy_and_keeps_original():
+    df = split_location_cols(_events_df())
+    assert df.loc[0, "location_x"] == 10.0
+    assert df.loc[0, "location_y"] == 20.0
+    assert df.loc[0, "pass_end_location_x"] == 30.0
+    # original list columns are preserved
+    assert df.loc[0, "location"] == [10.0, 20.0]
+
+
+def test_z_only_created_when_present():
+    df = split_location_cols(_events_df())
+    assert "shot_end_location_z" in df.columns
+    assert df.loc[1, "shot_end_location_z"] == 2.4
+    # 2d columns do not grow a z
+    assert "location_z" not in df.columns
+    assert "pass_end_location_z" not in df.columns
+
+
+def test_missing_values_yield_nan():
+    df = split_location_cols(_events_df())
+    assert np.isnan(df.loc[2, "location_x"])
+    assert np.isnan(df.loc[1, "pass_end_location_y"])
+
+
+def test_handles_numpy_arrays_from_parquet_roundtrip():
+    df = _events_df()
+    # simulate what pd.read_parquet gives back: arrays instead of lists
+    df["location"] = df["location"].map(lambda v: np.array(v) if v is not None else None)
+    df = split_location_cols(df)
+    assert df.loc[0, "location_x"] == 10.0
+    assert df.loc[3, "location_y"] == 60.0
+
+
+def test_absent_columns_are_ignored():
+    df = pd.DataFrame({"type": ["Pass"], "location": [[1.0, 2.0]]})
+    df = split_location_cols(df)
+    assert df.loc[0, "location_x"] == 1.0
+    assert "carry_end_location_x" not in df.columns
+
+
+def test_split_columns_are_numeric_and_hashable():
+    df = split_location_cols(_events_df())
+    assert df["location_x"].dtype.kind == "f"
+    # the flat columns support drop_duplicates, unlike the list column
+    df[["location_x", "location_y"]].drop_duplicates()