apache / apache/datafusion

Repartitioning by RoundRobinBatch(n) does not result in at least n batches for CSV sources

Open
#7,145 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Describe the bug

I've noticed that when reading from a CSV source table, explicit calls to `DataFrame.repartition(Partitioning::RoundRobinBatch(n))` behave differently vs. when reading from a JSON or Parquet source. For some reason when reading from CSV, the total number of batches generated is often less than `n`. This means we end up with empty partitions, which I think is surprising from a user perspective. If I explicitly request `n` partitions, I should expect `n` nonempty partitions if the source has sufficient data.

I would also expect to see consistent behavior regardless of the type of source table I am reading from.

### To Reproduce

First example, read from parquet and write to csv behaves as expected:
```rust
#[tokio::main]
async fn main() -> Result<()> {
// create local execution context
let ctx = SessionContext::new();

let local = Arc::new(LocalFileSystem::new());
let local_url = Url::parse("file://local").unwrap();
ctx.runtime_env().register_object_store(&local_url, local);

ctx.register_parquet(
"source",
"file://local/home/dev/dftest/test5.parquet",
ParquetReadOptions::default(),
)
.await?;
let path = "file://local/home/dev/dftest/testdata/";
let df = ctx
.sql("select * from source")
.await?
.repartition(Partitioning::RoundRobinBatch(10))?;
let in_path = "file://local/home/dev/dftest/csv_data/";
df.write_csv(in_path).await?;
```

```bash
ls -lh csv_data
total 70M
-rw-rw-r-- 1 dev dev 7.5M Jul 30 12:53 part-0.csv
-rw-rw-r-- 1 dev dev 7.5M Jul 30 12:53 part-1.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-2.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-3.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-4.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-5.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-6.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-7.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-8.csv
-rw-rw-r-- 1 dev dev 6.9M Jul 30 12:53 part-9.csv
```

Second example is nearly identical, except we make the source data CSV instead.
```rust
#[tokio::main]
async fn main() -> Result<()> {
// create local execution context
let ctx = SessionContext::new();

let local = Arc::new(LocalFileSystem::new());
let local_url = Url::parse("file://local").unwrap();
ctx.runtime_env().register_object_store(&local_url, local);

ctx.register_csv(
"source",
"file://local/home/dev/dftest/test5.csv",
CsvReadOptions::default(),
)
.await?;
let path = "file://local/home/dev/dftest/testdata/";
let df = ctx
.sql("select * from source")
.await?
.repartition(Partitioning::RoundRobinBatch(10))?;
let in_path = "file://local/home/dev/dftest/csv_data/";
df.write_csv(in_path).await?;
```

Now, we see many empty partitions written out.

```bash
ls -lh csv_data
total 70M
-rw-rw-r-- 1 dev dev 18M Jul 30 12:53 part-0.csv
-rw-rw-r-- 1 dev dev 17M Jul 30 12:53 part-1.csv
-rw-rw-r-- 1 dev dev 18M Jul 30 12:53 part-2.csv
-rw-rw-r-- 1 dev dev 16M Jul 30 12:53 part-3.csv
-rw-rw-r-- 1 dev dev 1.9M Jul 30 12:53 part-4.csv
-rw-rw-r-- 1 dev dev 0 Jul 30 12:53 part-5.csv
-rw-rw-r-- 1 dev dev 0 Jul 30 12:53 part-6.csv
-rw-rw-r-- 1 dev dev 0 Jul 30 12:53 part-7.csv
-rw-rw-r-- 1 dev dev 0 Jul 30 12:53 part-8.csv
-rw-rw-r-- 1 dev dev 0 Jul 30 12:53 part-9.csv
```

### Expected behavior

The examples above should behave the same regardless of if the source is CSV or not. To me, it makes sense for `RoundRobinBatch(n)` to ensure there is a minimum of `n` batches so there can be `n` nonempty partitions.

### Additional context

_No response_

Contributor guide

Open the contributing guide

Research direction

Reproduce the difference between register_csv and register_parquet using DataFrame.repartition(Partitioning::RoundRobinBatch(10)) and write_csv. Start by tracing the CSV source execution path and compare its produced batches with the Parquet path; done means a CSV source with sufficient data produces at least 10 nonempty output partitions, consistently with Parquet.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.