Encode rows into a single contiguous `BytesMut` and append row-end offsets. # Invariants `offsets` must be non-empty and seeded with `0` at index 0. `offsets.last()` must equal `out.len()` on entry. On success, exactly `batch.num_rows()` additional offsets are pushed, and `offsets.last()` equals the new `out.len()`.
(
&self,
batch: &RecordBatch,
row_capacity: usize,
out: &mut BytesMut,
offsets: &mut Vec<usize>,
)
| 940 | /// * On success, exactly `batch.num_rows()` additional offsets are pushed, and |
| 941 | /// `offsets.last()` equals the new `out.len()`. |
| 942 | pub(crate) fn encode_rows( |
| 943 | &self, |
| 944 | batch: &RecordBatch, |
| 945 | row_capacity: usize, |
| 946 | out: &mut BytesMut, |
| 947 | offsets: &mut Vec<usize>, |
| 948 | ) -> Result<(), AvroError> { |
| 949 | let out_len = out.len(); |
| 950 | if offsets.first() != Some(&0) || offsets.last() != Some(&out_len) { |
| 951 | return Err(AvroError::General( |
| 952 | "encode_rows requires offsets to start with 0 and end at out.len()".to_string(), |
| 953 | )); |
| 954 | } |
| 955 | let n = batch.num_rows(); |
| 956 | if n == 0 { |
| 957 | return Ok(()); |
| 958 | } |
| 959 | if offsets.len().checked_add(n).is_none() { |
| 960 | return Err(AvroError::General( |
| 961 | "encode_rows cannot append offsets: too many rows".to_string(), |
| 962 | )); |
| 963 | } |
| 964 | let mut column_encoders = self.prepare_for_batch(batch)?; |
| 965 | offsets.reserve(n); |
| 966 | let prefix_bytes = self.prefix.as_ref().map(|p| p.as_slice()); |
| 967 | let prefix_len = prefix_bytes.map_or(0, |p| p.len()); |
| 968 | let per_row_hint = row_capacity.max(prefix_len); |
| 969 | if let Some(additional) = n |
| 970 | .checked_mul(per_row_hint) |
| 971 | .filter(|&a| out_len.checked_add(a).is_some()) |
| 972 | { |
| 973 | out.reserve(additional); |
| 974 | } |
| 975 | let start_out_len = out.len(); |
| 976 | let start_offsets_len = offsets.len(); |
| 977 | let res = (|| -> Result<(), AvroError> { |
| 978 | let mut w = out.writer(); |
| 979 | if let [enc0] = column_encoders.as_mut_slice() { |
| 980 | for_rows_with_prefix!(n, prefix_bytes, w, |row| { |
| 981 | enc0.encode(&mut w, row)?; |
| 982 | offsets.push(w.get_ref().len()); |
| 983 | }); |
| 984 | } else { |
| 985 | for_rows_with_prefix!(n, prefix_bytes, w, |row| { |
| 986 | for enc in column_encoders.iter_mut() { |
| 987 | enc.encode(&mut w, row)?; |
| 988 | } |
| 989 | offsets.push(w.get_ref().len()); |
| 990 | }); |
| 991 | } |
| 992 | Ok(()) |
| 993 | })(); |
| 994 | if res.is_err() { |
| 995 | out.truncate(start_out_len); |
| 996 | offsets.truncate(start_offsets_len); |
| 997 | } else { |
| 998 | debug_assert_eq!( |
| 999 | *offsets.last().unwrap(), |