From 816cbacc9caa5fc820b356b611f93dd5a1ce4dee Mon Sep 17 00:00:00 2001
From: Ihor Andrianov <andrianoia@gmail.com>
Date: Mon, 31 Mar 2025 01:56:04 +0300
Subject: [PATCH] some smartie optimizations

---
 core/translate/group_by.rs  | 27 +++----------------
 core/translate/optimizer.rs | 53 +++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 23 deletions(-)

diff --git a/core/translate/group_by.rs b/core/translate/group_by.rs
index 86d6087e2..e00edf455 100644
--- a/core/translate/group_by.rs
+++ b/core/translate/group_by.rs
@@ -6,6 +6,7 @@ use crate::{
     function::AggFunc,
     schema::{Column, PseudoTable},
     types::{OwnedValue, Record},
+    util::exprs_are_equivalent,
     vdbe::{
         builder::{CursorType, ProgramBuilder},
         insn::Insn,
@@ -711,27 +712,7 @@ pub fn translate_aggregation_step_groupby(
 }
 
 pub fn is_column_in_group_by(expr: &ast::Expr, group_by_exprs: &[ast::Expr]) -> bool {
-    if let ast::Expr::Column {
-        database: _,
-        table: _,
-        column: col,
-        is_rowid_alias: _,
-    } = expr
-    {
-        group_by_exprs.iter().any(|ex| {
-            if let ast::Expr::Column {
-                database: _,
-                table: _,
-                column: group_col,
-                is_rowid_alias: _,
-            } = ex
-            {
-                col == group_col
-            } else {
-                false
-            }
-        })
-    } else {
-        false
-    }
+    group_by_exprs
+        .iter()
+        .any(|expr2| exprs_are_equivalent(expr, expr2))
 }
diff --git a/core/translate/optimizer.rs b/core/translate/optimizer.rs
index 6fa7f9619..dfd7ab16d 100644
--- a/core/translate/optimizer.rs
+++ b/core/translate/optimizer.rs
@@ -4,6 +4,7 @@ use limbo_sqlite3_parser::ast;
 
 use crate::{
     schema::{Index, Schema},
+    util::exprs_are_equivalent,
     Result,
 };
 
@@ -43,6 +44,8 @@ fn optimize_select_plan(plan: &mut SelectPlan, schema: &Schema) -> Result<()> {
 
     eliminate_unnecessary_orderby(plan, schema)?;
 
+    eliminate_orderby_like_groupby(plan)?;
+
     Ok(())
 }
 
@@ -117,6 +120,51 @@ fn query_is_already_ordered_by(
     }
 }
 
+fn eliminate_orderby_like_groupby(plan: &mut SelectPlan) -> Result<()> {
+    if plan.order_by.is_none() | plan.group_by.is_none() {
+        return Ok(());
+    }
+    if plan.table_references.len() == 0 {
+        return Ok(());
+    }
+
+    let o = plan.order_by.as_mut().unwrap();
+    let g = plan.group_by.as_mut().unwrap();
+
+    let mut insert_pos = 0;
+    let mut i = 0;
+
+    while i < o.len() {
+        let (key, order) = &o[i];
+
+        if matches!(order, Direction::Descending) {
+            i += 1;
+            continue;
+        }
+        if let Some(pos) = g
+            .exprs
+            .iter()
+            .position(|expr| exprs_are_equivalent(expr, key))
+        {
+            if pos != insert_pos {
+                let mut current_pos = pos;
+                while current_pos > insert_pos {
+                    g.exprs.swap(current_pos, current_pos - 1);
+                    current_pos -= 1;
+                }
+            }
+            insert_pos += 1;
+            o.remove(i);
+        } else {
+            i += 1;
+        }
+    }
+    if o.is_empty() {
+        plan.order_by = None
+    }
+    Ok(())
+}
+
 fn eliminate_unnecessary_orderby(plan: &mut SelectPlan, schema: &Schema) -> Result<()> {
     if plan.order_by.is_none() {
         return Ok(());
@@ -125,6 +173,11 @@ fn eliminate_unnecessary_orderby(plan: &mut SelectPlan, schema: &Schema) -> Resu
         return Ok(());
     }
 
+    // if pk will be removed later
+    if plan.group_by.is_some() {
+        return Ok(());
+    }
+
     let o = plan.order_by.as_mut().unwrap();
 
     if o.len() != 1 {