refactor: make ParquetSink tests a bit more readable

Author: wiedld

datafusion/core/src/datasource/file_format/parquet.rs

@@ -2335,42 +2335,46 @@ mod tests {
     async fn parquet_sink_write() -> Result<()> {
         let parquet_sink = create_written_parquet_sink("file:///").await?;
 
-        // assert written
-        let mut written = parquet_sink.written();
-        let written = written.drain();
-        assert_eq!(
-            written.len(),
-            1,
-            "expected a single parquet files to be written, instead found {}",
-            written.len()
-        );
+        // assert written to proper path
+        let (path, file_metadata) = get_written(parquet_sink)?;
+        let path_parts = path.parts().collect::<Vec<_>>();
+        assert_eq!(path_parts.len(), 1, "should not have path prefix");
 
         // check the file metadata
-        let (
-            path,
-            FileMetaData {
-                num_rows,
-                schema,
-                key_value_metadata,
-                ..
+        let expected_kv_meta = vec![
+            KeyValue {
+                key: "my-data".to_string(),
+                value: Some("stuff".to_string()),
+            },
+            KeyValue {
+                key: "my-data-bool-key".to_string(),
+                value: None,
             },
-        ) = written.take(1).next().unwrap();
+        ];
+        assert_file_metadata(file_metadata, expected_kv_meta);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn parquet_sink_parallel_write() -> Result<()> {
+        let opts = ParquetOptions {
+            allow_single_file_parallelism: true,
+            maximum_parallel_row_group_writers: 2,
+            maximum_buffered_record_batches_per_stream: 2,
+            ..Default::default()
+        };
+
+        let parquet_sink =
+            create_written_parquet_sink_using_config("file:///", opts).await?;
+
+        // assert written to proper path
+        let (path, file_metadata) = get_written(parquet_sink)?;
         let path_parts = path.parts().collect::<Vec<_>>();
         assert_eq!(path_parts.len(), 1, "should not have path prefix");
 
-        assert_eq!(num_rows, 2, "file metadata to have 2 rows");
-        assert!(
-            schema.iter().any(|col_schema| col_schema.name == "a"),
-            "output file metadata should contain col a"
-        );
-        assert!(
-            schema.iter().any(|col_schema| col_schema.name == "b"),
-            "output file metadata should contain col b"
-        );
-
-        let mut key_value_metadata = key_value_metadata.unwrap();
-        key_value_metadata.sort_by(|a, b| a.key.cmp(&b.key));
-        let expected_metadata = vec![
+        // check the file metadata
+        let expected_kv_meta = vec![
             KeyValue {
                 key: "my-data".to_string(),
                 value: Some("stuff".to_string()),
@@ -2380,7 +2384,7 @@ mod tests {
                 value: None,
             },
         ];
-        assert_eq!(key_value_metadata, expected_metadata);
+        assert_file_metadata(file_metadata, expected_kv_meta);
 
         Ok(())
     }
@@ -2391,18 +2395,8 @@ mod tests {
         let file_path = format!("file:///path/to/{}", filename);
         let parquet_sink = create_written_parquet_sink(file_path.as_str()).await?;
 
-        // assert written
-        let mut written = parquet_sink.written();
-        let written = written.drain();
-        assert_eq!(
-            written.len(),
-            1,
-            "expected a single parquet file to be written, instead found {}",
-            written.len()
-        );
-
-        let (path, ..) = written.take(1).next().unwrap();
-
+        // assert written to proper path
+        let (path, _) = get_written(parquet_sink)?;
         let path_parts = path.parts().collect::<Vec<_>>();
         assert_eq!(
             path_parts.len(),
@@ -2420,18 +2414,8 @@ mod tests {
         let file_path = "file:///path/to";
         let parquet_sink = create_written_parquet_sink(file_path).await?;
 
-        // assert written
-        let mut written = parquet_sink.written();
-        let written = written.drain();
-        assert_eq!(
-            written.len(),
-            1,
-            "expected a single parquet file to be written, instead found {}",
-            written.len()
-        );
-
-        let (path, ..) = written.take(1).next().unwrap();
-
+        // assert written to proper path
+        let (path, _) = get_written(parquet_sink)?;
         let path_parts = path.parts().collect::<Vec<_>>();
         assert_eq!(
             path_parts.len(),
@@ -2449,18 +2433,8 @@ mod tests {
         let file_path = "file:///path/to/";
         let parquet_sink = create_written_parquet_sink(file_path).await?;
 
-        // assert written
-        let mut written = parquet_sink.written();
-        let written = written.drain();
-        assert_eq!(
-            written.len(),
-            1,
-            "expected a single parquet file to be written, instead found {}",
-            written.len()
-        );
-
-        let (path, ..) = written.take(1).next().unwrap();
-
+        // assert written to proper path
+        let (path, _) = get_written(parquet_sink)?;
         let path_parts = path.parts().collect::<Vec<_>>();
         assert_eq!(
             path_parts.len(),
@@ -2474,6 +2448,14 @@ mod tests {
     }
 
     async fn create_written_parquet_sink(table_path: &str) -> Result<Arc<ParquetSink>> {
+        create_written_parquet_sink_using_config(table_path, ParquetOptions::default())
+            .await
+    }
+
+    async fn create_written_parquet_sink_using_config(
+        table_path: &str,
+        global: ParquetOptions,
+    ) -> Result<Arc<ParquetSink>> {
         let field_a = Field::new("a", DataType::Utf8, false);
         let field_b = Field::new("b", DataType::Utf8, false);
         let schema = Arc::new(Schema::new(vec![field_a, field_b]));
@@ -2495,6 +2477,7 @@ mod tests {
                     ("my-data".to_string(), Some("stuff".to_string())),
                     ("my-data-bool-key".to_string(), None),
                 ]),
+                global,
                 ..Default::default()
             },
         ));
@@ -2519,6 +2502,42 @@ mod tests {
         Ok(parquet_sink)
     }
 
+    fn get_written(parquet_sink: Arc<ParquetSink>) -> Result<(Path, FileMetaData)> {
+        let mut written = parquet_sink.written();
+        let written = written.drain();
+        assert_eq!(
+            written.len(),
+            1,
+            "expected a single parquet files to be written, instead found {}",
+            written.len()
+        );
+
+        let (path, file_metadata) = written.take(1).next().unwrap();
+        Ok((path, file_metadata))
+    }
+
+    fn assert_file_metadata(file_metadata: FileMetaData, expected_kv: Vec<KeyValue>) {
+        let FileMetaData {
+            num_rows,
+            schema,
+            key_value_metadata,
+            ..
+        } = file_metadata;
+        assert_eq!(num_rows, 2, "file metadata to have 2 rows");
+        assert!(
+            schema.iter().any(|col_schema| col_schema.name == "a"),
+            "output file metadata should contain col a"
+        );
+        assert!(
+            schema.iter().any(|col_schema| col_schema.name == "b"),
+            "output file metadata should contain col b"
+        );
+
+        let mut key_value_metadata = key_value_metadata.unwrap();
+        key_value_metadata.sort_by(|a, b| a.key.cmp(&b.key));
+        assert_eq!(key_value_metadata, expected_kv);
+    }
+
     #[tokio::test]
     async fn parquet_sink_write_partitions() -> Result<()> {
         let field_a = Field::new("a", DataType::Utf8, false);