diff --git a/src/ml4t/backtest/datafeed.py b/src/ml4t/backtest/datafeed.py index be3d6aad..08b78bcb 100644 --- a/src/ml4t/backtest/datafeed.py +++ b/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. diff --git a/tests/test_datafeed_memory.py b/tests/test_datafeed_memory.py index f815f73d..d31830c9 100644 --- a/tests/test_datafeed_memory.py +++ b/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