From 51fe3c08ab7f99045e6db1713dd88b755275c414 Mon Sep 17 00:00:00 2001
From: Sree Charan Manamala <sree.manamala@imply.io>
Date: Mon, 9 Sep 2024 12:07:54 +0530
Subject: [PATCH] Window Functions : Reject MVDs during window processing
 (#17002)

This commit aims to reject MVDs in window processing as we do not support them.
Earlier to this commit, query running a window aggregate partitioned by an MVD column would fail with ClassCastException
---
 .../rowsandcols/ArrayListRowsAndColumns.java  | 20 +++++++++++++++-
 .../sql/calcite/CalciteWindowQueryTest.java   | 23 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/ArrayListRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/ArrayListRowsAndColumns.java
index 883371fec65..7ce3df8e066 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/ArrayListRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/ArrayListRowsAndColumns.java
@@ -25,6 +25,7 @@ import it.unimi.dsi.fastutil.ints.IntComparator;
 import it.unimi.dsi.fastutil.ints.IntList;
 import org.apache.druid.common.semantic.SemanticCreator;
 import org.apache.druid.common.semantic.SemanticUtils;
+import org.apache.druid.error.InvalidInput;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.query.operator.ColumnWithDirection;
@@ -153,11 +154,28 @@ public class ArrayListRowsAndColumns<RowType> implements AppendableRowsAndColumn
       return new LimitedColumn(retVal, startOffset, endOffset);
     }
 
-    final Function<RowType, Object> adapterForValue = rowAdapter.columnFunction(name);
     final Optional<ColumnType> maybeColumnType = rowSignature.getColumnType(name);
     final ColumnType columnType = maybeColumnType.orElse(ColumnType.UNKNOWN_COMPLEX);
     final Comparator<Object> comparator = Comparator.nullsFirst(columnType.getStrategy());
 
+    final Function<RowType, Object> adapterForValue;
+    if (columnType.equals(ColumnType.STRING)) {
+      // special handling to reject MVDs
+      adapterForValue = f -> {
+        Object value = rowAdapter.columnFunction(name).apply(f);
+        if (value instanceof List) {
+          throw InvalidInput.exception(
+              "Encountered a multi value column [%s]. Window processing does not support MVDs. "
+              + "Consider using UNNEST or MV_TO_ARRAY.",
+              name
+          );
+        }
+        return value;
+      };
+    } else {
+      adapterForValue = rowAdapter.columnFunction(name);
+    }
+
     return new Column()
     {
       @Nonnull
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java
index 165b4aa3f63..ccf459e743e 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java
@@ -24,6 +24,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import org.apache.druid.error.DruidException;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.RE;
@@ -249,6 +250,28 @@ public class CalciteWindowQueryTest extends BaseCalciteQueryTest
         .run();
   }
 
+  @Test
+  public void testFailure_partitionByMVD()
+  {
+    final DruidException e = Assert.assertThrows(
+        DruidException.class,
+        () -> testBuilder()
+            .sql("select cityName, countryName, array_to_mv(array[1,length(cityName)]),\n"
+                 + "row_number() over (partition by  array_to_mv(array[1,length(cityName)]) order by countryName, cityName)\n"
+                 + "from wikipedia\n"
+                 + "where countryName in ('Austria', 'Republic of Korea') and cityName is not null\n"
+                 + "order by 1, 2, 3")
+            .queryContext(DEFAULT_QUERY_CONTEXT)
+            .run()
+    );
+
+    assertEquals(
+        "Encountered a multi value column [v0]. Window processing does not support MVDs. "
+        + "Consider using UNNEST or MV_TO_ARRAY.",
+        e.getMessage()
+    );
+  }
+
   private WindowOperatorQuery getWindowOperatorQuery(List<Query<?>> queries)
   {
     assertEquals(1, queries.size());