Skip to content

Commit

Permalink
cleaner api
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed May 22, 2024
1 parent 2daf742 commit ec6fbb1
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 10 deletions.
9 changes: 4 additions & 5 deletions datafusion-examples/examples/csv_opener.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

use std::{sync::Arc, vec};

use datafusion::common::Statistics;
use datafusion::{
assert_batches_eq,
datasource::{
Expand Down Expand Up @@ -58,11 +57,11 @@ async fn main() -> Result<()> {

let path = std::path::Path::new(&path).canonicalize()?;

let mut scan_config =
let scan_config =
FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema.clone())
.with_projection(vec![12, 0])
.with_limit(Some(5));
scan_config.add_file(PartitionedFile::new(path.display().to_string(), 10));
.with_projection(Some(vec![12, 0]))
.with_limit(Some(5))
.with_file(PartitionedFile::new(path.display().to_string(), 10));

let result =
FileStream::new(&scan_config, 0, opener, &ExecutionPlanMetricsSet::new())
Expand Down
8 changes: 3 additions & 5 deletions datafusion-examples/examples/json_opener.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ use datafusion::{
error::Result,
physical_plan::metrics::ExecutionPlanMetricsSet,
};
use datafusion_common::Statistics;

use futures::StreamExt;
use object_store::ObjectStore;
Expand Down Expand Up @@ -61,12 +60,11 @@ async fn main() -> Result<()> {
Arc::new(object_store),
);

let mut scan_config =
let scan_config =
FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema.clone())
.with_projection(Some(vec![1, 0]))
.with_limit(5);

scan_config.add_file(PartitionedFile::new(path.to_string(), 10));
.with_limit(Some(5))
.with_file(PartitionedFile::new(path.to_string(), 10));

let result =
FileStream::new(&scan_config, 0, opener, &ExecutionPlanMetricsSet::new())
Expand Down
16 changes: 16 additions & 0 deletions datafusion/core/src/datasource/physical_plan/file_scan_config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,22 @@ impl FileScanConfig {
self
}

/// Add a file as a single group
///
/// See [Self::file_groups] for more information.
pub fn with_file(mut self, file: PartitionedFile) -> Self {
self.add_file(file);
self
}

/// Add a new file group
///
/// See [Self::file_groups] for more information
pub fn with_file_group(mut self, file_group: Vec<PartitionedFile>) -> Self {
self.add_file_group(file_group);
self
}

/// Set the partitioning columns of the files
pub fn with_table_partition_cols(mut self, table_partition_cols: Vec<Field>) -> Self {
self.table_partition_cols = table_partition_cols;
Expand Down

0 comments on commit ec6fbb1

Please sign in to comment.