jayzhan211 commented on code in PR #10554:
URL: https://github.com/apache/datafusion/pull/10554#discussion_r1604216406
##########
datafusion/core/src/physical_optimizer/aggregate_statistics.rs:
##########
@@ -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));
Review Comment:
Since the optimized code does not have specific behavior on the non-indirect
plan, I think the test is not necessary
--
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]
