From 17a2558339556114c68258e560d9d4497a23ecd0 Mon Sep 17 00:00:00 2001 From: TPTBusiness Date: Sun, 26 Apr 2026 21:57:35 +0200 Subject: [PATCH] fix(auto-fixer): fix df['instrument'] KeyError on MultiIndex frames LLM-generated code often accesses df['instrument'] as a column, but 'instrument' is an index level (level 1) in the MultiIndex DataFrame. Replace with df.index.get_level_values(1) except when the variable was created via reset_index() (where the column actually exists). Co-Authored-By: Claude Sonnet 4.6 --- .../coder/factor_coder/auto_fixer.py | 51 +++++++++++++++---- test/qlib/test_auto_fixer.py | 21 ++++++++ 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/rdagent/components/coder/factor_coder/auto_fixer.py b/rdagent/components/coder/factor_coder/auto_fixer.py index 74551875..48f8ab24 100644 --- a/rdagent/components/coder/factor_coder/auto_fixer.py +++ b/rdagent/components/coder/factor_coder/auto_fixer.py @@ -56,15 +56,16 @@ class FactorAutoFixer: # match window size, which causes all-NaN output for intraday data with 96 bars/day # (window=240 > 96 means zero valid bars per day). The LLM sets its own min_periods. fix_methods = [ - self._fix_reset_index_groupby, # First: fix groupby(level=N) after reset_index() - self._fix_groupby_mixed_levels, # Second: fix groupby(level=[int, str]) - self._fix_groupby_column_on_multiindex, # Third: fix groupby(['instrument','date']) on MultiIndex - self._fix_chained_groupby, # Fourth: fix groupby(level=N).groupby('date') chain - self._fix_rolling_ddof, # Fifth: remove unsupported ddof kwarg - self._fix_groupby_apply_to_transform, # Sixth: fix groupby patterns - self._fix_inf_nan_handling, # Seventh: add inf/nan handling - self._fix_data_range_processing, # Eighth: ensure full data range - self._fix_multiindex_groupby, # Ninth: ensure groupby on MultiIndex + self._fix_instrument_column_access, # First: fix df['instrument'] on MultiIndex + self._fix_reset_index_groupby, # Second: fix groupby(level=N) after reset_index() + self._fix_groupby_mixed_levels, # Third: fix groupby(level=[int, str]) + self._fix_groupby_column_on_multiindex, # Fourth: fix groupby(['instrument','date']) on MultiIndex + self._fix_chained_groupby, # Fifth: fix groupby(level=N).groupby('date') chain + self._fix_rolling_ddof, # Sixth: remove unsupported ddof kwarg + self._fix_groupby_apply_to_transform, # Seventh: fix groupby patterns + self._fix_inf_nan_handling, # Eighth: add inf/nan handling + self._fix_data_range_processing, # Ninth: ensure full data range + self._fix_multiindex_groupby, # Tenth: ensure groupby on MultiIndex ] for fix_method in fix_methods: @@ -82,6 +83,38 @@ class FactorAutoFixer: return fixed_code + def _fix_instrument_column_access(self, code: str) -> str: + """ + Fix: df['instrument'] raises KeyError on a MultiIndex DataFrame because + 'instrument' is an index level (level 1), not a column. + + Replace df['instrument'] with df.index.get_level_values('instrument') + but only when the DataFrame has a MultiIndex (not after reset_index which + would have promoted it to a real column). + + Also fixes df.reset_index()['instrument'] correctly since after reset_index + the column exists. + """ + fixed_code = code + + # Skip if already fixed or if reset_index() is being used before the access + # We only fix bare df['instrument'] where df is the original MultiIndex frame. + # Heuristic: if the assignment lhs or context shows reset_index, leave it alone. + + # Pattern: ['instrument'] where varname is NOT a reset_index result + reset_vars = set(re.findall(r'(\w+)\s*=\s*\w[^=\n]*\.reset_index\(', fixed_code)) + + def _replace_instrument_access(m: re.Match) -> str: + var = m.group(1) + if var in reset_vars: + return m.group(0) # leave reset_index vars alone — column exists + self.fixes_applied.append(f"instrument_column: {var}['instrument'] → get_level_values(1)") + return f"{var}.index.get_level_values(1)" + + fixed_code = re.sub(r"(\w+)\['instrument'\]", _replace_instrument_access, fixed_code) + + return fixed_code + def _fix_reset_index_groupby(self, code: str) -> str: """ Fix: groupby(level=N) on a variable created by .reset_index() fails because diff --git a/test/qlib/test_auto_fixer.py b/test/qlib/test_auto_fixer.py index 0c7a5efe..ed03afb1 100644 --- a/test/qlib/test_auto_fixer.py +++ b/test/qlib/test_auto_fixer.py @@ -101,6 +101,27 @@ class TestMinPeriodsNotTouched: assert "min_periods=10" in result +class TestInstrumentColumnAccess: + def test_instrument_column_replaced(self, fixer): + code = "df['group_key'] = df['instrument'] + '_' + df['day_id'].astype(str)" + result = fixer.fix(code) + assert "df.index.get_level_values(1)" in result + assert "df['instrument']" not in result + + def test_reset_index_var_not_touched(self, fixer): + # After reset_index, 'instrument' IS a real column — must not be replaced + code = "df_r = df.reset_index()\nval = df_r['instrument'].unique()" + result = fixer.fix(code) + assert "df_r['instrument']" in result + assert "get_level_values" not in result + + def test_groupby_after_instrument_fix(self, fixer): + # Combined: df['instrument'] in a groupby context + code = "df['key'] = df['instrument']\nout = df.groupby(df['key'])[['$close']].mean()" + result = fixer.fix(code) + assert "df['instrument']" not in result + + class TestRollingDdof: def test_removes_ddof_from_rolling_args(self, fixer): result = fixer.fix("df.rolling(20, min_periods=1, ddof=1).std()")