spark_connect/
row.rs

1//! Row type mirroring `pyspark.sql.Row`.
2//!
3//! A Row represents a single record: an ordered collection of (field_name, value) pairs.
4//! Values are accessed by index or by field name. Supports all Spark SQL value types.
5
6use std::collections::BTreeMap;
7use std::fmt;
8
9/// A Value in a Row, supporting all Spark SQL types.
10#[derive(Debug, Clone, PartialEq)]
11pub enum Value {
12    /// NULL value
13    Null,
14    /// Boolean
15    Bool(bool),
16    /// Byte (i8)
17    Byte(i8),
18    /// Short (i16)
19    Short(i16),
20    /// Integer (i32)
21    Integer(i32),
22    /// Long (i64)
23    Long(i64),
24    /// Float (f32)
25    Float(f32),
26    /// Double (f64)
27    Double(f64),
28    /// String
29    String(String),
30    /// Binary (bytes)
31    Binary(Vec<u8>),
32    /// Date (days since epoch)
33    Date(i32),
34    /// Timestamp (microseconds since epoch)
35    Timestamp(i64),
36    /// Decimal (string value with optional precision and scale)
37    Decimal {
38        value: String,
39        precision: Option<i32>,
40        scale: Option<i32>,
41    },
42    /// Array of values
43    List(Vec<Value>),
44    /// Map of key-value pairs
45    Map(BTreeMap<String, Value>),
46    /// Struct (nested Row)
47    Struct(Vec<(String, Value)>),
48    /// A VARIANT value carried as its raw (value, metadata) binary components, matching
49    /// `pyspark.sql.types.VariantVal(value, metadata)`. Decoding to JSON/Python is done
50    /// lazily on the Python side (VariantVal.toJson/toPython via variant_utils).
51    Variant { value: Vec<u8>, metadata: Vec<u8> },
52}
53
54impl Value {
55    /// Get this value as a bool, or None if it's not a bool.
56    pub fn as_bool(&self) -> Option<bool> {
57        match self {
58            Value::Bool(b) => Some(*b),
59            _ => None,
60        }
61    }
62
63    /// Get this value as i64, or None if it's not an integer type.
64    pub fn as_i64(&self) -> Option<i64> {
65        match self {
66            Value::Byte(b) => Some(*b as i64),
67            Value::Short(s) => Some(*s as i64),
68            Value::Integer(i) => Some(*i as i64),
69            Value::Long(l) => Some(*l),
70            _ => None,
71        }
72    }
73
74    /// Get this value as f64, or None if it's not a float type.
75    pub fn as_f64(&self) -> Option<f64> {
76        match self {
77            Value::Float(f) => Some(*f as f64),
78            Value::Double(d) => Some(*d),
79            _ => None,
80        }
81    }
82
83    /// Get this value as a string reference, or None if it's not a string.
84    pub fn as_str(&self) -> Option<&str> {
85        match self {
86            Value::String(s) => Some(s),
87            _ => None,
88        }
89    }
90
91    /// Get this value as a bytes reference, or None if it's not binary.
92    pub fn as_bytes(&self) -> Option<&[u8]> {
93        match self {
94            Value::Binary(b) => Some(b),
95            _ => None,
96        }
97    }
98
99    /// Check if this value is null.
100    pub fn is_null(&self) -> bool {
101        matches!(self, Value::Null)
102    }
103}
104
105impl fmt::Display for Value {
106    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107        match self {
108            Value::Null => write!(f, "NULL"),
109            Value::Bool(b) => write!(f, "{}", b),
110            Value::Byte(b) => write!(f, "{}", b),
111            Value::Short(s) => write!(f, "{}", s),
112            Value::Integer(i) => write!(f, "{}", i),
113            Value::Long(l) => write!(f, "{}", l),
114            Value::Float(fl) => write!(f, "{}", fl),
115            Value::Double(d) => write!(f, "{}", d),
116            Value::String(s) => write!(f, "{}", s),
117            Value::Binary(b) => write!(f, "{:?}", b),
118            Value::Date(d) => write!(f, "{}", d),
119            Value::Timestamp(t) => write!(f, "{}", t),
120            Value::Decimal {
121                value,
122                precision,
123                scale,
124            } => {
125                write!(f, "Decimal({}", value)?;
126                if let Some(p) = precision {
127                    write!(f, ",{})", p)?;
128                    if let Some(s) = scale {
129                        write!(f, "s={}", s)?;
130                    }
131                } else {
132                    write!(f, ")")?;
133                }
134                Ok(())
135            }
136            Value::List(l) => {
137                write!(f, "[")?;
138                for (i, v) in l.iter().enumerate() {
139                    if i > 0 {
140                        write!(f, ", ")?;
141                    }
142                    write!(f, "{}", v)?;
143                }
144                write!(f, "]")
145            }
146            Value::Map(m) => {
147                write!(f, "{{")?;
148                for (i, (k, v)) in m.iter().enumerate() {
149                    if i > 0 {
150                        write!(f, ", ")?;
151                    }
152                    write!(f, "{}: {}", k, v)?;
153                }
154                write!(f, "}}")
155            }
156            Value::Struct(s) => {
157                write!(f, "(")?;
158                for (i, (k, v)) in s.iter().enumerate() {
159                    if i > 0 {
160                        write!(f, ", ")?;
161                    }
162                    write!(f, "{}={}", k, v)?;
163                }
164                write!(f, ")")
165            }
166            Value::Variant { value, metadata } => {
167                write!(
168                    f,
169                    "Variant(value={} bytes, metadata={} bytes)",
170                    value.len(),
171                    metadata.len()
172                )
173            }
174        }
175    }
176}
177
178/// A Row is an ordered collection of (field_name, value) pairs.
179/// Supports access by index or by field name.
180#[derive(Debug, Clone, PartialEq)]
181pub struct Row {
182    /// Field names (in order)
183    fields: Vec<String>,
184    /// Values corresponding to fields (in order)
185    values: Vec<Value>,
186}
187
188impl Row {
189    /// Create a new Row from field names and values.
190    /// Panics if lengths don't match.
191    pub fn new(fields: Vec<String>, values: Vec<Value>) -> Self {
192        assert_eq!(
193            fields.len(),
194            values.len(),
195            "field names and values must have the same length"
196        );
197        Row { fields, values }
198    }
199
200    /// Create an empty Row.
201    pub fn empty() -> Self {
202        Row {
203            fields: vec![],
204            values: vec![],
205        }
206    }
207
208    /// Get the number of fields in this Row.
209    pub fn len(&self) -> usize {
210        self.fields.len()
211    }
212
213    /// Check if this Row is empty.
214    pub fn is_empty(&self) -> bool {
215        self.fields.is_empty()
216    }
217
218    /// Get a field value by index.
219    pub fn get(&self, index: usize) -> Option<&Value> {
220        self.values.get(index)
221    }
222
223    /// Get a field value by index, or panic if out of bounds.
224    pub fn get_unchecked(&self, index: usize) -> &Value {
225        &self.values[index]
226    }
227
228    /// Get a field value by name.
229    pub fn get_by_name(&self, name: &str) -> Option<&Value> {
230        self.fields
231            .iter()
232            .position(|f| f == name)
233            .and_then(|i| self.values.get(i))
234    }
235
236    /// Get field names.
237    pub fn fields(&self) -> &[String] {
238        &self.fields
239    }
240
241    /// Get values.
242    pub fn values(&self) -> &[Value] {
243        &self.values
244    }
245
246    /// Convert into values.
247    pub fn into_values(self) -> Vec<Value> {
248        self.values
249    }
250}
251
252impl fmt::Display for Row {
253    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
254        write!(f, "[")?;
255        for (i, v) in self.values.iter().enumerate() {
256            if i > 0 {
257                write!(f, ", ")?;
258            }
259            write!(f, "{}", v)?;
260        }
261        write!(f, "]")
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn test_row_creation() {
271        let row = Row::new(
272            vec!["id".to_string(), "name".to_string()],
273            vec![Value::Long(1), Value::String("Alice".to_string())],
274        );
275
276        assert_eq!(row.len(), 2);
277        assert_eq!(row.get(0), Some(&Value::Long(1)));
278        assert_eq!(row.get(1), Some(&Value::String("Alice".to_string())));
279    }
280
281    #[test]
282    fn test_row_access_by_name() {
283        let row = Row::new(
284            vec!["id".to_string(), "name".to_string()],
285            vec![Value::Long(1), Value::String("Alice".to_string())],
286        );
287
288        assert_eq!(row.get_by_name("id"), Some(&Value::Long(1)));
289        assert_eq!(
290            row.get_by_name("name"),
291            Some(&Value::String("Alice".to_string()))
292        );
293        assert_eq!(row.get_by_name("nonexistent"), None);
294    }
295
296    #[test]
297    fn test_value_conversions() {
298        let b = Value::Bool(true);
299        assert_eq!(b.as_bool(), Some(true));
300
301        let i = Value::Integer(42);
302        assert_eq!(i.as_i64(), Some(42));
303
304        #[allow(clippy::approx_constant)]
305        let d = Value::Double(3.14);
306
307        assert_eq!(d.as_f64(), Some(3.14));
308
309        let s = Value::String("test".to_string());
310        assert_eq!(s.as_str(), Some("test"));
311    }
312
313    #[test]
314    fn test_value_date_timestamp_decimal() {
315        // Test Date
316        let date_val = Value::Date(18993); // days since epoch
317        assert_eq!(date_val, Value::Date(18993));
318
319        // Test Timestamp
320        let ts_val = Value::Timestamp(1693526400000000); // micros since epoch
321        assert_eq!(ts_val, Value::Timestamp(1693526400000000));
322
323        // Test Decimal
324        let dec_val = Value::Decimal {
325            value: "123.45".to_string(),
326            precision: Some(5),
327            scale: Some(2),
328        };
329        match dec_val {
330            Value::Decimal {
331                value,
332                precision,
333                scale,
334            } => {
335                assert_eq!(value, "123.45");
336                assert_eq!(precision, Some(5));
337                assert_eq!(scale, Some(2));
338            }
339            _ => panic!("Expected Decimal variant"),
340        }
341    }
342
343    #[test]
344    fn as_i64_covers_all_integer_widths_and_rejects_others() {
345        assert_eq!(Value::Byte(7).as_i64(), Some(7));
346        assert_eq!(Value::Short(300).as_i64(), Some(300));
347        assert_eq!(Value::Integer(70_000).as_i64(), Some(70_000));
348        assert_eq!(Value::Long(5_000_000_000).as_i64(), Some(5_000_000_000));
349        assert_eq!(Value::Double(1.0).as_i64(), None);
350        assert_eq!(Value::Null.as_i64(), None);
351    }
352
353    #[test]
354    fn as_f64_covers_float_and_double_and_rejects_others() {
355        assert_eq!(Value::Float(1.5).as_f64(), Some(1.5));
356        assert_eq!(Value::Double(2.5).as_f64(), Some(2.5));
357        assert_eq!(Value::Long(3).as_f64(), None);
358    }
359
360    #[test]
361    fn scalar_accessors_reject_wrong_types() {
362        assert_eq!(Value::Integer(1).as_bool(), None);
363        assert_eq!(Value::Bool(true).as_str(), None);
364        assert_eq!(Value::String("x".into()).as_bytes(), None);
365        assert_eq!(Value::Binary(vec![1, 2]).as_bytes(), Some(&[1u8, 2][..]));
366    }
367
368    #[test]
369    fn is_null_reflects_variant() {
370        assert!(Value::Null.is_null());
371        assert!(!Value::Integer(0).is_null());
372    }
373
374    #[test]
375    fn display_covers_every_value_variant() {
376        assert_eq!(Value::Null.to_string(), "NULL");
377        assert_eq!(Value::Bool(true).to_string(), "true");
378        assert_eq!(Value::Byte(1).to_string(), "1");
379        assert_eq!(Value::Short(2).to_string(), "2");
380        assert_eq!(Value::Integer(3).to_string(), "3");
381        assert_eq!(Value::Long(4).to_string(), "4");
382        assert_eq!(Value::Float(1.5).to_string(), "1.5");
383        assert_eq!(Value::Double(2.5).to_string(), "2.5");
384        assert_eq!(Value::String("hi".into()).to_string(), "hi");
385        assert_eq!(Value::Binary(vec![1, 2]).to_string(), "[1, 2]");
386        assert_eq!(Value::Date(19_000).to_string(), "19000");
387        assert_eq!(Value::Timestamp(123).to_string(), "123");
388        assert_eq!(
389            Value::List(vec![Value::Integer(1), Value::Integer(2)]).to_string(),
390            "[1, 2]"
391        );
392        let mut m = std::collections::BTreeMap::new();
393        m.insert("k".to_string(), Value::Integer(9));
394        assert_eq!(Value::Map(m).to_string(), "{k: 9}");
395        assert_eq!(
396            Value::Struct(vec![("a".to_string(), Value::Integer(1))]).to_string(),
397            "(a=1)"
398        );
399    }
400
401    #[test]
402    fn display_decimal_with_and_without_precision() {
403        assert_eq!(
404            Value::Decimal {
405                value: "123.45".into(),
406                precision: Some(5),
407                scale: Some(2),
408            }
409            .to_string(),
410            "Decimal(123.45,5)s=2"
411        );
412        assert_eq!(
413            Value::Decimal {
414                value: "7".into(),
415                precision: None,
416                scale: None,
417            }
418            .to_string(),
419            "Decimal(7)"
420        );
421    }
422
423    #[test]
424    fn row_helpers_and_display() {
425        let empty = Row::empty();
426        assert!(empty.is_empty());
427        assert_eq!(empty.len(), 0);
428
429        let row = Row::new(
430            vec!["a".to_string(), "b".to_string()],
431            vec![Value::Integer(1), Value::String("x".to_string())],
432        );
433        assert!(!row.is_empty());
434        assert_eq!(row.get(5), None);
435        assert_eq!(row.get_unchecked(0), &Value::Integer(1));
436        assert_eq!(row.fields(), &["a".to_string(), "b".to_string()]);
437        assert_eq!(row.values().len(), 2);
438        assert_eq!(row.to_string(), "[1, x]");
439        assert_eq!(
440            row.into_values(),
441            vec![Value::Integer(1), Value::String("x".to_string())]
442        );
443    }
444
445    #[test]
446    #[should_panic(expected = "same length")]
447    fn row_new_rejects_mismatched_lengths() {
448        let _ = Row::new(vec!["a".to_string()], vec![]);
449    }
450}
spark_connect/row.rs

spark_connect/
row.rs
