feat: more Aggregate functions

Signed-off-by: Andreas Reichel <andreas@manticore-projects.com>
starlake-ai · Apr 2, 2024 · 0c32a05 · 0c32a05
1 parent 57d7464
commit 0c32a05
Show file tree

Hide file tree

Showing 3 changed files with 188 additions and 28 deletions.
diff --git a/src/main/java/com/manticore/transpiler/ExpressionTranspiler.java b/src/main/java/com/manticore/transpiler/ExpressionTranspiler.java
@@ -23,7 +23,6 @@
 import net.sf.jsqlparser.expression.CastExpression;
 import net.sf.jsqlparser.expression.DateTimeLiteralExpression;
 import net.sf.jsqlparser.expression.Expression;
-import net.sf.jsqlparser.expression.ExpressionVisitor;
 import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
 import net.sf.jsqlparser.expression.ExtractExpression;
 import net.sf.jsqlparser.expression.Function;
@@ -55,7 +54,6 @@
 import net.sf.jsqlparser.statement.select.SelectVisitor;
 import net.sf.jsqlparser.util.deparser.ExpressionDeParser;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Map;
@@ -96,7 +94,8 @@ enum TranspiledFunction {
 
     , SAFE_ADD, SAFE_DIVIDE, SAFE_MULTIPLY, SAFE_NEGATE, SAFE_SUBTRACT, TRUNC
 
-    , ARRAY_CONCAT_AGG, COUNTIF
+    , ARRAY_CONCAT_AGG, COUNTIF, LOGICAL_AND, LOGICAL_OR
+
 
     , NVL, UNNEST;
     // @FORMATTER:ON
@@ -814,13 +813,10 @@ THEN Instr( source_value, Regexp_Extract( source_value, reg_exp ) )
           }
           break;
         case ARRAY_CONCAT_AGG:
-          //list_sort(flatten(list(x)), 'ASC', 'NULLS FIRST'))
+          // list_sort(flatten(list(x)), 'ASC', 'NULLS FIRST'))
           function.setName("List_Sort");
-          function.setParameters(
-                  new Function("Flatten", new Function("List", parameters.get(0)))
-                  , new StringValue("ASC")
-                  , new StringValue("NULLS FIRST")
-          );
+          function.setParameters(new Function("Flatten", new Function("List", parameters.get(0))),
+              new StringValue("ASC"), new StringValue("NULLS FIRST"));
           break;
         case COUNTIF:
           // COUNT(IF(x < 0, x, NULL))
@@ -832,16 +828,20 @@ public void visit(Column column) {
             }
           });
           // @todo: clarify if there can be only exactly 1 column
-          //        else, what do to on None or many?
-          assert expressionColumns.size()==1;
+          // else, what do to on None or many?
+          assert expressionColumns.size() == 1;
           Column column = expressionColumns.toArray(new Column[0])[0];
 
           warning("Different NULL handling");
           function.setName("Count");
-          function.setParameters(
-                  new Function( "If", parameters.get(0), column , new NullValue())
-          );
+          function.setParameters(new Function("If", parameters.get(0), column, new NullValue()));
 
+          break;
+        case LOGICAL_AND:
+          function.setName("Bool_And");
+          break;
+        case LOGICAL_OR:
+          function.setName("Bool_Or");
           break;
       }
     }
@@ -855,7 +855,7 @@ public void visit(AnalyticExpression function) {
 
     if (UnsupportedFunction.from(function) != null) {
       throw new RuntimeException(
-              "Unsupported: " + functionName + " is not supported by DuckDB (yet).");
+          "Unsupported: " + functionName + " is not supported by DuckDB (yet).");
     } else if (functionName.endsWith("$$")) {
       // work around for transpiling already transpiled functions twice
       // @todo: figure out a better way to achieve that
@@ -877,17 +877,22 @@ public void visit(Column column) {
             }
           });
           // @todo: clarify if there can be only exactly 1 column
-          //        else, what do to on None or many?
-          assert expressionColumns.size()==1;
+          // else, what do to on None or many?
+          assert expressionColumns.size() == 1;
           Column column = expressionColumns.toArray(new Column[0])[0];
 
           warning("Different NULL handling");
           function.setName("Count");
-          function.setExpression(
-                  new Function( "If", function.getExpression(), column , new NullValue())
-          );
+          function
+              .setExpression(new Function("If", function.getExpression(), column, new NullValue()));
 
           break;
