diff --git a/crates/core/src/sql/execute.rs b/crates/core/src/sql/execute.rs index 96093f0250d..2347e093bd2 100644 --- a/crates/core/src/sql/execute.rs +++ b/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"); diff --git a/crates/execution/src/pipelined.rs b/crates/execution/src/pipelined.rs index 8aca73aff15..f11d3bd0e90 100644 --- a/crates/execution/src/pipelined.rs +++ b/crates/execution/src/pipelined.rs @@ -33,7 +33,7 @@ impl From 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(()) diff --git a/crates/expr/src/expr.rs b/crates/expr/src/expr.rs index 8400b68052d..f3dedb95b88 100644 --- a/crates/expr/src/expr.rs +++ b/crates/expr/src/expr.rs @@ -75,12 +75,12 @@ pub enum ProjectList { Name(ProjectName), List(RelExpr, Vec<(Box, FieldProject)>), Limit(Box, u64), - Agg(RelExpr, AggType, Box, AlgebraicType), + Agg(RelExpr, AggType, AlgebraicType), } #[derive(Debug)] pub enum AggType { - Count, + Count { alias: Box }, } 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), } } } diff --git a/crates/expr/src/lib.rs b/crates/expr/src/lib.rs index 24b308a2ccc..21662dc4416 100644 --- a/crates/expr/src/lib.rs +++ b/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(); diff --git a/crates/physical-plan/src/plan.rs b/crates/physical-plan/src/plan.rs index 0dd03bb91c5..a92288a5294 100644 --- a/crates/physical-plan/src/plan.rs +++ b/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] diff --git a/crates/physical-plan/src/printer.rs b/crates/physical-plan/src/printer.rs index 104f3005f8a..0baeb6c58a1 100644 --- a/crates/physical-plan/src/printer.rs +++ b/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>), Fields(Vec>), + 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 {