From fbdaa58c9358c3a2a766093446f8bb3f2ae3a5fc Mon Sep 17 00:00:00 2001 From: Marijn Valk Date: Sun, 19 Feb 2023 11:02:31 +0100 Subject: [PATCH] add test for extract_partition_values Signed-off-by: Marijn Valk --- rust/src/writer/json.rs | 45 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/rust/src/writer/json.rs b/rust/src/writer/json.rs index c9fee9db77..5b3e2a0fed 100644 --- a/rust/src/writer/json.rs +++ b/rust/src/writer/json.rs @@ -437,11 +437,15 @@ fn extract_partition_values( #[cfg(test)] mod tests { + use super::*; + use crate::arrow::array::Int32Array; + use crate::arrow::datatypes::{ + DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, + }; use crate::writer::test_utils::get_delta_schema; use crate::writer::DeltaWriter; use crate::writer::JsonWriter; use crate::Schema; - use arrow::datatypes::Schema as ArrowSchema; use parquet::file::reader::FileReader; use parquet::file::serialized_reader::SerializedFileReader; use std::fs::File; @@ -488,4 +492,43 @@ mod tests { .collect::>(); assert_eq!(columns, vec!["id".to_string(), "value".to_string()]); } + + #[test] + fn test_extract_partition_values() { + let record_batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("col1", ArrowDataType::Int32, false), + ArrowField::new("col2", ArrowDataType::Int32, false), + ArrowField::new("col3", ArrowDataType::Int32, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![2, 1])), + Arc::new(Int32Array::from(vec![None, None])), + ], + ) + .unwrap(); + + assert_eq!( + extract_partition_values( + &[ + String::from("col1"), + String::from("col2"), + String::from("col3") + ], + &record_batch + ) + .unwrap(), + HashMap::from([ + (String::from("col1"), Some(String::from("1"))), + (String::from("col2"), Some(String::from("2"))), + (String::from("col3"), None), + ]) + ); + assert_eq!( + extract_partition_values(&[String::from("col1")], &record_batch).unwrap(), + HashMap::from([(String::from("col1"), Some(String::from("1"))),]) + ); + assert!(extract_partition_values(&[String::from("col4")], &record_batch).is_err()) + } }