diff --git a/src/question_sampling/_prepare.py b/src/question_sampling/_prepare.py index f004182a..85220cf8 100644 --- a/src/question_sampling/_prepare.py +++ b/src/question_sampling/_prepare.py @@ -165,6 +165,9 @@ def drop_questions_that_resolve_too_soon( is_na = dfq["forecast_horizons"] == "N/A" return dfq[~(empty_horizons | is_na)] + if dfq.empty: # `apply(axis=1)` on an empty frame returns a Series that drops every column + return dfq + forecast_due_date = forecast_datetime.replace(hour=23, minute=59, second=59, microsecond=999999) cutoff_date = forecast_due_date + timedelta(days=MIN_DAYS_TO_MARKET_RESOLUTION) resolves_too_soon = dfq.apply( @@ -226,7 +229,7 @@ def drop_questions_that_resolve_too_late( Returns dfq (pd.DataFrame): Questions that resolve within the cap """ - if not is_market_source(source): + if not is_market_source(source) or dfq.empty: return dfq cutoff_date = forecast_datetime + timedelta(days=MAX_DAYS_TO_MARKET_RESOLUTION) resolves_in_time = dfq.apply( diff --git a/src/question_sampling/create_question_set.py b/src/question_sampling/create_question_set.py index 8e400d0e..f29d0fb0 100644 --- a/src/question_sampling/create_question_set.py +++ b/src/question_sampling/create_question_set.py @@ -197,6 +197,9 @@ def create_question_set( forecast_datetime=forecast_datetime, ) num_questions = len(dfq) + if dfq.empty: + logger.warning(f"No questions from {source_name} are eligible for sampling.") + continue QUESTIONS[source_name] = { "dfq": dfq, "num_questions_available": num_questions, diff --git a/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_late.py b/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_late.py index 8b161538..34b22e6d 100644 --- a/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_late.py +++ b/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_late.py @@ -59,7 +59,7 @@ def test_empty_market_frame_stays_empty(): result = drop_questions_that_resolve_too_late( source="polymarket", dfq=dfq, forecast_datetime=FORECAST_DATETIME ) - assert result.empty + assert result.empty and list(result.columns) == list(dfq.columns) def test_data_questions_are_untouched(): diff --git a/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_soon.py b/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_soon.py index 687852a0..432b7402 100644 --- a/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_soon.py +++ b/src/tests/unit/question_sampling/test_drop_questions_that_resolve_too_soon.py @@ -63,7 +63,15 @@ def test_empty_market_frame_stays_empty(): result = drop_questions_that_resolve_too_soon( source="polymarket", dfq=dfq, forecast_datetime=FORECAST_DATETIME ) - assert result.empty + assert result.empty and list(result.columns) == list(dfq.columns) + + +def test_empty_data_frame_keeps_its_columns(): + dfq = pd.DataFrame({"id": [], "forecast_horizons": []}) + result = drop_questions_that_resolve_too_soon( + source="acled", dfq=dfq, forecast_datetime=FORECAST_DATETIME + ) + assert result.empty and list(result.columns) == list(dfq.columns) def test_data_questions_without_horizons_are_dropped(): diff --git a/src/tests/unit/question_sampling/test_question_sampling.py b/src/tests/unit/question_sampling/test_question_sampling.py index dad433db..109821c2 100644 --- a/src/tests/unit/question_sampling/test_question_sampling.py +++ b/src/tests/unit/question_sampling/test_question_sampling.py @@ -319,6 +319,16 @@ def resolution_dates(horizons: list[int]) -> list[str]: assert all(dates == resolution_dates([7, 30, 365]) for dates in by_source["fred"]) assert all(dates == resolution_dates([30, 365]) for dates in by_source["acled"]) + @pytest.mark.parametrize("emptied", ["polymarket", "fred"]) + def test_source_whose_questions_are_all_filtered_out_is_skipped(self, bank, emptied): + questions, dfmeta = bank + questions[emptied] = questions[emptied].assign(resolved=True) + + result = _create(questions, dfmeta) + + sources = {q["source"] for q in result.llm_question_set["questions"]} + assert sources == {"polymarket", "kalshi", "fred", "acled"} - {emptied} + def test_source_without_questions_is_skipped(self, bank): questions, dfmeta = bank questions["fred"] = pd.DataFrame(columns=constants.QUESTION_FILE_COLUMNS)