Window Functions : Context Parameter to Enable Transfer of RACs over wire (#17150) (#17182)

(cherry picked from commit 661614129e)

Co-authored-by: Sree Charan Manamala <sree.manamala@imply.io>

docs/querying/sql-window-functions.md

@@ -28,10 +28,6 @@ description: Reference for window functions
 Apache Druid supports two query languages: [Druid SQL](sql.md) and [native queries](querying.md).
 This document describes the SQL language.
 
-Window functions are an [experimental](../development/experimental.md) feature.
-Development and testing are still at early stage. Feel free to try window functions and provide your feedback.
-Windows functions are not currently supported by multi-stage-query engine so you cannot use them in SQL-based ingestion.
-
 :::
 
 Window functions in Apache Druid produce values based upon the relationship of one row within a window of rows to the other rows within the same window. A window is a group of related rows within a result set. For example, rows with the same value for a specific dimension.

sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java

@@ -90,6 +90,11 @@ public class PlannerContext
    */
   public static final String CTX_SQL_OUTER_LIMIT = "sqlOuterLimit";
 
+  /**
+   * Key to enable transfer of RACs over wire.
+   */
+  public static final String CTX_ENABLE_RAC_TRANSFER_OVER_WIRE = "enableRACOverWire";
+
   /**
    * Context key for {@link PlannerContext#isUseBoundsAndSelectors()}.
    */

sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java

@@ -1495,13 +1495,16 @@ public class DruidQuery
     // This would cause MSQ queries to plan as
     // Window over an inner scan and avoid
     // leaf operators
+    boolean pushLeafOperator = plannerContext.queryContext()
+                                             .getBoolean(PlannerContext.CTX_ENABLE_RAC_TRANSFER_OVER_WIRE, false)
+                               && !plannerContext.featureAvailable(EngineFeature.WINDOW_LEAF_OPERATOR);
     return new WindowOperatorQuery(
         dataSource,
         new LegacySegmentSpec(Intervals.ETERNITY),
         plannerContext.queryContextMap(),
         windowing.getSignature(),
         operators,
-        plannerContext.featureAvailable(EngineFeature.WINDOW_LEAF_OPERATOR) ? ImmutableList.of() : null
+        pushLeafOperator ? null : ImmutableList.of()
     );
   }
 

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

@@ -16045,7 +16045,6 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
         .run();
   }
 
-  @NotYetSupported(Modes.UNSUPPORTED_DATASOURCE)
   @Test
   public void testWindowingOverJoin()
   {

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

@@ -37,6 +37,7 @@ import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.sql.calcite.CalciteWindowQueryTest.WindowQueryTestInputClass.TestType;
 import org.apache.druid.sql.calcite.QueryTestRunner.QueryResults;
 import org.apache.druid.sql.calcite.QueryVerification.QueryResultsVerifier;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.junit.Assert;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -296,10 +297,32 @@ public class CalciteWindowQueryTest extends BaseCalciteQueryTest
     );
 
     assertEquals(
-        "Encountered a multi value column. Window processing does not support MVDs. "
+        "Encountered a multi value column [v0]. Window processing does not support MVDs. "
         + "Consider using UNNEST or MV_TO_ARRAY.",
         e.getMessage()
     );
+
+    final DruidException e1 = 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(ImmutableMap.of(
+                QueryContexts.ENABLE_DEBUG, true,
+                QueryContexts.CTX_SQL_STRINGIFY_ARRAYS, false,
+                PlannerContext.CTX_ENABLE_RAC_TRANSFER_OVER_WIRE, true
+            ))
+            .run()
+    );
+
+    assertEquals(
+        "Encountered a multi value column. Window processing does not support MVDs. "
+        + "Consider using UNNEST or MV_TO_ARRAY.",
+        e1.getMessage()
+    );
   }
 
   private WindowOperatorQuery getWindowOperatorQuery(List<Query<?>> queries)

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

@@ -89,7 +89,6 @@ public @interface NotYetSupported
     RESULT_MISMATCH(AssertionError.class, "(assertResulEquals|AssertionError: column content mismatch)"),
     LONG_CASTING(AssertionError.class, "expected: java.lang.Long"),
     UNSUPPORTED_NULL_ORDERING(DruidException.class, "(A|DE)SCENDING ordering with NULLS (LAST|FIRST)"),
-    UNSUPPORTED_DATASOURCE(DruidException.class, "WindowOperatorQuery must run on top of a query or inline data source"),
     UNION_WITH_COMPLEX_OPERAND(DruidException.class, "Only Table and Values are supported as inputs for Union"),
     UNION_MORE_STRICT_ROWTYPE_CHECK(DruidException.class, "Row signature mismatch in Union inputs"),
     JOIN_CONDITION_NOT_PUSHED_CONDITION(DruidException.class, "SQL requires a join with '.*' condition"),