feat: configurable entity column in DataFeed

src/ml4t/backtest/datafeed.py

@@ -44,6 +44,9 @@ class DataFeed:
             process(timestamp, assets_data)
     """
 
+    #: Column names checked (in order) when auto-detecting the entity column.
+    ENTITY_COL_CANDIDATES = ("symbol", "asset", "product", "ticker")
+
     def __init__(
         self,
         prices_path: str | None = None,
@@ -52,6 +55,8 @@ def __init__(
         prices_df: pl.DataFrame | None = None,
         signals_df: pl.DataFrame | None = None,
         context_df: pl.DataFrame | None = None,
+        *,
+        entity_col: str | None = None,
     ):
         self.prices = (
             prices_df
@@ -72,6 +77,9 @@ def __init__(
         if self.prices is None:
             raise ValueError("prices_path or prices_df required")
 
+        # Resolve entity column name
+        self._entity_col = self._resolve_entity_col(entity_col, self.prices.columns)
+
         # Pre-partition data by timestamp for O(1) lookups
         # Store DataFrames (memory efficient) instead of dicts (memory explosion)
         self._prices_by_ts = self._partition_by_timestamp(self.prices)
@@ -85,7 +93,7 @@ def __init__(
         self._timestamps = self._get_timestamps()
         self._idx = 0
         self._signal_columns = (
-            [c for c in self.signals.columns if c not in ("timestamp", "asset")]
+            [c for c in self.signals.columns if c not in ("timestamp", self._entity_col)]
             if self.signals is not None
             else []
         )
@@ -96,7 +104,7 @@ def __init__(
         )
 
         price_cols = self.prices.columns
-        self._price_asset_idx = price_cols.index("asset")
+        self._price_asset_idx = price_cols.index(self._entity_col)
         self._price_open_idx = price_cols.index("open") if "open" in price_cols else -1
         self._price_high_idx = price_cols.index("high") if "high" in price_cols else -1
         self._price_low_idx = price_cols.index("low") if "low" in price_cols else -1
@@ -105,7 +113,7 @@ def __init__(
 
         if self.signals is not None:
             signal_cols = self.signals.columns
-            self._signal_asset_idx = signal_cols.index("asset")
+            self._signal_asset_idx = signal_cols.index(self._entity_col)
             self._signal_col_indices = [signal_cols.index(c) for c in self._signal_columns]
         else:
             self._signal_asset_idx = -1
@@ -117,6 +125,27 @@ def __init__(
         else:
             self._context_col_indices = []
 
+    @classmethod
+    def _resolve_entity_col(cls, explicit: str | None, columns: list[str]) -> str:
+        """Determine the entity identifier column.
+
+        If *explicit* is given, validate it exists.  Otherwise auto-detect by
+        checking ``ENTITY_COL_CANDIDATES`` in order.
+        """
+        if explicit is not None:
+            if explicit not in columns:
+                raise ValueError(
+                    f"entity_col={explicit!r} not found in columns {columns}"
+                )
+            return explicit
+        for candidate in cls.ENTITY_COL_CANDIDATES:
+            if candidate in columns:
+                return candidate
+        raise ValueError(
+            f"Cannot detect entity column. Expected one of "
+            f"{cls.ENTITY_COL_CANDIDATES}, got columns {columns}"
+        )
+
     def _partition_by_timestamp(self, df: pl.DataFrame) -> dict[datetime, pl.DataFrame]:
         """Partition DataFrame into dict keyed by timestamp for O(1) access.
 

tests/test_datafeed_memory.py

@@ -249,3 +249,119 @@ def test_zero_close_is_kept_in_price_view(self):
         _ts, data, _ctx = next(iter(feed))
         assert data["AAPL"]["close"] == 0.0
         assert data._prices["AAPL"] == 0.0
+
+
+class TestDataFeedEntityColumn:
+    """Tests for configurable entity column detection."""
+
+    def test_auto_detect_symbol(self):
+        """DataFeed should auto-detect 'symbol' column."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "symbol": ["SPY"],
+                "close": [300.0],
+            }
+        )
+        feed = DataFeed(prices_df=prices)
+        assert feed._entity_col == "symbol"
+        _ts, data, _ctx = next(iter(feed))
+        assert "SPY" in data
+        assert data["SPY"]["close"] == 300.0
+
+    def test_auto_detect_asset(self):
+        """DataFeed should auto-detect 'asset' column (backward compat)."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "asset": ["AAPL"],
+                "close": [150.0],
+            }
+        )
+        feed = DataFeed(prices_df=prices)
+        assert feed._entity_col == "asset"
+
+    def test_auto_detect_product(self):
+        """DataFeed should auto-detect 'product' column (futures)."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "product": ["ES"],
+                "close": [4500.0],
+            }
+        )
+        feed = DataFeed(prices_df=prices)
+        assert feed._entity_col == "product"
+        _ts, data, _ctx = next(iter(feed))
+        assert "ES" in data
+
+    def test_symbol_preferred_over_asset(self):
+        """When both 'symbol' and 'asset' exist, prefer 'symbol'."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "symbol": ["SPY"],
+                "asset": ["SPY_LEGACY"],
+                "close": [300.0],
+            }
+        )
+        feed = DataFeed(prices_df=prices)
+        assert feed._entity_col == "symbol"
+
+    def test_explicit_entity_col(self):
+        """DataFeed should accept explicit entity_col parameter."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "ticker": ["MSFT"],
+                "close": [350.0],
+            }
+        )
+        feed = DataFeed(prices_df=prices, entity_col="ticker")
+        assert feed._entity_col == "ticker"
+        _ts, data, _ctx = next(iter(feed))
+        assert "MSFT" in data
+
+    def test_explicit_entity_col_not_found(self):
+        """DataFeed should raise if explicit entity_col doesn't exist."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "symbol": ["SPY"],
+                "close": [300.0],
+            }
+        )
+        with pytest.raises(ValueError, match="entity_col='isin'"):
+            DataFeed(prices_df=prices, entity_col="isin")
+
+    def test_no_entity_col_detected(self):
+        """DataFeed should raise if no entity column can be detected."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "identifier": ["SPY"],
+                "close": [300.0],
+            }
+        )
+        with pytest.raises(ValueError, match="Cannot detect entity column"):
+            DataFeed(prices_df=prices)
+
+    def test_symbol_with_signals(self):
+        """DataFeed should handle 'symbol' column in both prices and signals."""
+        prices = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "symbol": ["AAPL"],
+                "close": [150.0],
+            }
+        )
+        signals = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1)],
+                "symbol": ["AAPL"],
+                "momentum": [0.5],
+            }
+        )
+        feed = DataFeed(prices_df=prices, signals_df=signals)
+        _ts, data, _ctx = next(iter(feed))
+        assert data["AAPL"]["signals"]["momentum"] == 0.5