+        case LOGICAL_AND:
+          function.setName("Bool_And");
+          break;
+        case LOGICAL_OR:
+          function.setName("Bool_Or");
+          break;
       }
     }
     if (rewrittenExpression == null) {
@@ -1520,7 +1525,7 @@ public void visit(StringValue stringValue) {
   }
 
   public void visit(HexValue hexValue) {
-    super.visit( hexValue.getLongValue() );
+    super.visit(hexValue.getLongValue());
   }
 
   public static String convertUnicode(String input) {

diff --git a/src/main/resources/doc/JSQLTranspiler.ods b/src/main/resources/doc/JSQLTranspiler.ods
diff --git a/src/test/resources/com/manticore/transpiler/google_bigquery/aggregate_function.sql b/src/test/resources/com/manticore/transpiler/google_bigquery/aggregate_function.sql
@@ -253,16 +253,171 @@ SELECT
   product_name,
   SUM(product_count) AS product_sum,
   GROUPING(product_type) AS product_type_agg,
-  GROUPING(product_name) AS product_name_agg,
+  GROUPING(product_name) AS product_name_agg
 FROM Products
 GROUP BY GROUPING SETS(product_type, product_name, ())
 ORDER BY product_name, product_type;
 
 -- result
 "product_type","product_name","product_sum","product_type_agg","product_name_agg"
-"","jeans","6.00","1","0"
-"","polo","25.00","1","0"
-"","t-shirt","11.00","1","0"
-"pants","","6.00","0","1"
-"shirt","","36.00","0","1"
-"","","42.00","1","1"
+"","jeans","6","1","0"
+"","polo","25","1","0"
+"","t-shirt","11","1","0"
+"pants","","6","0","1"
+"shirt","","36","0","1"
+"","","42","1","1"
+
+
+-- provided
+SELECT LOGICAL_AND(x < 3) AS logical_and
+FROM UNNEST([1, 2, 4]) AS x;
+
+-- expected
+SELECT BOOL_AND(x < 3) AS logical_and
+FROM (SELECT UNNEST([1, 2, 4]) AS x) AS x;
+
+-- result
+"logical_and"
+"false"
+
+
+-- provided
+SELECT LOGICAL_OR(x < 3) AS logical_or FROM UNNEST([1, 2, 4]) AS x;
+
+-- expected
+SELECT BOOL_OR(x < 3) AS logical_or
+FROM (SELECT UNNEST([1, 2, 4]) AS x) as x;
+
+-- result
+"logical_or"
+"true"
+
+
+-- provided
+SELECT x, MAX(x) OVER (PARTITION BY MOD(x, 2)) AS max
+FROM UNNEST([8, NULL, 37, 55, NULL, 4]) AS x
+order by 1 nulls first;
+
+-- expected
+SELECT x, MAX(x) OVER (PARTITION BY MOD(x, 2)) AS max
+FROM (select UNNEST([8, NULL, 37, 55, NULL, 4]) AS x) as x
+order by 1 nulls first;
+
+-- result
+"x","max"
+"",""
+"",""
+"4","8"
+"8","8"
+"37","55"
+"55","55"
+
+
+-- provided
+WITH fruits AS (
+  SELECT 'apple'  fruit, 3.55 price UNION ALL
+  SELECT 'banana'  fruit, 2.10 price UNION ALL
+  SELECT 'pear'  fruit, 4.30 price
+)
+SELECT MAX_BY(fruit, price) as fruit
+FROM fruits;
+
+-- result
+"fruit"
+"pear"
+
+
+-- provided
+SELECT x, MIN(x) OVER (PARTITION BY MOD(x, 2)) AS min
+FROM UNNEST([8, NULL, 37, 4, NULL, 55]) AS x
+order by 1 nulls first;
+
+-- expected
+SELECT x, MIN(x) OVER (PARTITION BY MOD(x, 2)) AS min
+FROM (Select UNNEST([8, NULL, 37, 4, NULL, 55]) AS x) as x
+order by 1 nulls first;
+
+-- result
+"x","min"
+"",""
+"",""
+"4","4"
+"8","4"
+"37","37"
+"55","37"
+
+
+-- provided
+WITH fruits AS (
+  SELECT 'apple'  fruit, 3.55 price UNION ALL
+  SELECT 'banana'  fruit, 2.10 price UNION ALL
+  SELECT 'pear'  fruit, 4.30 price
+)
+SELECT MIN_BY(fruit, price) as fruit
+FROM fruits;
+
+-- result
+"fruit"
+"banana"
+
+
+-- provided
+SELECT STRING_AGG(fruit, ' & ' ORDER BY LENGTH(fruit)) AS string_agg
+FROM UNNEST(['apple', 'pear', 'banana', 'pear']) AS fruit;
+
+-- expected
+SELECT STRING_AGG(fruit, ' & ' ORDER BY CASE TYPEOF(FRUIT)
+                                          WHEN 'VARCHAR' THEN LENGTH(FRUIT::VARCHAR)
+                                          WHEN 'BLOB' THEN OCTET_LENGTH(FRUIT::BLOB)
+                                          ELSE -1
+                                          END) AS string_agg
+FROM (SELECT UNNEST(['apple', 'pear', 'banana', 'pear']) AS fruit) as fruit;
+
+-- result
+"string_agg"
+"pear & pear & apple & banana"
+
+
+-- provided
+SELECT
+  fruit,
+  STRING_AGG(fruit, ' & ') OVER (ORDER BY LENGTH(fruit)) AS string_agg
+FROM UNNEST(['apple', NULL, 'pear', 'banana', 'pear']) AS fruit;
+
+-- expected
+SELECT
+  fruit,
+  STRING_AGG(fruit, ' & ') OVER (ORDER BY CASE TYPEOF(FRUIT)
+                                          WHEN 'VARCHAR' THEN LENGTH(FRUIT::VARCHAR)
+                                          WHEN 'BLOB' THEN OCTET_LENGTH(FRUIT::BLOB)
+                                          ELSE -1
+                                          END) AS string_agg
+FROM (SELECT UNNEST(['apple', NULL, 'pear', 'banana', 'pear']) AS fruit) AS fruit;
+
+-- result
+"fruit","string_agg"
+"pear","pear & pear"
+"pear","pear & pear"
+"apple","pear & pear & apple"
+"banana","pear & pear & apple & banana"
+"","pear & pear & apple & banana"
+
+
+-- provided
+SELECT CORR(y, x) AS results
+FROM
+  UNNEST(
+    [
+      STRUCT(1.0 AS y, 5.0 AS x),
+      (3.0, 9.0),
+      (4.0, 7.0)]);
+
+-- expected
+SELECT CORR(Y,X)AS RESULTS
+FROM( SELECT UNNEST([{Y:1.0,X:5.0},(3.0,9.0),(4.0,7.0)], recursive=>true));
+
+-- result
+"results"
+"0.6546536707079772"
+
+