analyze_parquet_distribution.py

#!/usr/bin/env python3
"""
Analyze parquet dataset distribution and ordering.
"""

import argparse
import json
import pandas as pd
import numpy as np
from pathlib import Path

def analyze_parquet_distribution(parquet_path: str, num_samples: int = None, osl_column: str = None):
    """Analyze the distribution and ordering of samples in parquet file.

    Args:
        parquet_path: Path to the parquet file
        num_samples: Number of samples to analyze (None = all samples)
        osl_column: Column name for output sequence length (optional)
    """

    print("=" * 80)
    print(f"Analyzing: {parquet_path}")
    print("=" * 80)

    # Load the parquet file
    df_full = pd.read_parquet(parquet_path)

    # Optionally limit to first N samples
    if num_samples is not None and num_samples < len(df_full):
        df = df_full.head(num_samples)
        analyzing_subset = True
    else:
        df = df_full
        analyzing_subset = False

    print(f"\n📊 BASIC INFO")
    if analyzing_subset:
        print(f"{'Total samples in file:':<30} {len(df_full):,}")
        print(f"{'Analyzing first:':<30} {len(df):,} samples")
    else:
        print(f"{'Total samples:':<30} {len(df):,}")
    print(f"{'Columns:':<30} {list(df.columns)}")
    if 'dataset' in df.columns:
        print(f"{'Datasets:':<30} {df['dataset'].unique().tolist()}")

    # Detect which token column exists
    if 'num_tokens' in df.columns:
        token_col = 'num_tokens'
    elif 'tok_input_len' in df.columns:
        token_col = 'tok_input_len'
    else:
        print(f"\n❌ Error: No token length column found!")
        print(f"   Expected 'num_tokens' or 'tok_input_len' but found: {list(df.columns)}")
        return None

    print(f"{'Token column (ISL):':<30} {token_col}")

    # Detect or use provided OSL column
    osl_col = None
    if osl_column:
        if osl_column in df.columns:
            osl_col = osl_column
        else:
            print(f"\n⚠️  Warning: OSL column '{osl_column}' not found in dataset")
            print(f"   Available columns: {list(df.columns)}")
    else:
        # Auto-detect common OSL column names
        if 'tok_ref_output_len' in df.columns:
            osl_col = 'tok_ref_output_len'

    if osl_col:
        print(f"{'Token column (OSL):':<30} {osl_col}")

    # Analyze token distribution
    subset_label = f" (First {len(df):,})" if analyzing_subset else ""
    total = len(df)

    print(f"\n📈 ISL (INPUT SEQUENCE LENGTH) DISTRIBUTION{subset_label}")
    print(f"\n{df[token_col].describe()}")

    # Analyze OSL distribution if available
    if osl_col:
        print(f"\n📈 OSL (OUTPUT SEQUENCE LENGTH) DISTRIBUTION{subset_label}")
        print(f"\n{df[osl_col].describe()}")

        # OSL distribution by ranges
        print(f"\n📊 OSL DISTRIBUTION BY RANGES{subset_label}")
        osl_ranges = [
            (0, 1000, "0-1K"),
            (1000, 2000, "1K-2K"),
            (2000, 4000, "2K-4K"),
            (4000, 8000, "4K-8K"),
            (8000, 16000, "8K-16K"),
            (16000, 32000, "16K-32K"),
            (32000, float('inf'), "32K+")
        ]

        for min_tok, max_tok, label in osl_ranges:
            count = len(df[(df[osl_col] >= min_tok) & (df[osl_col] < max_tok)])
            pct = (count / total) * 100 if total > 0 else 0
            print(f"  {label:<10} {count:>6,} samples ({pct:>5.1f}%)")

    # ISL Distribution by ranges
    print(f"\n📊 ISL DISTRIBUTION BY RANGES{subset_label}")
    ranges = [
        (0, 1000, "0-1K"),
        (1000, 2000, "1K-2K"),
        (2000, 4000, "2K-4K"),
        (4000, 6000, "4K-6K"),
        (6000, 8000, "6K-8K"),
        (8000, 10000, "8K-10K"),
        (10000, float('inf'), "10K+")
    ]

    for min_tok, max_tok, label in ranges:
        count = len(df[(df[token_col] >= min_tok) & (df[token_col] < max_tok)])
        pct = (count / total) * 100
        print(f"  {label:<10} {count:>6,} samples ({pct:>5.1f}%)")

    # Check if samples are ordered
    print(f"\n🔍 ORDERING ANALYSIS")

    # Check if sorted by token column
    is_sorted_asc = df[token_col].is_monotonic_increasing
    is_sorted_desc = df[token_col].is_monotonic_decreasing

    print(f"{'Sorted ascending (by ' + token_col + '):':<35} {is_sorted_asc}")
    print(f"{'Sorted descending (by ' + token_col + '):':<35} {is_sorted_desc}")

    # Show first 10 and last 10 token counts
    print(f"\n📋 FIRST 10 SAMPLES (token counts):")
    print(df[token_col].head(10).tolist())

    last_label = f"LAST 10 of {len(df):,} SAMPLES" if analyzing_subset else "LAST 10 SAMPLES"
    print(f"\n📋 {last_label} (token counts):")
    print(df[token_col].tail(10).tolist())

    # Calculate differences to see if there's a pattern
    token_diffs = df[token_col].diff().dropna()

    print(f"\n📊 TOKEN LENGTH CHANGES")
    print(f"{'Always increasing:':<35} {(token_diffs > 0).all()}")
    print(f"{'Always decreasing:':<35} {(token_diffs < 0).all()}")
    print(f"{'Mixed (random-ish):':<35} {not ((token_diffs > 0).all() or (token_diffs < 0).all())}")

    # Show some statistics about the differences
    if not is_sorted_asc and not is_sorted_desc:
        print(f"\n{'Mean absolute difference:':<35} {token_diffs.abs().mean():.2f}")
        print(f"{'Std of differences:':<35} {token_diffs.std():.2f}")

    # Show a sample of data
    sample_label = f"first 3 of {len(df):,}" if analyzing_subset else "first 3"
    print(f"\n📄 SAMPLE RECORDS ({sample_label}):")
    print("=" * 80)
    for idx in range(min(3, len(df))):
        row = df.iloc[idx]
        print(f"\nSample {idx + 1}:")
        if 'dataset' in row:
            print(f"  Dataset:      {row['dataset']}")
        print(f"  Input tokens: {row[token_col]}")
        if 'prompt' in row:
            print(f"  Prompt:       {str(row['prompt'])[:100]}...")
        if 'text_input' in row:
            print(f"  Text input:   {str(row['text_input'])[:150]}...")

    # Show distribution visualization (text-based)
    hist_label = f"histogram - {subset_label.strip()}" if analyzing_subset else "histogram"
    print(f"\n📊 VISUAL DISTRIBUTION ({hist_label})")
    print("=" * 80)

    # Create histogram bins
    bins = [0, 1000, 2000, 4000, 6000, 8000, 10000, 16000]
    hist, bin_edges = np.histogram(df[token_col], bins=bins)

    max_count = hist.max()
    bar_width = 50

    for i, count in enumerate(hist):
        label = f"{int(bin_edges[i])}-{int(bin_edges[i+1])}"
        bar_len = int((count / max_count) * bar_width) if max_count > 0 else 0
        bar = "█" * bar_len
        print(f"{label:<12} {bar} {count:>6,}")

    print("=" * 80)

    # Check for specific token length patterns
    print(f"\n🎯 SPECIFIC TOKEN LENGTH TARGETS{subset_label}")
    targets = [512, 1000, 2048, 4096, 8192]
    for target in targets:
        exact = len(df[df[token_col] == target])
        within_50 = len(df[(df[token_col] >= target - 50) & (df[token_col] <= target + 50)])
        print(f"  ~{target} tokens: {exact:>4} exact, {within_50:>4} within ±50 tokens")

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Analyze parquet dataset distribution and ordering",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Analyze a parquet file
  %(prog)s dataset.parquet

  # Analyze only first 500 samples
  %(prog)s dataset.parquet --samples 500

  # Analyze with OSL (output sequence length) column
  %(prog)s dataset.parquet --osl-column tok_ref_output_len

  # Analyze first 1000 samples and save to JSON
  %(prog)s /path/to/perf_eval_ref.parquet --samples 1000 --output-json output.json

  # With full paths
  %(prog)s /Users/user/rhoai-install-stuff/real_datasets/gpt-oss/perf/perf_eval_ref.parquet --output-json /tmp/dataset.json
        """
    )

    parser.add_argument(
        "parquet_path",
        help="Path to the parquet file to analyze"
    )

    parser.add_argument(
        "--samples",
        type=int,
        metavar="N",
        help="Number of samples to analyze (default: all samples)"
    )

    parser.add_argument(
        "--output-json",
        metavar="FILE",
        help="Save the dataset to JSON format at the specified path"
    )

    parser.add_argument(
        "--osl-column",
        metavar="COLUMN",
        help="Column name for output sequence length (e.g., tok_ref_output_len)"
    )

    args = parser.parse_args()
    parquet_path = args.parquet_path

    if not Path(parquet_path).exists():
        print(f"❌ File not found: {parquet_path}")
        exit(1)

    df = analyze_parquet_distribution(parquet_path, num_samples=args.samples, osl_column=args.osl_column)

    print(f"\n✅ Analysis complete!")
    analyzed_label = f" (analyzed {len(df):,} of total)" if args.samples else ""
    print(f"\nDataFrame shape{analyzed_label}: {df.shape}")
    print(f"Memory usage: {df.memory_usage(deep=True).sum() / 1024**2:.2f} MB")

    # Save to JSON if requested
    if args.output_json:
        output_path = Path(args.output_json)
        print(f"\n💾 Saving dataset to JSON: {output_path}")

        # Use pandas to_json which handles numpy types properly
        df.to_json(output_path, orient='records', indent=2)

        file_size_mb = output_path.stat().st_size / 1024**2
        print(f"✅ Saved {len(df):,} records to {output_path}")
        print(f"   File size: {file_size_mb:.2f} MB")