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

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

@@ -264,6 +264,36 @@ 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 _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 +435,9 @@ 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 (
+            any(self.data_dir.rglob("metadata.jsonl")) or _has_complete_parquet_snapshot(self.data_dir)
+        ):
             return self.data_dir
 
         # Download data if not available

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

@@ -2,6 +2,8 @@
 
 """Tests for GAIA benchmark implementation."""
 
+import pytest
+
 from agent_framework_lab_gaia import gaia_scorer
 
 
@@ -34,3 +36,76 @@ 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_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()