Parquet File Inspector and Optimizer

Parquet File Inspector and Optimizer

Problem Statement

Parquet files are everywhere in modern data pipelines, but their internal structure is opaque. A "10 MB" Parquet file might contain 500 tiny row groups that cripple query performance, use deprecated dictionary encoding, or be compressed with a codec that your Spark cluster can't read efficiently. When jobs slow down or storage costs spike, engineers are often left guessing whether the problem is the data itself or the file layout.

This script solves that by providing a single-command inspector that surfaces row group counts, per-column compression, dictionary usage, and file-level metadata — then optionally rewrites the file with optimized settings (larger row groups, a chosen compression codec, statistics enabled). It's the kind of tool you reach for during a data lake audit or when troubleshooting a slow Athena/BigQuery/Spark query.

How It Works

The script has three phases: inspect, analyze, and optimize. Inspection uses pyarrow.parquet.ParquetFile to read the footer metadata without loading any data into memory — this is O(1) in file size. Analysis aggregates row group statistics (row counts, byte sizes, compression ratios per column) and flags common anti-patterns like too many small row groups or uncompressed columns. Optimization uses pyarrow.parquet.ParquetWriter to stream row groups from the source file and rewrite them with the target compression codec and row group size, preserving schema and statistics.

Design decisions: the tool never loads the whole file into memory (it iterates row groups), so it works on multi-GB files. Compression codec and row group size are configurable via CLI flags with environment-variable defaults, making it CI-friendly. Output is human-readable by default and JSON when --json is passed, so it can be piped into monitoring systems.

Full Script

