Restore context flag for window functions (#16229)

2024-04-03 01:27:13 -07:00 · 2024-04-03 01:27:13 -07:00 · 4bea865697
parent 218513ad55
commit 4bea865697
8 changed files with 58 additions and 3 deletions
--- a/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchSqlAggregatorTest.java
+++ b/extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchSqlAggregatorTest.java
@ -76,6 +76,7 @@ import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker;
 import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
 import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
 import org.apache.druid.sql.calcite.filtration.Filtration;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.util.CacheTestHelperModule.ResultCacheMode;
 import org.apache.druid.sql.calcite.util.CalciteTests;
 import org.apache.druid.sql.calcite.util.TestDataBuilder;
@ -1163,6 +1164,7 @@ public class HllSketchSqlAggregatorTest extends BaseCalciteQueryTest
  public void testHllWithOrderedWindowing()
  {
    testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql(
            "SELECT dim1,coalesce(cast(l1 as integer),-999),"
                + " HLL_SKETCH_ESTIMATE( DS_HLL(dim1) OVER ( ORDER BY l1 ), true)"
@ -1189,6 +1191,7 @@ public class HllSketchSqlAggregatorTest extends BaseCalciteQueryTest
    skipVectorize();
    for (int i = 0; i < 2; i++) {
      testBuilder()
+          .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
          .sql(
              "SELECT "
                  + " TIME_FLOOR(__time, 'P1D') as dayLvl,\n"
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
@ -168,6 +168,7 @@ import org.apache.druid.sql.calcite.external.LocalOperatorConversion;
 import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
 import org.apache.druid.sql.calcite.planner.CatalogResolver;
 import org.apache.druid.sql.calcite.planner.PlannerConfig;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.planner.PlannerFactory;
 import org.apache.druid.sql.calcite.rel.DruidQuery;
 import org.apache.druid.sql.calcite.run.SqlEngine;
@ -259,6 +260,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
                  .put(MultiStageQueryContext.CTX_MAX_NUM_TASKS, 2)
                  .put(MSQWarnings.CTX_MAX_PARSE_EXCEPTIONS_ALLOWED, 0)
                  .put(MSQTaskQueryMaker.USER_KEY, "allowAll")
+                  .put(PlannerContext.CTX_ENABLE_WINDOW_FNS, true)
                  .build();

  public static final Map<String, Object> DURABLE_STORAGE_MSQ_CONTEXT =
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java
@ -65,6 +65,7 @@ import org.apache.druid.sql.calcite.parser.DruidSqlIngest;
 import org.apache.druid.sql.calcite.parser.DruidSqlInsert;
 import org.apache.druid.sql.calcite.parser.DruidSqlParserUtils;
 import org.apache.druid.sql.calcite.parser.ExternalDestinationSqlIdentifier;
+import org.apache.druid.sql.calcite.run.EngineFeature;
 import org.apache.druid.sql.calcite.table.DatasourceTable;
 import org.checkerframework.checker.nullness.qual.Nullable;

@ -621,6 +622,15 @@ public class DruidSqlValidator extends BaseDruidSqlValidator
  @Override
  public void validateCall(SqlCall call, SqlValidatorScope scope)
  {
+    if (call.getKind() == SqlKind.OVER) {
+      if (!plannerContext.featureAvailable(EngineFeature.WINDOW_FUNCTIONS)) {
+        throw buildCalciteContextException(
+            StringUtils.format(
+                "The query contains window functions; To run these window functions, specify [%s] in query context.",
+                PlannerContext.CTX_ENABLE_WINDOW_FNS),
+            call);
+      }
+    }
    if (call.getKind() == SqlKind.NULLS_FIRST) {
      SqlNode op0 = call.getOperandList().get(0);
      if (op0.getKind() == SqlKind.DESCENDING) {
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
@ -84,6 +84,10 @@ public class PlannerContext
   */
  public static final String CTX_SQL_OUTER_LIMIT = "sqlOuterLimit";

+  /**
+   * Key to enable window functions.
+   */
+  public static final String CTX_ENABLE_WINDOW_FNS = "enableWindowing";

  /**
   * Context key for {@link PlannerContext#isUseBoundsAndSelectors()}.
@ -571,9 +575,15 @@ public class PlannerContext
   * Checks if the current {@link SqlEngine} supports a particular feature.
   *
   * When executing a specific query, use this method instead of {@link SqlEngine#featureAvailable(EngineFeature)}
+   * because it also verifies feature flags such as {@link #CTX_ENABLE_WINDOW_FNS}.
   */
  public boolean featureAvailable(final EngineFeature feature)
  {
+    if (feature == EngineFeature.WINDOW_FUNCTIONS &&
+        !QueryContexts.getAsBoolean(CTX_ENABLE_WINDOW_FNS, queryContext.get(CTX_ENABLE_WINDOW_FNS), false)) {
+      // Short-circuit: feature requires context flag.
+      return false;
+    }
    if (feature == EngineFeature.TIME_BOUNDARY_QUERY && !queryContext().isTimeBoundaryPlanningEnabled()) {
      // Short-circuit: feature requires context flag.
      return false;
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
@ -14927,10 +14927,22 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
        ));
  }

+  @Test
+  public void testWindowingErrorWithoutFeatureFlag()
+  {
+    DruidException e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, false))
+        .sql("SELECT dim1,ROW_NUMBER() OVER () from druid.foo")
+        .run());
+
+    assertThat(e, invalidSqlIs("The query contains window functions; To run these window functions, specify [enableWindowing] in query context. (line [1], column [13])"));
+  }
+
  @Test
  public void testUnSupportedNullsFirst()
  {
    DruidException e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 DESC NULLS FIRST) from druid.foo")
        .run());

@ -14941,6 +14953,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
  public void testUnSupportedNullsLast()
  {
    DruidException e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 NULLS LAST) from druid.foo")
        .run());
    assertThat(e, invalidSqlIs("ASCENDING ordering with NULLS LAST is not supported! (line [1], column [41])"));
@ -14952,6 +14965,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
    assumeFeatureAvailable(EngineFeature.WINDOW_FUNCTIONS);

    DruidException e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 RANGE BETWEEN 3 PRECEDING AND 2 FOLLOWING) from druid.foo")
        .run());
    assertThat(e, invalidSqlIs("The query contains a window frame which may return incorrect results. To disregard this warning, set [windowingStrictValidation] to false in the query context. (line [1], column [31])"));
