From b8a45028dc967d5018e822543c26ac51bdd63b35 Mon Sep 17 00:00:00 2001
From: Mihai Budiu <mbudiu@feldera.com>
Date: Tue, 30 Apr 2024 13:28:18 -0700
Subject: [PATCH] [CALCITE-6376] Selecting 6 columns with QUALIFY operation
 results in exception

Signed-off-by: Mihai Budiu <mbudiu@feldera.com>
---
 .../org/apache/calcite/runtime/FlatLists.java     |  5 +++++
 .../calcite/test/SqlToRelConverterTest.java       |  8 ++++++++
 .../apache/calcite/test/SqlToRelConverterTest.xml | 15 +++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/core/src/main/java/org/apache/calcite/runtime/FlatLists.java b/core/src/main/java/org/apache/calcite/runtime/FlatLists.java
index de7471eda6e9..544a8926f066 100644
--- a/core/src/main/java/org/apache/calcite/runtime/FlatLists.java
+++ b/core/src/main/java/org/apache/calcite/runtime/FlatLists.java
@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.checker.nullness.qual.PolyNull;
 
@@ -1323,6 +1324,10 @@ protected ComparableListImpl(List<T> list) {
       return list.size();
     }
 
+    @Override @NonNull public Object[] toArray(@NonNull ComparableListImpl<T> this) {
+      return this.list.toArray();
+    }
+
     @Override public int compareTo(List o) {
       return compare(list, o);
     }
diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
index e62eeb78b56f..4e5582199984 100644
--- a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
+++ b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
@@ -3490,6 +3490,14 @@ void checkCorrelatedMapSubQuery(boolean expand) {
     sql(sql).ok();
   }
 
+  /** Test case for <a href="https://issues.apache.org/jira/browse/CALCITE-6376">[CALCITE-6376]
+   * Selecting 6 columns with QUALIFY operation results in exception</a>. */
+  @Test void testQualifyWindow() {
+    sql("SELECT empno, ename, deptno, job, mgr, hiredate\n"
+        + "FROM emp\n"
+        + "QUALIFY ROW_NUMBER() over (partition by ename order by deptno) = 1")
+        .ok();
+  }
 
   @Test void testQualifyWithoutReferences() {
     sql("SELECT empno, ename, deptno\n"
diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
index 1bf1556aa7a9..2a029cfd09c2 100644
--- a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
+++ b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
@@ -6136,6 +6136,21 @@ LogicalProject(EMPNO=[$0], ENAME=[$1], DEPTNO=[$2])
       LogicalFilter(condition=[$3])
         LogicalProject(EMPNO=[$0], ENAME=[$1], DEPTNO=[$7], QualifyExpression=[=(RANK() OVER (PARTITION BY $1 ORDER BY $7 DESC), 1)])
           LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testQualifyWindow">
+    <Resource name="sql">
+      <![CDATA[SELECT empno, ename, deptno, job, mgr, hiredate
+FROM emp
+QUALIFY ROW_NUMBER() over (partition by ename order by deptno) = 1]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalProject(EMPNO=[$0], ENAME=[$1], DEPTNO=[$2], JOB=[$3], MGR=[$4], HIREDATE=[$5])
+  LogicalFilter(condition=[$6])
+    LogicalProject(EMPNO=[$0], ENAME=[$1], DEPTNO=[$7], JOB=[$2], MGR=[$3], HIREDATE=[$4], QualifyExpression=[=(ROW_NUMBER() OVER (PARTITION BY $1 ORDER BY $7), 1)])
+      LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>