@Override
  public List<IndexSegment> prune(List<IndexSegment> segments, QueryContext query) {
    if (segments.isEmpty()) {
      return segments;
    }

    // For LIMIT 0 case, keep one segment to create the schema
    int limit = query.getLimit();
    if (limit == 0) {
      return Collections.singletonList(segments.get(0));
    }

    // Skip pruning segments for upsert table because valid doc index is equivalent to a filter
    if (segments.get(0).getValidDocIds() != null) {
      return segments;
    }

    if (query.getOrderByExpressions() == null) {
      return pruneSelectionOnly(segments, query);
    } else {
      return pruneSelectionOrderBy(segments, query);
    }
  }
@Test
  public void testUpsertTable() {
    List<IndexSegment> indexSegments = Arrays
        .asList(getIndexSegment(0L, 10L, 10, true), getIndexSegment(20L, 30L, 10, true),
            getIndexSegment(40L, 50L, 10, true));

    // Should not prune any segment for upsert table
    QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable LIMIT 5");
    List<IndexSegment> result = _segmentPruner.prune(indexSegments, queryContext);
    assertEquals(result.size(), 3);

    queryContext =
        QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY testColumn LIMIT 5");
    result = _segmentPruner.prune(indexSegments, queryContext);
    assertEquals(result.size(), 3);
  }