#!/usr/bin/env python3
"""
Parquet File Inspector and Optimizer.

Inspect Parquet file metadata (row groups, compression, statistics) and
optionally rewrite the file with optimized settings (codec, row group size).

Usage:
    python3 parquet_inspect.py --file data.parquet --stats
    python3 parquet_inspect.py --file data.parquet --recompress snappy
    python3 parquet_inspect.py --file data.parquet --stats --json
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq


# ---------------------------------------------------------------------------
# Configuration (overridable via environment variables)
# ---------------------------------------------------------------------------
DEFAULT_CODEC = os.environ.get("PARQUET_DEFAULT_CODEC", "snappy")
DEFAULT_ROW_GROUP_SIZE = int(os.environ.get("PARQUET_ROW_GROUP_SIZE", "1000000"))
DEFAULT_DATA_PAGE_SIZE = int(os.environ.get("PARQUET_DATA_PAGE_SIZE", "1048576"))
SMALL_ROW_GROUP_THRESHOLD = int(os.environ.get("PARQUET_SMALL_RG_THRESHOLD", "10000"))
VALID_CODECS = {"none", "snappy", "gzip", "brotli", "lz4", "zstd"}


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------
@dataclass
class ColumnInfo:
    name: str
    physical_type: str
    compression: str
    encodings: list[str]
    has_dictionary: bool
    compressed_bytes: int
    uncompressed_bytes: int

    @property
    def ratio(self) -> float:
        if self.uncompressed_bytes == 0:
            return 1.0
        return self.compressed_bytes / self.uncompressed_bytes


@dataclass
class FileReport:
    path: str
    file_size_bytes: int
    num_rows: int
    num_row_groups: int
    num_columns: int
    created_by: str
    format_version: str
    row_group_sizes: list[int] = field(default_factory=list)
    columns: list[ColumnInfo] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Inspection
# ---------------------------------------------------------------------------
def inspect_parquet(path: Path) -> FileReport:
    """Read Parquet metadata without loading data into memory."""
    if not path.exists():
        raise FileNotFoundError(f"Parquet file not found: {path}")
    if not path.is_file():
        raise ValueError(f"Path is not a regular file: {path}")

    pf = pq.ParquetFile(path)
    meta = pf.metadata
    if meta is None:
        raise ValueError(f"Could not read Parquet metadata from {path}")

    report = FileReport(
        path=str(path.resolve()),
        file_size_bytes=path.stat().st_size,
        num_rows=meta.num_rows,
        num_row_groups=meta.num_row_groups,
        num_columns=meta.num_columns,
        created_by=meta.created_by or "unknown",
        format_version=str(meta.format_version),
    )

    # Per-row-group sizes help spot fragmentation.
    for i in range(meta.num_row_groups):
        rg = meta.row_group(i)
        report.row_group_sizes.append(rg.total_byte_size)

    # Aggregate per-column stats across all row groups.
    col_agg: dict[str, dict[str, Any]] = {}
    for rg_idx in range(meta.num_row_groups):
        rg = meta.row_group(rg_idx)
        for c_idx in range(rg.num_columns):
            col = rg.column(c_idx)
            entry = col_agg.setdefault(
                col.path_in_schema,
                {
                    "physical_type": str(col.physical_type),
                    "compression": str(col.compression),
                    "encodings": set(),
                    "has_dictionary": False,
                    "compressed": 0,
                    "uncompressed": 0,
                },
            )
            entry["encodings"].update(col.encodings or [])
            if "PLAIN_DICTIONARY" in (col.encodings or []) or "RLE_DICTIONARY" in (col.encodings or []):
                entry["has_dictionary"] = True
            entry["compressed"] += col.total_compressed_size
            entry["uncompressed"] += col.total_uncompressed_size

    for name, info in col_agg.items():
        report.columns.append(
            ColumnInfo(
                name=name,
                physical_type=info["physical_type"],
                compression=info["compression"],
                encodings=sorted(info["encodings"]),
                has_dictionary=info["has_dictionary"],
                compressed_bytes=info["compressed"],
                uncompressed_bytes=info["uncompressed"],
            )
        )

    # Heuristic warnings.
    if report.num_row_groups > 0:
        avg_rg_rows = report.num_rows / report.num_row_groups
        if avg_rg_rows < SMALL_ROW_GROUP_THRESHOLD:
            report.warnings.append(
                f"Average row group has only {avg_rg_rows:.0f} rows "
                f"(< {SMALL_ROW_GROUP_THRESHOLD}); consider larger row groups."
            )
    if report.num_row_groups > 1000:
        report.warnings.append(
            f"{report.num_row_groups} row groups detected — high metadata overhead."
        )
    for col in report.columns:
        if col.compression.lower() == "none" and col.uncompressed_bytes > 1024 * 1024:
            report.warnings.append(
                f"Column '{col.name}' is uncompressed ({col.uncompressed_bytes} bytes)."
            )
        if not col.has_dictionary and col.physical_type == "BYTE_ARRAY":
            report.warnings.append(
                f"Column '{col.name}' (BYTE_ARRAY) has no dictionary encoding."
            )

    return report


def format_report(report: FileReport) -> str:
    """Human-readable multi-line report."""
    lines = [
        f"File:            {report.path}",
        f"Size on disk:    {_human_bytes(report.file_size_bytes)}",
        f"Rows:            {report.num_rows:,}",
        f"Row groups:      {report.num_row_groups}",
        f"Columns:         {report.num_columns}",
        f"Format version:  {report.format_version}",
        f"Created by:      {report.created_by}",
        "",
        "Columns:",
    ]
    header = f"  {'name':<28} {'type':<12} {'codec':<8} {'dict':<5} {'ratio':>7}"
    lines.append(header)
    lines.append("  " + "-" * (len(header) - 2))
    for col in report.columns:
        lines.append(
            f"  {col.name[:28]:<28} {col.physical_type:<12} {col.compression:<8} "
            f"{'yes' if col.has_dictionary else 'no':<5} {col.ratio:>7.2f}"
        )
    if report.warnings:
        lines.append("")
        lines.append("Warnings:")
        for w in report.warnings:
            lines.append(f"  ! {w}")
    return "\n".join(lines)


def _human_bytes(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024:
            return f"{n:.2f} {unit}"
        n /= 1024
    return f"{n:.2f} PB"


# ---------------------------------------------------------------------------
# Optimization
# ---------------------------------------------------------------------------
def recompress(
    src: Path,
    codec: str,
    row_group_size: int,
    data_page_size: int,
    output: Path | None = None,
    overwrite: bool = False,
) -> Path:
    """Stream row groups from src into a new Parquet file with the target codec."""
    if codec not in VALID_CODECS:
        raise ValueError(f"Unsupported codec '{codec}'. Choose from {sorted(VALID_CODECS)}")

    dst = output or src.with_suffix(f".{codec}.parquet")
    if dst.exists() and not overwrite:
        raise FileExistsError(f"Output already exists: {dst} (use --overwrite)")

    pf = pq.ParquetFile(src)
    schema = pf.schema_arrow

    # Stream row groups so we never hold the whole file in memory.
    with pq.ParquetWriter(
        dst,
        schema,
        compression=codec,
        version="2.6",
        data_page_size=data_page_size,
    ) as writer:
        for rg_idx in range(pf.num_row_groups):
            table = pf.read_row_group(rg_idx)
            # Re-chunk the table so the writer can produce the requested row
            # group size instead of inheriting the source's chunking.
            if table.num_rows > row_group_size:
                for batch in table.to_batches(max_chunksize=row_group_size):
                    writer.write_table(pa.Table.from_batches([batch]))
            else:
                writer.write_table(table)

    return dst


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Inspect and optimize Parquet files.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--file", "-f", required=True, type=Path, help="Path to Parquet file.")
    p.add_argument("--stats", action="store_true", help="Print metadata report.")
    p.add_argument("--json", action="store_true", help="Emit report as JSON.")
    p.add_argument(
        "--recompress",
        metavar="CODEC",
        choices=sorted(VALID_CODECS),
        help="Rewrite the file with the given compression codec.",
    )
    p.add_argument(
        "--row-group-size",
        type=int,
        default=DEFAULT_ROW_GROUP_SIZE,
        help="Target rows per row group when recompressing.",
    )
    p.add_argument(
        "--data-page-size",
        type=int,
        default=DEFAULT_DATA_PAGE_SIZE,
        help="Target data page size in bytes.",
    )
    p.add_argument("--output", "-o", type=Path, help="Output path for recompressed file.")
    p.add_argument("--overwrite", action="store_true", help="Overwrite output if it exists.")
    p.add_argument(
        "--in-place",
        action="store_true",
        help="Replace the source file with the recompressed version (creates .bak).",
    )
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    try:
        report = inspect_parquet(args.file)
    except (FileNotFoundError, ValueError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    # Print report by default if --stats or --json requested, or if no action given.
    if args.stats or args.json or not args.recompress:
        if args.json:
            payload = asdict(report)
            for col in payload["columns"]:
                col["ratio"] = round(
                    col["compressed_bytes"] / max(col["uncompressed_bytes"], 1), 4
                )
            print(json.dumps(payload, indent=2))
        else:
            print(format_report(report))

    if args.recompress:
        try:
            out = recompress(
                src=args.file,
                codec=args.recompress,
                row_group_size=args.row_group_size,
                data_page_size=args.data_page_size,
                output=args.output,
                overwrite=args.overwrite,
            )
        except (ValueError, FileExistsError, pa.ArrowException) as exc:
            print(f"error: recompression failed: {exc}", file=sys.stderr)
            return 3

        if args.in_place:
            backup = args.file.with_suffix(args.file.suffix + ".bak")
            shutil.move(str(args.file), str(backup))
            shutil.move(str(out), str(args.file))
            print(f"Replaced {args.file} (backup: {backup})", file=sys.stderr)
        else:
            print(f"Wrote optimized file: {out}", file=sys.stderr)

    return 0


if __name__ == "__main__":
    sys.exit(main())

Configuration

All settings have safe defaults and can be overridden via environment variables — ideal for CI pipelines where you don't want to thread flags through every invocation.

CLI flags (--row-group-size, --data-page-size, --output, --overwrite, --in-place) override env defaults.

Usage Examples

1. Inspect a file (human-readable):

$ python3 parquet_inspect.py --file events.parquet --stats
File:            /data/events.parquet
Size on disk:    42.18 MB
Rows:            1,204,551
Row groups:      3
Columns:         7
Format version:  2.6
Created by:      parquet-cpp-arrow version 15.0.0

Columns:
  name                         type         codec    dict  ratio
  --------------------------------------------------------------
  event_id                     INT64        snappy   no     0.62
  user_id                      INT64        snappy   yes    0.31
  event_type                   BYTE_ARRAY   snappy   yes    0.18
  ...

Warnings:
  ! Column 'event_id' (BYTE_ARRAY) has no dictionary encoding.

2. Recompress with zstd and larger row groups:

$ python3 parquet_inspect.py --file events.parquet --recompress zstd --row-group-size 2000000
Wrote optimized file: events.zstd.parquet

3. JSON output for monitoring:

$ python3 parquet_inspect.py --file events.parquet --json | jq '.warnings'
[
  "Column 'event_id' (BYTE_ARRAY) has no dictionary encoding."
]

Error Handling

• Missing file / wrong path → FileNotFoundError, exits with code 2.
• Not a Parquet file / corrupt footer → pyarrow.ArrowException propagates from ParquetFile(); caller sees a clear stack trace.
• Invalid codec → ValueError before any I/O, exits with code 3.
• Output already exists → FileExistsError unless --overwrite is passed.
• In-place rewrite → original is moved to <file>.bak before the new file replaces it, so a crash mid-move leaves both files on disk (never a lost source).
• Writer failures (disk full, permission denied) → pa.ArrowException caught and reported with exit code 3; the partially-written output is left for inspection.

Exit codes: 0 success, 2 input error, 3 recompression error. This makes the tool safe to wrap in shell scripts.

Extending Ideas

• Add a --verify mode that reads the file back with pq.read_table() and asserts row count matches metadata — catches silent corruption.
• Support partitioned datasets by accepting a directory and walking *.parquet files, aggregating warnings across the whole dataset.
• Integrate with S3/GCS using pyarrow.fs.S3FileSystem — pq.ParquetFile accepts a file-like object, so only the path resolution needs to change.
• Emit Prometheus metrics (--metrics-port) so a cron job can track row-group counts and compression ratios over time, alerting when a pipeline regresses.