Bug fix for array type selector causing array aggregation over window frame fail (#16653)

processing/src/main/java/org/apache/druid/query/rowsandcols/semantic/DefaultColumnSelectorFactoryMaker.java

@@ -41,7 +41,6 @@ import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -209,7 +208,8 @@ public class DefaultColumnSelectorFactoryMaker implements ColumnSelectorFactoryM
             myClazz = float.class;
             break;
           case ARRAY:
-            myClazz = List.class;
+            myClazz = Object[].class;
+            break;
           default:
             throw DruidException.defensive("this class cannot handle type [%s]", columnAccessor.getType());
         }

processing/src/main/java/org/apache/druid/segment/virtual/ExpressionSelectors.java

@@ -475,15 +475,15 @@ public class ExpressionSelectors
     }
 
     final Class<?> clazz = selector.classOfObject();
-    if (Number.class.isAssignableFrom(clazz) || String.class.isAssignableFrom(clazz)) {
+    if (Number.class.isAssignableFrom(clazz) || String.class.isAssignableFrom(clazz) || Object[].class.isAssignableFrom(clazz)) {
-      // Number, String supported as-is.
+      // Number, String, Arrays supported as-is.
       return selector::getObject;
     } else if (clazz.isAssignableFrom(Number.class) || clazz.isAssignableFrom(String.class)) {
       // Might be Numbers and Strings. Use a selector that double-checks.
       return () -> {
         final Object val = selector.getObject();
         if (val instanceof List) {
-          NonnullPair<ExpressionType, Object[]> coerced = ExprEval.coerceListToArray((List) val, homogenizeMultiValue);
+          NonnullPair<ExpressionType, Object[]> coerced = ExprEval.coerceListToArray((List<?>) val, homogenizeMultiValue);
           if (coerced == null) {
             return null;
           }
@@ -496,7 +496,7 @@ public class ExpressionSelectors
       return () -> {
         final Object val = selector.getObject();
         if (val != null) {
-          NonnullPair<ExpressionType, Object[]> coerced = ExprEval.coerceListToArray((List) val, homogenizeMultiValue);
+          NonnullPair<ExpressionType, Object[]> coerced = ExprEval.coerceListToArray((List<?>) val, homogenizeMultiValue);
           if (coerced == null) {
             return null;
           }

processing/src/test/java/org/apache/druid/segment/virtual/ExpressionSelectorsTest.java

@@ -591,7 +591,22 @@ public class ExpressionSelectorsTest extends InitializedNullHandlingTest
 
     settableSupplier.set(ImmutableList.of("1", "2", "3"));
     Assert.assertArrayEquals(new String[]{"1", "2", "3"}, (Object[]) supplier.get());
+  }
 
+  @Test
+  public void test_supplierFromObjectSelector_onArray()
+  {
+    final SettableSupplier<Object[]> settableSupplier = new SettableSupplier<>();
+    final Supplier<Object> supplier = ExpressionSelectors.supplierFromObjectSelector(
+        objectSelectorFromSupplier(settableSupplier, Object[].class),
+        true
+    );
+
+    Assert.assertNotNull(supplier);
+    Assert.assertEquals(null, supplier.get());
+
+    settableSupplier.set(new String[]{"1", "2", "3"});
+    Assert.assertArrayEquals(new String[]{"1", "2", "3"}, (Object[]) supplier.get());
   }
 
   @Test

sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java

@@ -320,6 +320,36 @@ public class CalciteWindowQueryTest extends BaseCalciteQueryTest
         .run();
   }
 
+  @Test
+  public void testWithArrayConcat()
+  {
+    testBuilder()
+        .sql("select countryName, cityName, channel, "
+             + "array_concat_agg(ARRAY['abc', channel], 10000) over (partition by cityName order by countryName) as c\n"
+             + "from wikipedia\n"
+             + "where countryName in ('Austria', 'Republic of Korea') "
+             + "and (cityName in ('Vienna', 'Seoul') or cityName is null)\n"
+             + "group by countryName, cityName, channel")
+        .queryContext(ImmutableMap.of(
+            PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
+            QueryContexts.ENABLE_DEBUG, true
+        ))
+        .expectedResults(
+            ResultMatchMode.RELAX_NULLS,
+            ImmutableList.of(
+              new Object[]{"Austria", null, "#de.wikipedia", "[\"abc\",\"#de.wikipedia\"]"},
+              new Object[]{"Republic of Korea", null, "#en.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#en.wikipedia\",\"abc\",\"#ja.wikipedia\",\"abc\",\"#ko.wikipedia\"]"},
+              new Object[]{"Republic of Korea", null, "#ja.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#en.wikipedia\",\"abc\",\"#ja.wikipedia\",\"abc\",\"#ko.wikipedia\"]"},
+              new Object[]{"Republic of Korea", null, "#ko.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#en.wikipedia\",\"abc\",\"#ja.wikipedia\",\"abc\",\"#ko.wikipedia\"]"},
+              new Object[]{"Republic of Korea", "Seoul", "#ko.wikipedia", "[\"abc\",\"#ko.wikipedia\"]"},
+              new Object[]{"Austria", "Vienna", "#de.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#es.wikipedia\",\"abc\",\"#tr.wikipedia\"]"},
+              new Object[]{"Austria", "Vienna", "#es.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#es.wikipedia\",\"abc\",\"#tr.wikipedia\"]"},
+              new Object[]{"Austria", "Vienna", "#tr.wikipedia", "[\"abc\",\"#de.wikipedia\",\"abc\",\"#es.wikipedia\",\"abc\",\"#tr.wikipedia\"]"}
+            )
+        )
+        .run();
+  }
+
   private WindowOperatorQuery getWindowOperatorQuery(List<Query<?>> queries)
   {
     assertEquals(1, queries.size());