(
file_compression_type: FileCompressionType,
)
| 530 | #[cfg(feature = "compression")] |
| 531 | #[tokio::test] |
| 532 | async fn query_compress_data( |
| 533 | file_compression_type: FileCompressionType, |
| 534 | ) -> Result<()> { |
| 535 | use arrow_schema::{DataType, Field, Schema}; |
| 536 | use datafusion_common::DataFusionError; |
| 537 | use datafusion_datasource::file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD; |
| 538 | use futures::TryStreamExt; |
| 539 | |
| 540 | let mut cfg = SessionConfig::new(); |
| 541 | cfg.options_mut().catalog.has_header = true; |
| 542 | let session_state = SessionStateBuilder::new() |
| 543 | .with_config(cfg) |
| 544 | .with_default_features() |
| 545 | .build(); |
| 546 | let integration = LocalFileSystem::new_with_prefix(arrow_test_data()).unwrap(); |
| 547 | let path = Path::from("csv/aggregate_test_100.csv"); |
| 548 | let csv = CsvFormat::default().with_has_header(true); |
| 549 | let records_to_read = csv |
| 550 | .options() |
| 551 | .schema_infer_max_rec |
| 552 | .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); |
| 553 | let store = Arc::new(integration) as Arc<dyn ObjectStore>; |
| 554 | let original_stream = store.get(&path).await?; |
| 555 | |
| 556 | //convert original_stream to compressed_stream for next step |
| 557 | let compressed_stream = |
| 558 | file_compression_type.to_owned().convert_to_compress_stream( |
| 559 | original_stream |
| 560 | .into_stream() |
| 561 | .map_err(DataFusionError::from) |
| 562 | .boxed(), |
| 563 | ); |
| 564 | |
| 565 | //prepare expected schema for assert_eq |
| 566 | let expected = Schema::new(vec![ |
| 567 | Field::new("c1", DataType::Utf8, true), |
| 568 | Field::new("c2", DataType::Int64, true), |
| 569 | Field::new("c3", DataType::Int64, true), |
| 570 | Field::new("c4", DataType::Int64, true), |
| 571 | Field::new("c5", DataType::Int64, true), |
| 572 | Field::new("c6", DataType::Int64, true), |
| 573 | Field::new("c7", DataType::Int64, true), |
| 574 | Field::new("c8", DataType::Int64, true), |
| 575 | Field::new("c9", DataType::Int64, true), |
| 576 | Field::new("c10", DataType::Utf8, true), |
| 577 | Field::new("c11", DataType::Float64, true), |
| 578 | Field::new("c12", DataType::Float64, true), |
| 579 | Field::new("c13", DataType::Utf8, true), |
| 580 | ]); |
| 581 | |
| 582 | let compressed_csv = csv.with_file_compression_type(file_compression_type); |
| 583 | |
| 584 | //convert compressed_stream to decoded_stream |
| 585 | let decoded_stream = compressed_csv |
| 586 | .read_to_delimited_chunks_from_stream(compressed_stream.unwrap()) |
| 587 | .await; |
| 588 | let (schema, records_read) = compressed_csv |
| 589 | .infer_schema_from_stream(&session_state, records_to_read, decoded_stream) |
nothing calls this directly
no test coverage detected
searching dependent graphs…