You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/06/25 22:41:08 UTC

[GitHub] [arrow-datafusion] Dandandan opened a new pull request #620: [WIP] Optimize count(*) with table statistics

Dandandan opened a new pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620


   # Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Closes #.
   
    # Rationale for this change
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   
   # What changes are included in this PR?
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   # Are there any user-facing changes?
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->
   


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-870078343


   > Good stuff 👍 I recommend using logical plan builder to construct the new plans to make things more robust.
   
   Perhaps we can do this as a follow on PR


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659305261



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {
+                            Ok(LogicalPlan::Projection {
+                                expr: vec![Expr::Alias(
+                                    Box::new(Expr::Literal(ScalarValue::UInt64(Some(
+                                        num_rows as u64,
+                                    )))),
+                                    "COUNT(Uint8(1))".to_string(),
+                                )],
+                                input: Arc::new(LogicalPlan::EmptyRelation {
+                                    produce_one_row: true,
+                                    schema: Arc::new(DFSchema::empty()),
+                                }),
+                                schema: schema.clone(),
+                            })
+                        }
+                        _ => Ok(plan.clone()),
+                    };
+                }
+
+                Ok(plan.clone())
+            }
+            // Rest: recurse and find possible statistics
+            _ => {
+                let expr = plan.expressions();
+
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|plan| self.optimize(plan, execution_props))
+                    .collect::<Result<Vec<_>>>()?;
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+        }
+    }
+
+    fn name(&self) -> &str {
+        "count_statistics"
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use crate::error::Result;
+    use crate::execution::context::ExecutionProps;
+    use crate::logical_plan::LogicalPlan;
+    use crate::optimizer::aggregate_statistics::AggregateStatistics;
+    use crate::optimizer::optimizer::OptimizerRule;
+    use crate::{
+        datasource::{datasource::Statistics, TableProvider},
+        logical_plan::Expr,
+    };
+
+    struct TestTableProvider {
+        num_rows: usize,
+        is_exact: bool,
+    }
+
+    impl TableProvider for TestTableProvider {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
+        }
+
+        fn scan(
+            &self,
+            _projection: &Option<Vec<usize>>,
+            _batch_size: usize,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn statistics(&self) -> crate::datasource::datasource::Statistics {
+            Statistics {
+                num_rows: Some(self.num_rows),
+                total_byte_size: None,
+                column_statistics: None,
+            }
+        }
+        fn has_exact_statistics(&self) -> bool {
+            self.is_exact
+        }
+    }
+
+    #[test]
+    fn optimize_count_using_statistics() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Projection: UInt64(100) AS COUNT(Uint8(1))\
+            \n    EmptyRelation";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_not_exact() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: false,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+            \n    TableScan: test projection=None";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
+        let opt = AggregateStatistics::new();
+        let optimized_plan = opt.optimize(plan, &ExecutionProps::new()).unwrap();
+        let formatted_plan = format!("{:?}", optimized_plan);
+        assert_eq!(formatted_plan, expected);
+        assert_eq!(plan.schema(), plan.schema());
+    }

