Python: fix(lab-gaia): accept local parquet data in _ensure_data

python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py

@@ -264,6 +264,52 @@ def _read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
                 yield record
 
 
+_SHARDED_PARQUET_RE = re.compile(r"(metadata(?:-[^-]+)?)-(\d{5})-of-(\d{5})\.parquet")
+
+
+def _has_complete_parquet_snapshot(data_dir: Path) -> bool:
+    """Whether every sharded metadata parquet group holds all of its declared shards.
+
+    An interrupted copy can leave only part of a ``metadata-XXXXX-of-NNNNN.parquet``
+    family in the directory. Treating any matching filename as ready would run the
+    benchmark over a silently truncated task set, so an incomplete family forces the
+    normal download path instead. Plain unsharded metadata parquet files count as
+    complete on their own.
+    """
+    sharded: dict[tuple[Path, str], tuple[set[int], int]] = {}
+    unsharded_seen = False
+    for path in data_dir.rglob("metadata*.parquet"):
+        match = _SHARDED_PARQUET_RE.fullmatch(path.name)
+        if match is None:
+            unsharded_seen = True
+            continue
+        key = (path.parent, match.group(1))
+        indexes, total = sharded.get(key, (set(), int(match.group(3))))
+        indexes.add(int(match.group(2)))
+        sharded[key] = (indexes, max(total, int(match.group(3))))
+    if not sharded:
+        return unsharded_seen
+    # The loader reads every matching parquet file, so a single incomplete sharded
+    # family disqualifies the directory even when unsharded files also exist.
+    return all(len(indexes) == total for indexes, total in sharded.values())
+
+
+def _has_usable_local_data(data_dir: Path) -> bool:
+    """Whether the directory holds data that is safe to load without a download.
+
+    A JSONL-only layout is usable when ``metadata.jsonl`` is present, and a parquet
+    layout is usable when every sharded family is complete. A mixed layout must
+    not pass merely because JSONL exists: ``_load_gaia_local`` reads parquet files
+    first and only falls back to JSONL when no parquet file is present, so an
+    incomplete parquet family would silently truncate the task set even though a
+    complete ``metadata.jsonl`` sits beside it. When any metadata parquet file is
+    present, the parquet snapshot must therefore be complete.
+    """
+    if any(data_dir.rglob("metadata*.parquet")):
+        return _has_complete_parquet_snapshot(data_dir)
+    return any(data_dir.rglob("metadata.jsonl"))
+
+
 def _load_gaia_local(repo_dir: Path, wanted_levels: list[int] | None = None, max_n: int | None = None) -> list[Task]:
     """Load GAIA tasks from local repository directory."""
     tasks: list[Task] = []
@@ -405,7 +451,7 @@ async def _default_evaluator(self, task: Task, prediction: Prediction) -> Evalua
 
     def _ensure_data(self) -> Path:
         """Ensure GAIA data is available locally."""
-        if self.data_dir.exists() and any(self.data_dir.rglob("metadata.jsonl")):
+        if self.data_dir.exists() and _has_usable_local_data(self.data_dir):
             return self.data_dir
 
         # Download data if not available

python/packages/lab/gaia/tests/test_gaia.py

@@ -2,6 +2,7 @@
 
 """Tests for GAIA benchmark implementation."""
 
+import pytest
 from agent_framework_lab_gaia import gaia_scorer
 
 
@@ -34,3 +35,118 @@ def test_none_handling(self):
         """Test handling of None values."""
         assert gaia_scorer("None", "test") is False
         assert gaia_scorer("", "test") is False
