@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 AND the
      // input node would already return fewer than sort.offset + sort.fetch
      // rows (e.g. there is already an inner limit applied)
      return;
    }

    RelCollation collation = sort.getCollation();
    Preconditions.checkArgument(
        collation.equals(exchange.getCollation()),
        "Expected collation on exchange and sort to be the same"
    );

    final RexNode fetch;
    if (sort.fetch == null) {
      fetch = null;
    } else if (sort.offset == null) {
      fetch = sort.fetch;
    } else {
      int total = RexExpressionUtils.getValueAsInt(sort.fetch) + RexExpressionUtils.getValueAsInt(sort.offset);
      fetch = REX_BUILDER.makeLiteral(total, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
    }
    // do not transform sort-exchange copy when there's no fetch limit, or fetch amount is larger than threshold
    if (!collation.getFieldCollations().isEmpty()
        && (fetch == null || RexExpressionUtils.getValueAsInt(fetch) > DEFAULT_SORT_EXCHANGE_COPY_THRESHOLD)) {
      return;
    }

    final RelNode newExchangeInput = sort.copy(sort.getTraitSet(), exchange.getInput(), collation, null, fetch);
    final RelNode exchangeCopy = exchange.copy(exchange.getTraitSet(), newExchangeInput, exchange.getDistribution());
    final RelNode sortCopy = sort.copy(sort.getTraitSet(), exchangeCopy, collation,
        sort.offset == null ? REX_ZERO : sort.offset, sort.fetch);

    call.transformTo(sortCopy);
  }
@Test
  public void shouldMatchLimitNoOffsetYesSortOnSender() {
    // Given:
    RelCollation collation = RelCollations.of(1);
    SortExchange exchange = PinotLogicalSortExchange.create(_input, RelDistributions.SINGLETON, collation, true, false);
    Sort sort = LogicalSort.create(exchange, collation, 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());

    RelNode sortCopy = sortCopyCapture.getValue();
    Assert.assertTrue(sortCopy instanceof LogicalSort);
    Assert.assertTrue(((LogicalSort) sortCopy).getInput() instanceof PinotLogicalSortExchange);
    Assert.assertTrue(((LogicalSort) sortCopy).getInput().getInput(0) instanceof LogicalSort);

    LogicalSort innerSort = (LogicalSort) ((LogicalSort) sortCopy).getInput().getInput(0);
    Assert.assertEquals(innerSort.getCollation().getKeys().size(), 1);
    Assert.assertNull((innerSort).offset);
    Assert.assertEquals((innerSort).fetch, literal(1));
  }