foxhunt/ml/examples/create_small_parquet_files.rs

//! Create Small Test Parquet Files (1000 bars each)
//!
//! Agent-2: Extract first 1000 bars from existing Parquet files for fast testing.
//! This creates lightweight test files for rapid ML model validation.

use anyhow::{Context, Result};
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::ArrowWriter;
use parquet::file::properties::WriterProperties;
use std::fs::File;
use std::path::Path;
use std::sync::Arc;

struct FileMapping {
    input: &'static str,
    output: &'static str,
}

const FILE_MAPPINGS: &[FileMapping] = &[
    FileMapping {
        input: "test_data/ES_FUT_180d.parquet",
        output: "test_data/ES_FUT_small.parquet",
    },
    FileMapping {
        input: "test_data/NQ_FUT_180d.parquet",
        output: "test_data/NQ_FUT_small.parquet",
    },
    FileMapping {
        input: "test_data/6E_FUT_180d.parquet",
        output: "test_data/6E_FUT_small.parquet",
    },
    FileMapping {
        input: "test_data/ZN_FUT_90d.parquet",
        output: "test_data/ZN_FUT_small.parquet",
    },
];

fn create_small_parquet(
    input_path: &str,
    output_path: &str,
    num_rows: usize,
) -> Result<(usize, f64)> {
    println!("Processing {}...", input_path);

    // Check if input file exists
    if !Path::new(input_path).exists() {
        println!("  ⚠️  File not found, skipping...");
        println!();
        return Ok((0, 0.0));
    }

    // Open input Parquet file
    let input_file =
        File::open(input_path).with_context(|| format!("Failed to open {}", input_path))?;

    let builder = ParquetRecordBatchReaderBuilder::try_new(input_file)?;
    let original_rows = builder.metadata().file_metadata().num_rows() as usize;
    println!("  Original rows: {}", original_rows);

    // Create reader
    let mut reader = builder.build()?;

    // Read first batch and extract first 1000 rows
    let mut total_rows_extracted = 0;
    let mut batches_to_write = Vec::new();

    while let Some(Ok(batch)) = reader.next() {
        let rows_needed = num_rows.saturating_sub(total_rows_extracted);
        if rows_needed == 0 {
            break;
        }

        let rows_to_take = batch.num_rows().min(rows_needed);
        let sliced_batch = batch.slice(0, rows_to_take);
        batches_to_write.push(sliced_batch);
        total_rows_extracted += rows_to_take;

        if total_rows_extracted >= num_rows {
            break;
        }
    }

    println!("  Extracted rows: {}", total_rows_extracted);

    if batches_to_write.is_empty() {
        println!("  ⚠️  No data to write, skipping...");
        println!();
        return Ok((0, 0.0));
    }

    // Get schema from first batch
    let schema = batches_to_write[0].schema();

    // Write output Parquet file
    let output_file =
        File::create(output_path).with_context(|| format!("Failed to create {}", output_path))?;

    let props = WriterProperties::builder()
        .set_compression(parquet::basic::Compression::SNAPPY)
        .build();

    let mut writer = ArrowWriter::try_new(output_file, schema, Some(props))?;

    for batch in &batches_to_write {
        writer.write(batch)?;
    }
    writer.close()?;

    // Get file sizes
    let input_size = std::fs::metadata(input_path)?.len() as f64 / 1024.0; // KB
    let output_size = std::fs::metadata(output_path)?.len() as f64 / 1024.0; // KB

    println!("  Original size: {:.2} KB", input_size);
    println!("  Small file size: {:.2} KB", output_size);
    println!("  Compression ratio: {:.2}x", input_size / output_size);
    println!();

    Ok((total_rows_extracted, output_size))
}

#[tokio::main]
async fn main() -> Result<()> {
    println!("=".repeat(70));
    println!("Creating Small Test Parquet Files (1000 bars each)");
    println!("=".repeat(70));
    println!();

    let mut results = Vec::new();

    for mapping in FILE_MAPPINGS {
        match create_small_parquet(mapping.input, mapping.output, 1000) {
            Ok((rows, size_kb)) => {
                if rows > 0 {
                    let symbol = mapping
                        .output
                        .replace("test_data/", "")
                        .replace("_small.parquet", "");
                    results.push((symbol, rows, size_kb));
                }
            },
            Err(e) => {
                eprintln!("❌ Error processing {}: {}", mapping.input, e);
            },
        }
    }

    println!("=".repeat(70));
    println!("Summary");
    println!("=".repeat(70));
    println!("{:<15} {:<10} {:<15}", "Symbol", "Rows", "Size (KB)");
    println!("-".repeat(70));

    for (symbol, rows, size_kb) in &results {
        println!("{:<15} {:<10} {:<15.2}", symbol, rows, size_kb);
    }

    println!("=".repeat(70));

    let total_size: f64 = results.iter().map(|(_, _, size)| size).sum();
    println!(
        "\nTotal size: {:.2} KB ({:.2} MB)",
        total_size,
        total_size / 1024.0
    );
    println!("Files created: {}", results.len());

    println!("\n✅ Small test files created successfully!");

    Ok(())
}