From 2a28c9091f1ce2c88b400387312f087781ad4ef7 Mon Sep 17 00:00:00 2001
From: avinashgola <112467988+avinashgola@users.noreply.github.com>
Date: Sun, 6 Sep 2026 04:05:48 +0530
Subject: [PATCH] Fix oversized Parquet file validation

---
 src/together/lib/utils/files.py |  2 ++
 tests/unit/test_files_checks.py | 25 +++++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/together/lib/utils/files.py b/src/together/lib/utils/files.py
index c37c859b4..b4c07c862 100644
--- a/src/together/lib/utils/files.py
+++ b/src/together/lib/utils/files.py
@@ -171,7 +171,9 @@ def check_file(
         report_dict["message"] = (
             f"Maximum supported file size is {MAX_FILE_SIZE_GB} GB. Found file with size of {round(file_size / NUM_BYTES_IN_GB, 3)} GB."
         )
+        report_dict["file_size"] = file_size
         report_dict["is_check_passed"] = False
+        return report_dict
     elif file_size == 0:
         report_dict["message"] = "File is empty"
         report_dict["file_size"] = 0
diff --git a/tests/unit/test_files_checks.py b/tests/unit/test_files_checks.py
index 6aaedbd42..4f9538a1e 100644
--- a/tests/unit/test_files_checks.py
+++ b/tests/unit/test_files_checks.py
@@ -567,6 +567,31 @@ def test_check_parquet_reports_progress_callback(tmp_path: Path) -> None:
     assert parquet_events[-1].processed_bytes == parquet_events[-1].total_bytes == file.stat().st_size
 
 
+def test_check_file_rejects_oversized_parquet_before_format_checks(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    pyarrow = pytest.importorskip("pyarrow")
+    parquet = pytest.importorskip("pyarrow.parquet")
+
+    file = tmp_path / "oversized.parquet"
+    table = pyarrow.table(
+        {
+            "input_ids": [[1]],
+            "attention_mask": [[1]],
+            "labels": [[1]],
+            "position_ids": [[0]],
+        }
+    )
+    parquet.write_table(table, file)
+    monkeypatch.setattr("together.lib.utils.files.MAX_FILE_SIZE_GB", 0)
+
+    report = check_file(file)
+
+    assert not report["is_check_passed"]
+    assert report["file_size"] == file.stat().st_size
+    assert "Maximum supported file size" in report["message"]
+
+
 def test_check_progress_tracker_does_not_reset_across_phases(tmp_path: Path) -> None:
     from together.lib.cli.components.check_progress import CheckProgressTracker