You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by GitBox <gi...@apache.org> on 2022/11/18 00:05:05 UTC

[GitHub] [pinot] agavra opened a new pull request, #9832: [multistage] support sort push-down

agavra opened a new pull request, #9832:
URL: https://github.com/apache/pinot/pull/9832

   This replaces #9507
   
   **Review Notes**
   
   1. **!!NOTE!!**: there is a v1 change in this PR, if `SERVER_RETURN_FINAL_RESULT` is set, the server will now return the "proper" columns, ordered by the query projection instead of the internal format. You can see this change in the files for `*CombineOperator` and `SelectionOperatorService`
   2. [ImmutableSortExchangeCopyRule.java](https://github.com/apache/pinot/compare/master...agavra:pinot:sort_pushdown?expand=1#diff-8a22fa3cfc8f05c0581192f941dc92c76d3cf5ebf2250cac528256fa2db0764a) is a generated file that I'm checking in. I generated it within the Calcite code base and ran `./gradlew generateSources` and copied the contents over to Pinot. We could generate it as part of Pinot build, but that would require pulling in extra dependencies to do the codegen and add that to our build, and this file should change extremely infrequently.
   3. The "main" part of the change is the introduction of `PinotSortExchangeCopyRule`, which simply copies a `Sort` past an `Exchange` (dropping the offset, and setting the limit to limit + offset). This was inspired by the implementation of https://github.com/apache/calcite/blob/406c913b808b3234464d8c81d7352c4040dd281a/core/src/main/java/org/apache/calcite/rel/rules/SortJoinCopyRule.java 
   
   **What**
   
   This PR supports `SORT` operator push-down. For example, `SELECT * FROM basic ORDER BY col1 DESC LIMIT 2 OFFSET 1` will now generate the following calcite logical plan:
   ```
   LogicalSort(sort0=[$2], dir0=[DESC], offset=[1], fetch=[2])
     LogicalSortExchange(distribution=[hash], collation=[[2 DESC]])
       LogicalSort(sort0=[$2], dir0=[DESC], fetch=[+(2, 1)])
         LogicalTableScan(table=[[basic_order_by_basic]])
   ```
   
   which generates this pinot query plan:
   ```
   [0]@localhost:50762 MAIL_RECEIVE(RANDOM_DISTRIBUTED)
   ├── [1]@localhost:50761 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:50762} (Subtree Omitted)
   └── [1]@localhost:50760 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:50762}
      └── [1]@localhost:50760 SORT (LIMIT 2)
         └── [1]@localhost:50760 MAIL_RECEIVE(HASH_DISTRIBUTED)
            └── [2]@localhost:50760 MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:50761,[1]@localhost:50760}
               └── [2]@localhost:50760 SORT (LIMIT 3)
                  └── [2]@localhost:50760 TABLE SCAN (basic_order_by_basic) {OFFLINE=[basic_order_by_basic_OFFLINE_e625b8e5-ede7-40e3-8696-1992ecee3f9d]}
   ```
   
   Notice that there are now _two_ `SORT` operators, one right after the table scan and another right after the mail receive (the operator that consolidates responses from both localhost:50760 and localhost:50761).
   
   **Why**
   
   Without this change, the leaf nodes have no knowledge of how much data they should send to the multistage intermediate servers. You can imagine why this is a problem: if a query specifies `LIMIT 3` but queries a large data set with low selectivity, the server nodes will respond to the intermediate node with potentially gigabytes of data. This compounded with the observation that our current engine has no flow-control that pipes down to v1 (leaf server queries will complete in entirety before sending any data back) means that before this change our queries were likely to timeout (and cost a lot of money in network bandwidth).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1028365376


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   yes, I can try to add a unit test for it as well



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1030734248


##########
pinot-query-runtime/src/test/resources/queries/OrderBy.json:
##########
@@ -0,0 +1,221 @@
+{
+  "basic_order_by": {
+    "tables": {
+      "basic": {
+        "schema": [
+          {"name": "col0", "type": "INT"},
+          {"name": "col1", "type": "INT"},
+          {"name": "col2", "type": "STRING"}
+        ],
+        "inputs": [
+          [1, 2, "a"],
+          [2, 3, "b"],
+          [3, 1, "c"],
+          [4, 4, "d"],
+          [5, 5, "e"],
+          [6, 6, "f"]
+        ]
+      }
+    },
+    "queries": [
+      {"sql": "SELECT * FROM {basic} ORDER BY col1 LIMIT 2 OFFSET 1"},

Review Comment:
   I changed the `OFFSET 10` and `LIMIT 10` clauses to be 100 (they were meant to change the same thing) and added `OFFSET 0` and `LIMIT ALL`. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036173291


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, null, literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(1));
+  }
+
+  @Test
+  public void shouldMatchNoSortAndPushDownLimitPlusOffset() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, literal(2), literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldMatchSortOnly() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, null, null);
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertNull((innerSort).fetch);
+  }
+
+  @Test
+  public void shouldMatchLimitOffsetAndSort() {

Review Comment:
   missing `shouldMatchLimitNoOffsetSort`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1329951586

   @Jackie-Jiang - I removed the usage of this flag and v1 changes, this PR now only affects the multistage engine :) hopefully there's no concerns now


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1030729227


