feat(cubesql): Support Athena OFFSET and LIMIT push down

cube-js · Feb 12, 2024 · 00c2a6b · 00c2a6b
1 parent 311c5d7
commit 00c2a6b
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts b/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts
@@ -113,6 +113,12 @@ export class PrestodbQuery extends BaseQuery {
     const templates = super.sqlTemplates();
     templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
     templates.functions.DATEPART = 'DATE_PART({{ args_concat }})';
+    templates.statements.select = 'SELECT {{ select_concat | map(attribute=\'aliased\') | join(\', \') }} \n' +
+      'FROM (\n  {{ from }}\n) AS {{ from_alias }} \n' +
+      '{% if group_by %} GROUP BY {{ group_by | map(attribute=\'index\') | join(\', \') }}{% endif %}' +
+      '{% if order_by %} ORDER BY {{ order_by | map(attribute=\'expr\') | join(\', \') }}{% endif %}' +
+      '{% if offset %}\nOFFSET {{ offset }}{% endif %}' +
+      '{% if limit %}\nLIMIT {{ limit }}{% endif %}';
     templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})';
     templates.expressions.interval = 'INTERVAL \'{{ num }}\' {{ date_part }}';
     return templates;

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -21220,4 +21220,61 @@ limit
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_athena_offset_limit_push_down() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan_customized(
+            "
+            SELECT
+                SUM(s) AS s,
+                p AS p
+            FROM (
+                SELECT
+                    taxful_total_price AS p,
+                    CASE
+                        WHEN taxful_total_price = 1 THEN 0
+                        ELSE SUM(taxful_total_price)
+                    END AS s
+                FROM KibanaSampleDataEcommerce AS k
+                GROUP BY 1
+            ) AS t
+            GROUP BY 2
+            LIMIT 2
+            OFFSET 1
+            "
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            vec![
+                (
+                    "statements/select".to_string(),
+                    r#"SELECT {{ select_concat | map(attribute='aliased') | join(', ') }} 
+FROM ({{ from }}) AS {{ from_alias }} 
+{% if group_by %} GROUP BY {{ group_by | map(attribute='index') | join(', ') }}{% endif %}
+{% if order_by %} ORDER BY {{ order_by | map(attribute='expr') | join(', ') }}{% endif %}{% if offset %}
+OFFSET {{ offset }}{% endif %}{% if limit %}
+LIMIT {{ limit }}{% endif %}"#.to_string(),
+                ),
+            ]
+        )
+        .await;
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql;
+        assert!(sql.contains("OFFSET 1\nLIMIT 2"));
+    }
 }