Review comment:
       I suggest adding the following two negative tests:
   1. `SELECT count(*), a FROM test GROUP BY a`
   2 `SELECT count(*) FROM test WHERE a < 5`




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan edited a comment on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan edited a comment on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-869183633


   After this PR I also plan some follow up issues:
   
   * also supporting min/max statistics


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan edited a comment on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan edited a comment on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-869183633


   After this PR I also plan some follow up issues:
   
   * also supporting min/max statistics


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r660147059



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {

Review comment:
       oops -- sorry I missed that. Test is 👍 

##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,335 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::{sync::Arc, vec};
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{col, DFField, DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() => {
+                // aggregations that can not be replaced
+                // using statistics
+                let mut agg = vec![];
+                // expressions that can be replaced by constants
+                let mut projections = vec![];
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    for expr in aggr_expr {
+                        match expr {
+                            Expr::AggregateFunction {
+                                fun: AggregateFunction::Count,
+                                args,
+                                distinct: false,
+                            } if args
+                                == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] =>
+                            {
+                                projections.push(Expr::Alias(
+                                    Box::new(Expr::Literal(ScalarValue::UInt64(Some(
+                                        num_rows as u64,
+                                    )))),
+                                    "COUNT(Uint8(1))".to_string(),
+                                ));
+                            }
+                            _ => {
+                                agg.push(expr.clone());
+                            }
+                        }
+                    }
+
+                    return Ok(if agg.is_empty() {
+                        // table scan can be entirely removed
+
+                        LogicalPlan::Projection {
+                            expr: projections,
+                            input: Arc::new(LogicalPlan::EmptyRelation {
+                                produce_one_row: true,
+                                schema: Arc::new(DFSchema::empty()),
+                            }),
+                            schema: schema.clone(),
+                        }
+                    } else if projections.is_empty() {
+                        // no replacements -> return original plan
+                        plan.clone()
+                    } else {
+                        // Split into parts that can be supported and part that should stay in aggregate
+                        let agg_fields = agg
+                            .iter()
+                            .map(|x| x.to_field(input.schema()))
+                            .collect::<Result<Vec<DFField>>>()?;
+                        let agg_schema = DFSchema::new(agg_fields)?;
+                        let cols = agg
+                            .iter()
+                            .map(|e| e.name(&agg_schema))
+                            .collect::<Result<Vec<String>>>()?;
+                        projections.extend(cols.iter().map(|x| col(x)));
+                        LogicalPlan::Projection {
+                            expr: projections,
+                            schema: schema.clone(),
+                            input: Arc::new(LogicalPlan::Aggregate {
+                                input: input.clone(),
+                                group_expr: vec![],
+                                aggr_expr: agg,
+                                schema: Arc::new(agg_schema),
+                            }),
+                        }
+                    });
+                }
+                Ok(plan.clone())
+            }
+            // Rest: recurse and find possible statistics
+            _ => {
+                let expr = plan.expressions();
+
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|plan| self.optimize(plan, execution_props))
+                    .collect::<Result<Vec<_>>>()?;
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+        }
+    }
+
+    fn name(&self) -> &str {
+        "aggregate_statistics"
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use crate::error::Result;
+    use crate::execution::context::ExecutionProps;
+    use crate::logical_plan::LogicalPlan;
+    use crate::optimizer::aggregate_statistics::AggregateStatistics;
+    use crate::optimizer::optimizer::OptimizerRule;
+    use crate::{
+        datasource::{datasource::Statistics, TableProvider},
+        logical_plan::Expr,
+    };
+
+    struct TestTableProvider {
+        num_rows: usize,
+        is_exact: bool,
+    }
+
+    impl TableProvider for TestTableProvider {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
+        }
+
+        fn scan(
+            &self,
+            _projection: &Option<Vec<usize>>,
+            _batch_size: usize,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn statistics(&self) -> crate::datasource::datasource::Statistics {
+            Statistics {
+                num_rows: Some(self.num_rows),
+                total_byte_size: None,
+                column_statistics: None,
+            }
+        }
+        fn has_exact_statistics(&self) -> bool {
+            self.is_exact
+        }
+    }
+
+    #[test]
+    fn optimize_count_using_statistics() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Projection: UInt64(100) AS COUNT(Uint8(1))\
+            \n    EmptyRelation";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_not_exact() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: false,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+            \n    TableScan: test projection=None";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_sum() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select sum(a)/count(*) from test")

Review comment:
       This is a cool optimization 👍 




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-869183633


   After 
   
   * Optimizing out `count(*)` even when we have multiple aggregation expressions (e.g sum+ count).
   * also supporting min/max statistics


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb merged pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb merged pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620


   


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659305326



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {

Review comment:
       🤔  about this, don't we also need to check for empty `gby` as well?




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-870078343


   > Good stuff 👍 I recommend using logical plan builder to construct the new plans to make things more robust.
   
   Perhaps we can do this as a follow on PR


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan edited a comment on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan edited a comment on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-869183633


   After this PR I also plan some follow issues:
   
   * Optimizing out `count(*)` even when we have multiple aggregation expressions (e.g sum+ count).
   * also supporting min/max statistics


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb merged pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb merged pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620


   


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r660147832



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,335 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::{sync::Arc, vec};
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{col, DFField, DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() => {
+                // aggregations that can not be replaced
+                // using statistics
+                let mut agg = vec![];
+                // expressions that can be replaced by constants
+                let mut projections = vec![];
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    for expr in aggr_expr {
+                        match expr {
+                            Expr::AggregateFunction {
+                                fun: AggregateFunction::Count,
+                                args,
+                                distinct: false,
+                            } if args
+                                == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] =>
+                            {
+                                projections.push(Expr::Alias(
+                                    Box::new(Expr::Literal(ScalarValue::UInt64(Some(
+                                        num_rows as u64,
+                                    )))),
+                                    "COUNT(Uint8(1))".to_string(),
+                                ));
+                            }
+                            _ => {
+                                agg.push(expr.clone());
+                            }
+                        }
+                    }
+
+                    return Ok(if agg.is_empty() {
+                        // table scan can be entirely removed
+
+                        LogicalPlan::Projection {
+                            expr: projections,
+                            input: Arc::new(LogicalPlan::EmptyRelation {
+                                produce_one_row: true,
+                                schema: Arc::new(DFSchema::empty()),
+                            }),
+                            schema: schema.clone(),
+                        }
+                    } else if projections.is_empty() {
+                        // no replacements -> return original plan
+                        plan.clone()
+                    } else {
+                        // Split into parts that can be supported and part that should stay in aggregate
+                        let agg_fields = agg
+                            .iter()
+                            .map(|x| x.to_field(input.schema()))
+                            .collect::<Result<Vec<DFField>>>()?;
+                        let agg_schema = DFSchema::new(agg_fields)?;
+                        let cols = agg
+                            .iter()
+                            .map(|e| e.name(&agg_schema))
+                            .collect::<Result<Vec<String>>>()?;
+                        projections.extend(cols.iter().map(|x| col(x)));
+                        LogicalPlan::Projection {
+                            expr: projections,
+                            schema: schema.clone(),
+                            input: Arc::new(LogicalPlan::Aggregate {
+                                input: input.clone(),
+                                group_expr: vec![],
+                                aggr_expr: agg,
+                                schema: Arc::new(agg_schema),
+                            }),
+                        }
+                    });
+                }
+                Ok(plan.clone())
+            }
+            // Rest: recurse and find possible statistics
+            _ => {
+                let expr = plan.expressions();
+
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|plan| self.optimize(plan, execution_props))
+                    .collect::<Result<Vec<_>>>()?;
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+        }
+    }
+
+    fn name(&self) -> &str {
+        "aggregate_statistics"
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use crate::error::Result;
+    use crate::execution::context::ExecutionProps;
+    use crate::logical_plan::LogicalPlan;
+    use crate::optimizer::aggregate_statistics::AggregateStatistics;
+    use crate::optimizer::optimizer::OptimizerRule;
+    use crate::{
+        datasource::{datasource::Statistics, TableProvider},
+        logical_plan::Expr,
+    };
+
+    struct TestTableProvider {
+        num_rows: usize,
+        is_exact: bool,
+    }
+
+    impl TableProvider for TestTableProvider {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
+        }
+
+        fn scan(
+            &self,
+            _projection: &Option<Vec<usize>>,
+            _batch_size: usize,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn statistics(&self) -> crate::datasource::datasource::Statistics {
+            Statistics {
+                num_rows: Some(self.num_rows),
+                total_byte_size: None,
+                column_statistics: None,
+            }
+        }
+        fn has_exact_statistics(&self) -> bool {
+            self.is_exact
+        }
+    }
+
+    #[test]
+    fn optimize_count_using_statistics() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Projection: UInt64(100) AS COUNT(Uint8(1))\
+            \n    EmptyRelation";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_not_exact() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: false,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+            \n    TableScan: test projection=None";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_sum() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select sum(a)/count(*) from test")

Review comment:
       This is a cool optimization 👍 




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659339611



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {

Review comment:
       We have `group_expr.is_empty()`, but we should add a test for this as well 👍 




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan edited a comment on pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan edited a comment on pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#issuecomment-869183633


   After this PR I also plan some follow issues:
   
   * also supporting min/max statistics


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659494281



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {
+                            Ok(LogicalPlan::Projection {
+                                expr: vec![Expr::Alias(
+                                    Box::new(Expr::Literal(ScalarValue::UInt64(Some(
+                                        num_rows as u64,
+                                    )))),
+                                    "COUNT(Uint8(1))".to_string(),
+                                )],
+                                input: Arc::new(LogicalPlan::EmptyRelation {
+                                    produce_one_row: true,
+                                    schema: Arc::new(DFSchema::empty()),
+                                }),
+                                schema: schema.clone(),
+                            })
+                        }
+                        _ => Ok(plan.clone()),
+                    };
+                }
+
+                Ok(plan.clone())
+            }
+            // Rest: recurse and find possible statistics
+            _ => {
+                let expr = plan.expressions();
+
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|plan| self.optimize(plan, execution_props))
+                    .collect::<Result<Vec<_>>>()?;
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+        }
+    }
+
+    fn name(&self) -> &str {
+        "count_statistics"
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use crate::error::Result;
+    use crate::execution::context::ExecutionProps;
+    use crate::logical_plan::LogicalPlan;
+    use crate::optimizer::aggregate_statistics::AggregateStatistics;
+    use crate::optimizer::optimizer::OptimizerRule;
+    use crate::{
+        datasource::{datasource::Statistics, TableProvider},
+        logical_plan::Expr,
+    };
+
+    struct TestTableProvider {
+        num_rows: usize,
+        is_exact: bool,
+    }
+
+    impl TableProvider for TestTableProvider {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
+        }
+
+        fn scan(
+            &self,
+            _projection: &Option<Vec<usize>>,
+            _batch_size: usize,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn statistics(&self) -> crate::datasource::datasource::Statistics {
+            Statistics {
+                num_rows: Some(self.num_rows),
+                total_byte_size: None,
+                column_statistics: None,
+            }
+        }
+        fn has_exact_statistics(&self) -> bool {
+            self.is_exact
+        }
+    }
+
+    #[test]
+    fn optimize_count_using_statistics() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Projection: UInt64(100) AS COUNT(Uint8(1))\
+            \n    EmptyRelation";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn optimize_count_not_exact() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: false,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+            \n    TableScan: test projection=None";
+
+        assert_optimized_plan_eq(&plan, expected);
+        Ok(())
+    }
+
+    fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
+        let opt = AggregateStatistics::new();
+        let optimized_plan = opt.optimize(plan, &ExecutionProps::new()).unwrap();
+        let formatted_plan = format!("{:?}", optimized_plan);
+        assert_eq!(formatted_plan, expected);
+        assert_eq!(plan.schema(), plan.schema());
+    }

