Skip to content

Move aggregate statistic , combine aggregate and limited distinct aggregate test to slt #10554

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

Move aggregate statistic , combine aggregate and limited distinct aggregate test to slt #10554

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
62cfab2
move test to slt
jayzhan211 May 17, 2024
770be4d
remove more
jayzhan211 May 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
358 changes: 0 additions & 358 deletions datafusion/core/src/physical_optimizer/aggregate_statistics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,361 +304,3 @@ fn take_optimizable_max(
}
None
}

#[cfg(test)]
pub(crate) mod tests {

use super::*;
use crate::logical_expr::Operator;
use crate::physical_plan::aggregates::PhysicalGroupBy;
use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use crate::physical_plan::common;
use crate::physical_plan::expressions::Count;
use crate::physical_plan::filter::FilterExec;
use crate::physical_plan::memory::MemoryExec;
use crate::prelude::SessionContext;

use arrow::array::Int32Array;
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion_common::cast::as_int64_array;
use datafusion_physical_expr::expressions::cast;
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_plan::aggregates::AggregateMode;

/// Mock data using a MemoryExec which has an exact count statistic
fn mock_data() -> Result<Arc<MemoryExec>> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
]));

let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int32Array::from(vec![Some(1), Some(2), None])),
Arc::new(Int32Array::from(vec![Some(4), None, Some(6)])),
],
)?;

Ok(Arc::new(MemoryExec::try_new(
&[vec![batch]],
Arc::clone(&schema),
None,
)?))
}

/// Checks that the count optimization was applied and we still get the right result
async fn assert_count_optim_success(
plan: AggregateExec,
agg: TestAggregate,
) -> Result<()> {
let session_ctx = SessionContext::new();
let state = session_ctx.state();
let plan: Arc<dyn ExecutionPlan> = Arc::new(plan);

let optimized = AggregateStatistics::new()
.optimize(Arc::clone(&plan), state.config_options())?;

// A ProjectionExec is a sign that the count optimization was applied
assert!(optimized.as_any().is::<ProjectionExec>());

// run both the optimized and nonoptimized plan
let optimized_result =
common::collect(optimized.execute(0, session_ctx.task_ctx())?).await?;
let nonoptimized_result =
common::collect(plan.execute(0, session_ctx.task_ctx())?).await?;
assert_eq!(optimized_result.len(), nonoptimized_result.len());

// and validate the results are the same and expected
assert_eq!(optimized_result.len(), 1);
check_batch(optimized_result.into_iter().next().unwrap(), &agg);
// check the non optimized one too to ensure types and names remain the same
assert_eq!(nonoptimized_result.len(), 1);
check_batch(nonoptimized_result.into_iter().next().unwrap(), &agg);

Ok(())
}

fn check_batch(batch: RecordBatch, agg: &TestAggregate) {
let schema = batch.schema();
let fields = schema.fields();
assert_eq!(fields.len(), 1);

let field = &fields[0];
assert_eq!(field.name(), agg.column_name());
assert_eq!(field.data_type(), &DataType::Int64);
// note that nullabiolity differs

assert_eq!(
as_int64_array(batch.column(0)).unwrap().values(),
&[agg.expected_count()]
);
}

/// Describe the type of aggregate being tested
pub(crate) enum TestAggregate {
/// Testing COUNT(*) type aggregates
CountStar,

/// Testing for COUNT(column) aggregate
ColumnA(Arc<Schema>),
}

impl TestAggregate {
pub(crate) fn new_count_star() -> Self {
Self::CountStar
}

fn new_count_column(schema: &Arc<Schema>) -> Self {
Self::ColumnA(schema.clone())
}

/// Return appropriate expr depending if COUNT is for col or table (*)
pub(crate) fn count_expr(&self) -> Arc<dyn AggregateExpr> {
Arc::new(Count::new(
self.column(),
self.column_name(),
DataType::Int64,
))
}

/// what argument would this aggregate need in the plan?
fn column(&self) -> Arc<dyn PhysicalExpr> {
match self {
Self::CountStar => expressions::lit(COUNT_STAR_EXPANSION),
Self::ColumnA(s) => expressions::col("a", s).unwrap(),
}
}

/// What name would this aggregate produce in a plan?
fn column_name(&self) -> &'static str {
match self {
Self::CountStar => "COUNT(*)",
Self::ColumnA(_) => "COUNT(a)",
}
}

