apache-spark-connect 4.2.0

//! Server-free serialization checks for plan variants whose behavioral coverage would
//! otherwise only come from a live server: every JoinType maps to the correct proto
//! enum, and local/cached relations serialize to the right rel_type. These assert the
//! actual proto mapping (not a recorded snapshot) and run in CI without a server.

use spark_connect::column::col;
use spark_connect::plan::{self, JoinType, LogicalPlan};
use spark_connect::types::DataType;
use spark_connect_proto as proto;

fn base() -> LogicalPlan {
    plan::range(0, 5, 1)
}

#[test]
fn join_types_map_to_correct_proto() {
    let cases = [
        (JoinType::Inner, proto::join::JoinType::Inner),
        (JoinType::LeftOuter, proto::join::JoinType::LeftOuter),
        (JoinType::RightOuter, proto::join::JoinType::RightOuter),
        (JoinType::FullOuter, proto::join::JoinType::FullOuter),
        (JoinType::LeftSemi, proto::join::JoinType::LeftSemi),
        (JoinType::LeftAnti, proto::join::JoinType::LeftAnti),
        (JoinType::Cross, proto::join::JoinType::Cross),
    ];
    for (jt, expected) in cases {
        let rel = plan::join(base(), base(), jt, Some(col("id")), vec![]).to_proto();
        match rel.rel_type {
            Some(proto::relation::RelType::Join(j)) => {
                assert_eq!(
                    j.join_type, expected as i32,
                    "JoinType {jt:?} must serialize to proto {expected:?}"
                );
                assert!(
                    j.join_condition.is_some(),
                    "the `on` condition must be carried"
                );
            }
            other => panic!("expected Join relation, got {other:?}"),
        }
    }
    // using-columns form (no `on` condition)
    let rel = plan::join(
        base(),
        base(),
        JoinType::Inner,
        None,
        vec!["id".to_string()],
    )
    .to_proto();
    match rel.rel_type {
        Some(proto::relation::RelType::Join(j)) => {
            assert_eq!(j.using_columns, vec!["id".to_string()]);
        }
        other => panic!("expected Join, got {other:?}"),
    }
}

#[test]
fn local_and_cached_relations_serialize() {
    // LocalRelation with and without inline data.
    match plan::local_relation(DataType::Integer, None)
        .to_proto()
        .rel_type
    {
        Some(proto::relation::RelType::LocalRelation(_)) => {}
        other => panic!("expected LocalRelation, got {other:?}"),
    }
    match plan::local_relation(DataType::Integer, Some(vec![1, 2, 3]))
        .to_proto()
        .rel_type
    {
        Some(proto::relation::RelType::LocalRelation(lr)) => {
            assert_eq!(lr.data.as_deref(), Some(&[1u8, 2, 3][..]));
        }
        other => panic!("expected LocalRelation, got {other:?}"),
    }
    // CachedRemoteRelation carries the relation id.
    match plan::cached_remote_relation("rel-42").to_proto().rel_type {
        Some(proto::relation::RelType::CachedRemoteRelation(cr)) => {
            assert_eq!(cr.relation_id, "rel-42");
        }
        other => panic!("expected CachedRemoteRelation, got {other:?}"),
    }
}

#[test]
fn relation_changes_maps_to_correct_proto() {
    use std::collections::HashMap;
    let mut options = HashMap::new();
    options.insert("startingVersion".to_string(), "0".to_string());

    // Batch changes: is_streaming defaults to false.
    let batch = LogicalPlan::RelationChanges {
        table_name: "db.tbl".to_string(),
        options: options.clone(),
        is_streaming: None,
    }
    .to_proto();
    match batch.rel_type {
        Some(proto::relation::RelType::RelationChanges(rc)) => {
            assert_eq!(rc.unparsed_identifier, "db.tbl");
            assert!(!rc.is_streaming);
            assert_eq!(
                rc.options.get("startingVersion").map(String::as_str),
                Some("0")
            );
        }
        other => panic!("expected RelationChanges, got {other:?}"),
    }

    // Streaming changes: is_streaming = true.
    let streaming = LogicalPlan::RelationChanges {
        table_name: "db.tbl".to_string(),
        options: HashMap::new(),
        is_streaming: Some(true),
    }
    .to_proto();
    match streaming.rel_type {
        Some(proto::relation::RelType::RelationChanges(rc)) => {
            assert_eq!(rc.unparsed_identifier, "db.tbl");
            assert!(rc.is_streaming);
        }
        other => panic!("expected RelationChanges, got {other:?}"),
    }
}

#[test]
fn column_list_apis_accept_arrays_vecs_and_strings() {
    // The IntoIterator<Item = impl Into<Column>> APIs accept arrays of Columns, arrays
    // of &str (via From<&str> for Column), and Vec — all producing the same Project.
    use spark_connect_proto::relation::RelType;
    let by_cols = plan::project(base(), [col("a"), col("b")]).to_proto();
    let by_strs = plan::project(base(), ["a", "b"]).to_proto();
    let by_vec = plan::project(base(), vec![col("a")]).to_proto();
    for (rel, n) in [(by_cols, 2usize), (by_strs, 2), (by_vec, 1)] {
        match rel.rel_type {
            Some(RelType::Project(p)) => assert_eq!(p.expressions.len(), n),
            other => panic!("expected Project, got {other:?}"),
        }
    }
    // functions builders accept the same ergonomic forms.
    let _ = spark_connect::functions::array([col("a"), col("b")]);
    let _ = spark_connect::functions::array(["a", "b"]);
    let _ = col("x").isin([col("y"), col("z")]);
}