perf(common-crawl): parse WARC files with fastwarc instead of warcio

nemo_curator/stages/text/download/common_crawl/download.py

@@ -22,8 +22,8 @@
 
 import pandas as pd
 import requests
+from fastwarc.warc import ArchiveIterator, WarcRecordType
 from loguru import logger
-from warcio.archiveiterator import ArchiveIterator
 
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.text.download import DocumentDownloader
@@ -285,10 +285,11 @@ def _read_warc_record_s3(self, row: pd.Series) -> bytes | None:
 
             try:
                 stream = io.BytesIO(decompressed)
-                archive_iterator = ArchiveIterator(stream)
+                archive_iterator = ArchiveIterator(
+                    stream, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False
+                )
                 for record in archive_iterator:
-                    if record.rec_type == "response":
-                        return record.content_stream().read()
+                    return record.reader.read()
             except Exception as e:  # noqa: BLE001
                 logger.debug(f"Failed to parse WARC record {filename}: {e}, returning decompressed bytes")
                 return decompressed
@@ -302,13 +303,13 @@ def _read_warc_record_s3(self, row: pd.Series) -> bytes | None:
             logger.warning(f"S3 fetch failed for {filename}: {e}")
             return None
 
-    def _read_warc_record(self, row: pd.Series) -> bytes | None:  # noqa: C901, PLR0911
+    def _read_warc_record(self, row: pd.Series) -> bytes | None:  # noqa: PLR0911
         """Fetch a single WARC record using HTTPS range request.
 
         This method:
         1. Fetches gzip-compressed WARC record bytes via HTTP range request
         2. Decompresses the gzip content
-        3. Parses the WARC record format using warcio
+        3. Parses the WARC record format using fastwarc
         4. Extracts and returns the HTTP response body (the actual content)
         """
         filename = None
@@ -349,14 +350,15 @@ def _read_warc_record(self, row: pd.Series) -> bytes | None:  # noqa: C901, PLR0
                 # Content might not be gzip-compressed, use as-is
                 decompressed = raw_bytes
 
-            # Parse the WARC record using warcio to extract HTTP response body
+            # Parse the WARC record using fastwarc to extract HTTP response body
             try:
                 stream = io.BytesIO(decompressed)
-                archive_iterator = ArchiveIterator(stream)
+                archive_iterator = ArchiveIterator(
+                    stream, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False
+                )
                 for record in archive_iterator:
-                    if record.rec_type == "response":
-                        # Return the HTTP response body (content after HTTP headers)
-                        return record.content_stream().read()
+                    # Return the HTTP response body (content after HTTP headers)
+                    return record.reader.read()
             except Exception as e:  # noqa: BLE001
                 logger.debug(f"Failed to parse WARC record {filename}: {e}, returning decompressed bytes")
                 return decompressed

nemo_curator/stages/text/download/common_crawl/warc_iterator.py

@@ -16,12 +16,10 @@
 from pathlib import Path
 from typing import Any
 
+from fastwarc.warc import ArchiveIterator, WarcRecordType
 from fsspec.core import url_to_fs
 from loguru import logger
 
-# TODO: Consider using fastwarc https://github.com/NVIDIA-NeMo/Curator/issues/778
-from warcio.archiveiterator import ArchiveIterator
-
 from nemo_curator.stages.text.download import DocumentIterator
 
 
@@ -46,19 +44,44 @@ def iterate(self, file_path: str) -> Iterator[dict[str, Any]]:
         num_records = 0
         fs, fs_path = url_to_fs(file_path_str, **self.storage_options)
         with fs.open(fs_path, "rb") as file_pointer:
-            archive_iterator = ArchiveIterator(file_pointer, arc2warc=True)
+            # fastwarc wraps any file-like object and sniffs gzip itself, so the fsspec
+            # handle can be passed straight through. Non-response records are discarded
+            # in C++ before their headers reach Python, and auto_decode="all" keeps the
+            # HTTP body decoded from its Content-Encoding, as this iterator did before.
+            # strict_mode=False resynchronizes past a record with an unparseable WARC
+            # header instead of silently ending the file there, which is the default.
+            #
+            # Gaps that come with fastwarc 0.x, none of them reachable from Common Crawl's
+            # own files: a record resynchronized past is dropped silently -- the parser
+            # exposes no skip counter or callback and does not surface the record even
+            # with record_types=any_type, so there is nothing this loop can log, and a
+            # short record count is the only symptom; chunked transfer-encoding is not
+            # decoded even with auto_decode="all"; and the HTTP preamble is only stripped
+            # from records that declare Content-Type: application/http, which Common Crawl
+            # emits. ARC input, which the previous arc2warc=True accepted, is not supported.
+            archive_iterator = ArchiveIterator(
+                file_pointer, record_types=WarcRecordType.response, auto_decode="all", strict_mode=False
+            )
             while True:
                 try:
                     rec = next(archive_iterator)
