@Override
  public Expression resolveSelect(final int idx, final Expression expression) {
    for (int i = 0; i < partitionBys.size(); i++) {
      if (partitionBys.get(i).equals(expression)) {
        return new UnqualifiedColumnReferenceExp(getSchema().key().get(i).name());
      }
    }

    return expression;
  }
@Test
  public void shouldResolveSelectToPartitionByColumn() {
    // When:
    final Expression result = repartitionNode.resolveSelect(0, rewrittenPartitionBy);

    // Then:
    assertThat(result, is(new UnqualifiedColumnReferenceExp(K0)));
  }