Review comment:
       Done




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659553253



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,

Review comment:
       This is what it returns (I think this is the correct result), so I guess we should not do this for distinct.
   ```
   > select count(distinct 1) from X;
   +--------------------------+
   | COUNT(DISTINCT UInt8(1)) |
   +--------------------------+
   | 1                        |
   +--------------------------+
   1 row in set. Query took 0.332 seconds.
   ```
   




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659340234



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,

Review comment:
       Is this something we would see in practice?




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659339127



##########
File path: datafusion/src/datasource/datasource.rs
##########
@@ -108,6 +108,11 @@ pub trait TableProvider: Sync + Send {
     /// Statistics should be optional because not all data sources can provide statistics.
     fn statistics(&self) -> Statistics;
 
+    /// Returns whether statistics provided are exact values or estimates

Review comment:
       That's a similar thought process I had. 
   Maybe at some point it would also be nice to tell what parts of the statistics are exact (e.g. number of rows) and what estimated (such as distinct count).




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r659303495



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,

Review comment:
       Given the argument is a constant, you can probably also match `distinct: _` - so like `COUNT(distinct 1)` is the same as `COUNT(1)` 

##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {

Review comment:
       You could probably make this more general to match *any* literal value and optimize things like `SELECT count('foo') FROM bar`. However, given this is the rewrite that the DataFusion planner does for `SELECT * FROM bar` I think you have covered the most important case
   
   ```suggestion
                           } if args == &[Expr::Literal(_)] => {
   ```

##########
File path: datafusion/src/datasource/datasource.rs
##########
@@ -108,6 +108,11 @@ pub trait TableProvider: Sync + Send {
     /// Statistics should be optional because not all data sources can provide statistics.
     fn statistics(&self) -> Statistics;
 
+    /// Returns whether statistics provided are exact values or estimates

Review comment:
       The nice thing about adding a `has_exact_statistics` is that it is a backwards compatible API
   
   An alternate might be to encapsulate the "Exact statistics or not" into a field on `Statistics` itself, which feels to me like it keeps related things together more, but has the downside of changing Statistics / APIs

##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {
+                            Ok(LogicalPlan::Projection {
+                                expr: vec![Expr::Alias(
+                                    Box::new(Expr::Literal(ScalarValue::UInt64(Some(
+                                        num_rows as u64,
+                                    )))),
+                                    "COUNT(Uint8(1))".to_string(),
+                                )],
+                                input: Arc::new(LogicalPlan::EmptyRelation {
+                                    produce_one_row: true,
+                                    schema: Arc::new(DFSchema::empty()),
+                                }),
+                                schema: schema.clone(),
+                            })
+                        }
+                        _ => Ok(plan.clone()),
+                    };
+                }
+
+                Ok(plan.clone())
+            }
+            // Rest: recurse and find possible statistics
+            _ => {
+                let expr = plan.expressions();
+
+                // apply the optimization to all inputs of the plan
+                let inputs = plan.inputs();
+                let new_inputs = inputs
+                    .iter()
+                    .map(|plan| self.optimize(plan, execution_props))
+                    .collect::<Result<Vec<_>>>()?;
+
+                utils::from_plan(plan, &expr, &new_inputs)
+            }
+        }
+    }
+
+    fn name(&self) -> &str {
+        "count_statistics"
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::datatypes::{DataType, Field, Schema};
+
+    use crate::error::Result;
+    use crate::execution::context::ExecutionProps;
+    use crate::logical_plan::LogicalPlan;
+    use crate::optimizer::aggregate_statistics::AggregateStatistics;
+    use crate::optimizer::optimizer::OptimizerRule;
+    use crate::{
+        datasource::{datasource::Statistics, TableProvider},
+        logical_plan::Expr,
+    };
+
+    struct TestTableProvider {
+        num_rows: usize,
+        is_exact: bool,
+    }
+
+    impl TableProvider for TestTableProvider {
+        fn as_any(&self) -> &dyn std::any::Any {
+            unimplemented!()
+        }
+        fn schema(&self) -> arrow::datatypes::SchemaRef {
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]))
+        }
+
+        fn scan(
+            &self,
+            _projection: &Option<Vec<usize>>,
+            _batch_size: usize,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn statistics(&self) -> crate::datasource::datasource::Statistics {
+            Statistics {
+                num_rows: Some(self.num_rows),
+                total_byte_size: None,
+                column_statistics: None,
+            }
+        }
+        fn has_exact_statistics(&self) -> bool {
+            self.is_exact
+        }
+    }
+
+    #[test]
+    fn optimize_count_using_statistics() -> Result<()> {
+        use crate::execution::context::ExecutionContext;
+        let mut ctx = ExecutionContext::new();
+        ctx.register_table(
+            "test",
+            Arc::new(TestTableProvider {
+                num_rows: 100,
+                is_exact: true,
+            }),
+        )
+        .unwrap();
+
+        let plan = ctx
+            .create_logical_plan("select count(*) from test")
+            .unwrap();
+        let expected = "\
+            Projection: #COUNT(UInt8(1))\
+            \n  Projection: UInt64(100) AS COUNT(Uint8(1))\

Review comment:
       ❤️ 




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [arrow-datafusion] alamb commented on a change in pull request #620: Optimize count(*) with table statistics

Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #620:
URL: https://github.com/apache/arrow-datafusion/pull/620#discussion_r660147059



##########
File path: datafusion/src/optimizer/aggregate_statistics.rs
##########
@@ -0,0 +1,217 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilizing exact statistics from sources to avoid scanning data
+use std::sync::Arc;
+
+use crate::{
+    execution::context::ExecutionProps,
+    logical_plan::{DFSchema, Expr, LogicalPlan},
+    physical_plan::aggregates::AggregateFunction,
+    scalar::ScalarValue,
+};
+
+use super::{optimizer::OptimizerRule, utils};
+use crate::error::Result;
+
+/// Optimizer that uses available statistics for aggregate functions
+pub struct AggregateStatistics {}
+
+impl AggregateStatistics {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for AggregateStatistics {
+    fn optimize(
+        &self,
+        plan: &LogicalPlan,
+        execution_props: &ExecutionProps,
+    ) -> crate::error::Result<LogicalPlan> {
+        match plan {
+            // match only select count(*) from table_scan
+            LogicalPlan::Aggregate {
+                input,
+                group_expr,
+                aggr_expr,
+                schema,
+            } if group_expr.is_empty() && aggr_expr.len() == 1 => {
+                if let Some(num_rows) = match input.as_ref() {
+                    LogicalPlan::TableScan { source, .. }
+                        if source.has_exact_statistics() =>
+                    {
+                        source.statistics().num_rows
+                    }
+                    _ => None,
+                } {
+                    return match &aggr_expr[0] {
+                        Expr::AggregateFunction {
+                            fun: AggregateFunction::Count,
+                            args,
+                            distinct: false,
+                        } if args == &[Expr::Literal(ScalarValue::UInt8(Some(1)))] => {

Review comment:
       oops -- sorry I missed that. Test is 👍 




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org