-                    if rec.rec_type == "response":
-                        content = rec.content_stream().read()
-                        warc_id = rec.rec_headers.get_header("WARC-Record-ID")[10:-1]
-                        url = rec.rec_headers.get_header("WARC-Target-URI")
-                        yield {"url": url, "warc_id": warc_id, "source_id": filename, "content": content}
-                        num_records += 1
                 except StopIteration:
                     # End of file reached normally
                     break
+                except Exception as e:  # noqa: BLE001
+                    # next() has its own try because a stream the parser cannot open at all
+                    # fails here (fastwarc raises StreamError) rather than while reading a
+                    # record, and leaves nothing to resynchronize to. Report it once and stop
+                    # instead of calling next() again on a stream that is already dead.
+                    logger.error(f"Error processing record {num_records} in {filename}: {e!s}")
+                    break
+
+                try:
+                    content = rec.reader.read()
+                    warc_id = rec.headers.get("WARC-Record-ID")[10:-1]
+                    url = rec.headers.get("WARC-Target-URI")
+                    yield {"url": url, "warc_id": warc_id, "source_id": filename, "content": content}
+                    num_records += 1
                 except Exception as e:  # noqa: BLE001
                     # Handle corruption or other errors
                     logger.error(f"Error processing record {num_records} in {filename}: {e!s}")

pyproject.toml

