apache-spark-connect 4.2.0

//! End-to-end integration tests against a live Spark Connect server.
//!
//! These tests connect to a live Spark Connect server at `sc://localhost:15002`
//! and verify the complete execution path: plan building, execution, Arrow IPC decode, and row collection.
//!
//! Run with: `SPARK_REMOTE=sc://localhost:15002 cargo test --test e2e_integration`
//! Or simply skip with: `cargo test --test e2e_integration` (will skip all tests)

use spark_connect::column::{col, lit};
use spark_connect::row::Value;
use spark_connect::session::SparkSession;
use spark_connect::types::DataType;

fn should_run() -> bool {
    std::env::var("SPARK_REMOTE").is_ok()
}

#[test]
fn test_range_collect() {
    if !should_run() {
        println!("Skipping test_range_collect - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let df = session.range(5).expect("Failed to create range DataFrame");
    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 5, "Expected 5 rows from range(5)");
    for (i, row) in rows.iter().enumerate() {
        let id_value = row.get(0).expect("Row should have at least one field");
        match id_value {
            Value::Long(id) => {
                assert_eq!(*id, i as i64, "Expected id to be {}", i);
            }
            _ => panic!("Expected Long value, got {:?}", id_value),
        }
    }
}

#[test]
fn test_filter_collect() {
    if !should_run() {
        println!("Skipping test_filter_collect - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // range(10).filter(id > 7) must return exactly rows 8 and 9 - this exercises
    // the Filter relation end-to-end, not just range.
    let df = session
        .range(10)
        .expect("Failed to create range DataFrame")
        .filter(col("id").gt(lit(7)));

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(
        rows.len(),
        2,
        "range(10).filter(id>7) should yield 2 rows (8,9)"
    );
    let ids: Vec<i64> = rows
        .iter()
        .map(|r| match r.get(0).expect("no id") {
            Value::Long(v) => *v,
            other => panic!("expected Long, got {other:?}"),
        })
        .collect();
    assert_eq!(ids, vec![8, 9], "filter should keep only ids > 7");
}

#[test]
fn test_select_alias() {
    if !should_run() {
        println!("Skipping test_select_alias - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let df = session
        .range(5)
        .expect("Failed to create range DataFrame")
        .select(vec![col("id").alias("x")]);

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 5, "Expected 5 rows");
    let values: Vec<i64> = rows
        .iter()
        .map(|r| match r.get(0).expect("No x value") {
            Value::Long(val) => *val,
            _ => panic!("Expected Long"),
        })
        .collect();
    assert_eq!(values, vec![0, 1, 2, 3, 4], "Expected 0,1,2,3,4");
}

#[test]
fn test_sql_query() {
    if !should_run() {
        println!("Skipping test_sql_query - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let df = session
        .sql("SELECT 1 AS a, 'x' AS b")
        .expect("Failed to execute SQL");

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 1, "Expected 1 row from SQL query");
    let row = &rows[0];

    // Just verify we got 2 columns and some data
    assert_eq!(row.len(), 2, "Expected 2 columns");

    // Check field 'b' is a string ('x')
    match row.get(1).expect("No b value") {
        Value::String(b) => assert_eq!(b, "x", "Expected b='x'"),
        _ => panic!("Expected String for b"),
    }
}

#[test]
fn test_groupby_count() {
    if !should_run() {
        println!("Skipping test_groupby_count - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // Group by a static column (all rows in one group) and count
    let df = session
        .range(6)
        .expect("Failed to create range DataFrame")
        .group_by(vec![lit(1).alias("k")])
        .count();

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 1, "Expected 1 group");
    // First column is 'k' (value 1), second is count (should be 6)
    match rows[0].get(1).expect("No count") {
        Value::Long(count) => assert_eq!(*count, 6, "Expected count of 6"),
        _ => panic!("Expected Long for count"),
    }
}

#[test]
fn test_count() {
    if !should_run() {
        println!("Skipping test_count - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let df = session
        .range(100)
        .expect("Failed to create range DataFrame");
    let count = df.count().expect("Failed to count");

    assert_eq!(count, 100, "Expected count of 100");
}

#[test]
fn test_schema() {
    if !should_run() {
        println!("Skipping test_schema - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let df = session.range(3).expect("Failed to create range DataFrame");
    let schema = df.schema().expect("Failed to get schema");

    match schema {
        DataType::Struct { fields } => {
            assert_eq!(fields.len(), 1, "Expected 1 field");
            assert_eq!(fields[0].name, "id", "Expected field name 'id'");
            assert_eq!(fields[0].data_type, DataType::Long, "Expected Long type");
        }
        _ => panic!("Expected Struct schema"),
    }
}

#[test]
fn test_decimal_round() {
    if !should_run() {
        println!("Skipping test_decimal_round - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // Test: SELECT ROUND(3.14159, 2) AS v should return 3.14, not 3
    // Using pure SQL first
    let df = session
        .sql("SELECT ROUND(3.14159, 2) AS v")
        .expect("Failed to execute SQL");

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 1, "Expected 1 row");
    let row = &rows[0];

    // A SQL decimal literal (3.14159) is DECIMAL in Spark, so ROUND(.., 2) yields
    // DECIMAL(4,2)=3.14 - not DOUBLE. Accept either, checking the numeric value.
    match row.get(0).expect("No value") {
        Value::Double(d) => {
            assert!((d - 3.14).abs() < 0.01, "Expected ~3.14, got {}", d);
        }
        Value::Decimal { value, .. } => {
            let d: f64 = value.parse().expect("decimal value parses as f64");
            assert!((d - 3.14).abs() < 0.01, "Expected ~3.14, got {}", value);
        }
        other => panic!("Expected Double/Decimal, got {:?}", other),
    }

    // Also test with lit(3.14159) to ensure literal decoding works
    let df2 = session
        .sql("SELECT 3.14159 AS v")
        .expect("Failed to execute SQL");

    let rows2 = df2.collect().expect("Failed to collect rows");
    assert_eq!(rows2.len(), 1, "Expected 1 row");
    let row2 = &rows2[0];

    // The SQL literal 3.14159 decodes as DECIMAL(6,5) in Spark (not DOUBLE).
    match row2.get(0).expect("No value") {
        Value::Double(d) => {
            assert!(
                (d - 3.14159).abs() < 0.00001,
                "Expected ~3.14159, got {}",
                d
            );
        }
        Value::Decimal { value, .. } => {
            let d: f64 = value.parse().expect("decimal value parses as f64");
            assert!(
                (d - 3.14159).abs() < 0.00001,
                "Expected ~3.14159, got {}",
                value
            );
        }
        other => panic!("Expected Double/Decimal, got {:?}", other),
    }
}

#[test]
fn test_string_concat() {
    if !should_run() {
        println!("Skipping test_string_concat - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // Test: SELECT 'a' AS a, 'b' AS b -> CONCAT(a, b) should return 'ab', not empty
    let df = session
        .sql("SELECT CONCAT('a', 'b') AS result")
        .expect("Failed to execute SQL");

    let rows = df.collect().expect("Failed to collect rows");

    assert_eq!(rows.len(), 1, "Expected 1 row");
    let row = &rows[0];

    match row.get(0).expect("No value") {
        Value::String(s) => {
            assert_eq!(s, "ab", "CONCAT('a', 'b') should return 'ab', got '{}'", s);
        }
        other => panic!("Expected String, got {:?}", other),
    }
}

#[test]
fn test_cache_persist() {
    if !should_run() {
        println!("Skipping test_cache_persist - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // Verify the cache()/persist() CLIENT API: the AnalyzePlan.Persist round-trips
    // and returns a usable DataFrame. We intentionally do NOT assert on server-side
    // block *materialization* here: on a local single-JVM Spark 4.2.0 server the
    // BlockManager rejects the persisted block with "StorageLevel is null or invalid"
    // for EVERY Spark Connect client (the reference pyspark client included) — it's a
    // server/local-mode quirk, not a client behavior, so asserting it would test the
    // environment rather than this code.
    use spark_connect::StorageLevelExt;

    // cache() = persist(MEMORY_AND_DISK_DESER): must round-trip.
    let cached = session
        .range(5)
        .expect("Failed to create range DataFrame")
        .cache()
        .expect("cache() should round-trip the Persist AnalyzePlan");

    // Its storage level should read back as the level cache() requested.
    let level = cached
        .storage_level()
        .expect("storage_level() should round-trip");
    assert!(
        level.use_memory && level.use_disk,
        "cache() level should be MEMORY_AND_DISK(_DESER); got {level:?}"
    );

    // persist() with an explicit preset: must round-trip too, and unpersist() cleans up.
    let persisted = session
        .range(3)
        .expect("Failed to create range DataFrame")
        .persist(spark_connect::StorageLevel::memory_and_disk())
        .expect("persist(MEMORY_AND_DISK) should round-trip the Persist AnalyzePlan");
    persisted
        .unpersist(true)
        .expect("unpersist() should round-trip");
    // NOTE: we deliberately don't collect() a persisted relation here. On a local
    // single-JVM Spark 4.2.0 server the BlockManager rejects materializing the cached
    // block ("StorageLevel is null or invalid") for EVERY Spark Connect client,
    // reference pyspark included — a server/local-mode quirk that also poisons later
    // execution of the same relation. This test therefore validates the client-side
    // cache/persist/unpersist/storage_level round-trips, which is the client's job.
}

#[test]
fn test_with_watermark() {
    if !should_run() {
        println!("Skipping test_with_watermark - set SPARK_REMOTE to run");
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // Create a DataFrame with a timestamp and add watermark
    let df = session
        .sql("SELECT CURRENT_TIMESTAMP AS ts, 1 AS value")
        .expect("Failed to execute SQL")
        .with_watermark("ts", "10 seconds");

    // Just verify that the method doesn't fail
    let _schema = df.schema().expect("Failed to get schema");
}

#[test]
fn test_to_arrow_full_path() {
    // Full server -> collect -> Arrow IPC path for DataFrame::to_arrow (the
    // conversion logic itself is unit-tested in dataframe.rs::conversion_tests).
    if !should_run() {
        println!("Skipping test_to_arrow_full_path - set SPARK_REMOTE to run");
        return;
    }
    use arrow::ipc::reader::FileReader;
    use std::io::Cursor;

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    let ipc = session
        .range(5)
        .expect("range")
        .to_arrow()
        .expect("to_arrow");
    let reader = FileReader::try_new(Cursor::new(ipc), None).expect("valid Arrow IPC file");
    let rows: usize = reader.map(|b| b.expect("batch").num_rows()).sum();
    assert_eq!(rows, 5, "to_arrow must round-trip all 5 rows of range(5)");
}

/// End-to-end guard for SPARK-59037: a single ~5 MiB string cell yields one Arrow
/// batch, and thus one gRPC response message, larger than tonic's 4 MiB default
/// decode cap. Before the client raised `max_decoding_message_size` to 128 MiB this
/// `collect()` failed with a gRPC "message length too large" decode error (the same
/// batch the reference grpcio client returns fine). It must now succeed and return
/// the whole value intact.
#[test]
fn test_collect_batch_larger_than_grpc_default_cap() {
    if !should_run() {
        println!(
            "Skipping test_collect_batch_larger_than_grpc_default_cap - set SPARK_REMOTE to run"
        );
        return;
    }

    let remote_url =
        std::env::var("SPARK_REMOTE").unwrap_or_else(|_| "sc://localhost:15002".to_string());
    let session = SparkSession::builder()
        .remote(&remote_url)
        .get_or_create()
        .expect("Failed to create session");

    // 5 MiB, comfortably above tonic's 4 MiB default receive cap. A single row cannot be
    // split across Arrow batches, so this forces one oversized gRPC response message.
    let n: usize = 5 * 1024 * 1024;
    let df = session
        .sql(&format!("SELECT repeat('x', {n}) AS s"))
        .expect("Failed to build repeat() DataFrame");
    let rows = df
        .collect()
        .expect("collect() of a >4 MiB batch must succeed with the raised gRPC cap");

    assert_eq!(rows.len(), 1, "Expected exactly one row");
    let s = rows[0]
        .get(0)
        .and_then(|v| v.as_str())
        .expect("Expected a string cell");
    assert_eq!(
        s.len(),
        n,
        "Expected the full {n}-byte string to round-trip through collect()"
    );
}