+
+
+class TestEnsureData:
+    """Test local data detection in GAIA._ensure_data."""
+
+    def test_parquet_only_data_dir_is_accepted(self, tmp_path):
+        """A data dir in the parquet layout _load_gaia_local supports must not force a download (#9337)."""
+        pytest.importorskip("pyarrow")
+        from agent_framework_lab_gaia.gaia import GAIA
+
+        validation = tmp_path / "2023" / "validation"
+        validation.mkdir(parents=True)
+        rows = [
+            {
+                "task_id": "v1",
+                "Question": "What is 2+2?",
+                "Final answer": "4",
+                "Level": 1,
+                "file_name": "",
+                "Annotator Metadata": "{}",
+            }
+        ]
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00001.parquet")
+
+        runner = GAIA(data_dir=str(tmp_path))
+        assert runner._ensure_data() == tmp_path
+
+    def test_incomplete_sharded_parquet_snapshot_forces_download(self, tmp_path, monkeypatch):
+        """A partial metadata-...-of-N family must not count as ready local data (#9337)."""
+        pytest.importorskip("pyarrow")
+        monkeypatch.delenv("HF_TOKEN", raising=False)
+        from agent_framework_lab_gaia.gaia import GAIA
+
+        validation = tmp_path / "2023" / "validation"
+        validation.mkdir(parents=True)
+        rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}]
+
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        # interrupted copy: only shard 0 of 2 landed
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet")
+
+        runner = GAIA(data_dir=str(tmp_path))
+        with pytest.raises(RuntimeError, match="HF_TOKEN"):
+            runner._ensure_data()
+
+        # completing the family makes the directory ready again
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet")
+        assert runner._ensure_data() == tmp_path
+
+    def test_jsonl_does_not_rescue_an_incomplete_parquet_family(self, tmp_path, monkeypatch):
+        """A mixed layout must not pass on the JSONL shortcut (#9340 review).
+
+        _load_gaia_local reads parquet first and only falls back to JSONL when
+        no parquet file is present, so metadata.jsonl beside an incomplete
+        sharded family would silently truncate the task set.
+        """
+        pytest.importorskip("pyarrow")
+        monkeypatch.delenv("HF_TOKEN", raising=False)
+        from agent_framework_lab_gaia.gaia import GAIA
+
+        validation = tmp_path / "2023" / "validation"
+        validation.mkdir(parents=True)
+        rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}]
+
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        (validation / "metadata.jsonl").write_text("{}\n", encoding="utf-8")
+        # interrupted copy: only shard 0 of 2 landed beside the JSONL file
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet")
+
+        runner = GAIA(data_dir=str(tmp_path))
+        with pytest.raises(RuntimeError, match="HF_TOKEN"):
+            runner._ensure_data()
+
+        # completing the family makes the mixed directory ready again
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00001-of-00002.parquet")
+        assert runner._ensure_data() == tmp_path
+
+    def test_jsonl_only_layout_is_still_accepted(self, tmp_path, monkeypatch):
+        """A JSONL-only directory keeps working without any parquet files (#9340 review)."""
+        monkeypatch.delenv("HF_TOKEN", raising=False)
+        from agent_framework_lab_gaia.gaia import GAIA
+
+        validation = tmp_path / "2023" / "validation"
+        validation.mkdir(parents=True)
+        (validation / "metadata.jsonl").write_text("{}\n", encoding="utf-8")
+
+        runner = GAIA(data_dir=str(tmp_path))
+        assert runner._ensure_data() == tmp_path
+
+    def test_unsharded_file_does_not_rescue_an_incomplete_family(self, tmp_path, monkeypatch):
+        """An unsharded file must not mask a partial sharded family (#9337 review)."""
+        pytest.importorskip("pyarrow")
+        monkeypatch.delenv("HF_TOKEN", raising=False)
+        from agent_framework_lab_gaia.gaia import GAIA
+
+        validation = tmp_path / "2023" / "validation"
+        validation.mkdir(parents=True)
+        rows = [{"task_id": "v1", "Question": "q", "Final answer": "a", "Level": 1, "file_name": ""}]
+
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-00000-of-00002.parquet")
+        pq.write_table(pa.Table.from_pylist(rows), validation / "metadata-extra.parquet")
+
+        runner = GAIA(data_dir=str(tmp_path))
+        with pytest.raises(RuntimeError, match="HF_TOKEN"):
+            runner._ensure_data()