@@ -197,13 +197,16 @@ translation_all = [
 text_cpu = [
     # Download / Extract
     "beautifulsoup4",
+    # fastwarc and resiliparse pin each other exactly, so this cap also holds resiliparse
+    # below 1.x; both are used well beyond WARC parsing (see html_extractors/resiliparse.py).
+    # fastwarc 1.x also drops strict_mode's resync past a corrupt WARC header (warc_iterator.py).
+    "fastwarc<1",
     "justext",
     "lxml",
     "pycld2",
     "resiliparse",
     "s5cmd",
     "trafilatura==2.0.0",
-    "warcio",
     # Filters
     "fasttext==0.9.3",
     "sentencepiece",

tests/stages/text/download/common_crawl/test_warc_iterator.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gzip
+from collections.abc import Callable
 from pathlib import Path
 from unittest import mock
 
@@ -21,6 +23,24 @@
 
 from nemo_curator.stages.text.download.common_crawl.warc_iterator import CommonCrawlWarcIterator
 
+_OK_BODY = b"<html><body>ok</body></html>"
+_OK_HTTP = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n" + _OK_BODY
+
+
+def _response_record(record_id: str, http_payload: bytes, version: str = "WARC/1.0") -> bytes:
+    """Build a single WARC response record wrapping an already-serialized HTTP response."""
+    header = (
+        f"{version}\r\n"
+        f"WARC-Type: response\r\n"
+        f"WARC-Record-ID: <urn:uuid:{record_id}>\r\n"
+        f"WARC-Date: 2022-01-01T00:00:00Z\r\n"
+        f"WARC-Target-URI: http://example.com/{record_id}\r\n"
+        f"Content-Type: application/http;msgtype=response\r\n"
+        f"Content-Length: {len(http_payload)}\r\n"
+        f"\r\n"
+    ).encode()
+    return header + http_payload + b"\r\n\r\n"
+
 
 class TestCommonCrawlWarcIterator:
     """Test suite for CommonCrawlWarcIterator - focused on core logic correctness."""
@@ -63,7 +83,7 @@ def test_error_processing_record_continues(self, tmp_path: Path) -> None:
         raw_warc_path = tmp_path / "test.warc"
 
         # Create a WARC file with a response record that has no WARC-Record-ID header
-        # This will cause the get_header to return None, leading to the subscriptable error
+        # This makes headers.get return None, leading to the subscriptable error
         http_response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n<html><body>Test</body></html>\r\n"
         http_response_bytes = http_response.encode("utf-8")
         content_length = len(http_response_bytes)
@@ -119,6 +139,8 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa
                 "content": "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n<html><body><h1>Test Page</h1></body></html>\r\n",
                 "id": "response123",
                 "target_uri": "http://example.com/page",
+                # Marks the payload as an HTTP response, as Common Crawl's own WARC files do.
+                "content_type": "application/http;msgtype=response",
             },
             {
                 "type": "metadata",
@@ -147,6 +169,9 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa
             if config["target_uri"]:
                 header_parts.append(f"WARC-Target-URI: {config['target_uri']}\r\n")
 
+            if config.get("content_type"):
+                header_parts.append(f"Content-Type: {config['content_type']}\r\n")
+
             header_parts.append(f"Content-Length: {content_length}\r\n\r\n")
 
             warc_record = "".join(header_parts).encode() + content_bytes + b"\r\n\r\n"
@@ -166,12 +191,78 @@ def test_mixed_record_types_response_only_with_correct_values(self, tmp_path: Pa
         assert record["url"] == "http://example.com/page"
         assert record["warc_id"] == "response123"  # Stripped <urn:uuid: and >
         assert record["source_id"] == "mixed_types.warc"
-        # The content should be just the HTML body (warcio extracts body from HTTP response)
+        # The content should be just the HTML body, with the HTTP headers stripped
         assert record["content"] == html_content
 
         # Verify the content contains expected HTML
         assert b"<h1>Test Page</h1>" in record["content"]
 
+    @pytest.mark.parametrize(
+        ("content_encoding", "encode"),
+        [
+            (None, lambda body: body),
+            ("gzip", gzip.compress),
+        ],
+    )
+    def test_http_body_is_returned_decoded(
+        self,
+        tmp_path: Path,
+        content_encoding: str | None,
+        encode: Callable[[bytes], bytes],
+    ) -> None:
+        """The yielded content is the decoded HTTP body, whatever Content-Encoding the server used."""
+        body = b"<html><body>decoded</body></html>"
+        encoding_header = f"Content-Encoding: {content_encoding}\r\n" if content_encoding else ""
+        http_payload = (f"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n{encoding_header}\r\n").encode() + encode(body)
+
+        raw_warc_path = tmp_path / "encoded.warc"
+        raw_warc_path.write_bytes(_response_record("encoded123", http_payload))
+
+        records = list(CommonCrawlWarcIterator().iterate(str(raw_warc_path)))
+
+        assert len(records) == 1
+        assert records[0]["content"] == body
+
+    # expected_errors is 0 for both corrupt-record cases on purpose: fastwarc 0.x
+    # resynchronizes past the bad record and offers no way to observe that it did,
+    # so the drop cannot be logged (see the comment in CommonCrawlWarcIterator.iterate).
+    @pytest.mark.parametrize(
+        ("warc_bytes", "expected_ids", "expected_errors"),
+        [
+            (
+                _response_record("corrupt", _OK_HTTP, version="WARC/XX") + _response_record("good", _OK_HTTP),
+                ["good"],
+                0,
+            ),
+            (
+                _response_record("first", _OK_HTTP)
+                + _response_record("corrupt", _OK_HTTP, version="WARC/XX")
+                + _response_record("last", _OK_HTTP),
+                ["first", "last"],
+                0,
+            ),
+            (b"\x00\x01not-a-warc\r\n\r\n" + _response_record("good", _OK_HTTP), [], 1),
+        ],
+        ids=["corrupt-first", "corrupt-mid", "unreadable-stream"],
+    )
+    def test_corrupt_record_does_not_abandon_the_rest_of_the_file(
+        self,
+        tmp_path: Path,
+        warc_bytes: bytes,
+        expected_ids: list[str],
+        expected_errors: int,
+    ) -> None:
+        """A corrupt WARC header is resynchronized past; a stream the parser cannot open is logged once."""
+        raw_warc_path = tmp_path / "corrupt.warc"
+        raw_warc_path.write_bytes(warc_bytes)
+
+        with mock.patch.object(logger, "error") as mock_logger:
+            yielded = list(CommonCrawlWarcIterator().iterate(str(raw_warc_path)))
+
+        assert [record["warc_id"] for record in yielded] == expected_ids
+        assert [record["content"] for record in yielded] == [_OK_BODY] * len(expected_ids)
+        assert mock_logger.call_count == expected_errors
+
     def test_output_columns(self) -> None:
         """Test that output_columns returns the expected column names."""
         iterator = CommonCrawlWarcIterator()