From 4ce1e1651e8fea305fb27f743473926f9feeec23 Mon Sep 17 00:00:00 2001
From: Bruno Volpato <bruno.volpato@datadoghq.com>
Date: Fri, 5 Apr 2024 22:58:44 -0400
Subject: [PATCH] [CALCITE-6355] RelToSqlConverter[ORDER BY] generates an
 incorrect order by when NULLS LAST is used in non-projected field

---
 .../calcite/rel/rel2sql/SqlImplementor.java   | 36 +++++++++-------
 .../rel/rel2sql/RelToSqlConverterTest.java    | 42 +++++++++++++++++++
 2 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/rel/rel2sql/SqlImplementor.java b/core/src/main/java/org/apache/calcite/rel/rel2sql/SqlImplementor.java
index 242b526f890..9950674ea15 100644
--- a/core/src/main/java/org/apache/calcite/rel/rel2sql/SqlImplementor.java
+++ b/core/src/main/java/org/apache/calcite/rel/rel2sql/SqlImplementor.java
@@ -1880,7 +1880,7 @@ private boolean needNewSubQuery(
       if (rel instanceof Project
           && clauses.contains(Clause.ORDER_BY)
           && dialect.getConformance().isSortByOrdinal()
-          && hasSortByOrdinal()) {
+          && hasSortByOrdinal(node)) {
         // Cannot merge a Project that contains sort by ordinal under it.
         return true;
       }
@@ -1932,25 +1932,33 @@ && hasSortByOrdinal()) {
     }
 
     /**
-     * Return whether the current {@link SqlNode} in {@link Result} contains sort by column
-     * in ordinal format.
+     * Return whether the given {@link SqlNode} contains a sort by using an ordinal / numeric
+     * literal. Checks recursively if the node is a {@link SqlSelect} or a {@link SqlBasicCall}.
+     *
+     * @param sqlNode SqlNode to check
      */
-    private boolean hasSortByOrdinal(@UnknownInitialization Result this) {
-      if (node instanceof SqlSelect) {
-        final SqlNodeList orderList = ((SqlSelect) node).getOrderList();
+    private boolean hasSortByOrdinal(@UnknownInitialization Result this,
+                                     @Nullable SqlNode sqlNode) {
+      if (sqlNode == null) {
+        return false;
+      }
+      if (sqlNode instanceof SqlNumericLiteral) {
+        return true;
+      }
+      if (sqlNode instanceof SqlSelect) {
+        final SqlNodeList orderList = ((SqlSelect) sqlNode).getOrderList();
         if (orderList == null) {
           return false;
         }
-        for (SqlNode sqlNode : orderList) {
-          if (sqlNode instanceof SqlNumericLiteral) {
+        for (SqlNode child : orderList) {
+          if (hasSortByOrdinal(child)) {
             return true;
           }
-          if (sqlNode instanceof SqlBasicCall) {
-            for (SqlNode operand : ((SqlBasicCall) sqlNode).getOperandList()) {
-              if (operand instanceof SqlNumericLiteral) {
-                return true;
-              }
-            }
+        }
+      } else if (sqlNode instanceof SqlBasicCall) {
+        for (SqlNode operand : ((SqlBasicCall) sqlNode).getOperandList()) {
+          if (hasSortByOrdinal(operand)) {
+            return true;
           }
         }
       }
diff --git a/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java b/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
index 5e4e6b11bb7..6e55237bc9a 100644
--- a/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
+++ b/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
@@ -2162,6 +2162,48 @@ private SqlDialect nonOrdinalDialect() {
         .ok(prestoExpected);
   }
 
+  /**
+   * Test case for the base case of
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6355">[CALCITE-6355]
+   * RelToSqlConverter[ORDER BY] generates an incorrect order by when NULLS LAST is used in
+   * non-projected</a>.
+   */
+  @Test void testOrderByOrdinalDesc() {
+    String query = "select \"product_id\"\n"
+                   + "from \"product\"\n"
+                   + "where \"net_weight\" is not null\n"
+                   + "group by \"product_id\""
+                   + "order by MAX(\"net_weight\") desc";
+    final String expected = "SELECT \"product_id\"\n"
+                            + "FROM (SELECT \"product_id\", MAX(\"net_weight\") AS \"EXPR$1\"\n"
+                            + "FROM \"foodmart\".\"product\"\n"
+                            + "WHERE \"net_weight\" IS NOT NULL\n"
+                            + "GROUP BY \"product_id\"\n"
+                            + "ORDER BY 2 DESC) AS \"t3\"";
+    sql(query).ok(expected);
+  }
+
+  /**
+   * Test case for the problematic case of
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6355">[CALCITE-6355]
+   * RelToSqlConverter[ORDER BY] generates an incorrect order by when NULLS LAST is used in
+   * non-projected</a>.
+   */
+  @Test void testOrderByOrdinalDescNullsLast() {
+    String query = "select \"product_id\"\n"
+                   + "from \"product\"\n"
+                   + "where \"net_weight\" is not null\n"
+                   + "group by \"product_id\""
+                   + "order by MAX(\"net_weight\") desc nulls last";
+    final String expected = "SELECT \"product_id\"\n"
+                            + "FROM (SELECT \"product_id\", MAX(\"net_weight\") AS \"EXPR$1\"\n"
+                            + "FROM \"foodmart\".\"product\"\n"
+                            + "WHERE \"net_weight\" IS NOT NULL\n"
+                            + "GROUP BY \"product_id\"\n"
+                            + "ORDER BY 2 DESC NULLS LAST) AS \"t3\"";
+    sql(query).ok(expected);
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-3440">[CALCITE-3440]
    * RelToSqlConverter does not properly alias ambiguous ORDER BY</a>. */