From 06b11f67e0e530aeceded1326c156482849dc5c1 Mon Sep 17 00:00:00 2001 From: Rafael Fernandez Date: Sat, 28 Mar 2026 19:31:44 +0100 Subject: [PATCH] [SPARK-53316] Add Map and Struct literal types --- crates/connect/src/expressions.rs | 151 ++++++++++++++++++++++++++++++ 1 file changed, 151 insertions(+) diff --git a/crates/connect/src/expressions.rs b/crates/connect/src/expressions.rs index 488227a..c0265e9 100644 --- a/crates/connect/src/expressions.rs +++ b/crates/connect/src/expressions.rs @@ -249,6 +249,100 @@ where } } +/// Represents a Spark Map literal value. +/// +/// Used with [`lit`](crate::functions::lit) to create a map column literal. +#[derive(Clone, Debug)] +pub struct MapLiteral +where + K: Into + Clone, + V: Into + Clone, + spark::DataType: From, + spark::DataType: From, +{ + pub keys: Vec, + pub values: Vec, +} + +impl MapLiteral +where + K: Into + Clone, + V: Into + Clone, + spark::DataType: From, + spark::DataType: From, +{ + pub fn new(keys: Vec, values: Vec) -> Self { + assert_eq!( + keys.len(), + values.len(), + "Keys and values must have the same length" + ); + Self { keys, values } + } +} + +impl From> for spark::expression::Literal +where + K: Into + Clone, + V: Into + Clone, + spark::DataType: From, + spark::DataType: From, +{ + fn from(value: MapLiteral) -> Self { + let key_type = Some(spark::DataType::from( + value.keys.first().expect("Map cannot be empty").clone(), + )); + let value_type = Some(spark::DataType::from( + value.values.first().expect("Map cannot be empty").clone(), + )); + + let keys = value.keys.into_iter().map(|k| k.into()).collect(); + let values = value.values.into_iter().map(|v| v.into()).collect(); + + let map = spark::expression::literal::Map { + key_type, + value_type, + keys, + values, + }; + + spark::expression::Literal { + literal_type: Some(spark::expression::literal::LiteralType::Map(map)), + } + } +} + +/// Represents a Spark Struct literal value. +/// +/// Used with [`lit`](crate::functions::lit) to create a struct column literal. +#[derive(Clone, Debug)] +pub struct StructLiteral { + pub struct_type: spark::DataType, + pub elements: Vec, +} + +impl StructLiteral { + pub fn new(struct_type: spark::DataType, elements: Vec) -> Self { + Self { + struct_type, + elements, + } + } +} + +impl From for spark::expression::Literal { + fn from(value: StructLiteral) -> Self { + let struct_lit = spark::expression::literal::Struct { + struct_type: Some(value.struct_type), + elements: value.elements, + }; + + spark::expression::Literal { + literal_type: Some(spark::expression::literal::LiteralType::Struct(struct_lit)), + } + } +} + impl From<&str> for spark::expression::cast::CastToType { fn from(value: &str) -> Self { spark::expression::cast::CastToType::TypeStr(value.to_string()) @@ -266,3 +360,60 @@ impl From for spark::expression::cast::CastToType { spark::expression::cast::CastToType::Type(value.into()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_map_literal() { + let map = MapLiteral::new( + vec!["key1".to_string(), "key2".to_string()], + vec![1i32, 2i32], + ); + let literal: spark::expression::Literal = map.into(); + match literal.literal_type { + Some(spark::expression::literal::LiteralType::Map(m)) => { + assert_eq!(m.keys.len(), 2); + assert_eq!(m.values.len(), 2); + assert!(m.key_type.is_some()); + assert!(m.value_type.is_some()); + } + _ => panic!("Expected Map literal"), + } + } + + #[test] + fn test_struct_literal() { + let struct_type = + crate::types::DataType::Struct(Box::new(crate::types::StructType::new(vec![ + crate::types::StructField { + name: "name", + data_type: crate::types::DataType::String, + nullable: true, + metadata: None, + }, + crate::types::StructField { + name: "age", + data_type: crate::types::DataType::Integer, + nullable: true, + metadata: None, + }, + ]))); + + let elements = vec![ + spark::expression::Literal::from("Alice".to_string()), + spark::expression::Literal::from(30i32), + ]; + + let struct_lit = StructLiteral::new(struct_type.into(), elements); + let literal: spark::expression::Literal = struct_lit.into(); + match literal.literal_type { + Some(spark::expression::literal::LiteralType::Struct(s)) => { + assert_eq!(s.elements.len(), 2); + assert!(s.struct_type.is_some()); + } + _ => panic!("Expected Struct literal"), + } + } +}