Method with_fetch

datafusion/common/src/stats.rs:536–649

Calculates the statistics after applying `fetch` and `skip` operations. Here, `self` denotes per-partition statistics. Use the `n_partitions` parameter to compute global statistics in a multi-partition setting.

(
        mut self,
        fetch: Option<usize>,
        skip: usize,
        n_partitions: usize,
    )

    /// Here, `self` denotes per-partition statistics. Use the `n_partitions`
    /// parameter to compute global statistics in a multi-partition setting.
    pub fn with_fetch(
        mut self,
        fetch: Option<usize>,
        skip: usize,
        n_partitions: usize,
    ) -> Result<Self> {
        let fetch_val = fetch.unwrap_or(usize::MAX);

        // Get the ratio of rows after / rows before on a per-partition basis
        let num_rows_before = self.num_rows;

        self.num_rows = match self {
            Statistics {
                num_rows: Precision::Exact(nr),
                ..
            }
            | Statistics {
                num_rows: Precision::Inexact(nr),
                ..
            } => {
                // Here, the inexact case gives us an estimate of the number of rows.
                if nr <= skip {
                    // All input data will be skipped. Preserve the exactness of
                    // the input estimate: if the input was inexact, the
                    // resulting zero is also inexact.
                    check_num_rows(Some(0), self.num_rows.is_exact().unwrap())
                } else if nr <= fetch_val && skip == 0 {
                    // If the input does not reach the `fetch` globally, and `skip`
                    // is zero (meaning the input and output are identical), return
                    // input stats as is.
                    // TODO: Can input stats still be used, but adjusted, when `skip`
                    //       is non-zero?
                    return Ok(self);
                } else if nr - skip <= fetch_val {
                    // After `skip` input rows are skipped, the remaining rows are
                    // less than or equal to the `fetch` values, so `num_rows` must
                    // equal the remaining rows.
                    check_num_rows(
                        (nr - skip).checked_mul(n_partitions),
                        // We know that we have an estimate for the number of rows:
                        self.num_rows.is_exact().unwrap(),
                    )
                } else {
                    // At this point we know that we were given a `fetch` value
                    // as the `None` case would go into the branch above. Since
                    // the input has more rows than `fetch + skip`, the number
                    // of rows will be the `fetch`, other statistics will have to be downgraded to inexact.
                    check_num_rows(
                        fetch_val.checked_mul(n_partitions),
                        // We know that we have an estimate for the number of rows:
                        self.num_rows.is_exact().unwrap(),
                    )
                }
            }
            Statistics {
                num_rows: Precision::Absent,
                ..
            } => check_num_rows(fetch.and_then(|v| v.checked_mul(n_partitions)), false),