Skip to content

Commit f0b2d79

Browse files
2010YOUY012010YOUY01
committed
review: move old test to new impl + additional test
1 parent 3125d9a commit f0b2d79

3 files changed

Lines changed: 276 additions & 169 deletions

File tree

‎datafusion/physical-plan/src/aggregates/mod.rs‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4104,7 +4104,94 @@ mod tests {
41044104
}
41054105

41064106
#[tokio::test]
4107-
async fn test_migrated_hash_stream_skip_aggregation_after_threshold() -> Result<()> {
4107+
async fn test_partial_hash_stream_skip_aggregation_after_first_batch() -> Result<()> {
4108+
let schema = Arc::new(Schema::new(vec![
4109+
Field::new("key", DataType::Int32, true),
4110+
Field::new("val", DataType::Int32, true),
4111+
]));
4112+
4113+
let group_by =
4114+
PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]);
4115+
4116+
let aggr_expr = vec![
4117+
AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?])
4118+
.schema(Arc::clone(&schema))
4119+
.alias(String::from("COUNT(val)"))
4120+
.build()
4121+
.map(Arc::new)?,
4122+
];
4123+
4124+
let input_data = vec![
4125+
RecordBatch::try_new(
4126+
Arc::clone(&schema),
4127+
vec![
4128+
Arc::new(Int32Array::from(vec![1, 2, 3])),
4129+
Arc::new(Int32Array::from(vec![0, 0, 0])),
4130+
],
4131+
)
4132+
.unwrap(),
4133+
RecordBatch::try_new(
4134+
Arc::clone(&schema),
4135+
vec![
4136+
Arc::new(Int32Array::from(vec![2, 3, 4])),
4137+
Arc::new(Int32Array::from(vec![0, 0, 0])),
4138+
],
4139+
)
4140+
.unwrap(),
4141+
];
4142+
4143+
let input =
4144+
TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?;
4145+
let aggregate_exec = Arc::new(AggregateExec::try_new(
4146+
AggregateMode::Partial,
4147+
group_by,
4148+
aggr_expr,
4149+
vec![None],
4150+
Arc::clone(&input) as Arc<dyn ExecutionPlan>,
4151+
schema,
4152+
)?);
4153+
4154+
let session_config = SessionConfig::default()
4155+
.set_bool("datafusion.execution.enable_migration_aggregate", true)
4156+
.set(
4157+
"datafusion.execution.skip_partial_aggregation_probe_rows_threshold",
4158+
&ScalarValue::Int64(Some(2)),
4159+
)
4160+
.set(
4161+
"datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
4162+
&ScalarValue::Float64(Some(0.1)),
4163+
);
4164+
4165+
let ctx = Arc::new(TaskContext::default().with_session_config(session_config));
4166+
let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?;
4167+
4168+
allow_duplicates! {
4169+
assert_snapshot!(batches_to_sort_string(&output), @r"
4170+
+-----+-------------------+
4171+
| key | COUNT(val)[count] |
4172+
+-----+-------------------+
4173+
| 1 | 1 |
4174+
| 2 | 1 |
4175+
| 2 | 1 |
4176+
| 3 | 1 |
4177+
| 3 | 1 |
4178+
| 4 | 1 |
4179+
+-----+-------------------+
4180+
");
4181+
}
4182+
4183+
let metrics = aggregate_exec.metrics().unwrap();
4184+
let skipped_rows = metrics
4185+
.sum_by_name("skipped_aggregation_rows")
4186+
.map(|m| m.as_usize())
4187+
.unwrap_or(0);
4188+
assert_eq!(skipped_rows, 3);
4189+
4190+
Ok(())
4191+
}
4192+
4193+
#[tokio::test]
4194+
async fn test_partial_hash_stream_skip_aggregation_after_threshold() -> Result<()> {
41084195
let schema = Arc::new(Schema::new(vec![
41094196
Field::new("key", DataType::Int32, true),
41104197
Field::new("val", DataType::Int32, true),

‎datafusion/physical-plan/src/aggregates/row_hash.rs‎

Lines changed: 0 additions & 168 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,7 +1386,6 @@ impl GroupedHashAggregateStream {
13861386
mod tests {
13871387
use super::*;
13881388
use crate::InputOrderMode;
1389-
use crate::execution_plan::ExecutionPlan;
13901389
use crate::test::TestMemoryExec;
13911390
use arrow::array::{Int32Array, Int64Array};
13921391
use arrow::datatypes::{DataType, Field, Schema};
@@ -1501,152 +1500,6 @@ mod tests {
15011500
Ok(())
15021501
}
15031502

1504-
#[tokio::test]
1505-
async fn test_skip_aggregation_probe_not_locked_until_skip() -> Result<()> {
1506-
// Test that the probe is not locked until we actually decide to skip.
1507-
// This allows us to continue evaluating the skip condition across multiple batches.
1508-
//
1509-
// Scenario:
1510-
// - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip
1511-
// - Batch 2: Now hits ratio threshold (high cardinality) -> skip
1512-
//
1513-
// Without the fix, the probe would be locked after batch 1, preventing the skip
1514-
// decision from being made on batch 2.
1515-
1516-
let schema = Arc::new(Schema::new(vec![
1517-
Field::new("group_col", DataType::Int32, false),
1518-
Field::new("value_col", DataType::Int32, false),
1519-
]));
1520-
1521-
// Configure thresholds:
1522-
// - probe_rows_threshold: 100 rows
1523-
// - probe_ratio_threshold: 0.8 (80%)
1524-
let probe_rows_threshold = 100;
1525-
let probe_ratio_threshold = 0.8;
1526-
1527-
// Batch 1: 100 rows with only 10 unique groups
1528-
// Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip
1529-
// This will hit the rows threshold but not the ratio threshold
1530-
let batch1_rows = 100;
1531-
let batch1_groups = 10;
1532-
let mut group_ids_batch1 = Vec::new();
1533-
for i in 0..batch1_rows {
1534-
group_ids_batch1.push((i % batch1_groups) as i32);
1535-
}
1536-
let values_batch1: Vec<i32> = vec![1; batch1_rows];
1537-
1538-
let batch1 = RecordBatch::try_new(
1539-
Arc::clone(&schema),
1540-
vec![
1541-
Arc::new(Int32Array::from(group_ids_batch1)),
1542-
Arc::new(Int32Array::from(values_batch1)),
1543-
],
1544-
)?;
1545-
1546-
// Batch 2: 360 rows with 360 unique NEW groups (starting from group 10)
1547-
// After batch 2, total: 460 rows, 370 groups
1548-
// Ratio: 370/460 ≈ 0.804 (80.4%) > 0.8 -> SHOULD decide to skip
1549-
let batch2_rows = 360;
1550-
let batch2_groups = 360;
1551-
let group_ids_batch2: Vec<i32> = (batch1_groups..(batch1_groups + batch2_groups))
1552-
.map(|x| x as i32)
1553-
.collect();
1554-
let values_batch2: Vec<i32> = vec![1; batch2_rows];
1555-
1556-
let batch2 = RecordBatch::try_new(
1557-
Arc::clone(&schema),
1558-
vec![
1559-
Arc::new(Int32Array::from(group_ids_batch2)),
1560-
Arc::new(Int32Array::from(values_batch2)),
1561-
],
1562-
)?;
1563-
1564-
// Batch 3: This batch should be skipped since we decided to skip after batch 2
1565-
// 100 rows with 100 unique groups (continuing from where batch 2 left off)
1566-
let batch3_rows = 100;
1567-
let batch3_groups = 100;
1568-
let batch3_start_group = batch1_groups + batch2_groups;
1569-
let group_ids_batch3: Vec<i32> = (batch3_start_group
1570-
..(batch3_start_group + batch3_groups))
1571-
.map(|x| x as i32)
1572-
.collect();
1573-
let values_batch3: Vec<i32> = vec![1; batch3_rows];
1574-
1575-
let batch3 = RecordBatch::try_new(
1576-
Arc::clone(&schema),
1577-
vec![
1578-
Arc::new(Int32Array::from(group_ids_batch3)),
1579-
Arc::new(Int32Array::from(values_batch3)),
1580-
],
1581-
)?;
1582-
1583-
let input_partitions = vec![vec![batch1, batch2, batch3]];
1584-
1585-
let runtime = RuntimeEnvBuilder::default().build_arc()?;
1586-
let mut task_ctx = TaskContext::default().with_runtime(runtime);
1587-
1588-
// Configure skip aggregation settings
1589-
let mut session_config = task_ctx.session_config().clone();
1590-
session_config = session_config.set(
1591-
"datafusion.execution.skip_partial_aggregation_probe_rows_threshold",
1592-
&datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)),
1593-
);
1594-
session_config = session_config.set(
1595-
"datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
1596-
&datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)),
1597-
);
1598-
task_ctx = task_ctx.with_session_config(session_config);
1599-
let task_ctx = Arc::new(task_ctx);
1600-
1601-
// Create aggregate: COUNT(*) GROUP BY group_col
1602-
let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())];
1603-
let aggr_expr = vec![Arc::new(
1604-
AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?])
1605-
.schema(Arc::clone(&schema))
1606-
.alias("count_value")
1607-
.build()?,
1608-
)];
1609-
1610-
let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?;
1611-
let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec)));
1612-
1613-
// Use Partial mode
1614-
let aggregate_exec = AggregateExec::try_new(
1615-
AggregateMode::Partial,
1616-
PhysicalGroupBy::new_single(group_expr),
1617-
aggr_expr,
1618-
vec![None],
1619-
exec,
1620-
Arc::clone(&schema),
1621-
)?;
1622-
1623-
// Execute and collect results
1624-
let mut stream =
1625-
GroupedHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?;
1626-
let mut results = Vec::new();
1627-
1628-
while let Some(result) = stream.next().await {
1629-
let batch = result?;
1630-
results.push(batch);
1631-
}
1632-
1633-
// Check that skip aggregation actually happened
1634-
// The key metric is skipped_aggregation_rows
1635-
let metrics = aggregate_exec.metrics().unwrap();
1636-
let skipped_rows = metrics
1637-
.sum_by_name("skipped_aggregation_rows")
1638-
.map(|m| m.as_usize())
1639-
.unwrap_or(0);
1640-
1641-
// We expect batch 3's rows to be skipped (100 rows)
1642-
assert_eq!(
1643-
skipped_rows, batch3_rows,
1644-
"Expected batch 3's rows ({batch3_rows}) to be skipped",
1645-
);
1646-
1647-
Ok(())
1648-
}
1649-
16501503
#[tokio::test]
16511504
async fn test_emit_early_with_partially_sorted() -> Result<()> {
16521505
// Reproducer for #20445: EmitEarly with PartiallySorted panics in
@@ -1730,25 +1583,4 @@ mod tests {
17301583

17311584
Ok(())
17321585
}
1733-
1734-
#[test]
1735-
fn test_skip_aggregation_probe_equality_does_not_skip() {
1736-
// When num_groups / input_rows == probe_ratio_threshold, the `>` boundary
1737-
// means we must NOT skip — equality is not sufficient to trigger skip.
1738-
let threshold_ratio = 0.5_f64;
1739-
let threshold_rows = 10_usize;
1740-
let mut probe = SkipAggregationProbe::new(
1741-
threshold_rows,
1742-
threshold_ratio,
1743-
metrics::Count::new(),
1744-
);
1745-
1746-
// 10 rows, 5 groups → ratio = 5/10 = 0.5 exactly equals threshold
1747-
probe.update_state(10, 5);
1748-
1749-
assert!(
1750-
!probe.should_skip(),
1751-
"ratio == threshold should not trigger skip (boundary is exclusive)"
1752-
);
1753-
}
17541586
}

0 commit comments

Comments
 (0)