Coleta de Dados

A coleta de dados é o primeiro passo físico em nosso pipeline automatizado de ML. Na produção, esse passo extrai registros de registros de viagem brutos de lagos de dados externos ou lojas de fluxo.

Fonte de dados

Recuperamos registros de viagem de táxi amarelo publicados pela Comissão de Táxi e Limousine de Nova York (TLC). Esses registros são salvos como arquivos mensais em parquet contendo campos como:

  • tpep_pickup_datetime: O carimbo de data/hora em que o medidor estava envolvido.
  • tpep_dropoff_datetime: O carimbo de data/hora em que o medidor foi desvinculado.
  • PULocationID: TLC Taxi Zone ID onde o passageiro foi pego.
  • trip_distance: A distância de viagem decorrido em milhas.

Cache Raw Local

Nosso coletor baixa arquivos do endpoint da nuvem e os salva em um diretório local bruto: data/raw/yellow_tripdata_YYYY-MM.parquet

Se um arquivo já estiver presente localmente, o coletor pulará o download para economizar largura de banda e garantir execuções idempotentes e repetíveis.

Nota

Para forcar um novo download, exclua o arquivo parquet de destino em data/raw/ e execute o coletor novamente.

Execução CLI

Para coletar os conjuntos de dados de forma programática:

make collect PROJECT=nyc_taxi_demand_forecasting

Implementação Python

O código de download está estruturado como uma classe de coletor de responsabilidade única isolada:

Sourcesrc/nyc_taxi_demand_forecasting/data/collection.py
import logging
from pathlib import Path

import requests

from nyc_taxi_demand_forecasting.configuration import CollectionConfig

# Parquet frames its payload with a magic word at both ends: a missing
# trailer is exactly what a crashed write or a saved error page looks like.
_PARQUET_MAGIC = b"PAR1"
_MINIMUM_PARQUET_SIZE = 2 * len(_PARQUET_MAGIC)


class TlcYellowTaxiParquetCollector:
    """Collect immutable TLC parquet files into local raw storage.

    Example:
        collector.collect(config, Path("data/raw"))
    """

    def collect(
        self, config: CollectionConfig, output_directory: Path
    ) -> tuple[Path, ...]:
        output_directory.mkdir(parents=True, exist_ok=True)
        return tuple(
            self._download_month(config, month, output_directory)
            for month in config.months
        )

    def _download_month(
        self, config: CollectionConfig, month: int, output_directory: Path
    ) -> Path:
        output_path = output_directory / config.file_name(month)
        if output_path.exists():
            if self._is_cached_parquet(output_path):
                return output_path
            logging.getLogger(__name__).warning(
                "corrupt_cached_parquet_redownload",
                extra={"path": str(output_path)},
            )

        url = config.source_url(month)
        response = requests.get(url, timeout=60)
        response.raise_for_status()
        content = response.content
        if not self._is_parquet_content(content):
            raise ValueError(
                f"Invalid download from {url}: expected parquet "
                f"(PAR1 magic), got {content[:8]!r}"
            )
        self._write_atomically(output_path, content)
        return output_path

    def _is_cached_parquet(self, path: Path) -> bool:
        if not path.is_file() or path.stat().st_size < _MINIMUM_PARQUET_SIZE:
            return False
        with path.open("rb") as file:
            header = file.read(len(_PARQUET_MAGIC))
            file.seek(-len(_PARQUET_MAGIC), 2)
            trailer = file.read(len(_PARQUET_MAGIC))
        return header == _PARQUET_MAGIC and trailer == _PARQUET_MAGIC

    def _is_parquet_content(self, content: bytes) -> bool:
        return (
            len(content) >= _MINIMUM_PARQUET_SIZE
            and content[: len(_PARQUET_MAGIC)] == _PARQUET_MAGIC
            and content[-len(_PARQUET_MAGIC) :] == _PARQUET_MAGIC
        )

    def _write_atomically(self, output_path: Path, content: bytes) -> None:
        temporary_path = output_path.with_suffix(".tmp")
        temporary_path.write_bytes(content)
        temporary_path.replace(output_path)