public static Pair<CloudObjectIncrCheckpoint, Option<Dataset<Row>>> filterAndGenerateCheckpointBasedOnSourceLimit(Dataset<Row> sourceData,
                                                                                                                    long sourceLimit, QueryInfo queryInfo,
                                                                                                                    CloudObjectIncrCheckpoint cloudObjectIncrCheckpoint) {
    if (sourceData.isEmpty()) {
      // There is no file matching the prefix.
      CloudObjectIncrCheckpoint updatedCheckpoint =
              queryInfo.getEndInstant().equals(cloudObjectIncrCheckpoint.getCommit())
                      ? cloudObjectIncrCheckpoint
                      : new CloudObjectIncrCheckpoint(queryInfo.getEndInstant(), null);
      return Pair.of(updatedCheckpoint, Option.empty());
    }
    // Let's persist the dataset to avoid triggering the dag repeatedly
    sourceData.persist(StorageLevel.MEMORY_AND_DISK());
    // Set ordering in query to enable batching
    Dataset<Row> orderedDf = QueryRunner.applyOrdering(sourceData, queryInfo.getOrderByColumns());
    Option<String> lastCheckpoint = Option.of(cloudObjectIncrCheckpoint.getCommit());
    Option<String> lastCheckpointKey = Option.ofNullable(cloudObjectIncrCheckpoint.getKey());
    Option<String> concatenatedKey = lastCheckpoint.flatMap(checkpoint -> lastCheckpointKey.map(key -> checkpoint + key));

    // Filter until last checkpoint key
    if (concatenatedKey.isPresent()) {
      orderedDf = orderedDf.withColumn("commit_key",
          functions.concat(functions.col(queryInfo.getOrderColumn()), functions.col(queryInfo.getKeyColumn())));
      // Apply incremental filter
      orderedDf = orderedDf.filter(functions.col("commit_key").gt(concatenatedKey.get())).drop("commit_key");
      // If there are no more files where commit_key is greater than lastCheckpointCommit#lastCheckpointKey
      if (orderedDf.isEmpty()) {
        LOG.info("Empty ordered source, returning endpoint:" + queryInfo.getEndInstant());
        sourceData.unpersist();
        // queryInfo.getEndInstant() represents source table's last completed instant
        // If current checkpoint is c1#abc and queryInfo.getEndInstant() is c1, return c1#abc.
        // If current checkpoint is c1#abc and queryInfo.getEndInstant() is c2, return c2.
        CloudObjectIncrCheckpoint updatedCheckpoint =
            queryInfo.getEndInstant().equals(cloudObjectIncrCheckpoint.getCommit())
                ? cloudObjectIncrCheckpoint
                : new CloudObjectIncrCheckpoint(queryInfo.getEndInstant(), null);
        return Pair.of(updatedCheckpoint, Option.empty());
      }
    }

    // Limit based on sourceLimit
    WindowSpec windowSpec = Window.orderBy(col(queryInfo.getOrderColumn()), col(queryInfo.getKeyColumn()));
    // Add the 'cumulativeSize' column with running sum of 'limitColumn'
    Dataset<Row> aggregatedData = orderedDf.withColumn(CUMULATIVE_COLUMN_NAME,
        sum(col(queryInfo.getLimitColumn())).over(windowSpec));
    Dataset<Row> collectedRows = aggregatedData.filter(col(CUMULATIVE_COLUMN_NAME).leq(sourceLimit));

    Row row = null;
    if (collectedRows.isEmpty()) {
      // If the first element itself exceeds limits then return first element
      LOG.info("First object exceeding source limit: " + sourceLimit + " bytes");
      row = aggregatedData.select(queryInfo.getOrderColumn(), queryInfo.getKeyColumn(), CUMULATIVE_COLUMN_NAME).first();
      collectedRows = aggregatedData.limit(1);
    } else {
      // Get the last row and form composite key
      row = collectedRows.select(queryInfo.getOrderColumn(), queryInfo.getKeyColumn(), CUMULATIVE_COLUMN_NAME).orderBy(
          col(queryInfo.getOrderColumn()).desc(), col(queryInfo.getKeyColumn()).desc()).first();
    }
    LOG.info("Processed batch size: " + row.get(row.fieldIndex(CUMULATIVE_COLUMN_NAME)) + " bytes");
    sourceData.unpersist();
    return Pair.of(new CloudObjectIncrCheckpoint(row.getString(0), row.getString(1)), Option.of(collectedRows));
  }
@Test
  void testFileOrderingAcrossCommits() {
    List<Triple<String, Long, String>> filePathSizeAndCommitTime = new ArrayList<>();
    // Add file paths and sizes to the list
    filePathSizeAndCommitTime.add(Triple.of("path/to/file8.json", 100L, "commit3"));
    filePathSizeAndCommitTime.add(Triple.of("path/to/file6.json", 250L, "commit3"));
    filePathSizeAndCommitTime.add(Triple.of("path/to/file7.json", 50L, "commit3"));
    filePathSizeAndCommitTime.add(Triple.of("path/to/file0.json", 100L, "commit4"));
    filePathSizeAndCommitTime.add(Triple.of("path/to/file1.json", 50L, "commit4"));
    filePathSizeAndCommitTime.add(Triple.of("path/to/file2.json", 50L, "commit4"));
    Dataset<Row> inputDs = generateDataset(filePathSizeAndCommitTime);

    QueryInfo queryInfo = new QueryInfo(
        QUERY_TYPE_INCREMENTAL_OPT_VAL(), "commit3", "commit3",
        "commit4", "_hoodie_commit_time",
        "s3.object.key", "s3.object.size");
    Pair<CloudObjectIncrCheckpoint, Option<Dataset<Row>>> result = IncrSourceHelper.filterAndGenerateCheckpointBasedOnSourceLimit(
        inputDs, 50L, queryInfo, new CloudObjectIncrCheckpoint("commit3", "path/to/file8.json"));
    Row row = result.getRight().get().select("cumulativeSize").collectAsList().get((int) result.getRight().get().count() - 1);
    assertEquals("commit4#path/to/file0.json", result.getKey().toString());
    List<Row> rows = result.getRight().get().collectAsList();
    assertEquals(1, rows.size());
    assertEquals(100L, row.get(0));

    result = IncrSourceHelper.filterAndGenerateCheckpointBasedOnSourceLimit(
        inputDs, 350L, queryInfo, new CloudObjectIncrCheckpoint("commit3", "path/to/file8.json"));
    row = result.getRight().get().select("cumulativeSize").collectAsList().get((int) result.getRight().get().count() - 1);
    assertEquals("commit4#path/to/file2.json", result.getKey().toString());
    rows = result.getRight().get().collectAsList();
    assertEquals(3, rows.size());
    assertEquals(200L, row.get(0));
  }