Skip to main content

parquet/arrow/arrow_writer/
byte_array.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::basic::Encoding;
19use crate::bloom_filter::Sbbf;
20use crate::column::writer::encoder::{
21    ColumnValueEncoder, DataPageValues, DictionaryPage, create_bloom_filter,
22};
23use crate::data_type::{AsBytes, ByteArray, Int32Type};
24use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder};
25use crate::encodings::rle::RleEncoder;
26use crate::errors::{ParquetError, Result};
27use crate::file::properties::{
28    EnabledStatistics, ResolvedColumnProperties, WriterProperties, WriterVersion,
29};
30use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator};
31use crate::geospatial::statistics::GeospatialStatistics;
32use crate::schema::types::ColumnDescPtr;
33use crate::util::bit_util::num_required_bits;
34use crate::util::interner::{Interner, Storage};
35use crate::util::prefix::common_prefix_length;
36use arrow_array::types::ByteArrayType;
37use arrow_array::{
38    Array, ArrayAccessor, BinaryArray, BinaryViewArray, DictionaryArray, FixedSizeBinaryArray,
39    GenericByteArray, LargeBinaryArray, LargeStringArray, StringArray, StringViewArray,
40};
41use arrow_buffer::{ArrowNativeType, Buffer};
42use arrow_schema::DataType;
43
44macro_rules! downcast_dict_impl {
45    ($array:ident, $key:ident, $val:ident, $op:expr $(, $arg:expr)*) => {{
46        $op($array
47            .as_any()
48            .downcast_ref::<DictionaryArray<arrow_array::types::$key>>()
49            .unwrap()
50            .downcast_dict::<$val>()
51            .unwrap()$(, $arg)*)
52    }};
53}
54
55macro_rules! downcast_dict_op {
56    ($key_type:expr, $val:ident, $array:ident, $op:expr $(, $arg:expr)*) => {
57        match $key_type.as_ref() {
58            DataType::UInt8 => downcast_dict_impl!($array, UInt8Type, $val, $op$(, $arg)*),
59            DataType::UInt16 => downcast_dict_impl!($array, UInt16Type, $val, $op$(, $arg)*),
60            DataType::UInt32 => downcast_dict_impl!($array, UInt32Type, $val, $op$(, $arg)*),
61            DataType::UInt64 => downcast_dict_impl!($array, UInt64Type, $val, $op$(, $arg)*),
62            DataType::Int8 => downcast_dict_impl!($array, Int8Type, $val, $op$(, $arg)*),
63            DataType::Int16 => downcast_dict_impl!($array, Int16Type, $val, $op$(, $arg)*),
64            DataType::Int32 => downcast_dict_impl!($array, Int32Type, $val, $op$(, $arg)*),
65            DataType::Int64 => downcast_dict_impl!($array, Int64Type, $val, $op$(, $arg)*),
66            _ => unreachable!(),
67        }
68    };
69}
70
71macro_rules! downcast_op {
72    ($data_type:expr, $array:ident, $op:expr $(, $arg:expr)*) => {
73        match $data_type {
74            DataType::Utf8 => $op($array.as_any().downcast_ref::<StringArray>().unwrap()$(, $arg)*),
75            DataType::LargeUtf8 => {
76                $op($array.as_any().downcast_ref::<LargeStringArray>().unwrap()$(, $arg)*)
77            }
78            DataType::Utf8View => $op($array.as_any().downcast_ref::<StringViewArray>().unwrap()$(, $arg)*),
79            DataType::Binary => {
80                $op($array.as_any().downcast_ref::<BinaryArray>().unwrap()$(, $arg)*)
81            }
82            DataType::LargeBinary => {
83                $op($array.as_any().downcast_ref::<LargeBinaryArray>().unwrap()$(, $arg)*)
84            }
85            DataType::BinaryView => {
86                $op($array.as_any().downcast_ref::<BinaryViewArray>().unwrap()$(, $arg)*)
87            }
88            DataType::Dictionary(key, value) => match value.as_ref() {
89                DataType::Utf8 => downcast_dict_op!(key, StringArray, $array, $op$(, $arg)*),
90                DataType::LargeUtf8 => {
91                    downcast_dict_op!(key, LargeStringArray, $array, $op$(, $arg)*)
92                }
93                DataType::Utf8View => {
94                    downcast_dict_op!(key, StringViewArray, $array, $op$(, $arg)*)
95                }
96                DataType::Binary => downcast_dict_op!(key, BinaryArray, $array, $op$(, $arg)*),
97                DataType::LargeBinary => {
98                    downcast_dict_op!(key, LargeBinaryArray, $array, $op$(, $arg)*)
99                }
100                DataType::BinaryView => {
101                    downcast_dict_op!(key, BinaryViewArray, $array, $op$(, $arg)*)
102                }
103                DataType::FixedSizeBinary(_) => {
104                    downcast_dict_op!(key, FixedSizeBinaryArray, $array, $op$(, $arg)*)
105                }
106                d => unreachable!("cannot downcast {} dictionary value to byte array", d),
107            },
108            d => unreachable!("cannot downcast {} to byte array", d),
109        }
110    };
111}
112
113/// A fallback encoder, i.e. non-dictionary, for [`ByteArray`]
114struct FallbackEncoder {
115    encoder: FallbackEncoderImpl,
116    num_values: usize,
117    variable_length_bytes: i64,
118}
119
120/// The fallback encoder in use
121///
122/// Note: DeltaBitPackEncoder is boxed as it is rather large
123enum FallbackEncoderImpl {
124    Plain {
125        buffer: Vec<u8>,
126    },
127    DeltaLength {
128        buffer: Vec<u8>,
129        lengths: Box<DeltaBitPackEncoder<Int32Type>>,
130    },
131    Delta {
132        buffer: Vec<u8>,
133        last_value: Vec<u8>,
134        prefix_lengths: Box<DeltaBitPackEncoder<Int32Type>>,
135        suffix_lengths: Box<DeltaBitPackEncoder<Int32Type>>,
136    },
137}
138
139impl FallbackEncoder {
140    /// Create the fallback encoder for the given [`WriterProperties`] and the
141    /// column settings already resolved from them
142    fn new(props: &WriterProperties, column_props: &ResolvedColumnProperties) -> Result<Self> {
143        // Set either main encoder or fallback encoder.
144        let encoding = column_props
145            .encoding
146            .unwrap_or_else(|| match props.writer_version() {
147                WriterVersion::PARQUET_1_0 => Encoding::PLAIN,
148                WriterVersion::PARQUET_2_0 => Encoding::DELTA_BYTE_ARRAY,
149            });
150
151        let encoder = match encoding {
152            Encoding::PLAIN => FallbackEncoderImpl::Plain { buffer: vec![] },
153            Encoding::DELTA_LENGTH_BYTE_ARRAY => FallbackEncoderImpl::DeltaLength {
154                buffer: vec![],
155                lengths: Box::new(DeltaBitPackEncoder::new()),
156            },
157            Encoding::DELTA_BYTE_ARRAY => FallbackEncoderImpl::Delta {
158                buffer: vec![],
159                last_value: vec![],
160                prefix_lengths: Box::new(DeltaBitPackEncoder::new()),
161                suffix_lengths: Box::new(DeltaBitPackEncoder::new()),
162            },
163            _ => {
164                return Err(general_err!(
165                    "unsupported encoding {} for byte array",
166                    encoding
167                ));
168            }
169        };
170
171        Ok(Self {
172            encoder,
173            num_values: 0,
174            variable_length_bytes: 0,
175        })
176    }
177
178    /// Encode `values` to the in-progress page
179    fn encode<T>(&mut self, values: T, indices: impl ExactSizeIterator<Item = usize>)
180    where
181        T: ArrayAccessor + Copy,
182        T::Item: AsRef<[u8]>,
183    {
184        self.num_values += indices.len();
185        match &mut self.encoder {
186            FallbackEncoderImpl::Plain { buffer } => {
187                for idx in indices {
188                    let value = values.value(idx);
189                    let value = value.as_ref();
190                    buffer.extend_from_slice((value.len() as u32).as_bytes());
191                    buffer.extend_from_slice(value);
192                    self.variable_length_bytes += value.len() as i64;
193                }
194            }
195            FallbackEncoderImpl::DeltaLength { buffer, lengths } => {
196                for idx in indices {
197                    let value = values.value(idx);
198                    let value = value.as_ref();
199                    lengths.put(&[value.len() as i32]).unwrap();
200                    buffer.extend_from_slice(value);
201                    self.variable_length_bytes += value.len() as i64;
202                }
203            }
204            FallbackEncoderImpl::Delta {
205                buffer,
206                last_value,
207                prefix_lengths,
208                suffix_lengths,
209            } => {
210                for idx in indices {
211                    let value = values.value(idx);
212                    let value = value.as_ref();
213
214                    let prefix_length = common_prefix_length(last_value, value);
215                    let suffix_length = value.len() - prefix_length;
216
217                    last_value.clear();
218                    last_value.extend_from_slice(value);
219
220                    buffer.extend_from_slice(&value[prefix_length..]);
221                    prefix_lengths.put(&[prefix_length as i32]).unwrap();
222                    suffix_lengths.put(&[suffix_length as i32]).unwrap();
223                    self.variable_length_bytes += value.len() as i64;
224                }
225            }
226        }
227    }
228
229    /// Returns an estimate of the data page size in bytes
230    ///
231    /// This includes:
232    /// <already_written_encoded_byte_size> + <estimated_encoded_size_of_unflushed_bytes>
233    fn estimated_data_page_size(&self) -> usize {
234        match &self.encoder {
235            FallbackEncoderImpl::Plain { buffer, .. } => buffer.len(),
236            FallbackEncoderImpl::DeltaLength { buffer, lengths } => {
237                buffer.len() + lengths.estimated_data_encoded_size()
238            }
239            FallbackEncoderImpl::Delta {
240                buffer,
241                prefix_lengths,
242                suffix_lengths,
243                ..
244            } => {
245                buffer.len()
246                    + prefix_lengths.estimated_data_encoded_size()
247                    + suffix_lengths.estimated_data_encoded_size()
248            }
249        }
250    }
251
252    fn flush_data_page(
253        &mut self,
254        min_value: Option<ByteArray>,
255        max_value: Option<ByteArray>,
256    ) -> Result<DataPageValues<ByteArray>> {
257        let (buf, encoding) = match &mut self.encoder {
258            FallbackEncoderImpl::Plain { buffer } => (std::mem::take(buffer), Encoding::PLAIN),
259            FallbackEncoderImpl::DeltaLength { buffer, lengths } => {
260                let lengths = lengths.flush_buffer()?;
261
262                let mut out = Vec::with_capacity(lengths.len() + buffer.len());
263                out.extend_from_slice(&lengths);
264                out.extend_from_slice(buffer);
265                buffer.clear();
266                (out, Encoding::DELTA_LENGTH_BYTE_ARRAY)
267            }
268            FallbackEncoderImpl::Delta {
269                buffer,
270                prefix_lengths,
271                suffix_lengths,
272                last_value,
273            } => {
274                let prefix_lengths = prefix_lengths.flush_buffer()?;
275                let suffix_lengths = suffix_lengths.flush_buffer()?;
276
277                let mut out =
278                    Vec::with_capacity(prefix_lengths.len() + suffix_lengths.len() + buffer.len());
279                out.extend_from_slice(&prefix_lengths);
280                out.extend_from_slice(&suffix_lengths);
281                out.extend_from_slice(buffer);
282                buffer.clear();
283                last_value.clear();
284                (out, Encoding::DELTA_BYTE_ARRAY)
285            }
286        };
287
288        // Capture value of variable_length_bytes and reset for next page
289        let variable_length_bytes = Some(self.variable_length_bytes);
290        self.variable_length_bytes = 0;
291
292        Ok(DataPageValues {
293            buf: buf.into(),
294            num_values: std::mem::take(&mut self.num_values),
295            encoding,
296            min_value,
297            max_value,
298            nan_count: None,
299            variable_length_bytes,
300        })
301    }
302}
303
304/// [`Storage`] for the [`Interner`] used by [`DictEncoder`]
305#[derive(Debug, Default)]
306struct ByteArrayStorage {
307    /// Encoded dictionary data
308    page: Vec<u8>,
309
310    values: Vec<std::ops::Range<usize>>,
311}
312
313impl Storage for ByteArrayStorage {
314    type Key = u64;
315    type Value = [u8];
316
317    fn get(&self, idx: Self::Key) -> &Self::Value {
318        &self.page[self.values[idx as usize].clone()]
319    }
320
321    fn push(&mut self, value: &Self::Value) -> Self::Key {
322        let key = self.values.len();
323
324        self.page.reserve(4 + value.len());
325        self.page.extend_from_slice((value.len() as u32).as_bytes());
326
327        let start = self.page.len();
328        self.page.extend_from_slice(value);
329        self.values.push(start..self.page.len());
330
331        key as u64
332    }
333
334    fn estimated_memory_size(&self) -> usize {
335        self.page.capacity() * std::mem::size_of::<u8>()
336            + self.values.capacity() * std::mem::size_of::<std::ops::Range<usize>>()
337    }
338}
339
340/// A dictionary encoder for byte array data
341#[derive(Debug, Default)]
342struct DictEncoder {
343    interner: Interner<ByteArrayStorage>,
344    indices: Vec<u64>,
345    variable_length_bytes: i64,
346}
347
348impl DictEncoder {
349    /// Encode `values` to the in-progress page
350    fn encode<T>(&mut self, values: T, indices: impl ExactSizeIterator<Item = usize>)
351    where
352        T: ArrayAccessor + Copy,
353        T::Item: AsRef<[u8]>,
354    {
355        self.indices.reserve(indices.len());
356
357        for idx in indices {
358            let value = values.value(idx);
359            let interned = self.interner.intern(value.as_ref());
360            self.indices.push(interned);
361            self.variable_length_bytes += value.as_ref().len() as i64;
362        }
363    }
364
365    fn bit_width(&self) -> u8 {
366        let length = self.interner.storage().values.len();
367        num_required_bits(length.saturating_sub(1) as u64)
368    }
369
370    fn estimated_memory_size(&self) -> usize {
371        self.interner.estimated_memory_size() + self.indices.capacity() * std::mem::size_of::<u64>()
372    }
373
374    fn estimated_data_page_size(&self) -> usize {
375        let bit_width = self.bit_width();
376        1 + RleEncoder::max_buffer_size(bit_width, self.indices.len())
377    }
378
379    fn estimated_dict_page_size(&self) -> usize {
380        self.interner.storage().page.len()
381    }
382
383    fn flush_dict_page(self) -> DictionaryPage {
384        let storage = self.interner.into_inner();
385
386        DictionaryPage {
387            buf: storage.page.into(),
388            num_values: storage.values.len(),
389            is_sorted: false,
390        }
391    }
392
393    fn flush_data_page(
394        &mut self,
395        min_value: Option<ByteArray>,
396        max_value: Option<ByteArray>,
397    ) -> DataPageValues<ByteArray> {
398        let num_values = self.indices.len();
399        let buffer_len = self.estimated_data_page_size();
400        let mut buffer = Vec::with_capacity(buffer_len);
401        buffer.push(self.bit_width());
402
403        let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer);
404        encoder.put_batch(&self.indices);
405
406        self.indices.clear();
407
408        // Capture value of variable_length_bytes and reset for next page
409        let variable_length_bytes = Some(self.variable_length_bytes);
410        self.variable_length_bytes = 0;
411
412        DataPageValues {
413            buf: encoder.consume().into(),
414            num_values,
415            encoding: Encoding::RLE_DICTIONARY,
416            min_value,
417            max_value,
418            nan_count: None,
419            variable_length_bytes,
420        }
421    }
422}
423
424pub struct ByteArrayEncoder {
425    fallback: FallbackEncoder,
426    dict_encoder: Option<DictEncoder>,
427    statistics_enabled: EnabledStatistics,
428    min_value: Option<ByteArray>,
429    max_value: Option<ByteArray>,
430    bloom_filter: Option<Sbbf>,
431    bloom_filter_target_fpp: f64,
432    geo_stats_accumulator: Option<Box<dyn GeoStatsAccumulator>>,
433}
434
435impl ColumnValueEncoder for ByteArrayEncoder {
436    type T = ByteArray;
437    type Values = dyn Array;
438    fn flush_bloom_filter(&mut self) -> Option<Sbbf> {
439        let mut sbbf = self.bloom_filter.take()?;
440        sbbf.fold_to_target_fpp(self.bloom_filter_target_fpp);
441        Some(sbbf)
442    }
443
444    fn try_new(
445        descr: &ColumnDescPtr,
446        props: &WriterProperties,
447        column_props: &ResolvedColumnProperties,
448    ) -> Result<Self>
449    where
450        Self: Sized,
451    {
452        let dictionary = column_props.dictionary_enabled.then(DictEncoder::default);
453
454        let fallback = FallbackEncoder::new(props, column_props)?;
455
456        let (bloom_filter, bloom_filter_target_fpp) = create_bloom_filter(column_props)?;
457
458        let statistics_enabled = column_props.statistics_enabled;
459
460        let geo_stats_accumulator = try_new_geo_stats_accumulator(descr);
461
462        Ok(Self {
463            fallback,
464            statistics_enabled,
465            bloom_filter,
466            bloom_filter_target_fpp,
467            dict_encoder: dictionary,
468            min_value: None,
469            max_value: None,
470            geo_stats_accumulator,
471        })
472    }
473
474    fn write(&mut self, _values: &Self::Values, _offset: usize, _len: usize) -> Result<()> {
475        unreachable!("should call write_gather instead")
476    }
477
478    fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()> {
479        downcast_op!(
480            values.data_type(),
481            values,
482            encode,
483            indices.iter().copied(),
484            self
485        );
486        Ok(())
487    }
488
489    fn count_values_within_byte_budget_gather(
490        values: &Self::Values,
491        indices: &[usize],
492        byte_budget: usize,
493    ) -> Option<usize> {
494        // `ByteArrayEncoder` only ever writes via `write_gather`, so this
495        // is the relevant method.
496        //
497        // Two-stage walk for the simple offset-buffer byte array types:
498        //   1. If indices are contiguous, compute the total payload in
499        //      O(1) via a single subtraction on the offsets buffer.
500        //      When the total fits the budget — the overwhelmingly
501        //      common "small values" case — return immediately.
502        //   2. Otherwise, walk per-value byte sizes from the offsets
503        //      buffer (still cheap, no slice/UTF-8 construction) and
504        //      exit at the first value that pushes the cumulative sum
505        //      past the budget. This bounds skewed distributions: an
506        //      outlier value is caught wherever it lands in the chunk.
507        let count = match values.data_type() {
508            DataType::Utf8 => count_within_budget_offsets(
509                values.as_any().downcast_ref::<StringArray>().unwrap(),
510                indices,
511                byte_budget,
512            ),
513            DataType::LargeUtf8 => count_within_budget_offsets(
514                values.as_any().downcast_ref::<LargeStringArray>().unwrap(),
515                indices,
516                byte_budget,
517            ),
518            DataType::Binary => count_within_budget_offsets(
519                values.as_any().downcast_ref::<BinaryArray>().unwrap(),
520                indices,
521                byte_budget,
522            ),
523            DataType::LargeBinary => count_within_budget_offsets(
524                values.as_any().downcast_ref::<LargeBinaryArray>().unwrap(),
525                indices,
526                byte_budget,
527            ),
528            // View arrays carry each value's length in the low 32 bits of
529            // its u128 view word, so lengths are scannable without touching
530            // any data buffer — and the common small-value case skips even
531            // that scan via an O(1) conservative bound.
532            DataType::Utf8View => {
533                let array = values.as_any().downcast_ref::<StringViewArray>().unwrap();
534                count_within_budget_views(
535                    array.views(),
536                    indices,
537                    byte_budget,
538                    max_view_value_len(array.data_buffers()),
539                )
540            }
541            DataType::BinaryView => {
542                let array = values.as_any().downcast_ref::<BinaryViewArray>().unwrap();
543                count_within_budget_views(
544                    array.views(),
545                    indices,
546                    byte_budget,
547                    max_view_value_len(array.data_buffers()),
548                )
549            }
550            // The values in an arrow dictionary are already small and
551            // deduplicated, so there is nothing to bound — treat every
552            // chunk as fitting and stay on the batched path. (A per-value
553            // walk through dict keys on every chunk also measured ~+30-80%
554            // slower than `main`.)
555            DataType::Dictionary(_, _) => indices.len(),
556            // Every byte-array type `ByteArrayEncoder` is constructed for
557            // has an explicit arm above. A `Dictionary(value = FixedSizeBinary)`
558            // column hits the `Dictionary(_, _)` arm (its `values.data_type()`
559            // is `Dictionary`), and a bare `FixedSizeBinary` column is routed
560            // to the generic column writer, never this encoder — so no other
561            // type can reach here.
562            data_type => unreachable!("ByteArrayEncoder cannot be constructed for {data_type:?}"),
563        };
564        Some(count)
565    }
566
567    fn num_values(&self) -> usize {
568        match &self.dict_encoder {
569            Some(encoder) => encoder.indices.len(),
570            None => self.fallback.num_values,
571        }
572    }
573
574    fn has_dictionary(&self) -> bool {
575        self.dict_encoder.is_some()
576    }
577
578    fn compresses_against_previous_value(&self) -> bool {
579        // While dictionary encoding is active the data page holds RLE
580        // indices, which carry no cross-value state; only the DELTA_BYTE_ARRAY
581        // fallback shares prefixes with the preceding value.
582        self.dict_encoder.is_none()
583            && matches!(self.fallback.encoder, FallbackEncoderImpl::Delta { .. })
584    }
585
586    fn estimated_memory_size(&self) -> usize {
587        let encoder_size = match &self.dict_encoder {
588            Some(encoder) => encoder.estimated_memory_size(),
589            // For the FallbackEncoder, these unflushed bytes are already encoded.
590            // Therefore, the size should be the same as estimated_data_page_size.
591            None => self.fallback.estimated_data_page_size(),
592        };
593
594        let bloom_filter_size = self
595            .bloom_filter
596            .as_ref()
597            .map(|bf| bf.estimated_memory_size())
598            .unwrap_or_default();
599
600        let stats_size = self.min_value.as_ref().map(|v| v.len()).unwrap_or_default()
601            + self.max_value.as_ref().map(|v| v.len()).unwrap_or_default();
602
603        encoder_size + bloom_filter_size + stats_size
604    }
605
606    fn estimated_dict_page_size(&self) -> Option<usize> {
607        Some(self.dict_encoder.as_ref()?.estimated_dict_page_size())
608    }
609
610    /// Returns an estimate of the data page size in bytes
611    ///
612    /// This includes:
613    /// <already_written_encoded_byte_size> + <estimated_encoded_size_of_unflushed_bytes>
614    fn estimated_data_page_size(&self) -> usize {
615        match &self.dict_encoder {
616            Some(encoder) => encoder.estimated_data_page_size(),
617            None => self.fallback.estimated_data_page_size(),
618        }
619    }
620
621    fn flush_dict_page(&mut self) -> Result<Option<DictionaryPage>> {
622        match self.dict_encoder.take() {
623            Some(encoder) => {
624                if !encoder.indices.is_empty() {
625                    return Err(general_err!(
626                        "Must flush data pages before flushing dictionary"
627                    ));
628                }
629
630                if let Some(bloom_filter) = &mut self.bloom_filter {
631                    let storage = encoder.interner.storage();
632                    for range in &storage.values {
633                        bloom_filter.insert(&storage.page[range.clone()]);
634                    }
635                }
636
637                Ok(Some(encoder.flush_dict_page()))
638            }
639            _ => Ok(None),
640        }
641    }
642
643    fn flush_data_page(&mut self) -> Result<DataPageValues<ByteArray>> {
644        let min_value = self.min_value.take();
645        let max_value = self.max_value.take();
646
647        match &mut self.dict_encoder {
648            Some(encoder) => Ok(encoder.flush_data_page(min_value, max_value)),
649            _ => self.fallback.flush_data_page(min_value, max_value),
650        }
651    }
652
653    fn flush_geospatial_statistics(&mut self) -> Option<Box<GeospatialStatistics>> {
654        self.geo_stats_accumulator.as_mut().map(|a| a.finish())?
655    }
656}
657
658/// Encodes the provided `values` and `indices` to `encoder`
659///
660/// This is a free function so it can be used with `downcast_op!`
661fn encode<T, I>(values: T, indices: I, encoder: &mut ByteArrayEncoder)
662where
663    T: ArrayAccessor + Copy,
664    T::Item: Copy + Ord + AsRef<[u8]>,
665    I: ExactSizeIterator<Item = usize> + Clone,
666{
667    if encoder.statistics_enabled != EnabledStatistics::None {
668        if let Some(accumulator) = encoder.geo_stats_accumulator.as_mut() {
669            update_geo_stats_accumulator(accumulator.as_mut(), values, indices.clone());
670        } else if let Some((min, max)) = compute_min_max(values, indices.clone()) {
671            // Compare before copying: `write_gather` runs once per
672            // mini-batch, and a byte-budgeted mini-batch of large values can
673            // hold a single value, so an unconditional copy here would
674            // duplicate every value once for `min` and once for `max`.
675            let min = min.as_ref();
676            if encoder.min_value.as_ref().is_none_or(|m| m.data() > min) {
677                encoder.min_value = Some(min.to_vec().into());
678            }
679
680            let max = max.as_ref();
681            if encoder.max_value.as_ref().is_none_or(|m| m.data() < max) {
682                encoder.max_value = Some(max.to_vec().into());
683            }
684        }
685    }
686
687    // While a dictionary is in use the filter is populated from its distinct values in
688    // `flush_dict_page`, so each value is hashed once rather than once per row.
689    match &mut encoder.dict_encoder {
690        Some(dict_encoder) => dict_encoder.encode(values, indices),
691        None => {
692            if let Some(bloom_filter) = &mut encoder.bloom_filter {
693                for idx in indices.clone() {
694                    bloom_filter.insert(values.value(idx).as_ref());
695                }
696            }
697            encoder.fallback.encode(values, indices)
698        }
699    }
700}
701
702/// Upper bound on any single value's byte length in a view array.
703fn max_view_value_len(buffers: &[Buffer]) -> usize {
704    /// Bytes that fit inline in a u128 view word (the rest is len + prefix).
705    const MAX_INLINE_VIEW_LEN: usize = 12;
706    // An out-of-line view's data is a contiguous slice of exactly one data
707    // buffer, so it cannot exceed the largest buffer; inline views hold at
708    // most `MAX_INLINE_VIEW_LEN`. Loose (a value is usually far smaller than
709    // a whole buffer) but O(number of buffers) and always sound.
710    buffers
711        .iter()
712        .map(|b| b.len())
713        .max()
714        .unwrap_or(0)
715        .max(MAX_INLINE_VIEW_LEN)
716}
717
718/// Number of leading `indices` whose cumulative plain-encoded size fits
719/// `byte_budget` (boundary value included), for view arrays (`Utf8View`,
720/// `BinaryView`).
721fn count_within_budget_views(
722    views: &[u128],
723    indices: &[usize],
724    byte_budget: usize,
725    max_value_len: usize,
726) -> usize {
727    // Each plain-encoded BYTE_ARRAY value carries a 4-byte length prefix, so
728    // the budget is compared against `value_len + size_of::<u32>()` — the
729    // bytes actually written to the page, not just the payload.
730    //
731    // Stage 1: O(1) conservative bound. View arrays have no prefix-sum
732    // offsets buffer, so the exact span subtraction used by
733    // `count_within_budget_offsets` is unavailable; instead bound every
734    // value by `max_value_len`. Skips the walk for the common small-value
735    // case (what view arrays are built for, and where there is nothing to
736    // bound).
737    let per_value = max_value_len + std::mem::size_of::<u32>();
738    if indices.len().saturating_mul(per_value) <= byte_budget {
739        return indices.len();
740    }
741    // Stage 2: exact per-value scan, reading each length from the low 32
742    // bits of its u128 view word (no data-buffer dereference).
743    let mut cum: usize = 0;
744    for (i, idx) in indices.iter().enumerate() {
745        let len = (views[*idx] as u32) as usize;
746        cum = cum.saturating_add(len + std::mem::size_of::<u32>());
747        if cum > byte_budget {
748            return i + 1;
749        }
750    }
751    indices.len()
752}
753
754/// Number of leading `indices` whose cumulative plain-encoded size fits
755/// `byte_budget` (boundary value included), for offset-buffer byte arrays
756/// (`Utf8`/`LargeUtf8`/`Binary`/`LargeBinary`).
757///
758/// `indices` are assumed sorted ascending — they always are here, since
759/// they come from `non_null_indices`, which is built in array order.
760fn count_within_budget_offsets<T: ByteArrayType>(
761    values: &GenericByteArray<T>,
762    indices: &[usize],
763    byte_budget: usize,
764) -> usize {
765    if indices.is_empty() {
766        return 0;
767    }
768    let n = indices.len();
769    let first = indices[0];
770    let last = indices[n - 1];
771    let offsets = values.value_offsets();
772    // Each plain-encoded value carries a 4-byte length prefix on the page.
773    let prefix_overhead = std::mem::size_of::<u32>();
774
775    // Stage 1: O(1) span upper bound. The span `offsets[last+1] -
776    // offsets[first]` covers every array position in `[first, last]`, a
777    // superset of `indices` — and the skipped positions in a nullable
778    // column are nulls with zero offset delta, so the span still equals the
779    // exact payload. If it fits the budget, every value fits. Covers the
780    // common small-value case for both non-null and (sparse) nullable
781    // columns.
782    if last >= first {
783        let payload = (offsets[last + 1] - offsets[first]).as_usize();
784        if payload + n * prefix_overhead <= byte_budget {
785            return n;
786        }
787    }
788
789    // Stage 2: scan per-index lengths from the offsets buffer.
790    let mut cum: usize = 0;
791    for (i, idx) in indices.iter().enumerate() {
792        let len = (offsets[idx + 1] - offsets[*idx]).as_usize() + prefix_overhead;
793        cum = cum.saturating_add(len);
794        if cum > byte_budget {
795            return i + 1;
796        }
797    }
798    n
799}
800
801/// Computes the min and max for the provided array and indices
802///
803/// This is a free function so it can be used with `downcast_op!`
804fn compute_min_max<T>(
805    array: T,
806    mut valid: impl Iterator<Item = usize>,
807) -> Option<(T::Item, T::Item)>
808where
809    T: ArrayAccessor,
810    T::Item: Copy + Ord + AsRef<[u8]>,
811{
812    let first_idx = valid.next()?;
813
814    let first_val = array.value(first_idx);
815    let mut min = first_val;
816    let mut max = first_val;
817    for idx in valid {
818        let val = array.value(idx);
819        min = min.min(val);
820        max = max.max(val);
821    }
822    Some((min, max))
823}
824
825/// Updates geospatial statistics for the provided array and indices
826fn update_geo_stats_accumulator<T>(
827    bounder: &mut dyn GeoStatsAccumulator,
828    array: T,
829    valid: impl Iterator<Item = usize>,
830) where
831    T: ArrayAccessor,
832    T::Item: Copy + Ord + AsRef<[u8]>,
833{
834    if bounder.is_valid() {
835        for idx in valid {
836            let val = array.value(idx);
837            bounder.update_wkb(val.as_ref());
838        }
839    }
840}