From eb677a7ea817c072c5d5807ed8692269fdd898f6 Mon Sep 17 00:00:00 2001 From: Samar Date: Thu, 30 Jul 2026 05:17:04 +0300 Subject: [PATCH 1/2] Fix conditional retrieval loading for non-retrieval settings Setting 1 (option=2) and Setting 2 (option=3) do not require table retrieval files, but get_retrieved_tables() was called unconditionally, causing an AssertionError when retrieval files were absent. This fix moves the get_retrieved_tables() call inside the preprocessing_option == 1 guard, so it is only invoked when retrieval is actually needed (Setting 0). Added test verifying that preprocessing_option 2 and 3 succeed without retrieval files. --- eval/reforce/convert_beaver_to_reforce.py | 4 +- test_bug/test_retrieval_conditional.py | 72 +++++++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) create mode 100644 test_bug/test_retrieval_conditional.py diff --git a/eval/reforce/convert_beaver_to_reforce.py b/eval/reforce/convert_beaver_to_reforce.py index 72d661f..1dce934 100644 --- a/eval/reforce/convert_beaver_to_reforce.py +++ b/eval/reforce/convert_beaver_to_reforce.py @@ -325,7 +325,9 @@ def get_table_key(t_name): new_join_keys.append([c.lower() for c in jk]) item['join_keys'] = new_join_keys - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if preprocessing_option == 1: + retrieved_tables = get_retrieved_tables(dataset) for idx, item in enumerate(beaver_questions): if (idx + 1) % 10 == 0: diff --git a/test_bug/test_retrieval_conditional.py b/test_bug/test_retrieval_conditional.py new file mode 100644 index 0000000..6d36e3d --- /dev/null +++ b/test_bug/test_retrieval_conditional.py @@ -0,0 +1,72 @@ +""" +Test: get_retrieved_tables is only called when retrieval is actually needed. + +Before the fix: + get_retrieved_tables(dataset) was called unconditionally at line 328, + BEFORE the preprocessing_option check. This caused AssertionError + for Setting 1 (option=2) and Setting 2 (option=3) even though + they don't need retrieval files. + +After the fix: + get_retrieved_tables(dataset) is only called when preprocessing_option == 1, + which corresponds to Setting 0 (end-to-end mode requiring retrieval). +""" +import sys +import os +import json + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "eval", "reforce")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +def test_setting_without_retrieval(): + """Setting 1 (option=2) should succeed without retrieval files.""" + from convert_beaver_to_reforce import convert_beaver_to_reforce + + data_dir = os.path.join(os.path.dirname(__file__), "..", "data") + output_path = os.path.join(os.path.dirname(__file__), "test_opt2_output.json") + + convert_beaver_to_reforce( + dataset="dw", + beaver_questions_path=os.path.join(data_dir, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(data_dir, "dw", "dev_tables.json"), + output_path=output_path, + sampled_questions_path=output_path.replace(".json", "_sampled.json"), + preprocessing_option=2, # = Setting 1: gold tables, no retrieval needed + db_id="dw", + ) + + assert os.path.exists(output_path), f"Output file {output_path} not created" + with open(output_path) as f: + data = json.load(f) + assert len(data) > 0, "Output data is empty" + print(f"[PASS] Setting 1 (option=2): {len(data)} entries processed successfully") + os.remove(output_path) + +def test_setting_without_retrieval_option3(): + """Setting 2 (option=3) should also succeed without retrieval files.""" + from convert_beaver_to_reforce import convert_beaver_to_reforce + + data_dir = os.path.join(os.path.dirname(__file__), "..", "data") + output_path = os.path.join(os.path.dirname(__file__), "test_opt3_output.json") + + convert_beaver_to_reforce( + dataset="dw", + beaver_questions_path=os.path.join(data_dir, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(data_dir, "dw", "dev_tables.json"), + output_path=output_path, + sampled_questions_path=output_path.replace(".json", "_sampled.json"), + preprocessing_option=3, # = Setting 2: all hints, no retrieval needed + db_id="dw", + ) + + assert os.path.exists(output_path) + with open(output_path) as f: + data = json.load(f) + assert len(data) > 0 + print(f"[PASS] Setting 2 (option=3): {len(data)} entries processed successfully") + os.remove(output_path) + +if __name__ == "__main__": + test_setting_without_retrieval() + test_setting_without_retrieval_option3() + print("\nAll tests passed!") From 427ac9f06a9631629878d66b1922ccf330c210c2 Mon Sep 17 00:00:00 2001 From: Samar Date: Thu, 30 Jul 2026 05:24:02 +0300 Subject: [PATCH 2/2] Extend retrieval-conditional fix to DIN-SQL and DAIL-SQL pipelines - Same fix as the ReFoRCE converter: moved get_retrieved_tables() inside the option == 1 guard for DIN-SQL and DAIL-SQL. - Updated regression test to cover all three pipelines. --- .../preprocessed_data/beaver_preprocess.py | 6 +- .../preprocessed_data/beaver_preprocess_v2.py | 6 +- test_bug/test_retrieval_conditional.py | 143 ++++++++++++------ 3 files changed, 106 insertions(+), 49 deletions(-) diff --git a/eval/dailsql/preprocessed_data/beaver_preprocess.py b/eval/dailsql/preprocessed_data/beaver_preprocess.py index 85c4036..1d4547e 100644 --- a/eval/dailsql/preprocessed_data/beaver_preprocess.py +++ b/eval/dailsql/preprocessed_data/beaver_preprocess.py @@ -281,7 +281,9 @@ def convert_beaver_questions_to_dailsql_format(dataset, beaver_questions_path, o # Load spacy for tokenization nlp = spacy.load("en_core_web_sm") - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if option == 1: + retrieved_tables = get_retrieved_tables(dataset) dailsql_format = [] for question_info in tqdm(beaver_questions): @@ -289,7 +291,7 @@ def convert_beaver_questions_to_dailsql_format(dataset, beaver_questions_path, o # Option 1+: Add gold tables to question if option >= 1: - if option == 1: + if option == 1 and retrieved_tables is not None: gold_tables = retrieved_tables[question_info['id']] else: gold_tables = question_info.get('tables', []) diff --git a/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py b/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py index 096bb3c..7c6fd3b 100644 --- a/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py +++ b/eval/dinsql/preprocessed_data/beaver_preprocess_v2.py @@ -234,7 +234,9 @@ def convert_beaver_questions_to_dinsql_format(dataset, beaver_questions_path, ou with open(beaver_questions_path, 'r') as f: beaver_questions = json.load(f) - retrieved_tables = get_retrieved_tables(dataset) + retrieved_tables = None + if option == 1: + retrieved_tables = get_retrieved_tables(dataset) dinsql_format = [] for question_info in beaver_questions: @@ -247,7 +249,7 @@ def convert_beaver_questions_to_dinsql_format(dataset, beaver_questions_path, ou # Option 1+: Store gold_tables for filtering if option >= 1: - if option == 1: + if option == 1 and retrieved_tables is not None: base_item['tables'] = retrieved_tables[question_info['id']] else: base_item['tables'] = question_info.get('tables', []) diff --git a/test_bug/test_retrieval_conditional.py b/test_bug/test_retrieval_conditional.py index 6d36e3d..f89eacb 100644 --- a/test_bug/test_retrieval_conditional.py +++ b/test_bug/test_retrieval_conditional.py @@ -1,72 +1,125 @@ """ Test: get_retrieved_tables is only called when retrieval is actually needed. +Affected files: + - eval/reforce/convert_beaver_to_reforce.py + - eval/dinsql/preprocessed_data/beaver_preprocess_v2.py + - eval/dailsql/preprocessed_data/beaver_preprocess.py + Before the fix: - get_retrieved_tables(dataset) was called unconditionally at line 328, - BEFORE the preprocessing_option check. This caused AssertionError - for Setting 1 (option=2) and Setting 2 (option=3) even though - they don't need retrieval files. + get_retrieved_tables(dataset) was called unconditionally before + the option/ setting check, causing AssertionError for Setting 1 (option=2) + and Setting 2 (option=3) even though they use gold tables from data. After the fix: - get_retrieved_tables(dataset) is only called when preprocessing_option == 1, - which corresponds to Setting 0 (end-to-end mode requiring retrieval). + get_retrieved_tables(dataset) is only called when option == 1 + (Setting 0 / end-to-end mode requiring retrieval). """ import sys import os import json -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "eval", "reforce")) -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_DIR) -def test_setting_without_retrieval(): - """Setting 1 (option=2) should succeed without retrieval files.""" - from convert_beaver_to_reforce import convert_beaver_to_reforce - - data_dir = os.path.join(os.path.dirname(__file__), "..", "data") - output_path = os.path.join(os.path.dirname(__file__), "test_opt2_output.json") +DATA_DIR = os.path.join(REPO_DIR, "data") +OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def test_reforce_setting_1(): + """ReFoRCE: Setting 1 (option=2) without retrieval.""" + from eval.reforce.convert_beaver_to_reforce import convert_beaver_to_reforce + out = os.path.join(OUTPUT_DIR, "reforce_opt2.json") + sampled = out.replace(".json", "_sampled.json") convert_beaver_to_reforce( dataset="dw", - beaver_questions_path=os.path.join(data_dir, "dw", "dev_sampled.json"), - beaver_tables_path=os.path.join(data_dir, "dw", "dev_tables.json"), - output_path=output_path, - sampled_questions_path=output_path.replace(".json", "_sampled.json"), - preprocessing_option=2, # = Setting 1: gold tables, no retrieval needed + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(DATA_DIR, "dw", "dev_tables.json"), + output_path=out, + sampled_questions_path=sampled, + preprocessing_option=2, db_id="dw", ) - - assert os.path.exists(output_path), f"Output file {output_path} not created" - with open(output_path) as f: + assert os.path.exists(out) + with open(out) as f: data = json.load(f) - assert len(data) > 0, "Output data is empty" - print(f"[PASS] Setting 1 (option=2): {len(data)} entries processed successfully") - os.remove(output_path) + assert len(data) > 0 + os.remove(out) + if os.path.exists(sampled): + os.remove(sampled) + print(f" [PASS] ReFoRCE option=2: {len(data)} entries") -def test_setting_without_retrieval_option3(): - """Setting 2 (option=3) should also succeed without retrieval files.""" - from convert_beaver_to_reforce import convert_beaver_to_reforce - - data_dir = os.path.join(os.path.dirname(__file__), "..", "data") - output_path = os.path.join(os.path.dirname(__file__), "test_opt3_output.json") + +def test_dinsql_setting_1(): + """DIN-SQL: Setting 1 (option=2) without retrieval.""" + from eval.dinsql.preprocessed_data.beaver_preprocess_v2 import convert_beaver_questions_to_dinsql_format - convert_beaver_to_reforce( + out = os.path.join(OUTPUT_DIR, "dinsql_opt2.json") + convert_beaver_questions_to_dinsql_format( dataset="dw", - beaver_questions_path=os.path.join(data_dir, "dw", "dev_sampled.json"), - beaver_tables_path=os.path.join(data_dir, "dw", "dev_tables.json"), - output_path=output_path, - sampled_questions_path=output_path.replace(".json", "_sampled.json"), - preprocessing_option=3, # = Setting 2: all hints, no retrieval needed - db_id="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + output_path=out, + option=2, ) + assert os.path.exists(out) + with open(out) as f: + data = json.load(f) + assert len(data) > 0 + os.remove(out) + print(f" [PASS] DIN-SQL option=2: {len(data)} entries") + + +def test_dailsql_setting_1(): + """DAIL-SQL: Setting 1 (option=2) without retrieval.""" + from eval.dailsql.preprocessed_data.beaver_preprocess import convert_beaver_questions_to_dailsql_format - assert os.path.exists(output_path) - with open(output_path) as f: + out = os.path.join(OUTPUT_DIR, "dailsql_opt2.json") + convert_beaver_questions_to_dailsql_format( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + output_path=out, + option=2, + ) + assert os.path.exists(out) + with open(out) as f: data = json.load(f) assert len(data) > 0 - print(f"[PASS] Setting 2 (option=3): {len(data)} entries processed successfully") - os.remove(output_path) + os.remove(out) + print(f" [PASS] DAIL-SQL option=2: {len(data)} entries") + + +def test_setting_0_still_requires_retrieval(): + """Setting 0 (option=1) should still fail if retrieval files are missing.""" + from eval.reforce.convert_beaver_to_reforce import convert_beaver_to_reforce + + out = os.path.join(OUTPUT_DIR, "setting0_should_fail.json") + sampled = out.replace(".json", "_sampled.json") + try: + convert_beaver_to_reforce( + dataset="dw", + beaver_questions_path=os.path.join(DATA_DIR, "dw", "dev_sampled.json"), + beaver_tables_path=os.path.join(DATA_DIR, "dw", "dev_tables.json"), + output_path=out, + sampled_questions_path=sampled, + preprocessing_option=1, + db_id="dw", + ) + print(" [FAIL] Setting 0 should have raised an error (no retrieval files)") + except AssertionError: + print(" [PASS] Setting 0 raises AssertionError as expected (no retrieval files)") + + if os.path.exists(out): + os.remove(out) + if os.path.exists(sampled): + os.remove(sampled) + if __name__ == "__main__": - test_setting_without_retrieval() - test_setting_without_retrieval_option3() - print("\nAll tests passed!") + print("=== Test: Non-retrieval settings without retrieval files ===\n") + print("ReFoRCE:"); test_reforce_setting_1() + print("DIN-SQL:"); test_dinsql_setting_1() + print("DAIL-SQL:"); test_dailsql_setting_1() + print("\n=== Test: Setting 0 still requires retrieval ===") + test_setting_0_still_requires_retrieval() + print("\n=== All tests completed! ===")