jayzhan211 commented on code in PR #26953:
URL: https://github.com/apache/datafusion/pull/26953#discussion_r4538684594
##########
datafusion/physical-plan/src/sorts/sort/spill_tests.rs:
##########
@@ -679,6 +683,202 @@ async fn test_sort_preparation_keeps_output_reserved() ->
Result<()> {
Ok(())
}
+/// Merging sorted runs in the row format holds the keys of every run's current
+/// batch. Two `Int32` keys encode to 18 bytes per row against 8 bytes in the
+/// batches, and a single `FixedSizeBinary(16)` key and two long `Utf8View`
keys
+/// also encode to more than their runs, so merging full-size runs needs more
+/// memory than the sort reserves while buffering them. The payloads make runs
+/// larger than their keys or keep buffers shared after sorting, so these runs
+/// keep `batch_size` and must not be charged once per batch either.
+async fn check_merge_keys_larger_than_batches() -> Result<()> {
+ let rows = 8192;
+ let int_keys = |batch: i32| -> Vec<ArrayRef> {
+ vec![
+ Arc::new(Int32Array::from_iter_values(
+ (0..rows).map(|i| (i * 31 + batch) % 1000),
+ )),
+ Arc::new(Int32Array::from_iter_values((0..rows).rev())),
+ ]
+ };
+ let int_fields = [("a", DataType::Int32), ("b", DataType::Int32)];
+ check_merge_keys(&int_fields, 2, 16, true, |batch|
Ok(int_keys(batch))).await?;
+ // Longer than 12 bytes, so views keep them in data buffers.
+ let strings = |batch: i32| -> StringArray {
+ (0..rows)
+ .map(|i| Some(format!("{batch}-{i:0>64}")))
+ .collect()
+ };
+ let dictionary =
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8));
+ for payload in [DataType::Utf8View, dictionary] {
+ let fields = [&int_fields[..], &[("c", payload.clone())]].concat();
+ // Runs of this `Utf8View` payload are not split, as on `main`
+ let within_pool = payload != DataType::Utf8View;
+ check_merge_keys(&fields, 2, 16, within_pool, |batch| {
+ let payload = cast(&strings(batch), &payload)?;
+ Ok([int_keys(batch), vec![payload]].concat())
+ })
+ .await?;
+ }
+ // Views into 1024 dictionary values share their bytes, as a
+ // dictionary-encoded Parquet column produces. Sorting the buffered batches
+ // before a spill grows past the pool, on `main` too.
+ let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
+ check_merge_keys(&fields, 2, 16, false, |batch| {
Review Comment:
Since the key-size check went in, no test reaches `views_share_bytes`:
removing `&& !views_share_bytes(run)` from `merge_run_batch_size` still passes
every `sorts::` and `spill::` test. It's still needed for computed keys, which
always count as larger than the run. With round 2's `shared.parquet`, this
passes on `main` and here at `-m 256M`, `512M` and `1G`, and fails without the
gate with `Not enough memory to continue external sort`:
```sql
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'shared.parquet' ORDER BY a + 1, b;
```
Could you add a computed-key case? The diff below passes here and fails
without the gate.
<details>
<summary>Test diff for spill_tests.rs</summary>
```diff
-use datafusion_physical_expr::expressions::Column;
-use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_expr::expressions::{CastExpr, Column};
+use datafusion_physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
@@ async fn check_merge_keys_larger_than_batches
let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
- check_merge_keys(&fields, 2, 16, false, |batch| {
+ let shared_views = |batch: i32| -> Result<Vec<ArrayRef>> {
let values: StringArray = (0..1024)
@@
Ok([int_keys(batch), vec![views]].concat())
- })
- .await?;
+ };
+ check_merge_keys(&fields, 2, 16, false, shared_views).await?;
+ // A computed key counts as larger than the run, so only the shared
views
+ // keep this run at `batch_size`.
+ let a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let keys = vec![
+ Arc::new(CastExpr::new(a, DataType::Int64, None)) as Arc<dyn
PhysicalExpr>,
+ Arc::new(Column::new("b", 1)),
+ ];
+ check_merge_exprs(&fields, keys, 16, false, shared_views).await?;
@@ async fn check_merge_keys(
columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
+) -> Result<()> {
+ let keys = (0..keys)
+ .map(|i| Arc::new(Column::new(fields[i].0, i)) as Arc<dyn
PhysicalExpr>)
+ .collect();
+ check_merge_exprs(fields, keys, buffered, within_pool, columns).await
+}
+
+/// Like [`check_merge_keys`], sorting by `keys`.
+async fn check_merge_exprs(
+ fields: &[(&str, DataType)],
+ keys: Vec<Arc<dyn PhysicalExpr>>,
+ buffered: usize,
+ within_pool: bool,
+ columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
) -> Result<()> {
@@
- let ordering =
- LexOrdering::new((0..keys).map(|i| {
- PhysicalSortExpr::new_default(Arc::new(Column::new(fields[i].0,
i)))
- }))
+ let ordering =
LexOrdering::new(keys.into_iter().map(PhysicalSortExpr::new_default))
.expect("at least one sort key");
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
