aprender-db 0.31.1

//! Top-K Selection API demonstration
//!
//! This example demonstrates the high-performance Top-K selection algorithm
//! that achieves 28.75x speedup over full sorting.
//!
//! Algorithm: O(N log K) heap-based selection vs O(N log N) full sort
//!
//! Run with: cargo run --example `topk_selection` --release

use arrow::array::{Float64Array, Int32Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
use std::time::Instant;
use trueno_db::topk::{SortOrder, TopKSelection};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Trueno-DB Top-K Selection Example ===\n");

    // Create large dataset (1M rows) for performance demonstration
    println!("Creating sample dataset (1M rows)...");
    let batch = create_sample_batch(1_000_000)?;
    println!("  ✓ Created batch: {} rows, {} columns\n", batch.num_rows(), batch.num_columns());

    // Demonstrate Top-K selection (descending - largest values)
    println!("=== Top-10 Highest Scores ===");
    let start = Instant::now();
    let top10_high = batch.top_k(1, 10, SortOrder::Descending)?;
    let duration = start.elapsed();

    println!("  Algorithm: O(N log K) heap-based selection");
    println!("  Time: {duration:?}");
    println!("  Results:");

    let score_col = top10_high
        .column(1)
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("Example should work with valid test data");
    let id_col = top10_high
        .column(0)
        .as_any()
        .downcast_ref::<Int32Array>()
        .expect("Example should work with valid test data");

    for i in 0..top10_high.num_rows() {
        println!("    #{}: user_id={}, score={:.2}", i + 1, id_col.value(i), score_col.value(i));
    }
    println!();

    // Demonstrate Top-K selection (ascending - smallest values)
    println!("=== Top-10 Lowest Scores ===");
    let start = Instant::now();
    let top10_low = batch.top_k(1, 10, SortOrder::Ascending)?;
    let duration = start.elapsed();

    println!("  Algorithm: O(N log K) heap-based selection");
    println!("  Time: {duration:?}");
    println!("  Results:");

    let score_col = top10_low
        .column(1)
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("Example should work with valid test data");
    let id_col = top10_low
        .column(0)
        .as_any()
        .downcast_ref::<Int32Array>()
        .expect("Example should work with valid test data");

    for i in 0..top10_low.num_rows() {
        println!("    #{}: user_id={}, score={:.2}", i + 1, id_col.value(i), score_col.value(i));
    }
    println!();

    // Demonstrate larger K values
    println!("=== Performance Comparison (K=100) ===");
    let start = Instant::now();
    let top100 = batch.top_k(1, 100, SortOrder::Descending)?;
    let duration = start.elapsed();

    println!("  Dataset: 1M rows");
    println!("  K: 100");
    println!("  Time: {duration:?}");
    println!("  Complexity: O(N log K) = O(1M * log(100)) ≈ O(6.6M operations)");
    println!("  vs Full Sort: O(N log N) = O(1M * log(1M)) ≈ O(20M operations)");
    println!("  Speedup: ~{}x\n", (20_000_000 / 6_600_000));

    // Show actual results from top 100
    println!("  Top 5 from Top-100 results:");
    let score_col = top100
        .column(1)
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("Example should work with valid test data");
    let id_col = top100
        .column(0)
        .as_any()
        .downcast_ref::<Int32Array>()
        .expect("Example should work with valid test data");

    for i in 0..5.min(top100.num_rows()) {
        println!("    #{}: user_id={}, score={:.2}", i + 1, id_col.value(i), score_col.value(i));
    }
    println!("    ... ({} more results)", top100.num_rows() - 5);
    println!();

    // Explain the algorithm
    println!("=== Algorithm Explanation ===");
    println!("Descending (largest K values):");
    println!("  1. Use min-heap of size K");
    println!("  2. Keep smallest value at heap top");
    println!("  3. When we see larger value, replace top");
    println!("  4. Final heap contains K largest values\n");

    println!("Ascending (smallest K values):");
    println!("  1. Use max-heap of size K");
    println!("  2. Keep largest value at heap top");
    println!("  3. When we see smaller value, replace top");
    println!("  4. Final heap contains K smallest values\n");

    println!("=== Performance Benefits ===");
    println!("✓ Memory: O(K) vs O(N) for full sort");
    println!("✓ Time: O(N log K) vs O(N log N)");
    println!("✓ Measured speedup: 28.75x for K=10, N=1M (release build)");
    println!("✓ Use case: ORDER BY ... LIMIT queries\n");

    Ok(())
}

fn create_sample_batch(num_rows: usize) -> Result<RecordBatch, Box<dyn std::error::Error>> {
    use rand::Rng;

    let schema = Schema::new(vec![
        Field::new("user_id", DataType::Int32, false),
        Field::new("score", DataType::Float64, false),
    ]);

    let mut rng = rand::thread_rng();

    // Generate random scores for realistic Top-K selection
    let user_ids: Vec<i32> = (0..num_rows).map(|i| i as i32).collect();
    let scores: Vec<f64> = (0..num_rows).map(|_| rng.gen_range(0.0..1000.0)).collect();

    let batch = RecordBatch::try_new(
        Arc::new(schema),
        vec![Arc::new(Int32Array::from(user_ids)), Arc::new(Float64Array::from(scores))],
    )?;

    Ok(batch)
}