| 3720 | } |
| 3721 | |
| 3722 | fn finish_incremental_inner( |
| 3723 | &mut self, |
| 3724 | rows: RowCollection, |
| 3725 | max_result_size: u64, |
| 3726 | ) -> Result<RowCollectionIter, String> { |
| 3727 | // How much additional memory is required to make a sorted view. |
| 3728 | let sorted_view_mem = rows.entries().saturating_mul(std::mem::size_of::<usize>()); |
| 3729 | let required_memory = rows.byte_len().saturating_add(sorted_view_mem); |
| 3730 | |
| 3731 | // Bail if creating the sorted view would require us to use too much memory. |
| 3732 | if required_memory > usize::cast_from(max_result_size) { |
| 3733 | let max_bytes = ByteSize::b(max_result_size); |
| 3734 | return Err(format!("total result exceeds max size of {max_bytes}",)); |
| 3735 | } |
| 3736 | |
| 3737 | let batch_num_rows = rows.count(); |
| 3738 | |
| 3739 | let sorted_view = rows; |
| 3740 | let mut iter = sorted_view |
| 3741 | .into_row_iter() |
| 3742 | .apply_offset(self.remaining_offset) |
| 3743 | .with_projection(self.project.clone()); |
| 3744 | |
| 3745 | if let Some(limit) = self.remaining_limit { |
| 3746 | iter = iter.with_limit(limit); |
| 3747 | }; |
| 3748 | |
| 3749 | self.remaining_offset = self.remaining_offset.saturating_sub(batch_num_rows); |
| 3750 | if let Some(remaining_limit) = self.remaining_limit.as_mut() { |
| 3751 | *remaining_limit -= iter.count(); |
| 3752 | } |
| 3753 | |
| 3754 | // TODO(parkmycar): Re-think how we can calculate the total response size without |
| 3755 | // having to iterate through the entire collection of Rows, while still |
| 3756 | // respecting the LIMIT, OFFSET, and projections. |
| 3757 | // |
| 3758 | // Note: It feels a bit bad always calculating the response size, but we almost |
| 3759 | // always need it to either check the `max_returned_query_size`, or for reporting |
| 3760 | // in the query history. |
| 3761 | let response_size: usize = iter.clone().map(|row| row.data().len()).sum(); |
| 3762 | |
| 3763 | // Bail if we would end up returning more data to the client than they can support. |
| 3764 | if let Some(max) = &mut self.remaining_max_returned_query_size { |
| 3765 | if let Some(remaining) = max.checked_sub(response_size.cast_into()) { |
| 3766 | *max = remaining; |
| 3767 | } else { |
| 3768 | let max_bytes = ByteSize::b(self.max_returned_query_size.expect("known to exist")); |
| 3769 | return Err(format!("total result exceeds max size of {max_bytes}")); |
| 3770 | } |
| 3771 | } |
| 3772 | |
| 3773 | Ok(iter) |
| 3774 | } |
| 3775 | } |
| 3776 | |
| 3777 | /// Compares two rows columnwise, using [compare_columns]. |