@VisibleForTesting
  static Optional<String> findBestColumnForPartitioning(List<Pair<String, Double>> columnNameToWeightPairs,
      Function<String, Double> cardinalityExtractor, double topCandidateRatio, int numPartitions) {
    return columnNameToWeightPairs.stream().filter(colToWeight -> cardinalityExtractor.apply(colToWeight.getLeft())
            > numPartitions * PartitionRule.ACCEPTABLE_CARDINALITY_TO_NUM_PARTITIONS_RATIO)
        .max(Comparator.comparingDouble(Pair::getRight)).map(Pair::getRight).flatMap(maxWeight -> {
          double topCandidatesThreshold = maxWeight * topCandidateRatio;
          return columnNameToWeightPairs.stream().filter(colToWeight -> colToWeight.getRight() > topCandidatesThreshold)
              .max(Comparator.comparingDouble(colToWeight -> cardinalityExtractor.apply(colToWeight.getLeft())))
              .map(Pair::getLeft);
        });
  }
@Test
  public void testFindBestColumnForPartitioning() {

    // d has max weight, but c has more cardinality
    String[] columnNames = {"a", "b", "c", "d", "e"};
    double[] cardinalities = {100, 90, 80, 70, 60};
    double[] weights = {0.2, 0.3, 0.7, 0.8, 0.5};
    assertEquals(Optional.of("c"), findBestColumn(columnNames, cardinalities, weights));

    // d has max weight and max cardinality
    columnNames = new String[]{"a", "b", "c", "d", "e"};
    cardinalities = new double[]{100, 90, 80, 200, 60};
    weights = new double[]{0.2, 0.3, 0.7, 0.8, 0.5};
    assertEquals(Optional.of("d"), findBestColumn(columnNames, cardinalities, weights));

    // d has max weight, but its cardinality compared to numPartition is lower than threshold;
    columnNames = new String[]{"a", "b", "c", "d", "e"};
    cardinalities = new double[]{100, 90, 80, 10, 60};
    weights = new double[]{0.1, 0.1, 0.4, 0.8, 0.2};
    assertEquals(Optional.of("c"), findBestColumn(columnNames, cardinalities, weights));
  }