Weijun-H commented on code in PR #24889:
URL: https://github.com/apache/datafusion/pull/24889#discussion_r3967560263
##########
datafusion/core/tests/memory_limit/mod.rs:
##########
@@ -223,6 +224,117 @@ mod count_distinct_spill {
}
}
+/// `GROUP BY` on a single nested key under a memory limit, on both the legacy
+/// `GroupedHashAggregateStream` and the migrated streams. The legacy stream
used
+/// to emit duplicate groups after spilling.
+#[tokio::test]
+async fn nested_key_spill_keeps_groups_unique() {
+ const NESTED_KEY_ROWS: usize = 200_000;
+ const NESTED_KEY_GROUPS: i64 = 16;
+ const NESTED_KEY_BATCH_ROWS: usize = 8_192;
+
+ /// Small enough that the final stages must spill their `count(distinct)`
+ /// state, large enough that the migrated final stream can hold one merged
+ /// batch under a `FairSpillPool` shared by four partitions.
+ const NESTED_KEY_MEMORY_LIMIT: usize = 8 * 1024 * 1024;
+
+ fn nested_key_struct_fields() -> Fields {
+ Fields::from(vec![
+ Field::new("list", DataType::new_list(DataType::Int64, true),
true),
+ Field::new("num", DataType::Int64, true),
+ ])
+ }
+
+ fn nested_key_table() -> MemTable {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new_struct("st", nested_key_struct_fields(), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batches = (0..NESTED_KEY_ROWS)
+ .step_by(NESTED_KEY_BATCH_ROWS)
+ .map(|start| {
+ let rows = start..(start +
NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS);
+ let mut list = ListBuilder::new(Int64Builder::new());
+ let mut num = Vec::with_capacity(rows.len());
+ let mut valid = Vec::with_capacity(rows.len());
+ for row in rows.clone() {
+ let group = row as i64 % NESTED_KEY_GROUPS;
+ match row % 37 {
+ 0 => list.append_null(),
+ 1 => list.append(true),
+ _ => {
+ list.values().append_value(group);
+ list.values().append_value(group + 1);
+ list.append(true);
+ }
+ }
+ num.push((row % 41 != 0).then_some(group));
+ valid.push(row % 43 != 0);
+ }
+ let st = StructArray::new(
+ nested_key_struct_fields(),
+ vec![Arc::new(list.finish()),
Arc::new(Int64Array::from(num))],
+ Some(NullBuffer::from(valid)),
+ );
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(st),
+ Arc::new(Int64Array::from_iter_values(
+ rows.map(|row| row as i64),
+ )),
+ ],
+ )
+ .unwrap()
+ })
+ .collect();
+ MemTable::try_new(schema, vec![batches]).unwrap()
+ }
+
+ async fn run_nested_key_query(memory_limit: Option<usize>, legacy: bool)
-> String {
+ let mut runtime = RuntimeEnvBuilder::new()
+ .with_disk_manager_builder(DiskManagerBuilder::default());
+ if let Some(limit) = memory_limit {
+ runtime =
runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit)));
+ }
+ let config = SessionConfig::new()
+ .with_target_partitions(4)
+ // small batches: the merged spill stream arrives in many batches
and
+ // groups span batch boundaries
+ .with_batch_size(64)
+ .set_bool("datafusion.execution.enable_migration_aggregate",
!legacy);
+ let ctx =
+ SessionContext::new_with_config_rt(config,
runtime.build_arc().unwrap());
+ ctx.register_table("t", Arc::new(nested_key_table()))
+ .unwrap();
+ let batches = ctx
+ .sql(
+ "select st, count(v), count(distinct v), sum(v), avg(v),
min(v), max(v) \
+ from t group by st",
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ batches_to_sort_string(&batches)
+ }
+
+ let expected = run_nested_key_query(None, false).await;
+ for legacy in [true, false] {
+ assert_eq!(
+ run_nested_key_query(None, legacy).await,
+ expected,
+ "unbounded, legacy={legacy}"
+ );
+ assert_eq!(
+ run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT), legacy).await,
Review Comment:
Could we assert that the aggregate actually spills in the memory-limited
runs? Otherwise this can keep passing after a plan or memory change stops
exercising the fix.
##########
datafusion/core/tests/memory_limit/mod.rs:
##########
@@ -223,6 +224,117 @@ mod count_distinct_spill {
}
}
+/// `GROUP BY` on a single nested key under a memory limit, on both the legacy
+/// `GroupedHashAggregateStream` and the migrated streams. The legacy stream
used
+/// to emit duplicate groups after spilling.
+#[tokio::test]
+async fn nested_key_spill_keeps_groups_unique() {
+ const NESTED_KEY_ROWS: usize = 200_000;
+ const NESTED_KEY_GROUPS: i64 = 16;
+ const NESTED_KEY_BATCH_ROWS: usize = 8_192;
+
+ /// Small enough that the final stages must spill their `count(distinct)`
+ /// state, large enough that the migrated final stream can hold one merged
+ /// batch under a `FairSpillPool` shared by four partitions.
+ const NESTED_KEY_MEMORY_LIMIT: usize = 8 * 1024 * 1024;
+
+ fn nested_key_struct_fields() -> Fields {
+ Fields::from(vec![
+ Field::new("list", DataType::new_list(DataType::Int64, true),
true),
+ Field::new("num", DataType::Int64, true),
+ ])
+ }
+
+ fn nested_key_table() -> MemTable {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new_struct("st", nested_key_struct_fields(), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batches = (0..NESTED_KEY_ROWS)
+ .step_by(NESTED_KEY_BATCH_ROWS)
+ .map(|start| {
+ let rows = start..(start +
NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS);
+ let mut list = ListBuilder::new(Int64Builder::new());
+ let mut num = Vec::with_capacity(rows.len());
+ let mut valid = Vec::with_capacity(rows.len());
+ for row in rows.clone() {
+ let group = row as i64 % NESTED_KEY_GROUPS;
+ match row % 37 {
+ 0 => list.append_null(),
+ 1 => list.append(true),
+ _ => {
+ list.values().append_value(group);
+ list.values().append_value(group + 1);
+ list.append(true);
+ }
+ }
+ num.push((row % 41 != 0).then_some(group));
+ valid.push(row % 43 != 0);
+ }
+ let st = StructArray::new(
+ nested_key_struct_fields(),
+ vec![Arc::new(list.finish()),
Arc::new(Int64Array::from(num))],
+ Some(NullBuffer::from(valid)),
+ );
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(st),
+ Arc::new(Int64Array::from_iter_values(
+ rows.map(|row| row as i64),
+ )),
+ ],
+ )
+ .unwrap()
+ })
+ .collect();
+ MemTable::try_new(schema, vec![batches]).unwrap()
+ }
+
+ async fn run_nested_key_query(memory_limit: Option<usize>, legacy: bool)
-> String {
+ let mut runtime = RuntimeEnvBuilder::new()
+ .with_disk_manager_builder(DiskManagerBuilder::default());
+ if let Some(limit) = memory_limit {
+ runtime =
runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit)));
+ }
+ let config = SessionConfig::new()
+ .with_target_partitions(4)
+ // small batches: the merged spill stream arrives in many batches
and
+ // groups span batch boundaries
+ .with_batch_size(64)
+ .set_bool("datafusion.execution.enable_migration_aggregate",
!legacy);
+ let ctx =
+ SessionContext::new_with_config_rt(config,
runtime.build_arc().unwrap());
+ ctx.register_table("t", Arc::new(nested_key_table()))
+ .unwrap();
+ let batches = ctx
+ .sql(
+ "select st, count(v), count(distinct v), sum(v), avg(v),
min(v), max(v) \
+ from t group by st",
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ batches_to_sort_string(&batches)
+ }
+
+ let expected = run_nested_key_query(None, false).await;
+ for legacy in [true, false] {
+ assert_eq!(
+ run_nested_key_query(None, legacy).await,
Review Comment:
The unlimited migrated run already produces `expected`. We can skip
repeating it here and keep only the unlimited legacy check plus both limited
runs.
--
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]