@ -14963,6 +14977,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
    assumeFeatureAvailable(EngineFeature.WINDOW_FUNCTIONS);

    DruidException e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 ROWS BETWEEN dim1 PRECEDING AND dim1 FOLLOWING) from druid.foo")
        .run());
    assertThat(e, invalidSqlIs("Window frames with expression based lower/upper bounds are not supported. (line [1], column [31])"));
@ -14976,11 +14991,13 @@ public class CalciteQueryTest extends BaseCalciteQueryTest

    DruidException e;
    e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) from druid.foo")
        .run());
    assertThat(e, invalidSqlIs("Query bounds with both lower and upper bounds as PRECEDING or FOLLOWING is not supported. (line [1], column [31])"));

    e = assertThrows(DruidException.class, () -> testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("SELECT dim1,ROW_NUMBER() OVER (ORDER BY dim1 ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) from druid.foo")
        .run());
    assertThat(e, invalidSqlIs("Query bounds with both lower and upper bounds as PRECEDING or FOLLOWING is not supported. (line [1], column [31])"));
@ -14994,6 +15011,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
    DruidException e = assertThrows(
        DruidException.class,
        () -> testBuilder()
+            .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
            .sql("SELECT ntile(4) OVER (ORDER BY dim1 ROWS BETWEEN 1 FOLLOWING AND CURRENT ROW) from druid.foo")
            .run()
    );
@ -15201,6 +15219,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest

    testBuilder()
        .sql(sql)
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .expectedQuery(
            WindowOperatorQueryBuilder.builder()
                .setDataSource(
@ -15288,6 +15307,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
        )
        .queryContext(
            ImmutableMap.of(
+                PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
                QueryContexts.ENABLE_DEBUG, true
            )
        )
@ -15388,6 +15408,7 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
        )
        .queryContext(
            ImmutableMap.of(
+                PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
                QueryContexts.ENABLE_DEBUG, true
            )
        )
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java
@ -20,8 +20,10 @@
 package org.apache.druid.sql.calcite;

 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import org.apache.druid.sql.calcite.NotYetSupported.Modes;
 import org.apache.druid.sql.calcite.NotYetSupported.NotYetSupportedProcessor;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;

@ -52,6 +54,7 @@ public class CalciteSysQueryTest extends BaseCalciteQueryTest
    msqIncompatible();

    testBuilder()
+        .queryContext(ImmutableMap.of(PlannerContext.CTX_ENABLE_WINDOW_FNS, true))
        .sql("select datasource, sum(duration) over () from sys.tasks group by datasource")
        .expectedResults(ImmutableList.of(
            new Object[]{"foo", 11L},
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteWindowQueryTest.java
@ -35,6 +35,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.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
@ -197,6 +198,7 @@ public class CalciteWindowQueryTest extends BaseCalciteQueryTest
          .skipVectorize(true)
          .sql(testCase.getSql())
          .queryContext(ImmutableMap.of(
+              PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
              QueryContexts.ENABLE_DEBUG, true,
              QueryContexts.WINDOWING_STRICT_VALIDATION, false
              ))
@ -219,6 +221,7 @@ public class CalciteWindowQueryTest extends BaseCalciteQueryTest
          .skipVectorize(true)
          .sql(testCase.getSql())
          .queryContext(ImmutableMap.of(QueryContexts.ENABLE_DEBUG, true,
+                                        PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
                                        QueryContexts.MAX_SUBQUERY_BYTES_KEY, "100000",
                                        QueryContexts.WINDOWING_STRICT_VALIDATION, false
                        )
--- a/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java
@ -57,6 +57,7 @@ import org.apache.druid.sql.calcite.NotYetSupported.Modes;
 import org.apache.druid.sql.calcite.NotYetSupported.NotYetSupportedProcessor;
 import org.apache.druid.sql.calcite.QueryTestRunner.QueryResults;
 import org.apache.druid.sql.calcite.planner.PlannerCaptureHook;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.partition.NumberedShardSpec;
 import org.joda.time.DateTime;
@ -480,9 +481,11 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
      testBuilder()
          .skipVectorize(true)
          .queryContext(ImmutableMap.of(
-              PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
-              QueryContexts.ENABLE_DEBUG, true)
-              )
+                            PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
+                            PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
+                            QueryContexts.ENABLE_DEBUG, true
+                        )
+          )
          .sql(testCase.getQueryString())
          .expectedResults(new TextualResultsVerifier(testCase.getExpectedResults(), null))
          .run();