alamb commented on code in PR #3051:
URL: https://github.com/apache/arrow-datafusion/pull/3051#discussion_r941486353
##########
datafusion/core/src/physical_plan/file_format/parquet.rs:
##########
@@ -856,6 +857,107 @@ mod tests {
)
}
+ const EXPECTED_USER_DEFINED_METADATA: &str = "some-user-defined-metadata";
+
+ #[tokio::test]
+ async fn route_data_access_ops_to_parquet_file_reader_factory() {
+ let c1: ArrayRef =
+ Arc::new(StringArray::from(vec![Some("Foo"), None, Some("bar")]));
+ let c2: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), Some(2),
None]));
+ let c3: ArrayRef = Arc::new(Int8Array::from(vec![Some(10), Some(20),
None]));
+
+ let batch = create_batch(vec![
+ ("c1", c1.clone()),
+ ("c2", c2.clone()),
+ ("c3", c3.clone()),
+ ]);
+
+ let file_schema = batch.schema().clone();
+ let (in_memory_object_store, parquet_files_meta) =
+ store_parquet_in_memory(vec![batch]).await;
+ let file_groups = parquet_files_meta
+ .into_iter()
+ .map(|meta| PartitionedFile {
+ object_meta: meta,
+ partition_values: vec![],
+ range: None,
+ extensions:
Some(Arc::new(String::from(EXPECTED_USER_DEFINED_METADATA))),
+ })
+ .collect();
+
+ // prepare the scan
+ let parquet_exec = ParquetExec::new(
+ FileScanConfig {
+ // just any url that doesn't point to in memory object store
+ object_store_url: ObjectStoreUrl::local_filesystem(),
+ file_groups: vec![file_groups],
+ file_schema,
+ statistics: Statistics::default(),
+ projection: None,
+ limit: None,
+ table_partition_cols: vec![],
+ },
+ None,
+ None,
+ )
+ .with_parquet_file_reader_factory(Arc::new(
+
InMemoryParquetFileReaderFactory(Arc::clone(&in_memory_object_store)),
+ ));
+
+ let session_ctx = SessionContext::new();
+
+ let task_ctx = session_ctx.task_ctx();
+ let read = collect(Arc::new(parquet_exec), task_ctx).await.unwrap();
+
+ let expected = vec![
+ "+-----+----+----+",
+ "| c1 | c2 | c3 |",
+ "+-----+----+----+",
+ "| Foo | 1 | 10 |",
+ "| | 2 | 20 |",
+ "| bar | | |",
+ "+-----+----+----+",
+ ];
+
+ assert_batches_sorted_eq!(expected, &read);
+ }
+
+ #[derive(Debug)]
+ struct InMemoryParquetFileReaderFactory(Arc<dyn ObjectStore>);
Review Comment:
Thanks @Cheappie -- I think the comments are good enough
In fact it is perhaps good to see what was required to be made `pub` to use
these new APIs. Hopefully it makes downstream integrations easier
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
