tool_nan_inf_check.py

import os
import polars as pl
import numpy as np

# Configuration
OUTPUT_BASE = "datasets/refined/unified_2_final"
INTERVALS = ['5m', '15m', '30m', '1h', '2h', '4h', '6h', '8h', '12h']
FINAL_OUTPUT = "unified_training_dataset_softclipped.parquet"


def check_parquet_file(file_path: str) -> dict:
    """Check a single parquet file for NaNs, Infs, and nulls."""
    try:
        df = pl.read_parquet(file_path)

        # Get basic info
        total_rows = df.height
        total_cols = len(df.columns)
        total_cells = total_rows * total_cols

        # Check for nulls
        null_counts = df.null_count()
        total_nulls = sum(null_counts.row(0))

        # Check for NaNs and Infs in numeric columns
        numeric_cols = [col for col, dtype in zip(df.columns, df.dtypes)
                        if dtype in [pl.Float32, pl.Float64, pl.Int32, pl.Int64]]

        nan_counts = {}
        inf_counts = {}
        total_nans = 0
        total_infs = 0

        for col in numeric_cols:
            # Check for NaNs
            nan_count = df.select(pl.col(col).is_nan().sum()).item()
            if nan_count > 0:
                nan_counts[col] = nan_count
                total_nans += nan_count

            # Check for Infs
            inf_count = df.select(pl.col(col).is_infinite().sum()).item()
            if inf_count > 0:
                inf_counts[col] = inf_count
                total_infs += inf_count

        return {
            'file_path': file_path,
            'total_rows': total_rows,
            'total_cols': total_cols,
            'total_cells': total_cells,
            'total_nulls': total_nulls,
            'total_nans': total_nans,
            'total_infs': total_infs,
            'null_counts': null_counts,
            'nan_counts': nan_counts,
            'inf_counts': inf_counts,
            'has_issues': total_nulls > 0 or total_nans > 0 or total_infs > 0
        }

    except Exception as e:
        return {
            'file_path': file_path,
            'error': str(e),
            'has_issues': True
        }


def main():
    print("🔍 Checking parquet files for NaNs, Infs, and missing values...\n")

    all_results = []
    issues_found = False

    for interval in INTERVALS:
        file_path = os.path.join(OUTPUT_BASE, interval, FINAL_OUTPUT)

        if not os.path.exists(file_path):
            print(f"⚠️  {interval}: File not found - {file_path}")
            continue

        print(f"📊 Checking {interval}...")
        result = check_parquet_file(file_path)
        all_results.append(result)

        if 'error' in result:
            print(f"❌ {interval}: ERROR - {result['error']}")
            issues_found = True
        elif result['has_issues']:
            print(f"❌ {interval}: ISSUES FOUND")
            print(f"   • Nulls: {result['total_nulls']:,}")
            print(f"   • NaNs: {result['total_nans']:,}")
            print(f"   • Infs: {result['total_infs']:,}")

            # Show problematic columns
            if result['nan_counts']:
                print(f"   • NaN columns: {list(result['nan_counts'].keys())}")
            if result['inf_counts']:
                print(f"   • Inf columns: {list(result['inf_counts'].keys())}")

            issues_found = True
        else:
            print(f"✅ {interval}: Clean ({result['total_rows']:,} rows, {result['total_cols']} cols)")

    # Summary
    print("\n" + "=" * 50)
    if issues_found:
        print("❌ ISSUES FOUND in one or more files!")
        print("\nDetailed breakdown:")
        for result in all_results:
            if result.get('has_issues', False):
                print(f"\n📁 {os.path.basename(os.path.dirname(result['file_path']))}:")
                if 'error' in result:
                    print(f"   ERROR: {result['error']}")
                else:
                    print(f"   Total cells: {result['total_cells']:,}")
                    print(f"   Nulls: {result['total_nulls']:,}")
                    print(f"   NaNs: {result['total_nans']:,}")
                    print(f"   Infs: {result['total_infs']:,}")
    else:
        print("✅ ALL FILES ARE CLEAN!")
        total_files = len([r for r in all_results if 'error' not in r])
        print(f"Checked {total_files} files successfully.")


if __name__ == "__main__":
    main()