Table Writes¶

def write_table(
    self,
    data: pl.DataFrame | pl.LazyFrame,
    location: str | Path,
    metadata: dict[str, str] | None = None,
    row_group_size: int | None = None,
) -> None:
    """Write one parquet table and fail if the output already exists.

    Args:
        data: Polars dataframe or lazy frame to write.
        location: Relative or absolute output table location.
        metadata: Optional parquet key-value footer metadata.
        row_group_size: Optional parquet row-group size.

    Raises:
        FileExistsError: If `location` already exists.

    Examples:
        Write a prepared scratch table during ingestion:

        >>> store.write_table(data, temp_path, row_group_size=config.row_group_size)

        Write a manifest with parquet footer metadata:

        >>> store.write_table(manifest, "manifest.parquet", metadata={"stage": "ingested"})
    """
    out_path = Path(self.resolve(location))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    if out_path.exists():
        raise FileExistsError(f"File exists: {out_path}")
    if isinstance(data, pl.LazyFrame):
        data.sink_parquet(out_path, metadata=metadata, row_group_size=row_group_size)
        return
    data.write_parquet(out_path, metadata=metadata, row_group_size=row_group_size)

def open_table_writer(
    self,
    location: str | Path,
    schema: pa.Schema,
    compression: str,
    *,
    use_content_defined_chunking: bool = False,
) -> TableWriter:
    """Open a parquet writer for incrementally writing one table.

    Args:
        location: Relative or absolute output table location.
        schema: Arrow schema for all chunks written through the returned writer.
        compression: Parquet compression codec passed to PyArrow.
        use_content_defined_chunking: Whether PyArrow should use content-defined
            chunking when writing parquet data.

    Returns:
        Table writer that keeps the parquet file open until closed.

    Raises:
        FileExistsError: If `location` already exists.

    Examples:
        Write bounded normalization output as chunks become available:

        >>> writer = store.open_table_writer(temp_path, chunk.to_arrow().schema, "zstd")
        >>> writer.write_table(chunk, row_group_size=config.row_group_size)
        >>> writer.close()
    """
    return LocalTableWriter(
        Path(self.resolve(location)),
        schema,
        compression,
        use_content_defined_chunking=use_content_defined_chunking,
    )

def table_size_bytes(self, location: str | Path) -> int | None:
    """Return the table file size in bytes.

    This is used by sharding to decide when an open shard should roll over to
    a new file.

    Args:
        location: Relative or absolute table location.

    Returns:
        File size in bytes, or `None` if the file does not exist.
    """
    path = Path(self.resolve(location))
    if not path.exists():
        return None
    return path.stat().st_size