This is an automated email from the ASF dual-hosted git repository.

Jackie-Jiang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new 9bacd8af0f9 Handle literal ORDER BY keys in single-stage queries 
(#20134)
9bacd8af0f9 is described below

commit 9bacd8af0f9424267662170ca8cbc21e1fd938e0
Author: Xiaotian (Jackie) Jiang <[email protected]>
AuthorDate: Mon Nov 2 21:04:42 2026 -0800

    Handle literal ORDER BY keys in single-stage queries (#20134)
---
 .../parsers/rewriter/OrderByLiteralRemover.java    |  45 ++++++
 .../sql/parsers/rewriter/QueryRewriterFactory.java |   3 +-
 .../rewriter/OrderByLiteralRemoverTest.java        |  75 +++++++++
 .../parsers/rewriter/QueryRewriterFactoryTest.java |  55 ++++---
 .../SelectionPartiallyOrderedByDescOperation.java  |   9 +-
 .../SelectionPartiallyOrderedByLinearOperator.java |  14 +-
 .../query/reduce/SelectionDataTableReducer.java    |   2 +
 .../org/apache/pinot/queries/BaseQueriesTest.java  |   2 +-
 ...nnerSegmentSelectionSingleValueQueriesTest.java | 176 +++++++++++++++++++++
 9 files changed, 334 insertions(+), 47 deletions(-)

diff --git 
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemover.java
 
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemover.java
new file mode 100644
index 00000000000..0d50bfd18e4
--- /dev/null
+++ 
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemover.java
@@ -0,0 +1,45 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.sql.parsers.rewriter;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+
+
+/// Removes literal ORDER BY keys, which impose no ordering constraint. Run 
after constant folding and alias/ordinal
+/// resolution. This rewriter is stateless and thread-safe; each invocation 
mutates its own query.
+public class OrderByLiteralRemover implements QueryRewriter {
+  @Override
+  public PinotQuery rewrite(PinotQuery pinotQuery) {
+    List<Expression> orderByList = pinotQuery.getOrderByList();
+    if (orderByList != null) {
+      List<Expression> orderByExpressions = new 
ArrayList<>(orderByList.size());
+      for (Expression expression : orderByList) {
+        if 
(!CalciteSqlParser.removeOrderByFunctions(expression).isSetLiteral()) {
+          orderByExpressions.add(expression);
+        }
+      }
+      pinotQuery.setOrderByList(orderByExpressions.isEmpty() ? null : 
orderByExpressions);
+    }
+    return pinotQuery;
+  }
+}
diff --git 
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactory.java
 
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactory.java
index 132f1bf9c9a..977a1d3bd41 100644
--- 
a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactory.java
+++ 
b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactory.java
@@ -38,10 +38,11 @@ public class QueryRewriterFactory {
   //   (reference) into the group-by list, but the alias should not be applied 
to the reference.
   //   E.g. SELECT a + 1 AS a FROM table GROUP BY 1
   //   AggregationOptimizer should run early to optimize aggregation patterns 
before other rewrites.
+  //   OrderByLiteralRemover must run after OrdinalsUpdater so integer 
ordinals are resolved before removing literals.
   public static final List<String> DEFAULT_QUERY_REWRITERS_CLASS_NAMES =
       List.of(CompileTimeFunctionsInvoker.class.getName(), 
AggregationOptimizer.class.getName(),
           SelectionsRewriter.class.getName(), 
PredicateComparisonRewriter.class.getName(),
-          AliasApplier.class.getName(), OrdinalsUpdater.class.getName(),
+          AliasApplier.class.getName(), OrdinalsUpdater.class.getName(), 
OrderByLiteralRemover.class.getName(),
           NonAggregationGroupByToDistinctQueryRewriter.class.getName(), 
RlsFiltersRewriter.class.getName(),
           CastTypeAliasRewriter.class.getName());
 
diff --git 
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemoverTest.java
 
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemoverTest.java
new file mode 100644
index 00000000000..689f5500f6a
--- /dev/null
+++ 
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/OrderByLiteralRemoverTest.java
@@ -0,0 +1,75 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.sql.parsers.rewriter;
+
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.apache.pinot.sql.parsers.SqlCompilationException;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Tests literal removal through the full rewrite pipeline. Test cases have 
no shared mutable state.
+public class OrderByLiteralRemoverTest {
+  @DataProvider
+  public Object[][] queries() {
+    return new Object[][]{
+        {"SELECT a FROM t", "SELECT a FROM t"},
+        {"SELECT a FROM t ORDER BY a DESC NULLS FIRST", "SELECT a FROM t ORDER 
BY a DESC NULLS FIRST"},
+        {"SELECT a FROM t ORDER BY 'x'", "SELECT a FROM t"},
+        {"SELECT a FROM t ORDER BY NULL DESC NULLS FIRST, true, 1.5", "SELECT 
a FROM t"},
+        {"SELECT a FROM t ORDER BY now() DESC", "SELECT a FROM t"},
+        {"SELECT a FROM t ORDER BY concat('a', 'b')", "SELECT a FROM t"},
+        {"SELECT a, 'x' FROM t ORDER BY 2 DESC", "SELECT a, 'x' FROM t"},
+        {"SELECT a, 42 AS constant FROM t ORDER BY 2", "SELECT a, 42 AS 
constant FROM t"},
+        {"SELECT a, 'x' AS constant FROM t ORDER BY constant", "SELECT a, 'x' 
AS constant FROM t"},
+        {"SELECT a, now() AS ts FROM t ORDER BY 2 DESC, ts", "SELECT a, now() 
AS ts FROM t"},
+        {"SELECT a FROM t ORDER BY 'x' DESC, a ASC NULLS FIRST, 'y' ASC",
+            "SELECT a FROM t ORDER BY a ASC NULLS FIRST"},
+        {"SELECT a, b FROM t ORDER BY 2 DESC, 'x', 1 ASC", "SELECT a, b FROM t 
ORDER BY b DESC, a ASC"},
+        {"SELECT a AS x FROM t ORDER BY 'constant', x DESC", "SELECT a AS x 
FROM t ORDER BY a DESC"},
+        {"SELECT a FROM t ORDER BY 'x', add(a, 1) DESC NULLS LAST",
+            "SELECT a FROM t ORDER BY add(a, 1) DESC NULLS LAST"},
+        {"SELECT DISTINCT a FROM t ORDER BY 'x', a", "SELECT DISTINCT a FROM t 
ORDER BY a"},
+        {"SELECT a, COUNT(*) FROM t GROUP BY a ORDER BY 'x' DESC, COUNT(*) 
DESC",
+            "SELECT a, COUNT(*) FROM t GROUP BY a ORDER BY COUNT(*) DESC"},
+        {"SELECT a FROM t GROUP BY a ORDER BY 'x', a", "SELECT a FROM t GROUP 
BY a ORDER BY a"}
+    };
+  }
+
+  @Test(dataProvider = "queries")
+  public void testRemoveLiterals(String query, String expectedQuery) {
+    PinotQuery actual = CalciteSqlParser.compileToPinotQuery(query + " LIMIT 5 
OFFSET 3");
+    PinotQuery expected = CalciteSqlParser.compileToPinotQuery(expectedQuery + 
" LIMIT 5 OFFSET 3");
+    assertEquals(actual.getOrderByList(), expected.getOrderByList());
+    assertEquals(actual.getLimit(), 5);
+    assertEquals(actual.getOffset(), 3);
+  }
+
+  @Test
+  public void testInvalidOrdinalsAreNotIgnored() {
+    for (int ordinal : new int[]{0, -1, 2}) {
+      assertThrows(SqlCompilationException.class,
+          () -> CalciteSqlParser.compileToPinotQuery("SELECT a FROM t ORDER BY 
" + ordinal));
+    }
+  }
+}
diff --git 
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactoryTest.java
 
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactoryTest.java
index 50323ecd591..c7a70c66a34 100644
--- 
a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactoryTest.java
+++ 
b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/QueryRewriterFactoryTest.java
@@ -18,47 +18,52 @@
  */
 package org.apache.pinot.sql.parsers.rewriter;
 
-import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import static org.apache.pinot.sql.parsers.CalciteSqlParser.QUERY_REWRITERS;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
 public class QueryRewriterFactoryTest {
 
   @Test
   public void testQueryRewriters() {
     // Default behavior
     QueryRewriterFactory.init(null);
-    Assert.assertEquals(QUERY_REWRITERS.size(), 9);
-    Assert.assertTrue(QUERY_REWRITERS.get(0) instanceof 
CompileTimeFunctionsInvoker);
-    Assert.assertTrue(QUERY_REWRITERS.get(1) instanceof AggregationOptimizer);
-    Assert.assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(3) instanceof 
PredicateComparisonRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(4) instanceof AliasApplier);
-    Assert.assertTrue(QUERY_REWRITERS.get(5) instanceof OrdinalsUpdater);
-    Assert.assertTrue(QUERY_REWRITERS.get(6) instanceof 
NonAggregationGroupByToDistinctQueryRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(7) instanceof RlsFiltersRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(8) instanceof CastTypeAliasRewriter);
+    assertEquals(QUERY_REWRITERS.size(), 10);
+    assertTrue(QUERY_REWRITERS.get(0) instanceof CompileTimeFunctionsInvoker);
+    assertTrue(QUERY_REWRITERS.get(1) instanceof AggregationOptimizer);
+    assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
+    assertTrue(QUERY_REWRITERS.get(3) instanceof PredicateComparisonRewriter);
+    assertTrue(QUERY_REWRITERS.get(4) instanceof AliasApplier);
+    assertTrue(QUERY_REWRITERS.get(5) instanceof OrdinalsUpdater);
+    assertTrue(QUERY_REWRITERS.get(6) instanceof OrderByLiteralRemover);
+    assertTrue(QUERY_REWRITERS.get(7) instanceof 
NonAggregationGroupByToDistinctQueryRewriter);
+    assertTrue(QUERY_REWRITERS.get(8) instanceof RlsFiltersRewriter);
+    assertTrue(QUERY_REWRITERS.get(9) instanceof CastTypeAliasRewriter);
 
     // Check init with other configs
     
QueryRewriterFactory.init("org.apache.pinot.sql.parsers.rewriter.PredicateComparisonRewriter,"
         + "org.apache.pinot.sql.parsers.rewriter.CompileTimeFunctionsInvoker,"
         + "org.apache.pinot.sql.parsers.rewriter.SelectionsRewriter");
-    Assert.assertEquals(QUERY_REWRITERS.size(), 3);
-    Assert.assertTrue(QUERY_REWRITERS.get(0) instanceof 
PredicateComparisonRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(1) instanceof 
CompileTimeFunctionsInvoker);
-    Assert.assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
+    assertEquals(QUERY_REWRITERS.size(), 3);
+    assertTrue(QUERY_REWRITERS.get(0) instanceof PredicateComparisonRewriter);
+    assertTrue(QUERY_REWRITERS.get(1) instanceof CompileTimeFunctionsInvoker);
+    assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
 
     // Revert back to default behavior
     QueryRewriterFactory.init(null);
-    Assert.assertEquals(QUERY_REWRITERS.size(), 9);
-    Assert.assertTrue(QUERY_REWRITERS.get(0) instanceof 
CompileTimeFunctionsInvoker);
-    Assert.assertTrue(QUERY_REWRITERS.get(1) instanceof AggregationOptimizer);
-    Assert.assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(3) instanceof 
PredicateComparisonRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(4) instanceof AliasApplier);
-    Assert.assertTrue(QUERY_REWRITERS.get(5) instanceof OrdinalsUpdater);
-    Assert.assertTrue(QUERY_REWRITERS.get(6) instanceof 
NonAggregationGroupByToDistinctQueryRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(7) instanceof RlsFiltersRewriter);
-    Assert.assertTrue(QUERY_REWRITERS.get(8) instanceof CastTypeAliasRewriter);
+    assertEquals(QUERY_REWRITERS.size(), 10);
+    assertTrue(QUERY_REWRITERS.get(0) instanceof CompileTimeFunctionsInvoker);
+    assertTrue(QUERY_REWRITERS.get(1) instanceof AggregationOptimizer);
+    assertTrue(QUERY_REWRITERS.get(2) instanceof SelectionsRewriter);
+    assertTrue(QUERY_REWRITERS.get(3) instanceof PredicateComparisonRewriter);
+    assertTrue(QUERY_REWRITERS.get(4) instanceof AliasApplier);
+    assertTrue(QUERY_REWRITERS.get(5) instanceof OrdinalsUpdater);
+    assertTrue(QUERY_REWRITERS.get(6) instanceof OrderByLiteralRemover);
+    assertTrue(QUERY_REWRITERS.get(7) instanceof 
NonAggregationGroupByToDistinctQueryRewriter);
+    assertTrue(QUERY_REWRITERS.get(8) instanceof RlsFiltersRewriter);
+    assertTrue(QUERY_REWRITERS.get(9) instanceof CastTypeAliasRewriter);
   }
 }
diff --git 
a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByDescOperation.java
 
b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByDescOperation.java
index 90ae5fc5cf3..3df89b696ca 100644
--- 
a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByDescOperation.java
+++ 
b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByDescOperation.java
@@ -18,7 +18,6 @@
  */
 package org.apache.pinot.core.operator.query;
 
-import com.google.common.base.Preconditions;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -34,6 +33,7 @@ import org.apache.pinot.spi.query.QueryScanCostContext;
 
 
 /// An operator for order-by queries DESC that are partially sorted over the 
sorting keys.
+/// The planner guarantees the sorted prefix can be read in reverse; a 
literal-only prefix imposes no direction.
 ///
 /// @see LinearSelectionOrderByOperator
 public class SelectionPartiallyOrderedByDescOperation extends 
LinearSelectionOrderByOperator {
@@ -45,13 +45,6 @@ public class SelectionPartiallyOrderedByDescOperation 
extends LinearSelectionOrd
   public SelectionPartiallyOrderedByDescOperation(IndexSegment indexSegment, 
QueryContext queryContext,
       List<ExpressionContext> expressions, BaseProjectOperator<?> 
projectOperator, int numSortedExpressions) {
     super(indexSegment, queryContext, expressions, projectOperator, 
numSortedExpressions);
-    assert queryContext.getOrderByExpressions() != null;
-    Preconditions.checkArgument(!queryContext.getOrderByExpressions().stream()
-            .filter(expr -> expr.getExpression().getType() == 
ExpressionContext.Type.IDENTIFIER)
-            .findFirst()
-            .orElseThrow(() -> new IllegalArgumentException("The query is not 
order by identifiers"))
-            .isAsc(),
-        "%s can only be used when the first column in order by is DESC", 
EXPLAIN_NAME);
   }
 
   @Override
diff --git 
a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByLinearOperator.java
 
b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByLinearOperator.java
index 845d65e0181..9245f56e77a 100644
--- 
a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByLinearOperator.java
+++ 
b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/SelectionPartiallyOrderedByLinearOperator.java
@@ -24,7 +24,6 @@ import java.util.function.Supplier;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.core.common.BlockValSet;
 import org.apache.pinot.core.operator.BaseProjectOperator;
-import org.apache.pinot.core.operator.DocIdOrderedOperator;
 import org.apache.pinot.core.operator.blocks.ValueBlock;
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.segment.spi.IndexSegment;
@@ -33,7 +32,8 @@ import org.apache.pinot.spi.query.QueryScanCostContext;
 
 /// The operator used when selecting with order-by on a way that the segment 
layout can be used to prune results.
 ///
-/// This requires that the first order-by expression is on a column that is 
sorted in the segment and:
+/// The planner guarantees the sorted prefix is compatible with the input 
order. A literal-only prefix accepts any
+/// input order. When the prefix contains a column, it must be sorted in the 
segment and:
 /// 1. The order-by is ASC
 /// 2. The order-by is DESC and the input operators support descending 
iteration (e.g. full scan)
 ///
@@ -47,16 +47,6 @@ public class SelectionPartiallyOrderedByLinearOperator 
extends LinearSelectionOr
   public SelectionPartiallyOrderedByLinearOperator(IndexSegment indexSegment, 
QueryContext queryContext,
       List<ExpressionContext> expressions, BaseProjectOperator<?> 
projectOperator, int numSortedExpressions) {
     super(indexSegment, queryContext, expressions, projectOperator, 
numSortedExpressions);
-    boolean firstColIsAsc = queryContext.getOrderByExpressions().stream()
-        .filter(expr -> expr.getExpression().getType() == 
ExpressionContext.Type.IDENTIFIER)
-        .findFirst()
-        .orElseThrow(() -> new IllegalArgumentException("The query is not 
order by identifiers"))
-        .isAsc();
-    DocIdOrderedOperator.DocIdOrder docIdOrder = 
DocIdOrderedOperator.DocIdOrder.fromAsc(firstColIsAsc);
-    if (!projectOperator.isCompatibleWith(docIdOrder)) {
-      throw new IllegalStateException(EXPLAIN_NAME + " requires the input 
operator to be compatible with order: "
-          + docIdOrder + ", but found: " + projectOperator.toExplainString());
-    }
   }
 
   @Override
diff --git 
a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/SelectionDataTableReducer.java
 
b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/SelectionDataTableReducer.java
index 0b1884e6568..f53d852915b 100644
--- 
a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/SelectionDataTableReducer.java
+++ 
b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/SelectionDataTableReducer.java
@@ -54,6 +54,8 @@ public class SelectionDataTableReducer implements 
DataTableReducer {
     }
     if (_queryContext.getOrderByExpressions() == null) {
       // Selection only
+      // TODO: Honor OFFSET without ORDER BY, as MSE does. SSE currently 
ignores it. This requires servers to retain
+      //       LIMIT + OFFSET rows and the broker to skip OFFSET rows.
       List<Object[]> reducedRows = 
SelectionOperatorUtils.reduceWithoutOrdering(dataTableMap.values(), limit,
           _queryContext.isNullHandlingEnabled());
       brokerResponseNative.setResultTable(
diff --git 
a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java 
b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java
index 3767e4c07e0..099c579aa1a 100644
--- a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java
@@ -194,7 +194,7 @@ public abstract class BaseQueriesTest {
   /// different index segments in the test and overriding getDistinctInstances.
   /// This can be particularly useful to test statistical aggregation 
functions.
   /// @see StatisticalQueriesTest for an example use case.
-  private BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, 
PlanMaker planMaker) {
+  protected BrokerResponseNative getBrokerResponse(PinotQuery pinotQuery, 
PlanMaker planMaker) {
     List<List<IndexSegment>> instances = getDistinctInstances();
     if (instances.size() == 2) {
       return getBrokerResponseDistinctInstances(pinotQuery, planMaker);
diff --git 
a/pinot-core/src/test/java/org/apache/pinot/queries/InnerSegmentSelectionSingleValueQueriesTest.java
 
b/pinot-core/src/test/java/org/apache/pinot/queries/InnerSegmentSelectionSingleValueQueriesTest.java
index b333f786e9d..b3d94f2980e 100644
--- 
a/pinot-core/src/test/java/org/apache/pinot/queries/InnerSegmentSelectionSingleValueQueriesTest.java
+++ 
b/pinot-core/src/test/java/org/apache/pinot/queries/InnerSegmentSelectionSingleValueQueriesTest.java
@@ -21,6 +21,7 @@ package org.apache.pinot.queries;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.apache.pinot.common.request.PinotQuery;
 import org.apache.pinot.common.response.broker.BrokerResponseNative;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
@@ -28,9 +29,19 @@ import org.apache.pinot.core.operator.BaseOperator;
 import org.apache.pinot.core.operator.ExecutionStatistics;
 import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
 import org.apache.pinot.core.operator.query.EmptySelectionOperator;
+import org.apache.pinot.core.operator.query.SelectionOnlyOperator;
+import org.apache.pinot.core.operator.query.SelectionOrderByOperator;
 import 
org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByDescOperation;
 import 
org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByLinearOperator;
+import org.apache.pinot.core.plan.SelectionPlanNode;
+import 
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.segment.spi.SegmentContext;
 import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.apache.pinot.sql.parsers.rewriter.OrderByLiteralRemover;
+import org.apache.pinot.sql.parsers.rewriter.QueryRewriter;
+import org.apache.pinot.sql.parsers.rewriter.QueryRewriterFactory;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
@@ -43,6 +54,171 @@ public class InnerSegmentSelectionSingleValueQueriesTest 
extends BaseSingleValue
   private static final String SELECTION_QUERY = "SELECT column1, column5, 
column11 FROM testTable";
   private static final String ORDER_BY = " ORDER BY column6, column1";
 
+  @DataProvider(name = "literalOrderByOptions")
+  public Object[][] literalOrderByOptions() {
+    return new Object[][]{{false, false}, {false, true}, {true, false}, {true, 
true}};
+  }
+
+  @Test(dataProvider = "literalOrderByOptions")
+  public void testSelectionOrderByLiterals(boolean allowReverseOrder, boolean 
enableNullHandling) {
+    String options = "SET allowReverseOrder = " + allowReverseOrder + "; SET 
enableNullHandling = "
+        + enableNullHandling + "; ";
+    for (String orderBy : List.of("2 DESC", "2 ASC", "2", "now() DESC", "NULL 
DESC", "'a' DESC, 'b' ASC")) {
+      String query = options + "SELECT column1, now() AS t FROM testTable 
ORDER BY " + orderBy + " LIMIT 5 OFFSET 3";
+      BaseOperator<SelectionResultsBlock> operator = getOperator(query);
+      assertTrue(operator instanceof SelectionOnlyOperator, query);
+      assertEquals(operator.nextBlock().getRows().size(), 5, query);
+      assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), 5L, 
query);
+      BrokerResponseNative response = getBrokerResponse(query);
+      assertTrue(response.getExceptions().isEmpty(), query);
+      assertEquals(response.getResultTable().getRows().size(), 5, query);
+      assertEquals(response.getResultTable().getDataSchema().getColumnNames(), 
new String[]{"column1", "t"});
+    }
+
+    // A literal referenced by a select-list ordinal is also removed from 
ORDER BY.
+    BaseOperator<SelectionResultsBlock> operator =
+        getOperator(options + "SELECT now() AS t FROM testTable ORDER BY 1 
DESC LIMIT 5 OFFSET 3");
+    assertTrue(operator instanceof SelectionOnlyOperator);
+    assertEquals(operator.nextBlock().getRows().size(), 5);
+    assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), 5L);
+
+    // With no remaining sort keys, preserve selection-only semantics: OFFSET 
is ignored.
+    BrokerResponseNative response = getBrokerResponse(options
+        + "SELECT 'constant' AS value FROM testTable ORDER BY 1 DESC LIMIT 10 
OFFSET 119997");
+    assertTrue(response.getExceptions().isEmpty());
+    assertEquals(response.getResultTable().getRows().size(), 10);
+    for (Object[] row : response.getResultTable().getRows()) {
+      assertEquals(row, new Object[]{"constant"});
+    }
+    response = getBrokerResponse(options
+        + "SELECT column1 FROM testTable WHERE column1 < 0 ORDER BY 'constant' 
DESC LIMIT 5 OFFSET 3");
+    assertTrue(response.getExceptions().isEmpty());
+    assertTrue(response.getResultTable().getRows().isEmpty());
+  }
+
+  @Test(dataProvider = "literalOrderByOptions")
+  public void testSelectionOrderByLiteralsAndSortedColumn(boolean 
allowReverseOrder, boolean enableNullHandling) {
+    String options = "SET allowReverseOrder = " + allowReverseOrder + "; SET 
enableNullHandling = "
+        + enableNullHandling + "; ";
+    for (boolean asc : List.of(true, false)) {
+      String direction = asc ? " ASC" : " DESC";
+      String literalDirection = asc ? " DESC" : " ASC";
+      String query = options + "SELECT daysSinceEpoch FROM testTable ORDER BY 
'leading'" + literalDirection
+          + ", daysSinceEpoch" + direction + ", 'trailing'" + literalDirection 
+ " LIMIT 5 OFFSET 3";
+      BaseOperator<SelectionResultsBlock> operator = getOperator(query);
+      if (asc || allowReverseOrder) {
+        assertTrue(operator instanceof 
SelectionPartiallyOrderedByLinearOperator, query);
+      } else {
+        assertTrue(operator instanceof 
SelectionPartiallyOrderedByDescOperation, query);
+      }
+      assertEquals(operator.nextBlock().getRows().size(), 8, query);
+      assertEquals(operator.getExecutionStatistics().getNumDocsScanned(),
+          asc || allowReverseOrder ? 8L : 30000L, query);
+      BrokerResponseNative response = getBrokerResponse(query);
+      assertTrue(response.getExceptions().isEmpty(), query);
+      List<Object[]> expectedRows = getBrokerResponse(options
+          + "SELECT daysSinceEpoch FROM testTable ORDER BY daysSinceEpoch" + 
direction + " LIMIT 5 OFFSET 3")
+          .getResultTable().getRows();
+      List<Object[]> rows = response.getResultTable().getRows();
+      assertEquals(rows.size(), expectedRows.size(), query);
+      for (int i = 0; i < rows.size(); i++) {
+        assertEquals(rows.get(i), expectedRows.get(i), query);
+      }
+    }
+  }
+
+  @Test(dataProvider = "literalOrderByOptions")
+  public void testOrderByLiteralPrefixAndUnsortedSuffix(boolean 
allowReverseOrder, boolean enableNullHandling) {
+    String options = "SET allowReverseOrder = " + allowReverseOrder + "; SET 
enableNullHandling = "
+        + enableNullHandling + "; ";
+    for (String suffix : List.of("column1 DESC, column6 ASC", "add(column1, 1) 
ASC, column6 DESC")) {
+      String query = options + "SELECT column1, column6 FROM testTable ORDER 
BY 'constant' ASC, " + suffix
+          + " LIMIT 5 OFFSET 3";
+      BaseOperator<SelectionResultsBlock> operator = getOperator(query);
+      assertTrue(operator instanceof SelectionOrderByOperator, query);
+      assertEquals(operator.nextBlock().getRows().size(), 8, query);
+      assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), 
30000L, query);
+      BrokerResponseNative response = getBrokerResponse(query);
+      assertTrue(response.getExceptions().isEmpty(), query);
+      List<Object[]> expectedRows = getBrokerResponse(options
+          + "SELECT column1, column6 FROM testTable ORDER BY " + suffix + " 
LIMIT 5 OFFSET 3")
+          .getResultTable().getRows();
+      List<Object[]> rows = response.getResultTable().getRows();
+      assertEquals(rows.size(), expectedRows.size(), query);
+      for (int i = 0; i < rows.size(); i++) {
+        assertEquals(rows.get(i), expectedRows.get(i), query);
+      }
+    }
+  }
+
+  @Test(dataProvider = "literalOrderByOptions")
+  public void testLiteralOrderByWithoutRewriter(boolean allowReverseOrder, 
boolean enableNullHandling) {
+    // Custom rewriter lists and requests from older brokers can retain 
literal ORDER BY keys.
+    for (String orderBy : List.of("2 DESC", "2 ASC", "now() DESC", "'a' DESC, 
'b' ASC", "'a' ASC, 'b' DESC")) {
+      String sql = "SELECT column1, now() AS t FROM testTable ORDER BY " + 
orderBy + " LIMIT 5 OFFSET 3";
+      PinotQuery query = compileWithoutLiteralRemoval(sql, allowReverseOrder, 
enableNullHandling);
+      SelectionResultsBlock block = new SelectionPlanNode(new 
SegmentContext(getIndexSegment()),
+          QueryContextConverterUtils.getQueryContext(query)).run().nextBlock();
+      assertEquals(block.getRows().size(), 8, sql);
+      BrokerResponseNative response = getBrokerResponse(query, PLAN_MAKER);
+      assertTrue(response.getExceptions().isEmpty(), sql);
+      assertEquals(response.getResultTable().getRows().size(), 5, sql);
+      assertEquals(response.getResultTable().getDataSchema().getColumnNames(), 
new String[]{"column1", "t"});
+    }
+
+    // Retaining ORDER BY also retains its OFFSET handling at the broker.
+    PinotQuery query = compileWithoutLiteralRemoval(
+        "SELECT 'constant' AS value FROM testTable ORDER BY 1 DESC LIMIT 10 
OFFSET 119997",
+        allowReverseOrder, enableNullHandling);
+    BrokerResponseNative response = getBrokerResponse(query, PLAN_MAKER);
+    assertTrue(response.getExceptions().isEmpty());
+    assertEquals(response.getResultTable().getRows().size(), 3);
+    for (Object[] row : response.getResultTable().getRows()) {
+      assertEquals(row, new Object[]{"constant"});
+    }
+  }
+
+  @Test(dataProvider = "literalOrderByOptions")
+  public void testLiteralPrefixWithoutRewriter(boolean allowReverseOrder, 
boolean enableNullHandling) {
+    String options = "SET allowReverseOrder = " + allowReverseOrder + "; SET 
enableNullHandling = "
+        + enableNullHandling + "; ";
+    for (String direction : List.of("ASC", "DESC")) {
+      for (String[] expressions : List.of(
+          new String[]{"daysSinceEpoch", "daysSinceEpoch ASC"},
+          new String[]{"daysSinceEpoch", "daysSinceEpoch DESC"},
+          new String[]{"column1, column6", "column1 DESC, column6 ASC"},
+          new String[]{"column1, column6", "add(column1, 1) ASC, column6 
DESC"})) {
+        String select = "SELECT " + expressions[0] + " FROM testTable ORDER BY 
";
+        String suffix = expressions[1] + " LIMIT 5 OFFSET 3";
+        String sql = select + "'constant' " + direction + ", " + suffix;
+        PinotQuery query = compileWithoutLiteralRemoval(sql, 
allowReverseOrder, enableNullHandling);
+        BrokerResponseNative response = getBrokerResponse(query, PLAN_MAKER);
+        assertTrue(response.getExceptions().isEmpty(), sql);
+        List<Object[]> expectedRows = getBrokerResponse(options + select + 
suffix).getResultTable().getRows();
+        List<Object[]> rows = response.getResultTable().getRows();
+        assertEquals(rows.size(), expectedRows.size(), sql);
+        for (int i = 0; i < rows.size(); i++) {
+          assertEquals(rows.get(i), expectedRows.get(i), sql);
+        }
+      }
+    }
+  }
+
+  private static PinotQuery compileWithoutLiteralRemoval(String sql, boolean 
allowReverseOrder,
+      boolean enableNullHandling) {
+    PinotQuery query = 
CalciteSqlParser.compileToPinotQueryWithoutRewrites(sql);
+    // Apply a custom pipeline without modifying the global parser 
configuration used by other tests.
+    for (QueryRewriter rewriter : QueryRewriterFactory.getQueryRewriters()) {
+      if (!(rewriter instanceof OrderByLiteralRemover)) {
+        query = rewriter.rewrite(query);
+      }
+    }
+    query.setQueryOptions(new HashMap<>(Map.of("allowReverseOrder", 
Boolean.toString(allowReverseOrder),
+        "enableNullHandling", Boolean.toString(enableNullHandling))));
+    assertTrue(query.getOrderByListSize() > 0, sql);
+    return query;
+  }
+
   @Test
   public void testSelectLimitZero() {
     String limit = " LIMIT 0";


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to