##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/IntExprRexVisitor.java:
##########
@@ -0,0 +1,126 @@
+/**
+ * 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.query.planner.logical;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexFieldAccess;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexLocalRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
+import org.apache.calcite.rex.RexPatternFieldRef;
+import org.apache.calcite.rex.RexRangeRef;
+import org.apache.calcite.rex.RexSubQuery;
+import org.apache.calcite.rex.RexTableInputRef;
+import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+
+
+/**
+ * {@code IntExprRexVisitor} will visit a {@link RexNode} that
+ * contains only integer literals and uses the {@link SqlStdOperatorTable#PLUS}
+ * operator and return the Integer that the expression represents.
+ */
+public class IntExprRexVisitor implements RexVisitor<Integer> {
+
+  public static final IntExprRexVisitor INSTANCE = new IntExprRexVisitor();
+
+  private IntExprRexVisitor() {
+  }
+
+  public Integer visit(RexNode in) {
+    if (in == null) {
+      return null;
+    }

Review Comment:
   good catch, this was leftover from a previous refactoring but returning 0 is actually the correct behavior now



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1028436129


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   added a test for this, PTAL



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1030730187


##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/IntExprRexVisitor.java:
##########
@@ -0,0 +1,126 @@
+/**
+ * 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.query.planner.logical;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexFieldAccess;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexLocalRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
+import org.apache.calcite.rex.RexPatternFieldRef;
+import org.apache.calcite.rex.RexRangeRef;
+import org.apache.calcite.rex.RexSubQuery;
+import org.apache.calcite.rex.RexTableInputRef;
+import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+
+
+/**
+ * {@code IntExprRexVisitor} will visit a {@link RexNode} that
+ * contains only integer literals and uses the {@link SqlStdOperatorTable#PLUS}
+ * operator and return the Integer that the expression represents.
+ */
+public class IntExprRexVisitor implements RexVisitor<Integer> {
+
+  public static final IntExprRexVisitor INSTANCE = new IntExprRexVisitor();
+
+  private IntExprRexVisitor() {
+  }
+
+  public Integer visit(RexNode in) {
+    if (in == null) {
+      return null;
+    }
+
+    return in.accept(this);
+  }
+
+  @Override
+  public Integer visitInputRef(RexInputRef inputRef) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Integer visitLocalRef(RexLocalRef localRef) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Integer visitLiteral(RexLiteral literal) {
+    return literal.getValueAs(Integer.class);
+  }
+
+  @Override
+  public Integer visitCall(RexCall call) {
+    SqlOperator operator = call.getOperator();
+    if (!(operator == SqlStdOperatorTable.PLUS)) {

Review Comment:
   same as the above, this was actually left over from a refactor where I used the PLUS call to add the offset and fetch, but now I'm just resolving it in place. I'll just delete this class, it doesn't make much sense the way it is now.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1027002234


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   Good catch! I'll add a test for this in a follow-up PR (I need to introduce a new testing framework to do this well) but in the meantime here's proof that the new code works:
   ```
   Generated Plan for 'SELECT * FROM {basic} LIMIT 2':
   
   [0]@localhost:52943 MAIL_RECEIVE(RANDOM_DISTRIBUTED)
   ├── [1]@localhost:52941 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:52943} (Subtree Omitted)
   └── [1]@localhost:52942 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:52943}
      └── [1]@localhost:52942 SORT LIMIT 2
         └── [1]@localhost:52942 MAIL_RECEIVE(HASH_DISTRIBUTED)
            └── [2]@localhost:52941 MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:52941,[1]@localhost:52942}
               └── [2]@localhost:52941 SORT LIMIT 2
                  └── [2]@localhost:52941 TABLE SCAN (basic_order_by_basic) {OFFLINE=[basic_order_by_basic_OFFLINE_3d103a31-4bea-46a0-b83a-cbfa305ac403]}
   
   
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1027164349


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   Is it possible we test this rule in unit test somehow? 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1324147001

   > @Jackie-Jiang - yes, the V2 engine requires that the server returns columns in the order that would be the result of the logical plan (and also that it doesn't have extra or missing columns, which could be the case of the order by column wasn't selected or a column was selected multiple times).
   > 
   > Calcite makes assumptions that each stage will return results that are "correct" for that logical stage. The alternative would be, instead of building this into v1, wrapping the v2 leaf stage with an operator that rearranges columns but it felt like this flag would be the exact thing we should leverage.
   
   Is it possible we re-arrange columns in intermediate stage instead if we pass data schema to intermediate stage? This way we can leave the leaf stage untouched. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036210452


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, null, literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(1));
+  }
+
+  @Test
+  public void shouldMatchNoSortAndPushDownLimitPlusOffset() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, literal(2), literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldMatchSortOnly() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, null, null);
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertNull((innerSort).fetch);
+  }
+
+  @Test
+  public void shouldMatchLimitOffsetAndSort() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, literal(1), literal(2));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldNotMatchOnlySortAlreadySorted() {

Review Comment:
   this is much harder to test, essentially the "alreadySmaller" could see that there is already an inner sort with a limit that limits the rows to fewer than what we'd need to satisfy the sort in question - alternatively, we could have table stats to show that the input is already smaller than the limit in question.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] Jackie-Jiang commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
Jackie-Jiang commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1329930985

   Currently `serverReturnFinalResult` is used by v1 engine to return the final aggregation result instead of the intermediate aggregation result.
   The purpose for the flag in this PR is to order the columns based on the selection sequence (IIUC, the query constructed from v2 will always have order-by columns included in the selection sequence). This flag can also be applied to group-by queries which currently put group-by columns in from of other columns. We don't want to overload the same flag, so I'd suggest adding a separate flag for it.
   Alternatively, is it possible to construct the leaf stage query in a way that the v1 result always match the logical stage order?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] Jackie-Jiang commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
Jackie-Jiang commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1035078758


##########
pinot-core/src/main/java/org/apache/pinot/core/query/selection/SelectionOperatorService.java:
##########
@@ -129,36 +128,23 @@ public void reduceWithOrdering(Collection<DataTable> dataTables, boolean nullHan
    * <p>Should be called after method "reduceWithOrdering()".
    */
   public ResultTable renderResultTableWithOrdering() {
-    int[] columnIndices = SelectionOperatorUtils.getColumnIndices(_selectionColumns, _dataSchema);
-    int numColumns = columnIndices.length;
-
-    // Construct the result data schema
-    String[] columnNames = _dataSchema.getColumnNames();
-    ColumnDataType[] columnDataTypes = _dataSchema.getColumnDataTypes();
-    String[] resultColumnNames = new String[numColumns];
-    ColumnDataType[] resultColumnDataTypes = new ColumnDataType[numColumns];
-    for (int i = 0; i < numColumns; i++) {
-      int columnIndex = columnIndices[i];
-      resultColumnNames[i] = columnNames[columnIndex];
-      resultColumnDataTypes[i] = columnDataTypes[columnIndex];
-    }
-    DataSchema resultDataSchema = new DataSchema(resultColumnNames, resultColumnDataTypes);
+    Iterator<Object[]> rows = new Iterator<Object[]>() {
+      @Override
+      public boolean hasNext() {
+        return _rows.size() > _offset;
+      }
 
-    // Extract the result rows
-    LinkedList<Object[]> rowsInSelectionResults = new LinkedList<>();
-    while (_rows.size() > _offset) {
-      Object[] row = _rows.poll();
-      assert row != null;
-      Object[] extractedRow = new Object[numColumns];
-      for (int i = 0; i < numColumns; i++) {
-        Object value = row[columnIndices[i]];
-        if (value != null) {
-          extractedRow[i] = resultColumnDataTypes[i].convertAndFormat(value);
-        }
+      @Override
+      public Object[] next() {
+        return _rows.poll();
       }
-      rowsInSelectionResults.addFirst(extractedRow);
-    }
+    };
 
-    return new ResultTable(resultDataSchema, rowsInSelectionResults);
+    return SelectionOperatorUtils.arrangeColumnsToMatchProjection(

Review Comment:
   (format) Reformat this to match the [Pinot Style](https://docs.pinot.apache.org/developers/developers-and-contributors/code-setup#setup-ide)



##########
pinot-core/src/main/java/org/apache/pinot/core/query/selection/SelectionOperatorUtils.java:
##########
@@ -647,4 +649,42 @@ public static <T> void addToPriorityQueue(T value, PriorityQueue<T> queue, int m
       queue.offer(value);
     }
   }
+
+  public static <T> T arrangeColumnsToMatchProjection(DataSchema dataSchema, Iterator<Object[]> rows,

Review Comment:
   Suggest not extracting the whole method because it is not as readable, and not very reusable. It is like forcefully stitching 2 methods together, with potential performance overhead of using iterator, passing function and extra per-record if check.
   We can extract the part of creating `resultDataSchema`: `public static DataSchema arrangeColumns(DataSchema dataSchema, int[] columnIndices)`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1026267981


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   IIUC, this doesn't handle the case where collation is empty. Is this intended to be out of scope of this PR?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1030910935


##########
pinot-query-runtime/src/test/resources/queries/OrderBy.json:
##########
@@ -0,0 +1,221 @@
+{
+  "basic_order_by": {
+    "tables": {
+      "basic": {
+        "schema": [
+          {"name": "col0", "type": "INT"},
+          {"name": "col1", "type": "INT"},
+          {"name": "col2", "type": "STRING"}
+        ],
+        "inputs": [
+          [1, 2, "a"],
+          [2, 3, "b"],
+          [3, 1, "c"],
+          [4, 4, "d"],
+          [5, 5, "e"],
+          [6, 6, "f"]
+        ]
+      }
+    },
+    "queries": [
+      {"sql": "SELECT * FROM {basic} ORDER BY col1 LIMIT 2 OFFSET 1"},

Review Comment:
   Sorry, didn't do the math. yeah. the 10 is doing the same thing for 100.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036169120


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);

Review Comment:
   it seems RelDistributions type doesn't matter but is there a specific reason why all are created against SINGLETON?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036208974


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, null, literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(1));
+  }
+
+  @Test
+  public void shouldMatchNoSortAndPushDownLimitPlusOffset() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, literal(2), literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldMatchSortOnly() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, null, null);
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertNull((innerSort).fetch);
+  }
+
+  @Test
+  public void shouldMatchLimitOffsetAndSort() {

Review Comment:
   correct! ^



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036206249


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);

Review Comment:
   no specific reason, that's just the easiest to create and this test doesn't really care about which type it is



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1030733303


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.planner.logical.IntExprRexVisitor;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  private static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(
+        metadataQuery,
+        exchange.getInput(),
+        sort.getCollation(),
+        sort.offset,
+        sort.fetch)) {
+      // Don't rewrite anything if the input is already sorted
+      return;
+    }
+
+    RelCollation collation = sort.getCollation();
+    Preconditions.checkArgument(
+        collation.equals(exchange.getCollation()),
+        "Expected collation on exchange and sort to be the same"
+    );
+
+    // for now, we push down the sort with the offset + limit as the

Review Comment:
   I actually just deleted this comment altogether, there's no reason not to push down the offset/limit



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] codecov-commenter commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
codecov-commenter commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1320920704

   # [Codecov](https://codecov.io/gh/apache/pinot/pull/9832?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) Report
   > Merging [#9832](https://codecov.io/gh/apache/pinot/pull/9832?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (adb3106) into [master](https://codecov.io/gh/apache/pinot/commit/16704a3cb1fc04958483fb565cef44f69e51cb17?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (16704a3) will **decrease** coverage by `48.35%`.
   > The diff coverage is `12.32%`.
   
   ```diff
   @@              Coverage Diff              @@
   ##             master    #9832       +/-   ##
   =============================================
   - Coverage     64.18%   15.82%   -48.36%     
   + Complexity     4977      175     -4802     
   =============================================
     Files          1911     1913        +2     
     Lines        102620   102738      +118     
     Branches      15612    15627       +15     
   =============================================
   - Hits          65862    16263    -49599     
   - Misses        31983    85280    +53297     
   + Partials       4775     1195     -3580     
   ```
   
   | Flag | Coverage Δ | |
   |---|---|---|
   | unittests1 | `?` | |
   | unittests2 | `15.82% <12.32%> (+0.07%)` | :arrow_up: |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#carryforward-flags-in-the-pull-request-comment) to find out more.
   
   | [Impacted Files](https://codecov.io/gh/apache/pinot/pull/9832?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | Coverage Δ | |
   |---|---|---|
   | [...not/core/operator/combine/BaseCombineOperator.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9vcGVyYXRvci9jb21iaW5lL0Jhc2VDb21iaW5lT3BlcmF0b3IuamF2YQ==) | `0.00% <0.00%> (-86.91%)` | :arrow_down: |
   | [...nMaxValueBasedSelectionOrderByCombineOperator.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9vcGVyYXRvci9jb21iaW5lL01pbk1heFZhbHVlQmFzZWRTZWxlY3Rpb25PcmRlckJ5Q29tYmluZU9wZXJhdG9yLmphdmE=) | `0.00% <0.00%> (-62.42%)` | :arrow_down: |
   | [...rator/combine/SelectionOrderByCombineOperator.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9vcGVyYXRvci9jb21iaW5lL1NlbGVjdGlvbk9yZGVyQnlDb21iaW5lT3BlcmF0b3IuamF2YQ==) | `0.00% <0.00%> (-65.63%)` | :arrow_down: |
   | [...core/query/selection/SelectionOperatorService.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9xdWVyeS9zZWxlY3Rpb24vU2VsZWN0aW9uT3BlcmF0b3JTZXJ2aWNlLmphdmE=) | `0.00% <0.00%> (-94.45%)` | :arrow_down: |
   | [...t/core/query/selection/SelectionOperatorUtils.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9xdWVyeS9zZWxlY3Rpb24vU2VsZWN0aW9uT3BlcmF0b3JVdGlscy5qYXZh) | `0.00% <0.00%> (-92.26%)` | :arrow_down: |
   | [...org/apache/pinot/query/planner/stage/SortNode.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcXVlcnktcGxhbm5lci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcGxhbm5lci9zdGFnZS9Tb3J0Tm9kZS5qYXZh) | `85.00% <0.00%> (-9.45%)` | :arrow_down: |
   | [...pinot/query/planner/logical/IntExprRexVisitor.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcXVlcnktcGxhbm5lci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcGxhbm5lci9sb2dpY2FsL0ludEV4cHJSZXhWaXNpdG9yLmphdmE=) | `14.28% <14.28%> (ø)` | |
   | [...not/query/planner/logical/RelToStageConverter.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcXVlcnktcGxhbm5lci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcGxhbm5lci9sb2dpY2FsL1JlbFRvU3RhZ2VDb252ZXJ0ZXIuamF2YQ==) | `83.33% <100.00%> (ø)` | |
   | [...ache/pinot/query/planner/logical/StagePlanner.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcXVlcnktcGxhbm5lci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcGxhbm5lci9sb2dpY2FsL1N0YWdlUGxhbm5lci5qYXZh) | `100.00% <100.00%> (ø)` | |
   | [...che/pinot/query/runtime/operator/SortOperator.java](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcXVlcnktcnVudGltZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcnVudGltZS9vcGVyYXRvci9Tb3J0T3BlcmF0b3IuamF2YQ==) | `93.33% <100.00%> (ø)` | |
   | ... and [1228 more](https://codecov.io/gh/apache/pinot/pull/9832/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | |
   
   :mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation)
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr merged pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr merged PR #9832:
URL: https://github.com/apache/pinot/pull/9832


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1033909356


##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/ServerRequestPlanVisitor.java:
##########
@@ -90,6 +92,9 @@ public static ServerPlanRequestContext build(MailboxService<TransferableBlock> m
     PinotQuery pinotQuery = new PinotQuery();
     pinotQuery.setLimit(DEFAULT_LEAF_NODE_LIMIT);
     pinotQuery.setExplain(false);
+    pinotQuery.setQueryOptions(ImmutableMap.of(
+        CommonConstants.Broker.Request.QueryOptionKey.SERVER_RETURN_FINAL_RESULT,

Review Comment:
   I removed this flag altogether.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1029892679


##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/IntExprRexVisitor.java:
##########
@@ -0,0 +1,126 @@
+/**
+ * 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.query.planner.logical;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexFieldAccess;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexLocalRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
+import org.apache.calcite.rex.RexPatternFieldRef;
+import org.apache.calcite.rex.RexRangeRef;
+import org.apache.calcite.rex.RexSubQuery;
+import org.apache.calcite.rex.RexTableInputRef;
+import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+
+
+/**
+ * {@code IntExprRexVisitor} will visit a {@link RexNode} that
+ * contains only integer literals and uses the {@link SqlStdOperatorTable#PLUS}
+ * operator and return the Integer that the expression represents.
+ */
+public class IntExprRexVisitor implements RexVisitor<Integer> {
+
+  public static final IntExprRexVisitor INSTANCE = new IntExprRexVisitor();
+
+  private IntExprRexVisitor() {
+  }
+
+  public Integer visit(RexNode in) {
+    if (in == null) {
+      return null;
+    }
+
+    return in.accept(this);
+  }
+
+  @Override
+  public Integer visitInputRef(RexInputRef inputRef) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Integer visitLocalRef(RexLocalRef localRef) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Integer visitLiteral(RexLiteral literal) {
+    return literal.getValueAs(Integer.class);
+  }
+
+  @Override
+  public Integer visitCall(RexCall call) {
+    SqlOperator operator = call.getOperator();
+    if (!(operator == SqlStdOperatorTable.PLUS)) {

Review Comment:
   Can you explain why only plus is supported?



##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.planner.logical.IntExprRexVisitor;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  private static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(
+        metadataQuery,
+        exchange.getInput(),
+        sort.getCollation(),
+        sort.offset,
+        sort.fetch)) {
+      // Don't rewrite anything if the input is already sorted

Review Comment:
   Add comment saying don't rewrite anything if the input is already sorted and the fetch number greater than row count?



##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/IntExprRexVisitor.java:
##########
@@ -0,0 +1,126 @@
+/**
+ * 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.query.planner.logical;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexFieldAccess;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexLocalRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
+import org.apache.calcite.rex.RexPatternFieldRef;
+import org.apache.calcite.rex.RexRangeRef;
+import org.apache.calcite.rex.RexSubQuery;
+import org.apache.calcite.rex.RexTableInputRef;
+import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+
+
+/**
+ * {@code IntExprRexVisitor} will visit a {@link RexNode} that
+ * contains only integer literals and uses the {@link SqlStdOperatorTable#PLUS}
+ * operator and return the Integer that the expression represents.
+ */
+public class IntExprRexVisitor implements RexVisitor<Integer> {
+
+  public static final IntExprRexVisitor INSTANCE = new IntExprRexVisitor();
+
+  private IntExprRexVisitor() {
+  }
+
+  public Integer visit(RexNode in) {
+    if (in == null) {
+      return null;
+    }

Review Comment:
   Should this just return 0 or we can say do Precheck(in !=null)? I feel it is not good to return null here.



##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   Thanks a lot for doing this!



##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.planner.logical.IntExprRexVisitor;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  private static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(
+        metadataQuery,
+        exchange.getInput(),
+        sort.getCollation(),
+        sort.offset,
+        sort.fetch)) {
+      // Don't rewrite anything if the input is already sorted
+      return;
+    }
+
+    RelCollation collation = sort.getCollation();
+    Preconditions.checkArgument(
+        collation.equals(exchange.getCollation()),
+        "Expected collation on exchange and sort to be the same"
+    );
+
+    // for now, we push down the sort with the offset + limit as the

Review Comment:
   Can we put a comment block saying something like /***customized code for pinot***/.xxxxx/**********/



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1324124148

   @Jackie-Jiang - yes, the V2 engine requires that the server returns columns in the order that would be the result of the logical plan (and also that it doesn't have extra or missing columns, which could be the case of the order by column wasn't selected or a column was selected multiple times). 
   
   Calcite makes assumptions that each stage will return results that are "correct" for that logical stage. The alternative would be, instead of building this into v1, wrapping the v2 leaf stage with an operator that rearranges columns but it felt like this flag would be the exact thing we should leverage.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on PR #9832:
URL: https://github.com/apache/pinot/pull/9832#issuecomment-1324225162

   > Is it possible we re-arrange columns in intermediate stage instead if we pass data schema to intermediate stage? This way we can leave the leaf stage untouched.
   
   @61yao no it isn't, I tried that for the initial PR but we generate data schemas based on the result of the calcite logical plan. I tried hacking calcite to allow that behavior and it gets really brittle (it complains during the validation phase that the expected schema and the operator schema don't match up)


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1029748399


##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/ServerRequestPlanVisitor.java:
##########
@@ -90,6 +92,9 @@ public static ServerPlanRequestContext build(MailboxService<TransferableBlock> m
     PinotQuery pinotQuery = new PinotQuery();
     pinotQuery.setLimit(DEFAULT_LEAF_NODE_LIMIT);
     pinotQuery.setExplain(false);
+    pinotQuery.setQueryOptions(ImmutableMap.of(
+        CommonConstants.Broker.Request.QueryOptionKey.SERVER_RETURN_FINAL_RESULT,

Review Comment:
   If we want to re-use the code, let's use a different config option that is invisible to user. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] 61yao commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
61yao commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1029934624


##########
pinot-query-runtime/src/test/resources/queries/OrderBy.json:
##########
@@ -0,0 +1,221 @@
+{
+  "basic_order_by": {
+    "tables": {
+      "basic": {
+        "schema": [
+          {"name": "col0", "type": "INT"},
+          {"name": "col1", "type": "INT"},
+          {"name": "col2", "type": "STRING"}
+        ],
+        "inputs": [
+          [1, 2, "a"],
+          [2, 3, "b"],
+          [3, 1, "c"],
+          [4, 4, "d"],
+          [5, 5, "e"],
+          [6, 6, "f"]
+        ]
+      }
+    },
+    "queries": [
+      {"sql": "SELECT * FROM {basic} ORDER BY col1 LIMIT 2 OFFSET 1"},

Review Comment:
   Can we add a test for offest 0, offset 100, limit 100 and Limit all (psql supports this). as well as limit null statement 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036156995


##########
pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRule.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.base.Preconditions;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+
+
+public class PinotSortExchangeCopyRule extends RelRule<RelRule.Config> {
+
+  public static final PinotSortExchangeCopyRule SORT_EXCHANGE_COPY =
+      PinotSortExchangeCopyRule.Config.DEFAULT.toRule();
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+
+  /**
+   * Creates a PinotSortExchangeCopyRule.
+   */
+  protected PinotSortExchangeCopyRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final SortExchange exchange = call.rel(1);
+    final RelMetadataQuery metadataQuery = call.getMetadataQuery();
+
+    if (RelMdUtil.checkInputForCollationAndLimit(

Review Comment:
   not sure I understand the comment here. isn't this also the behavior of the default immutablesortexchangecopyrule? anything specific we want to add in the future?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] walterddr commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
walterddr commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036177859


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, null, literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(1));
+  }
+
+  @Test
+  public void shouldMatchNoSortAndPushDownLimitPlusOffset() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, literal(2), literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldMatchSortOnly() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, null, null);
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertNull((innerSort).fetch);
+  }
+
+  @Test
+  public void shouldMatchLimitOffsetAndSort() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, literal(1), literal(2));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldNotMatchOnlySortAlreadySorted() {

Review Comment:
   seems like alreadySorted is tested but  alreadySmaller is not. from 
   ```
   public static boolean checkInputForCollationAndLimit( 
    ... {
     return alreadySorted(mq, input, collation) && alreadySmaller(mq, input, offset, fetch);
   }
   ```
   but given we are dealing with exchange / sort, i assume already smaller is always satisfied. yes?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

[GitHub] [pinot] agavra commented on a diff in pull request #9832: [multistage] support sort push-down

Posted by GitBox <gi...@apache.org>.
agavra commented on code in PR #9832:
URL: https://github.com/apache/pinot/pull/9832#discussion_r1036208045


##########
pinot-query-planner/src/test/java/org/apache/calcite/rel/rules/PinotSortExchangeCopyRuleTest.java:
##########
@@ -0,0 +1,222 @@
+/**
+ * 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.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalSortExchange;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.query.type.TypeSystem;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class PinotSortExchangeCopyRuleTest {
+
+  public static final TypeFactory TYPE_FACTORY = new TypeFactory(new TypeSystem());
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+
+  private AutoCloseable _mocks;
+
+  @Mock
+  private RelOptRuleCall _call;
+  @Mock
+  private RelNode _input;
+  @Mock
+  private RelOptCluster _cluster;
+  @Mock
+  private RelMetadataQuery _query;
+
+  @BeforeMethod
+  public void setUp() {
+    _mocks = MockitoAnnotations.openMocks(this);
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    Mockito.when(_input.getTraitSet()).thenReturn(traits);
+    Mockito.when(_input.getCluster()).thenReturn(_cluster);
+    Mockito.when(_call.getMetadataQuery()).thenReturn(_query);
+    Mockito.when(_query.getMaxRowCount(Mockito.any())).thenReturn(null);
+  }
+
+  @AfterMethod
+  public void tearDown()
+      throws Exception {
+    _mocks.close();
+  }
+
+  @Test
+  public void shouldMatchLimitNoOffsetNoSort() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, null, literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(1));
+  }
+
+  @Test
+  public void shouldMatchNoSortAndPushDownLimitPlusOffset() {
+    // Given:
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, RelCollations.EMPTY);
+    Sort sort = LogicalSort.create(exchange, RelCollations.EMPTY, literal(2), literal(1));
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 0);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertEquals((innerSort).fetch, literal(3));
+  }
+
+  @Test
+  public void shouldMatchSortOnly() {
+    // Given:
+    RelCollation collation = RelCollations.of(1);
+    SortExchange exchange = LogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation);
+    Sort sort = LogicalSort.create(exchange, collation, null, null);
+    Mockito.when(_call.rel(0)).thenReturn(sort);
+    Mockito.when(_call.rel(1)).thenReturn(exchange);
+
+    // When:
+    PinotSortExchangeCopyRule.SORT_EXCHANGE_COPY.onMatch(_call);
+
+    // Then:
+    ArgumentCaptor<RelNode> sortCopyCapture = ArgumentCaptor.forClass(LogicalSort.class);
+    Mockito.verify(_call, Mockito.times(1)).transformTo(sortCopyCapture.capture(), Mockito.anyMap());
+
+    RelNode sortCopy = sortCopyCapture.getValue();
+    Assert.assertTrue(sortCopy instanceof LogicalSort);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof LogicalSortExchange);
+    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);
+
+    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
+    Assert.assertEquals(innerSort.getCollation(), collation);
+    Assert.assertNull((innerSort).offset);
+    Assert.assertNull((innerSort).fetch);
+  }
+
+  @Test
+  public void shouldMatchLimitOffsetAndSort() {

Review Comment:
   the first test is `shouldMatchLimitNoOffsetNoSort`, or do you mean `yes limit, no offset, yes sort`? I can add that



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org