feat(dataset): add collect_all_keys option and keep CSV export columns aligned

docs/examples/export_entire_dataset_to_file.mdx

@@ -11,7 +11,7 @@ import RunnableCodeBlock from '@site/src/components/RunnableCodeBlock';
 import JsonExample from '!!raw-loader!roa-loader!./code_examples/export_entire_dataset_to_file_json.py';
 import CsvExample from '!!raw-loader!roa-loader!./code_examples/export_entire_dataset_to_file_csv.py';
 
-This example demonstrates how to use the <ApiLink to="class/BasicCrawler#export_data">`BasicCrawler.export_data`</ApiLink> method of the crawler to export the entire default dataset to a single file. This method supports exporting data in either CSV or JSON format and also accepts additional keyword arguments so you can fine-tune the underlying `json.dump` or `csv.writer` behavior.
+This example demonstrates how to use the <ApiLink to="class/BasicCrawler#export_data">`BasicCrawler.export_data`</ApiLink> method of the crawler to export the entire default dataset to a single file. This method supports exporting data in either CSV or JSON format and also accepts additional keyword arguments so you can fine-tune the underlying `json.dump` or `csv.DictWriter` behavior.
 
 :::note
 
@@ -31,3 +31,17 @@ For these examples, we are using the <ApiLink to="class/BeautifulSoupCrawler">`B
         </RunnableCodeBlock>
     </TabItem>
 </Tabs>
+
+## CSV columns
+
+Dataset items don't have to share a schema. Different handlers can push different fields, so the items of one dataset often have different keys.
+
+By default, the CSV columns are the keys of the first non-empty item. When a later item has a key the first one doesn't, its value isn't written, and Crawlee logs a warning naming the dropped keys. To use the keys of all items as columns instead, pass `collect_all_keys=True`:
+
+```python
+await crawler.export_data(path='results.csv', collect_all_keys=True)
+```
+
+No value is dropped in that mode. Collecting all keys means reading the whole dataset before the first row can be written, so only the default mode writes rows as it goes.
+
+In both modes, cells for columns an item doesn't have stay empty, or hold the `restval` value if you pass one.

src/crawlee/_types.py

@@ -737,7 +737,10 @@ class ExportToKwargs(TypedDict):
 
 
 class ExportDataJsonKwargs(TypedDict):
-    """Keyword arguments for dataset's `export_data_json` method."""
+    """Keyword arguments forwarded to `json.dump` by the JSON export.
+
+    Mirrors the keyword arguments of `json.dump`, so consult its documentation for their exact meaning.
+    """
 
     skipkeys: NotRequired[bool]
     """If True (default: False), dict keys that are not of a basic type (str, int, float, bool, None) will be skipped
@@ -773,7 +776,12 @@ class ExportDataJsonKwargs(TypedDict):
 
 
 class ExportDataCsvKwargs(TypedDict):
-    """Keyword arguments for dataset's `export_data_csv` method."""
+    """Keyword arguments forwarded to `csv.DictWriter` by the CSV export.
+
+    Mirrors the keyword arguments of `csv.DictWriter`, so consult its documentation for their exact meaning.
+    The `fieldnames` and `extrasaction` arguments are omitted, because the export derives the columns from the
+    exported items and manages how keys outside them are handled.
+    """
 
     dialect: NotRequired[str]
     """Specifies a dialect to be used in CSV parsing and writing."""
@@ -800,6 +808,9 @@ class ExportDataCsvKwargs(TypedDict):
     """Controls when quotes should be generated by the writer and recognized by the reader. Can take any of
     the `QUOTE_*` constants, with a default of `QUOTE_MINIMAL`."""
 
+    restval: NotRequired[str]
+    """The value written for columns an item does not contain. Defaults to an empty string."""
+
     skipinitialspace: NotRequired[bool]
     """When True, spaces immediately following the delimiter are ignored. Defaults to False."""
 

src/crawlee/_utils/file.py

@@ -6,6 +6,7 @@
 import os
 import sys
 import tempfile
+from logging import getLogger
 from pathlib import Path
 from typing import TYPE_CHECKING, overload
 
@@ -17,6 +18,8 @@
 
     from crawlee._types import ExportDataCsvKwargs, ExportDataJsonKwargs, JsonSerializable
 
+logger = getLogger(__name__)
+
 if sys.platform == 'win32':
 
     def _write_file(path: Path, data: str | bytes) -> None:
@@ -194,25 +197,59 @@ async def export_json_to_stream(
 async def export_csv_to_stream(
     iterator: AsyncIterator[Mapping[str, JsonSerializable]],
     dst: TextIO,
-    **kwargs: Unpack[ExportDataCsvKwargs],
+    *,
+    collect_all_keys: bool = False,
+    **writer_kwargs: Unpack[ExportDataCsvKwargs],
 ) -> None:
-    # Set lineterminator to '\n' if not explicitly provided. This prevents double line endings on Windows.
-    # The csv.writer default is '\r\n', which when written to a file in text mode on Windows gets converted
-    # to '\r\r\n' due to newline translation. By using '\n', we let the platform handle the line ending
-    # conversion: '\n' stays as '\n' on Unix, and becomes '\r\n' on Windows.
-    if 'lineterminator' not in kwargs:
-        kwargs['lineterminator'] = '\n'
-
-    writer = csv.writer(dst, **kwargs)
-    write_header = True
+    # Set lineterminator to '\n' if not explicitly provided. This prevents double line endings on Windows. The
+    # `csv.DictWriter` default is '\r\n', which when written to a file in text mode on Windows gets converted to
+    # '\r\r\n' due to newline translation. By using '\n', we let the platform handle the line ending conversion:
+    # '\n' stays as '\n' on Unix, and becomes '\r\n' on Windows.
+    if 'lineterminator' not in writer_kwargs:
+        writer_kwargs['lineterminator'] = '\n'
+
+    # The real writer cannot be built until the first item's keys are known, so construct a throwaway one up front
+    # to validate the options. Without this, an export configured with e.g. `delimiter='ab'` would raise only when
+    # the dataset happens to contain an item, and silently succeed on a run that scraped nothing.
+    csv.DictWriter(dst, fieldnames=[], extrasaction='ignore', **writer_kwargs)
+
+    if collect_all_keys:
+        items = [item async for item in iterator if item]
+        if not items:
+            return
+
+        fieldnames = list(dict.fromkeys(key for item in items for key in item))
+        writer = csv.DictWriter(dst, fieldnames=fieldnames, extrasaction='ignore', **writer_kwargs)
+        writer.writeheader()
+        for item in items:
+            writer.writerow(item)
+        return
+
+    writer = None
+    header_keys: set[str] = set()
+    dropped_keys: set[str] = set()
 
-    # Iterate over the dataset and write to CSV.
     async for item in iterator:
         if not item:
             continue
 
-        if write_header:
-            writer.writerow(item.keys())
-            write_header = False
-
-        writer.writerow(item.values())
+        if writer is None:
+            writer = csv.DictWriter(dst, fieldnames=list(item), extrasaction='ignore', **writer_kwargs)
+            writer.writeheader()
+            header_keys = set(writer.fieldnames)
+
+        dropped_keys.update(item.keys() - header_keys)
+        writer.writerow(item)
+
+    # The header comes from the first item, so any key introduced by a later item is dropped. Dropping them is
+    # intentional and matches Crawlee for JS. Doing it silently is not, so report it once for the whole export.
+    if dropped_keys:
+        logger.warning(
+            'CSV export dropped %d key(s) not present in the first item: %s. '
+            'Pass collect_all_keys=True to include keys from all items as columns.',
+            len(dropped_keys),
+            # The keys are annotated as `str`, but nothing enforces that at runtime, and a storage client that
+            # does not round-trip items through JSON hands them over as pushed. Stringify them so building the
+            # message cannot fail on a key that is not a string, or on a set mixing several key types.
+            ', '.join(sorted(str(key) for key in dropped_keys)),
+        )

src/crawlee/crawlers/_basic/_basic_crawler.py

@@ -913,6 +913,8 @@ async def export_data(
         dataset_id: str | None = None,
         dataset_name: str | None = None,
         dataset_alias: str | None = None,
+        *,
+        collect_all_keys: bool = False,
         **additional_kwargs: Unpack[ExportDataKwargs],
     ) -> None:
         """Export all items from a Dataset to a JSON or CSV file.
@@ -926,6 +928,9 @@ async def export_data(
             dataset_id: The ID of the Dataset to export from.
             dataset_name: The name of the Dataset to export from (global scope, named storage).
             dataset_alias: The alias of the Dataset to export from (run scope, unnamed storage).
+            collect_all_keys: CSV only. When True, the columns are the keys of all items combined, so no value is
+                dropped. When False, the columns are the keys of the first non-empty item, and keys introduced by
+                later items are dropped with a warning.
             additional_kwargs: Extra keyword arguments forwarded to the JSON/CSV exporter depending on the file format.
         """
         dataset = await Dataset.open(
@@ -941,7 +946,7 @@ async def export_data(
         if path.suffix == '.csv':
             dst = StringIO()
             csv_kwargs = cast('ExportDataCsvKwargs', additional_kwargs)
-            await export_csv_to_stream(dataset.iterate_items(), dst, **csv_kwargs)
+            await export_csv_to_stream(dataset.iterate_items(), dst, collect_all_keys=collect_all_keys, **csv_kwargs)
             await atomic_write(path, dst.getvalue())
         elif path.suffix == '.json':
             dst = StringIO()

src/crawlee/storages/_dataset.py

@@ -317,6 +317,8 @@ async def export_to(
         to_kvs_name: str | None = None,
         to_kvs_storage_client: StorageClient | None = None,
         to_kvs_configuration: Configuration | None = None,
+        *,
+        collect_all_keys: bool = False,
         **kwargs: Unpack[ExportDataCsvKwargs],
     ) -> None: ...
 
@@ -328,6 +330,8 @@ async def export_to(  # noqa: PLR0917
         to_kvs_name: str | None = None,
         to_kvs_storage_client: StorageClient | None = None,
         to_kvs_configuration: Configuration | None = None,
+        *,
+        collect_all_keys: bool = False,
         **kwargs: Any,
     ) -> None:
         """Export the entire dataset into a specified file stored under a key in a key-value store.
@@ -346,6 +350,9 @@ async def export_to(  # noqa: PLR0917
                 Specify only one of ID or name.
             to_kvs_storage_client: Storage client to use for the key-value store.
             to_kvs_configuration: Configuration for the key-value store.
+            collect_all_keys: CSV only. When True, the columns are the keys of all items combined, so no value is
+                dropped. When False, the columns are the keys of the first non-empty item, and keys introduced by
+                later items are dropped with a warning.
             kwargs: Additional parameters for the export operation, specific to the chosen content type.
         """
         kvs = await KeyValueStore.open(
@@ -357,7 +364,7 @@ async def export_to(  # noqa: PLR0917
         dst = StringIO()
 
         if content_type == 'csv':
-            await export_csv_to_stream(self.iterate_items(), dst, **kwargs)
+            await export_csv_to_stream(self.iterate_items(), dst, collect_all_keys=collect_all_keys, **kwargs)
             await kvs.set_value(key, dst.getvalue(), 'text/csv')
         elif content_type == 'json':
             await export_json_to_stream(self.iterate_items(), dst, **kwargs)