chore(cubestore): Upgrade DF: fix limit pushdown for LastRowByKey

cube-js · Dec 2, 2024 · faf7d75 · faf7d75
1 parent 019c5f2
commit faf7d75
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 19 deletions.
diff --git a/rust/cubestore/cubestore-sql-tests/src/tests.rs b/rust/cubestore/cubestore-sql-tests/src/tests.rs
@@ -8647,12 +8647,12 @@ async fn limit_pushdown_unique_key(service: Box<dyn SqlClient>) {
     //===========================
     let res = assert_limit_pushdown(
         &service,
-        "SELECT a, b, SUM(c) FROM (
+        "SELECT a FROM (SELECT a, b, SUM(c) FROM (
                 SELECT * FROM foo.pushdown_where_group1
                 union all
                 SELECT * FROM foo.pushdown_where_group2
                 ) as `tb`
-                GROUP BY 1, 2 ORDER BY 1 LIMIT 3",
+                GROUP BY 1, 2 ORDER BY 1 LIMIT 3) x",
         Some("ind1"),
         true,
         false,
@@ -8665,18 +8665,18 @@ async fn limit_pushdown_unique_key(service: Box<dyn SqlClient>) {
         vec![
             Row::new(vec![
                 TableValue::Int(11),
-                TableValue::Int(18),
-                TableValue::Int(3)
+                // TableValue::Int(18),
+                // TableValue::Int(3)
             ]),
             Row::new(vec![
                 TableValue::Int(11),
-                TableValue::Int(45),
-                TableValue::Int(1)
+                // TableValue::Int(45),
+                // TableValue::Int(1)
             ]),
             Row::new(vec![
                 TableValue::Int(12),
-                TableValue::Int(20),
-                TableValue::Int(4)
+                // TableValue::Int(20),
+                // TableValue::Int(4)
             ]),
         ]
     );

diff --git a/rust/cubestore/cubestore/src/queryplanner/merge_sort.rs b/rust/cubestore/cubestore/src/queryplanner/merge_sort.rs
@@ -41,15 +41,11 @@ impl LastRowByUniqueKeyExec {
                 "Empty unique_key passed for LastRowByUniqueKeyExec".to_string(),
             ));
         }
-        let schema = input.schema();
+        let properties = input.properties().clone();
         Ok(Self {
             input,
             unique_key,
-            properties: PlanProperties::new(
-                EquivalenceProperties::new(schema),
-                Partitioning::UnknownPartitioning(1),
-                ExecutionMode::Bounded,
-            ),
+            properties,
         })
     }
 
@@ -83,6 +79,10 @@ impl ExecutionPlan for LastRowByUniqueKeyExec {
         &self.properties
     }
 
+    fn maintains_input_order(&self) -> Vec<bool> {
+        vec![true]
+    }
+
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         vec![&self.input]
     }

diff --git a/rust/cubestore/cubestore/src/queryplanner/query_executor.rs b/rust/cubestore/cubestore/src/queryplanner/query_executor.rs
@@ -691,11 +691,19 @@ impl CubeTable {
                             )));
                         }
                     }
-                    Arc::new(MemoryExec::try_new(
-                        &[record_batches.clone()],
-                        index_projection_schema.clone(),
-                        index_projection_or_none_on_schema_match.clone(),
-                    )?)
+                    Arc::new(
+                        MemoryExec::try_new(
+                            &[record_batches.clone()],
+                            index_projection_schema.clone(),
+                            index_projection_or_none_on_schema_match.clone(),
+                        )?
+                        .with_sort_information(vec![
+                            lex_ordering_for_index(
+                                self.index_snapshot.index.get_row(),
+                                &index_projection_schema,
+                            )?,
+                        ]),
+                    )
                 } else {
                     let remote_path = chunk.get_row().get_full_name(chunk.get_id());
                     let local_path = self