Skip to content

Commit c7a5a93

Browse files
authored
Implement test for parquet pruning disabling (#754)
1 parent e4df37a commit c7a5a93

File tree

1 file changed

+51
-7
lines changed

1 file changed

+51
-7
lines changed

datafusion/tests/parquet_pruning.rs

Lines changed: 51 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,10 @@ use arrow::{
3030
};
3131
use chrono::{Datelike, Duration};
3232
use datafusion::{
33-
datasource::{parquet::ParquetTable, TableProvider},
33+
datasource::TableProvider,
3434
logical_plan::{col, lit, Expr, LogicalPlan, LogicalPlanBuilder},
3535
physical_plan::{plan_metrics, SQLMetric},
36-
prelude::ExecutionContext,
36+
prelude::{ExecutionConfig, ExecutionContext},
3737
scalar::ScalarValue,
3838
};
3939
use hashbrown::HashMap;
@@ -136,6 +136,47 @@ async fn prune_date64() {
136136
assert_eq!(output.result_rows, 1, "{}", output.description());
137137
}
138138

139+
#[tokio::test]
140+
async fn prune_disabled() {
141+
let query = "SELECT * FROM t where nanos < to_timestamp('2020-01-02 01:01:11Z')";
142+
let expected_rows = 10;
143+
144+
// with pruning
145+
let output = ContextWithParquet::new(Scenario::Timestamps)
146+
.await
147+
.query(query)
148+
.await;
149+
150+
// This should prune one without error
151+
assert_eq!(output.predicate_evaluation_errors(), Some(0));
152+
assert_eq!(output.row_groups_pruned(), Some(1));
153+
assert_eq!(
154+
output.result_rows,
155+
expected_rows,
156+
"{}",
157+
output.description()
158+
);
159+
160+
// same query, without pruning
161+
let config = ExecutionConfig::new().with_parquet_pruning(false);
162+
163+
let output = ContextWithParquet::with_config(Scenario::Timestamps, config)
164+
.await
165+
.query(query)
166+
.await;
167+
println!("{}", output.description());
168+
169+
// This should not prune any
170+
assert_eq!(output.predicate_evaluation_errors(), Some(0));
171+
assert_eq!(output.row_groups_pruned(), Some(0));
172+
assert_eq!(
173+
output.result_rows,
174+
expected_rows,
175+
"{}",
176+
output.description()
177+
);
178+
}
179+
139180
// ----------------------
140181
// Begin test fixture
141182
// ----------------------
@@ -207,15 +248,18 @@ impl TestOutput {
207248
/// and the appropriate scenario
208249
impl ContextWithParquet {
209250
async fn new(scenario: Scenario) -> Self {
210-
let file = make_test_file(scenario).await;
251+
Self::with_config(scenario, ExecutionConfig::new()).await
252+
}
211253

212-
// now, setup a the file as a data source and run a query against it
213-
let mut ctx = ExecutionContext::new();
254+
async fn with_config(scenario: Scenario, config: ExecutionConfig) -> Self {
255+
let file = make_test_file(scenario).await;
214256
let parquet_path = file.path().to_string_lossy();
215257

216-
let table = ParquetTable::try_new(parquet_path, 4).unwrap();
258+
// now, setup a the file as a data source and run a query against it
259+
let mut ctx = ExecutionContext::with_config(config);
217260

218-
let provider = Arc::new(table);
261+
ctx.register_parquet("t", &parquet_path).unwrap();
262+
let provider = ctx.deregister_table("t").unwrap().unwrap();
219263
ctx.register_table("t", provider.clone()).unwrap();
220264

221265
Self {

0 commit comments

Comments
 (0)