proof_of_sql/sql/postprocessing/
slice_postprocessing.rs

1use super::{PostprocessingError, PostprocessingResult, PostprocessingStep};
2use crate::base::{database::OwnedTable, scalar::Scalar};
3use serde::{Deserialize, Serialize};
4
5/// A `SlicePostprocessing` represents a slice of an `OwnedTable`.
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct SlicePostprocessing {
8    /// number of rows to return
9    ///
10    /// - if None, specify all rows
11    number_rows: Option<u64>,
12
13    /// number of rows to skip
14    ///
15    /// - if None, specify the first row as starting point
16    /// - if Some(nonnegative), specify the offset from the beginning
17    /// - if Some(negative), specify the offset from the end
18    ///   (e.g. -1 is the last row, -2 is the second to last row, etc.)
19    offset_value: Option<i64>,
20}
21
22impl SlicePostprocessing {
23    /// Create a new `SlicePostprocessing` with the given `number_rows` and `offset`.
24    #[must_use]
25    pub fn new(number_rows: Option<u64>, offset_value: Option<i64>) -> Self {
26        Self {
27            number_rows,
28            offset_value,
29        }
30    }
31}
32
33impl<S: Scalar> PostprocessingStep<S> for SlicePostprocessing {
34    /// Apply the slice transformation to the given `OwnedTable`.
35    fn apply(&self, owned_table: OwnedTable<S>) -> PostprocessingResult<OwnedTable<S>> {
36        let num_rows = owned_table.num_rows();
37        let limit = self.number_rows.unwrap_or(num_rows as u64);
38        let offset = self.offset_value.unwrap_or(0);
39        // Be permissive with data types at first so that computation can be done.
40        // If the conversion fails, we will return None.
41        let possible_starting_row = if offset < 0 {
42            num_rows as i128 + i128::from(offset)
43        } else {
44            i128::from(offset)
45        };
46        // The `possible_ending_row` is NOT inclusive.
47        let possible_ending_row = (possible_starting_row + i128::from(limit)).min(num_rows as i128);
48        let starting_row = usize::try_from(possible_starting_row).map_err(|_| {
49            PostprocessingError::InvalidSliceIndex {
50                index: possible_starting_row,
51            }
52        })?;
53        let ending_row = usize::try_from(possible_ending_row).map_err(|_| {
54            PostprocessingError::InvalidSliceIndex {
55                index: possible_ending_row,
56            }
57        })?;
58        Ok(OwnedTable::<S>::try_from_iter(
59            owned_table
60                .into_inner()
61                .into_iter()
62                .map(|(identifier, column)| (identifier, column.slice(starting_row, ending_row))),
63        )
64        .expect("Sliced columns of an existing table should have equal length"))
65    }
66}