From f7432dabcf05a0cbea546026b27ea606a5b36653 Mon Sep 17 00:00:00 2001
From: Fernando Velasquez <fevelasquez@google.com>
Date: Wed, 16 Mar 2022 23:05:17 -0400
Subject: [PATCH] Added support for ANY filters when building deduplicate
 aggregations

---
 .../builder/BigQueryBaseSQLBuilder.java       |  3 ++
 .../BigQueryDeduplicateSQLBuilder.java        | 31 ++++++++++++-------
 .../BigQueryDeduplicateSQLBuilderTest.java    | 30 ++++++++++++------
 3 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryBaseSQLBuilder.java b/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryBaseSQLBuilder.java
index d06fd108b9..e82f6933ec 100644
--- a/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryBaseSQLBuilder.java
+++ b/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryBaseSQLBuilder.java
@@ -50,6 +50,9 @@ public abstract class BigQueryBaseSQLBuilder {
   public static final String ROW_NUMBER_PARTITION_COLUMN =
     "ROW_NUMBER() OVER ( PARTITION BY %s ORDER BY %s ) AS `%s`";
   public static final String NULLS_LAST = "NULLS LAST";
+  public static final String IFNULL_FUNCTION = "IFNULL";
+  public static final String ZERO = "0";
+  public static final String ONE = "1";
 
   /**
    * Builds SQL statement
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilder.java b/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilder.java
index 8d37bbc71b..cfb6f54288 100644
--- a/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilder.java
+++ b/src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilder.java
@@ -134,18 +134,27 @@ protected String getOrderByFields(List<DeduplicateAggregationDefinition.FilterEx
    * @return Order by SQL expression
    */
   protected String getOrderByField(DeduplicateAggregationDefinition.FilterExpression filterExpression) {
-    String order;
-
-    // MAX of a value means ORDER DESCENDING and selecting the first result.
-    // MIN of a value means ORDER ASCENDING and selecting the first result.
-    if (filterExpression.getFilterFunction() == DeduplicateAggregationDefinition.FilterFunction.MAX) {
-      order = ORDER_DESC;
-    } else {
-      order = ORDER_ASC;
+    String exp = ((SQLExpression) filterExpression.getExpression()).extract();
+
+    switch (filterExpression.getFilterFunction()) {
+      case MIN:
+        // MIN of a value means ORDER ASCENDING and selecting the first result.
+        // ...[ORDER BY] exp ASC NULLS LAST
+        return exp + SPACE + ORDER_ASC + SPACE + NULLS_LAST;
+      case MAX:
+        // MAX of a value means ORDER DESCENDING and selecting the first result.
+        // ...[ORDER BY] exp DESC NULLS LAST
+        return exp + SPACE + ORDER_DESC + SPACE + NULLS_LAST;
+      case ANY_NULLS_FIRST:
+        // ANY_NULLS_FIRST means order with null values first and pick the first.
+        // ...[ORDER BY] IFNULL(exp , 0 , 1) ASC
+        return IFNULL_FUNCTION + OPEN_GROUP + exp + COMMA + ZERO + COMMA + ONE + CLOSE_GROUP + SPACE + ORDER_ASC;
+      case ANY_NULLS_LAST:
+      default:
+        // ANY_NULLS_LAST means order with null values first and pick the first.
+        // ...[ORDER BY] IFNULL(exp , 0 , 1) DESC
+        return IFNULL_FUNCTION + OPEN_GROUP + exp + COMMA + ZERO + COMMA + ONE + CLOSE_GROUP + SPACE + ORDER_DESC;
     }
-
-    // some_field ASC/DESC
-    return ((SQLExpression) filterExpression.getExpression()).extract() + SPACE + order + SPACE + NULLS_LAST;
   }
 
 }
diff --git a/src/test/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilderTest.java b/src/test/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilderTest.java
index bdefc38cc9..f445cd2f69 100644
--- a/src/test/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilderTest.java
+++ b/src/test/java/io/cdap/plugin/gcp/bigquery/sqlengine/builder/BigQueryDeduplicateSQLBuilderTest.java
@@ -105,20 +105,22 @@ public void testGetInnerSelect() {
 
   @Test
   public void testGetSelectedFields() {
-    Assert.assertEquals("a AS alias_a , "
-                          + "b AS alias_b , "
-                          + "c AS c , "
-                          + "d AS d , "
-                          + "e AS e , "
-                          + "f AS f , "
-                    + "ROW_NUMBER() OVER ( PARTITION BY c , d , e ORDER BY e DESC NULLS LAST , f ASC NULLS LAST ) AS" +
-                    " `the_row_number`", helper.getSelectedFields(def));
+    Assert.assertEquals(
+      "a AS alias_a , "
+        + "b AS alias_b , "
+        + "c AS c , "
+        + "d AS d , "
+        + "e AS e , "
+        + "f AS f , "
+        + "ROW_NUMBER() OVER ( PARTITION BY c , d , e ORDER BY e DESC NULLS LAST , f ASC NULLS LAST ) AS" +
+        " `the_row_number`",
+      helper.getSelectedFields(def));
   }
 
   @Test
   public void testGetRowNumColumn() {
     Assert.assertEquals("ROW_NUMBER() OVER ( PARTITION BY c , d , e ORDER BY e DESC NULLS LAST , " +
-            "f ASC NULLS LAST ) AS `the_row_number`", helper.getRowNumColumn(def));
+                          "f ASC NULLS LAST ) AS `the_row_number`", helper.getRowNumColumn(def));
   }
 
   @Test
@@ -139,9 +141,19 @@ public void testGetOrderByField() {
     DeduplicateAggregationDefinition.FilterExpression minField2 =
       new DeduplicateAggregationDefinition.FilterExpression(factory.compile("field2"),
                                                             DeduplicateAggregationDefinition.FilterFunction.MIN);
+    DeduplicateAggregationDefinition.FilterExpression minField3 =
+      new DeduplicateAggregationDefinition
+        .FilterExpression(factory.compile("field3"),
+                          DeduplicateAggregationDefinition.FilterFunction.ANY_NULLS_FIRST);
+    DeduplicateAggregationDefinition.FilterExpression minField4 =
+      new DeduplicateAggregationDefinition
+        .FilterExpression(factory.compile("field4"),
+                          DeduplicateAggregationDefinition.FilterFunction.ANY_NULLS_LAST);
 
     Assert.assertEquals("field1 DESC NULLS LAST", helper.getOrderByField(maxField1));
     Assert.assertEquals("field2 ASC NULLS LAST", helper.getOrderByField(minField2));
+    Assert.assertEquals("IFNULL(field3 , 0 , 1) ASC", helper.getOrderByField(minField3));
+    Assert.assertEquals("IFNULL(field4 , 0 , 1) DESC", helper.getOrderByField(minField4));
   }