-
|
Consider this code: Note: better example in the comments below use std::sync::Arc;
use datafusion::{
config::ConfigOptions,
error::DataFusionError,
physical_plan::{ExecutionPlan, displayable},
prelude::*,
};
fn repartition_subtree(
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
let new_children = plan
.children()
.into_iter()
.map(|child| repartition_subtree(Arc::clone(child), config))
.collect::<Result<_, _>>()?;
let plan = plan.with_new_children(new_children)?;
// change partitions to 4
let repartitioned_plan = plan.repartitioned(4, config)?;
if let Some(new_plan) = repartitioned_plan {
Ok(new_plan)
} else {
Ok(plan)
}
}
#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
let config = SessionConfig::new().with_target_partitions(2);
// register the table
let ctx = SessionContext::new_with_config(config);
ctx.register_parquet("example", "<path to datafusion repo>/datafusion/core/tests/data/test_statistics_per_partition/", ParquetReadOptions::new())
.await?;
// create a plan to run a SQL query
let df = ctx
.sql("SELECT id, MIN(0) FROM example GROUP BY id LIMIT 100")
.await?;
let plan = df.create_physical_plan().await.unwrap();
println!("plan: {}", displayable(plan.as_ref()).indent(false));
let state = ctx.state();
let new_plan = repartition_subtree(plan, state.config_options()).unwrap();
println!("new_plan: {}", displayable(new_plan.as_ref()).indent(false));
Ok(())
}The output is this: I plan the query with a Notice that the number of It seems to be because of this code: datafusion/datafusion/datasource/src/file_groups.rs Lines 192 to 199 in d24eb4a Even though, I'm scanning entire files in this case, the range is set in these file groups and repartitioning is skipped! |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 6 replies
-
|
I don't think it's that ranges check, from what I can see it's due to this condition: datafusion/datafusion/datasource/src/file_groups.rs Lines 219 to 225 in 6ea305e Specifically, the datafusion/datafusion/datasource/src/source.rs Lines 250 to 263 in 6ea305e
Which is from this config: datafusion/datafusion/common/src/config.rs Lines 931 to 932 in 6ea305e If we set it to something smaller than let config = SessionConfig::new()
.with_target_partitions(2)
.with_repartition_file_min_size(1024); // HereOutput: |
Beta Was this translation helpful? Give feedback.
-
|
You're right! In this example, it is due to the file sizes.
But even in this example, there is one problem. With target_partitions set to 4, I get 4 file groups as expected. let config = SessionConfig::new()
.with_target_partitions(4);I would have expected setting |
Beta Was this translation helpful? Give feedback.
-
|
The way I see it, they are governed by separate configs; it looks like |
Beta Was this translation helpful? Give feedback.
-
|
Here's a better example: use std::sync::Arc;
use datafusion::{
config::ConfigOptions,
error::DataFusionError,
physical_plan::{ExecutionPlan, displayable},
prelude::*,
};
fn repartition_subtree(
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
let new_children = plan
.children()
.into_iter()
.map(|child| repartition_subtree(Arc::clone(child), config))
.collect::<Result<_, _>>()?;
let plan = plan.with_new_children(new_children)?;
// change partitions to 4
let repartitioned_plan = plan.repartitioned(4, config)?;
if let Some(new_plan) = repartitioned_plan {
Ok(new_plan)
} else {
Ok(plan)
}
}
#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
let config = SessionConfig::new()
.with_target_partitions(2);
// register the table
let ctx = SessionContext::new_with_config(config);
ctx.register_parquet("example", "./partitioned_output", ParquetReadOptions::new())
.await?;
// create a plan to run a SQL query
let df = ctx
.sql("SELECT b, MIN(0) FROM example GROUP BY b LIMIT 100")
.await?;
let plan = df.create_physical_plan().await.unwrap();
println!("plan: {}", displayable(plan.as_ref()).indent(false));
let state = ctx.state();
let new_plan = repartition_subtree(plan, state.config_options()).unwrap();
println!("new_plan: {}", displayable(new_plan.as_ref()).indent(false));
Ok(())
}Here's the code to generate the data (excuse the code quality, it was written with AI): use std::sync::Arc;
use arrow::array::{Int32Array, Int64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::dataframe::DataFrameWriteOptions;
use datafusion::prelude::*;
use rand::Rng;
const NUM_ROWS: usize = 10_000_000;
const BATCH_SIZE: usize = 1_000_000; // Process in batches to manage memory
#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
println!("Creating DataFusion context...");
let ctx = SessionContext::new();
// Define schema
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false), // 0 to 1000
Field::new("b", DataType::Int64, false), // 0 to 1,000,000
]));
println!("Generating {} random rows...", NUM_ROWS);
// Generate data in batches to avoid memory issues
let mut batches = Vec::new();
let mut rng = rand::thread_rng();
for batch_num in 0..(NUM_ROWS / BATCH_SIZE) {
let start = batch_num * BATCH_SIZE;
println!(
" Generating batch {} ({} - {} rows)...",
batch_num + 1,
start,
start + BATCH_SIZE
);
// Generate random data for column 'a' (0 to 1000)
let col_a: Vec<i32> = (0..BATCH_SIZE).map(|_| rng.gen_range(0..=1000)).collect();
// Generate random data for column 'b' (0 to 1,000,000)
let col_b: Vec<i64> = (0..BATCH_SIZE)
.map(|_| rng.gen_range(0..=1_000_000))
.collect();
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(col_a)),
Arc::new(Int64Array::from(col_b)),
],
)?;
batches.push(batch);
}
println!("Creating DataFrame from {} batches...", batches.len());
// Create a DataFrame from the record batches
let df = ctx.read_batches(batches)?;
// Show sample data
println!("\nSample data (first 10 rows):");
df.clone().limit(0, Some(10))?.show().await?;
// Show row count
let row_count = df.clone().count().await?;
println!("\nTotal rows: {}", row_count);
// Write as partitioned parquet files (partitioned by column 'a')
let output_path = "./partitioned_output";
println!(
"\nWriting partitioned parquet files to '{}'...",
output_path
);
println!("Partitioning by column 'a'...");
df.write_parquet(
output_path,
DataFrameWriteOptions::new().with_partition_by(vec!["a".to_string()]),
None, // Use default parquet writer options
)
.await?;
println!("\n✓ Successfully wrote partitioned parquet files!");
println!(" Output directory: {}", output_path);
Ok(())
}With the latest main of datafusion (commit The partitioning remains the same. Now, if I remove the range check, we get this: The partitioning increases as expected. Now my question is this: is this something that is just lacking an implementation? I don't see a good reason to avoid repartitioning when ranges exist for a file. I don't have enough context to decide this. I'm willing to implement this if there's no objections to it. |
Beta Was this translation helpful? Give feedback.
-
|
From what I can tell, I think it was mainly to simplify implementation; I can't see a specific reason we disallow repartitioning if there are existing ranges in terms of correctness. For reference, this PR introduced that logic so it is quite old: #5057 Feel free to raise an issue & PR if you feel like tackling it |
Beta Was this translation helpful? Give feedback.
-
|
Alright, that makes sense. I'll open an issue (and probably a PR too). Thanks! |
Beta Was this translation helpful? Give feedback.
-
|
Opened an issue: #18940 |
Beta Was this translation helpful? Give feedback.
From what I can tell, I think it was mainly to simplify implementation; I can't see a specific reason we disallow repartitioning if there are existing ranges in terms of correctness. For reference, this PR introduced that logic so it is quite old: #5057
Feel free to raise an issue & PR if you feel like tackling it