EXPLAIN for COUNT and LIMIT

crates/core/src/sql/execute.rs

@@ -364,6 +364,7 @@ pub(crate) mod tests {
         let result = run_for_testing(&db, sql)?;
         assert_eq!(result, vec![product![5u64]], "Inventory");
 
+        // TODO: This should return product![2u64]
         let sql = "SELECT count(*) as n FROM inventory limit 2";
         let result = run_for_testing(&db, sql)?;
         assert_eq!(result, vec![product![5u64]], "Inventory");

crates/execution/src/pipelined.rs

@@ -33,7 +33,7 @@ impl From<ProjectListPlan> for ProjectListExecutor {
             ProjectListPlan::Name(plan) => Self::Name(plan.into()),
             ProjectListPlan::List(plan, fields) => Self::List(plan.into(), fields),
             ProjectListPlan::Limit(plan, n) => Self::Limit(Box::new((*plan).into()), n),
-            ProjectListPlan::Agg(plan, AggType::Count) => Self::Agg(plan.into(), AggType::Count),
+            ProjectListPlan::Agg(plan, agg) => Self::Agg(plan.into(), agg),
         }
     }
 }
@@ -78,10 +78,10 @@ impl ProjectListExecutor {
             // and if so, we retrieve the count from table metadata.
             // It's a valid optimization but one that should be done by the optimizer.
             // There should be no optimizations performed during execution.
-            Self::Agg(PipelinedExecutor::TableScan(table_scan), AggType::Count) => {
+            Self::Agg(PipelinedExecutor::TableScan(table_scan), AggType::Count { alias: _ }) => {
                 f(product![tx.table_or_err(table_scan.table)?.num_rows()])?;
             }
-            Self::Agg(plan, AggType::Count) => {
+            Self::Agg(plan, AggType::Count { alias: _ }) => {
                 plan.execute(tx, metrics, &mut |_| {
                     n += 1;
                     Ok(())

crates/expr/src/expr.rs

@@ -75,12 +75,12 @@ pub enum ProjectList {
     Name(ProjectName),
     List(RelExpr, Vec<(Box<str>, FieldProject)>),
     Limit(Box<ProjectList>, u64),
-    Agg(RelExpr, AggType, Box<str>, AlgebraicType),
+    Agg(RelExpr, AggType, AlgebraicType),
 }
 
 #[derive(Debug)]
 pub enum AggType {
-    Count,
+    Count { alias: Box<str> },
 }
 
 impl ProjectList {
@@ -120,7 +120,7 @@ impl ProjectList {
                     f(name, ty);
                 }
             }
-            Self::Agg(_, _, name, ty) => f(name, ty),
+            Self::Agg(_, AggType::Count { alias }, ty) => f(alias, ty),
         }
     }
 }

crates/expr/src/lib.rs

@@ -52,7 +52,9 @@ pub(crate) fn type_proj(input: RelExpr, proj: ast::Project, vars: &Relvars) -> T
             Ok(ProjectList::Name(ProjectName::Some(input, var)))
         }
         ast::Project::Star(Some(SqlIdent(var))) => Err(Unresolved::var(&var).into()),
-        ast::Project::Count(SqlIdent(alias)) => Ok(ProjectList::Agg(input, AggType::Count, alias, AlgebraicType::U64)),
+        ast::Project::Count(SqlIdent(alias)) => {
+            Ok(ProjectList::Agg(input, AggType::Count { alias }, AlgebraicType::U64))
+        }
         ast::Project::Exprs(elems) => {
             let mut projections = vec![];
             let mut names = HashSet::new();

crates/physical-plan/src/plan.rs

@@ -1152,11 +1152,9 @@ pub mod tests_utils {
 mod tests {
     use super::*;
 
-    use crate::compile::compile_select_list;
     use crate::printer::ExplainOptions;
     use expect_test::{expect, Expect};
     use spacetimedb_expr::check::SchemaView;
-    use spacetimedb_expr::statement::{parse_and_type_sql, Statement};
     use spacetimedb_lib::{
         db::auth::{StAccess, StTableType},
         AlgebraicType,
@@ -1858,37 +1856,60 @@ Index Join: Rhs on p
 
         let db = SchemaViewer::new(vec![t.clone()]);
 
-        let compile = |sql| {
-            let stmt = parse_and_type_sql(sql, &db).unwrap();
-            let Statement::Select(select) = stmt else {
-                unreachable!()
-            };
-            compile_select_list(select).optimize().unwrap()
-        };
+        check_query(
+            &db,
+            "SELECT * FROM t LIMIT 5",
+            expect![[r#"
+Limit: 5
+  -> Seq Scan on t
+  Output: t.x, t.y"#]],
+        );
 
-        let plan = compile("select * from t limit 5");
+        check_query(
+            &db,
+            "SELECT * FROM t WHERE x = 1 LIMIT 5",
+            expect![[r#"
+Limit: 5
+  -> Index Scan using Index id 0 on t
+      Index Cond: (t.x = U8(1))
+  Output: t.x, t.y"#]],
+        );
 
-        assert!(matches!(
-            plan,
-            ProjectListPlan::Name(ProjectPlan::None(PhysicalPlan::TableScan(
-                TableScan { limit: Some(5), .. },
-                _
-            )))
-        ));
+        check_query(
+            &db,
+            "SELECT * FROM t WHERE y = 1 LIMIT 5",
+            expect![[r#"
+Limit: 5
+  -> Seq Scan on t
+    Filter: (t.y = U8(1))
+  Output: t.x, t.y"#]],
+        );
+    }
 
-        let plan = compile("select * from t where x = 1 limit 5");
+    #[test]
+    fn count() {
+        let db = data();
 
-        assert!(matches!(
-            plan,
-            ProjectListPlan::Name(ProjectPlan::None(PhysicalPlan::IxScan(
-                IxScan { limit: Some(5), .. },
-                _
-            )))
-        ));
+        check_query(
+            &db,
+            "SELECT COUNT(*) AS n FROM p",
+            expect![[r#"
+Count
+  -> Seq Scan on p
+  Output: p.id, p.name"#]],
+        );
 
-        let plan = compile("select * from t where y = 1 limit 5");
+        check_query(
+            &db,
+            "SELECT COUNT(*) AS n FROM p WHERE id = 1",
+            expect![[r#"
+Count
+  -> Seq Scan on p
+    Filter: (p.id = U64(1))
+  Output: p.id, p.name"#]],
+        );
 
-        assert!(matches!(plan, ProjectListPlan::Limit(_, 5)) && plan.has_filter() && plan.has_table_scan(None));
+        // TODO: Is not yet possible to combine correctly `COUNT` with `LIMIT`
     }
 
     #[test]

crates/physical-plan/src/printer.rs

@@ -2,6 +2,7 @@ use crate::dml::MutationPlan;
 use crate::plan::{IxScan, Label, PhysicalExpr, PhysicalPlan, ProjectListPlan, ProjectPlan, Sarg, Semi, TupleField};
 use crate::{PhysicalCtx, PlanCtx};
 use itertools::Itertools;
+use spacetimedb_expr::expr::AggType;
 use spacetimedb_expr::StatementSource;
 use spacetimedb_lib::AlgebraicValue;
 use spacetimedb_primitives::{ColId, ConstraintId, IndexId};
@@ -236,6 +237,13 @@ pub enum Line<'a> {
         rhs: Field<'a>,
         ident: u16,
     },
+    Limit {
+        limit: u64,
+        ident: u16,
+    },
+    Count {
+        ident: u16,
+    },
     Insert {
         table_name: &'a str,
         ident: u16,
@@ -262,6 +270,8 @@ impl Line<'_> {
             Line::HashJoin { ident, .. } => *ident,
             Line::NlJoin { ident, .. } => *ident,
             Line::JoinExpr { ident, .. } => *ident,
+            Line::Limit { ident, .. } => *ident,
+            Line::Count { ident, .. } => *ident,
             Line::Insert { ident, .. } => *ident,
             Line::Update { ident, .. } => *ident,
             Line::Delete { ident, .. } => *ident,
@@ -282,6 +292,7 @@ enum Output<'a> {
     Unknown,
     Star(Vec<Field<'a>>),
     Fields(Vec<Field<'a>>),
+    Alias(&'a str),
     Empty,
 }
 
@@ -381,11 +392,18 @@ fn eval_plan<'a>(lines: &mut Lines<'a>, plan: &'a PhysicalPlan, ident: u16) {
             lines.add_table(*label, &scan.schema);
 
             let schema = lines.labels.table_by_label(label).unwrap();
-
+            let table = schema.name;
             lines.output = Output::Star(Output::fields(schema));
 
+            let ident = if let Some(limit) = scan.limit {
+                lines.add(Line::Limit { limit, ident });
+                ident + 2
+            } else {
+                ident
+            };
+
             lines.add(Line::TableScan {
-                table: schema.name,
+                table,
                 label: label.0,
                 ident,
             });
@@ -397,6 +415,13 @@ fn eval_plan<'a>(lines: &mut Lines<'a>, plan: &'a PhysicalPlan, ident: u16) {
 
             let index = idx.schema.indexes.iter().find(|x| x.index_id == idx.index_id).unwrap();
 
+            let ident = if let Some(limit) = idx.limit {
+                lines.add(Line::Limit { limit, ident });
+                ident + 2
+            } else {
+                ident
+            };
+
             lines.add(Line::IxScan {
                 table_name: &idx.schema.table_name,
                 index: PrintName::index(idx.index_id, &index.index_name),
@@ -549,7 +574,22 @@ impl<'a> Explain<'a> {
                     eval_plan(&mut lines, plan, 0);
                     lines.output = Output::Fields(Output::tuples(fields, &lines));
                 }
-                &ProjectListPlan::Limit(_, _) | &ProjectListPlan::Agg(_, _) => todo!(),
+                ProjectListPlan::Limit(plan, limit) => {
+                    lines.add(Line::Limit {
+                        limit: *limit,
+                        ident: 0,
+                    });
+                    eval_plan(&mut lines, plan, 2);
+                }
+                ProjectListPlan::Agg(plan, agg) => {
+                    match agg {
+                        AggType::Count { alias } => {
+                            lines.add(Line::Count { ident: 0 });
+                            lines.output = Output::Alias(alias)
+                        }
+                    }
+                    eval_plan(&mut lines, plan, 2);
+                }
             },
             PlanCtx::DML(plan) => {
                 eval_dml_plan(&mut lines, plan, 0);
@@ -735,6 +775,12 @@ impl fmt::Display for Explain<'_> {
                     writeln!(f, "{:ident$}Inner Unique: {unique}", "")?;
                     write!(f, "{:ident$}Join Cond: ({} = {})", "", lhs, rhs)?;
                 }
+                Line::Limit { limit, ident: _ } => {
+                    write!(f, "{:ident$}{arrow}Limit: {limit}", "")?;
+                }
+                Line::Count { ident: _ } => {
+                    write!(f, "{:ident$}{arrow}Count", "")?;
+                }
                 Line::Insert { table_name, ident: _ } => {
                     write!(f, "{:ident$}{arrow}Insert on {table_name}", "")?;
                 }
@@ -766,6 +812,7 @@ impl fmt::Display for Explain<'_> {
                 let columns = fields.iter().map(|x| format!("{}", x)).join(", ");
                 Some(columns)
             }
+            Output::Alias(alias) => Some(alias.to_string()),
             Output::Empty => Some("void".to_string()),
         };
         let end = if self.options.show_timings || self.options.show_schema {