/// What is the expected count?
fn expected_count(&self) -> i64 {
match self {
TestAggregate::CountStar => 3,
TestAggregate::ColumnA(_) => 2,
}
}
}

#[tokio::test]
async fn test_count_partial_direct_child() -> Result<()> {
// basic test case with the aggregation applied on a source with exact statistics
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_star();

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
source,
Arc::clone(&schema),
)?;

let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(partial_agg),
Arc::clone(&schema),
)?;

assert_count_optim_success(final_agg, agg).await?;

Ok(())
}

#[tokio::test]
async fn test_count_partial_with_nulls_direct_child() -> Result<()> {
// basic test case with the aggregation applied on a source with exact statistics
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_column(&schema);

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
source,
Arc::clone(&schema),
)?;

let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(partial_agg),
Arc::clone(&schema),
)?;

assert_count_optim_success(final_agg, agg).await?;

Ok(())
}

#[tokio::test]
async fn test_count_partial_indirect_child() -> Result<()> {
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_star();

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
source,
Arc::clone(&schema),
)?;

// We introduce an intermediate optimization step between the partial and final aggregtator
let coalesce = CoalescePartitionsExec::new(Arc::new(partial_agg));
Copy link
Contributor Author

@jayzhan211 jayzhan211 May 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the optimized code does not have specific behavior on the non-indirect plan, I think the test is not necessary


let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(coalesce),
Arc::clone(&schema),
)?;

assert_count_optim_success(final_agg, agg).await?;

Ok(())
}

#[tokio::test]
async fn test_count_partial_with_nulls_indirect_child() -> Result<()> {
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_column(&schema);

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
source,
Arc::clone(&schema),
)?;

// We introduce an intermediate optimization step between the partial and final aggregtator
let coalesce = CoalescePartitionsExec::new(Arc::new(partial_agg));

let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(coalesce),
Arc::clone(&schema),
)?;

assert_count_optim_success(final_agg, agg).await?;

Ok(())
}

#[tokio::test]
async fn test_count_inexact_stat() -> Result<()> {
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_star();

// adding a filter makes the statistics inexact
let filter = Arc::new(FilterExec::try_new(
expressions::binary(
expressions::col("a", &schema)?,
Operator::Gt,
cast(expressions::lit(1u32), &schema, DataType::Int32)?,
&schema,
)?,
source,
)?);

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
filter,
Arc::clone(&schema),
)?;

let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(partial_agg),
Arc::clone(&schema),
)?;

let conf = ConfigOptions::new();
let optimized =
AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;

// check that the original ExecutionPlan was not replaced
assert!(optimized.as_any().is::<AggregateExec>());

Ok(())
}

#[tokio::test]
async fn test_count_with_nulls_inexact_stat() -> Result<()> {
let source = mock_data()?;
let schema = source.schema();
let agg = TestAggregate::new_count_column(&schema);

// adding a filter makes the statistics inexact
let filter = Arc::new(FilterExec::try_new(
expressions::binary(
expressions::col("a", &schema)?,
Operator::Gt,
cast(expressions::lit(1u32), &schema, DataType::Int32)?,
&schema,
)?,
source,
)?);

let partial_agg = AggregateExec::try_new(
AggregateMode::Partial,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
filter,
Arc::clone(&schema),
)?;

let final_agg = AggregateExec::try_new(
AggregateMode::Final,
PhysicalGroupBy::default(),
vec![agg.count_expr()],
vec![None],
Arc::new(partial_agg),
Arc::clone(&schema),
)?;

let conf = ConfigOptions::new();
let optimized =
AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?;

// check that the original ExecutionPlan was not replaced
assert!(optimized.as_any().is::<AggregateExec>());

Ok(())
}
}
Loading