modosaic.providers.adapters.parquet

modosaic.providers.adapters.parquet

ParquetAdapter

ParquetAdapter(parquet_path, image_column='image', id_column=None, extension_column=None, metadata_columns=None, batch_size=512)

Bases: DatasetAdapter

Dataset adapter that streams images from parquet rows.

Initialize the adapter.

Parameters:

Name	Type	Description	Default
parquet_path	str | Path	Parquet file or directory containing parquet files.	required
image_column	str	Column or nested path containing image bytes or image paths.	'image'
id_column	str | None	Optional column or nested path used as the sample ID.	None
extension_column	str | None	Optional column or nested path containing file extensions.	None
metadata_columns	Sequence[str] | None	Optional column paths copied into record metadata.	None
batch_size	int	Number of rows read from parquet at a time.	512

Source code in modosaic/providers/adapters/parquet.py

def __init__(
        self,
        parquet_path: str | Path,
        image_column: str = "image",
        id_column: str | None = None,
        extension_column: str | None = None,
        metadata_columns: Sequence[str] | None = None,
        batch_size: int = 512,
):
    """Initialize the adapter.

    Args:
        parquet_path: Parquet file or directory containing parquet files.
        image_column: Column or nested path containing image bytes or image
            paths.
        id_column: Optional column or nested path used as the sample ID.
        extension_column: Optional column or nested path containing file
            extensions.
        metadata_columns: Optional column paths copied into record metadata.
        batch_size: Number of rows read from parquet at a time.
    """
    self.parquet_path = Path(parquet_path)
    self.image_column = image_column
    self.id_column = id_column
    self.extension_column = extension_column
    self.metadata_columns = tuple(metadata_columns or ())
    self.batch_size = batch_size

iter_samples

iter_samples()

Yield image records from parquet rows.

Returns:

Type	Description
Iterator[ImageRecord]	Iterator of image records.

Source code in modosaic/providers/adapters/parquet.py

@override
def iter_samples(self) -> Iterator[ImageRecord]:
    """Yield image records from parquet rows.

    Returns:
        Iterator of image records.
    """
    read_columns = {self._base_column(self.image_column)}
    if self.id_column:
        read_columns.add(self._base_column(self.id_column))
    if self.extension_column:
        read_columns.add(self._base_column(self.extension_column))
    for metadata_column in self.metadata_columns:
        read_columns.add(self._base_column(metadata_column))

    seen_rows = 0
    for parquet_path in self._resolve_parquet_files():
        logger.debug(f"Reading parquet samples from {parquet_path}")
        parquet_file = pq.ParquetFile(parquet_path)
        for batch in parquet_file.iter_batches(batch_size=self.batch_size, columns=sorted(read_columns)):
            rows = batch.to_pydict()
            for idx in range(batch.num_rows):
                image_value = self._value_from_column(rows, self.image_column, idx)
                id_value = self._value_from_column(rows, self.id_column, idx) if self.id_column else None
                extension_value = (
                    self._value_from_column(rows, self.extension_column, idx)
                    if self.extension_column
                    else None
                )

                image_bytes, inferred_extension, source_path = self._coerce_image_value(
                    image_value,
                    parquet_source=parquet_path,
                )

                sample_id = (
                    str(id_value)
                    if self.id_column and id_value is not None
                    else f"row_{seen_rows}"
                )

                extension = (
                    str(extension_value).strip()
                    if self.extension_column and extension_value
                    else inferred_extension
                )
                extension = ExtensionService.normalize_extension(extension, fallback=".bin")

                metadata = {col: self._value_from_column(rows, col, idx) for col in self.metadata_columns}
                metadata["source_parquet"] = str(parquet_path)
                if source_path:
                    metadata["resolved_image_path"] = source_path

                logger.debug(f"Loading parquet image sample {sample_id} from {parquet_path}")
                yield ImageRecord(
                    sample_id=sample_id,
                    image_bytes=image_bytes,
                    extension=extension,
                    metadata=metadata,
                )
                seen_rows += 1