Skip to main content

parquet/arrow/arrow_writer/
mod.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
18//! Contains writer which writes arrow data into parquet data.
19
20use crate::column::chunker::ContentDefinedChunker;
21
22use bytes::Bytes;
23use std::io::Write;
24use std::slice::Iter;
25use std::sync::{Arc, Mutex};
26use std::vec::IntoIter;
27
28use arrow_array::cast::AsArray;
29use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter};
30use arrow_array::{PrimitiveArray, types::*};
31use arrow_schema::{
32    ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit,
33};
34
35use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_precision};
36
37use crate::arrow::ArrowSchemaConverter;
38use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder;
39use crate::basic::PageType;
40use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter};
41use crate::column::page_encryption::PageEncryptor;
42use crate::column::writer::encoder::ColumnValueEncoder;
43use crate::column::writer::{
44    ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer,
45};
46use crate::data_type::{ByteArray, FixedLenByteArray};
47#[cfg(feature = "encryption")]
48use crate::encryption::encrypt::FileEncryptor;
49use crate::errors::{ParquetError, Result};
50use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData};
51use crate::file::properties::{WriterProperties, WriterPropertiesPtr};
52use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
53use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
54use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor};
55use levels::{ArrayLevels, calculate_array_levels};
56
57mod byte_array;
58mod levels;
59
60#[doc(inline)]
61pub use crate::column::page_store::{
62    InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs,
63    PageStoreFactory,
64};
65
66/// Encodes [`RecordBatch`] to parquet
67///
68/// Writes Arrow `RecordBatch`es to a Parquet writer. Multiple [`RecordBatch`] will be encoded
69/// to the same row group, up to `max_row_group_size` rows. Any remaining rows will be
70/// flushed on close, leading the final row group in the output file to potentially
71/// contain fewer than `max_row_group_size` rows
72///
73/// # Example: Writing `RecordBatch`es
74/// ```
75/// # use std::sync::Arc;
76/// # use bytes::Bytes;
77/// # use arrow_array::{ArrayRef, Int64Array};
78/// # use arrow_array::RecordBatch;
79/// # use parquet::arrow::arrow_writer::ArrowWriter;
80/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
81/// let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
82/// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
83///
84/// let mut buffer = Vec::new();
85/// let mut writer = ArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap();
86/// writer.write(&to_write).unwrap();
87/// writer.close().unwrap();
88///
89/// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap();
90/// let read = reader.next().unwrap().unwrap();
91///
92/// assert_eq!(to_write, read);
93/// ```
94///
95/// # Memory Usage and Limiting
96///
97/// The nature of Parquet requires buffering of an entire row group before it can
98/// be flushed to the underlying writer. Data is mostly buffered in its encoded
99/// form, reducing memory usage. However, some data such as dictionary keys,
100/// large strings or very nested data may still result in non-trivial memory
101/// usage.
102///
103/// See Also:
104/// * [`ArrowWriter::memory_size`]: the current memory usage of the writer.
105/// * [`ArrowWriter::in_progress_size`]: Estimated size of the buffered row group,
106///
107/// Call [`Self::flush`] to trigger an early flush of a row group based on a
108/// memory threshold and/or global memory pressure. However,  smaller row groups
109/// result in higher metadata overheads, and thus may worsen compression ratios
110/// and query performance.
111///
112/// ```no_run
113/// # use std::io::Write;
114/// # use arrow_array::RecordBatch;
115/// # use parquet::arrow::ArrowWriter;
116/// # let mut writer: ArrowWriter<Vec<u8>> = todo!();
117/// # let batch: RecordBatch = todo!();
118/// writer.write(&batch).unwrap();
119/// // Trigger an early flush if anticipated size exceeds 1_000_000
120/// if writer.in_progress_size() > 1_000_000 {
121///     writer.flush().unwrap();
122/// }
123/// ```
124///
125/// ## Type Support
126///
127/// The writer supports writing all Arrow [`DataType`]s that have a direct mapping to
128/// Parquet types including  [`StructArray`] and [`ListArray`].
129///
130/// The following are not supported:
131///
132/// * [`IntervalMonthDayNanoArray`]: Parquet does not [support nanosecond intervals].
133///
134/// [`DataType`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html
135/// [`StructArray`]: https://docs.rs/arrow/latest/arrow/array/struct.StructArray.html
136/// [`ListArray`]: https://docs.rs/arrow/latest/arrow/array/type.ListArray.html
137/// [`IntervalMonthDayNanoArray`]: https://docs.rs/arrow/latest/arrow/array/type.IntervalMonthDayNanoArray.html
138/// [support nanosecond intervals]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#interval
139///
140/// ## Type Compatibility
141/// The writer can write Arrow [`RecordBatch`]s that are logically equivalent. This means that for
142/// a  given column, the writer can accept multiple Arrow [`DataType`]s that contain the same
143/// value type.
144///
145/// For example, the following [`DataType`]s are all logically equivalent and can be written
146/// to the same column:
147/// * String, LargeString, StringView
148/// * Binary, LargeBinary, BinaryView
149///
150/// The writer can will also accept both native and dictionary encoded arrays if the dictionaries
151/// contain compatible values.
152/// ```
153/// # use std::sync::Arc;
154/// # use arrow_array::{DictionaryArray, LargeStringArray, RecordBatch, StringArray, UInt8Array};
155/// # use arrow_schema::{DataType, Field, Schema};
156/// # use parquet::arrow::arrow_writer::ArrowWriter;
157/// let record_batch1 = RecordBatch::try_new(
158///    Arc::new(Schema::new(vec![Field::new("col", DataType::LargeUtf8, false)])),
159///    vec![Arc::new(LargeStringArray::from_iter_values(vec!["a", "b"]))]
160///  )
161/// .unwrap();
162///
163/// let mut buffer = Vec::new();
164/// let mut writer = ArrowWriter::try_new(&mut buffer, record_batch1.schema(), None).unwrap();
165/// writer.write(&record_batch1).unwrap();
166///
167/// let record_batch2 = RecordBatch::try_new(
168///     Arc::new(Schema::new(vec![Field::new(
169///         "col",
170///         DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
171///          false,
172///     )])),
173///     vec![Arc::new(DictionaryArray::new(
174///          UInt8Array::from_iter_values(vec![0, 1]),
175///          Arc::new(StringArray::from_iter_values(vec!["b", "c"])),
176///      ))],
177///  )
178///  .unwrap();
179///  writer.write(&record_batch2).unwrap();
180///  writer.close();
181/// ```
182pub struct ArrowWriter<W: Write> {
183    /// Underlying Parquet writer
184    writer: SerializedFileWriter<W>,
185
186    /// The in-progress row group if any
187    in_progress: Option<ArrowRowGroupWriter>,
188
189    /// A copy of the Arrow schema.
190    ///
191    /// The schema is used to verify that each record batch written has the correct schema
192    arrow_schema: SchemaRef,
193
194    /// Creates new [`ArrowRowGroupWriter`] instances as required
195    row_group_writer_factory: ArrowRowGroupWriterFactory,
196
197    /// The maximum number of rows to write to each row group, or None for unlimited
198    max_row_group_row_count: Option<usize>,
199
200    /// The maximum size in bytes for a row group, or None for unlimited
201    max_row_group_bytes: Option<usize>,
202
203    /// CDC chunkers persisted across row groups (one per leaf column).
204    cdc_chunkers: Option<Vec<ContentDefinedChunker>>,
205}
206
207impl<W: Write + Send> std::fmt::Debug for ArrowWriter<W> {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        let buffered_memory = self.in_progress_size();
210        f.debug_struct("ArrowWriter")
211            .field("writer", &self.writer)
212            .field("in_progress_size", &format_args!("{buffered_memory} bytes"))
213            .field("in_progress_rows", &self.in_progress_rows())
214            .field("arrow_schema", &self.arrow_schema)
215            .field("max_row_group_row_count", &self.max_row_group_row_count)
216            .field("max_row_group_bytes", &self.max_row_group_bytes)
217            .finish()
218    }
219}
220
221impl<W: Write + Send> ArrowWriter<W> {
222    /// Try to create a new Arrow writer
223    ///
224    /// The writer will fail if:
225    ///  * a `SerializedFileWriter` cannot be created from the ParquetWriter
226    ///  * the Arrow schema contains unsupported datatypes such as Unions
227    pub fn try_new(
228        writer: W,
229        arrow_schema: SchemaRef,
230        props: Option<WriterProperties>,
231    ) -> Result<Self> {
232        let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
233        Self::try_new_with_options(writer, arrow_schema, options)
234    }
235
236    /// Try to create a new Arrow writer with [`ArrowWriterOptions`].
237    ///
238    /// The writer will fail if:
239    ///  * a `SerializedFileWriter` cannot be created from the ParquetWriter
240    ///  * the Arrow schema contains unsupported datatypes such as Unions
241    pub fn try_new_with_options(
242        writer: W,
243        arrow_schema: SchemaRef,
244        options: ArrowWriterOptions,
245    ) -> Result<Self> {
246        let mut props = options.properties;
247
248        let schema = if let Some(parquet_schema) = options.schema_descr {
249            parquet_schema.clone()
250        } else {
251            let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types());
252            if let Some(schema_root) = &options.schema_root {
253                converter = converter.schema_root(schema_root);
254            }
255
256            converter.convert(&arrow_schema)?
257        };
258
259        if !options.skip_arrow_metadata {
260            // add serialized arrow schema
261            add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
262        }
263
264        let max_row_group_row_count = props.max_row_group_row_count();
265        let max_row_group_bytes = props.max_row_group_bytes();
266
267        let props_ptr = Arc::new(props);
268        let file_writer =
269            SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
270
271        let mut row_group_writer_factory =
272            ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone());
273        if let Some(page_store_factory) = options.page_store_factory {
274            row_group_writer_factory =
275                row_group_writer_factory.with_page_store_factory(page_store_factory);
276        }
277
278        let cdc_chunkers = props_ptr
279            .content_defined_chunking()
280            .map(|opts| {
281                file_writer
282                    .schema_descr()
283                    .columns()
284                    .iter()
285                    .map(|desc| ContentDefinedChunker::new(desc, opts))
286                    .collect::<Result<Vec<_>>>()
287            })
288            .transpose()?;
289
290        Ok(Self {
291            writer: file_writer,
292            in_progress: None,
293            arrow_schema,
294            row_group_writer_factory,
295            max_row_group_row_count,
296            max_row_group_bytes,
297            cdc_chunkers,
298        })
299    }
300
301    /// Returns metadata for any flushed row groups
302    pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
303        self.writer.flushed_row_groups()
304    }
305
306    /// Estimated memory usage, in bytes, of this `ArrowWriter`
307    ///
308    /// This estimate is formed bu summing the values of
309    /// [`ArrowColumnWriter::memory_size`] all in progress columns.
310    pub fn memory_size(&self) -> usize {
311        match &self.in_progress {
312            Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(),
313            None => 0,
314        }
315    }
316
317    /// Anticipated encoded size of the in progress row group.
318    ///
319    /// This estimate the row group size after being completely encoded is,
320    /// formed by summing the values of
321    /// [`ArrowColumnWriter::get_estimated_total_bytes`] for all in progress
322    /// columns.
323    pub fn in_progress_size(&self) -> usize {
324        match &self.in_progress {
325            Some(in_progress) => in_progress
326                .writers
327                .iter()
328                .map(|x| x.get_estimated_total_bytes())
329                .sum(),
330            None => 0,
331        }
332    }
333
334    /// Returns the number of rows buffered in the in progress row group
335    pub fn in_progress_rows(&self) -> usize {
336        self.in_progress
337            .as_ref()
338            .map(|x| x.buffered_rows)
339            .unwrap_or_default()
340    }
341
342    /// Returns the number of bytes written by this instance
343    pub fn bytes_written(&self) -> usize {
344        self.writer.bytes_written()
345    }
346
347    /// Encodes the provided [`RecordBatch`]
348    ///
349    /// If this would cause the current row group to exceed [`WriterProperties::max_row_group_row_count`]
350    /// rows or [`WriterProperties::max_row_group_bytes`] bytes, the contents of `batch` will be
351    /// written to one or more row groups such that limits are respected.
352    ///
353    /// If both limits are `None`, all data is written to a single row group.
354    /// If one limit is set, that limit is respected.
355    /// If both limits are set, the lower bound (whichever triggers first) is respected.
356    ///
357    /// This will fail if the `batch`'s schema does not match the writer's schema.
358    pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
359        if batch.num_rows() == 0 {
360            return Ok(());
361        }
362
363        let in_progress = match &mut self.in_progress {
364            Some(in_progress) => in_progress,
365            x => x.insert(
366                self.row_group_writer_factory
367                    .create_row_group_writer(self.writer.flushed_row_groups().len())?,
368            ),
369        };
370
371        if let Some(max_rows) = self.max_row_group_row_count
372            && in_progress.buffered_rows + batch.num_rows() > max_rows
373        {
374            let to_write = max_rows - in_progress.buffered_rows;
375            let a = batch.slice(0, to_write);
376            let b = batch.slice(to_write, batch.num_rows() - to_write);
377            self.write(&a)?;
378            return self.write(&b);
379        }
380
381        // Check byte limit: if we have buffered data, use measured average row size
382        // to split batch proactively before exceeding byte limit
383        if let Some(max_bytes) = self.max_row_group_bytes
384            && in_progress.buffered_rows > 0
385        {
386            let current_bytes = in_progress.get_estimated_total_bytes();
387
388            if current_bytes >= max_bytes {
389                self.flush()?;
390                return self.write(batch);
391            }
392
393            if let Some(avg_row_bytes) = current_bytes
394                .checked_div(in_progress.buffered_rows)
395                .filter(|avg_row_bytes| *avg_row_bytes > 0)
396            {
397                // At this point, `current_bytes < max_bytes` (checked above)
398                let remaining_bytes = max_bytes - current_bytes;
399                let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
400
401                if batch.num_rows() > rows_that_fit {
402                    if rows_that_fit > 0 {
403                        let a = batch.slice(0, rows_that_fit);
404                        let b = batch.slice(rows_that_fit, batch.num_rows() - rows_that_fit);
405                        self.write(&a)?;
406                        return self.write(&b);
407                    } else {
408                        self.flush()?;
409                        return self.write(batch);
410                    }
411                }
412            }
413        }
414
415        match self.cdc_chunkers.as_mut() {
416            Some(chunkers) => in_progress.write_with_chunkers(batch, chunkers)?,
417            None => in_progress.write(batch)?,
418        }
419
420        let should_flush = self
421            .max_row_group_row_count
422            .is_some_and(|max| in_progress.buffered_rows >= max)
423            || self
424                .max_row_group_bytes
425                .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
426
427        if should_flush {
428            self.flush()?
429        }
430        Ok(())
431    }
432
433    /// Writes the given buf bytes to the internal buffer.
434    ///
435    /// It's safe to use this method to write data to the underlying writer,
436    /// because it will ensure that the buffering and byte‐counting layers are used.
437    pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
438        self.writer.write_all(buf)
439    }
440
441    /// Flushes underlying writer
442    pub fn sync(&mut self) -> std::io::Result<()> {
443        self.writer.flush()
444    }
445
446    /// Flushes all buffered rows into a new row group
447    ///
448    /// Note the underlying writer is not flushed with this call.
449    /// If this is a desired behavior, please call [`ArrowWriter::sync`].
450    pub fn flush(&mut self) -> Result<()> {
451        let in_progress = match self.in_progress.take() {
452            Some(in_progress) => in_progress,
453            None => return Ok(()),
454        };
455
456        let mut row_group_writer = self.writer.next_row_group()?;
457        for chunk in in_progress.close()? {
458            chunk.append_to_row_group(&mut row_group_writer)?;
459        }
460        row_group_writer.close()?;
461        Ok(())
462    }
463
464    /// Additional [`KeyValue`] metadata to be written in addition to those from [`WriterProperties`]
465    ///
466    /// This method provide a way to append kv_metadata after write RecordBatch
467    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
468        self.writer.append_key_value_metadata(kv_metadata)
469    }
470
471    /// Returns a reference to the underlying writer.
472    pub fn inner(&self) -> &W {
473        self.writer.inner()
474    }
475
476    /// Returns a mutable reference to the underlying writer.
477    ///
478    /// **Warning**: if you write directly to this writer, you will skip
479    /// the `TrackedWrite` buffering and byte‐counting layers. That’ll cause
480    /// the file footer’s recorded offsets and sizes to diverge from reality,
481    /// resulting in an unreadable or corrupted Parquet file.
482    ///
483    /// If you want to write safely to the underlying writer, use [`Self::write_all`].
484    pub fn inner_mut(&mut self) -> &mut W {
485        self.writer.inner_mut()
486    }
487
488    /// Flushes any outstanding data and returns the underlying writer.
489    pub fn into_inner(mut self) -> Result<W> {
490        self.flush()?;
491        self.writer.into_inner()
492    }
493
494    /// Close and finalize the underlying Parquet writer
495    ///
496    /// Unlike [`Self::close`] this does not consume self
497    ///
498    /// Attempting to write after calling finish will result in an error
499    pub fn finish(&mut self) -> Result<ParquetMetaData> {
500        self.flush()?;
501        self.writer.finish()
502    }
503
504    /// Close and finalize the underlying Parquet writer
505    pub fn close(mut self) -> Result<ParquetMetaData> {
506        self.finish()
507    }
508
509    /// Create a new row group writer and return its column writers.
510    #[deprecated(
511        since = "56.2.0",
512        note = "Use `ArrowRowGroupWriterFactory` instead, see `ArrowColumnWriter` for an example"
513    )]
514    pub fn get_column_writers(&mut self) -> Result<Vec<ArrowColumnWriter>> {
515        self.flush()?;
516        let in_progress = self
517            .row_group_writer_factory
518            .create_row_group_writer(self.writer.flushed_row_groups().len())?;
519        Ok(in_progress.writers)
520    }
521
522    /// Append the given column chunks to the file as a new row group.
523    #[deprecated(
524        since = "56.2.0",
525        note = "Use `SerializedFileWriter` directly instead, see `ArrowColumnWriter` for an example"
526    )]
527    pub fn append_row_group(&mut self, chunks: Vec<ArrowColumnChunk>) -> Result<()> {
528        let mut row_group_writer = self.writer.next_row_group()?;
529        for chunk in chunks {
530            chunk.append_to_row_group(&mut row_group_writer)?;
531        }
532        row_group_writer.close()?;
533        Ok(())
534    }
535
536    /// Converts this writer into a lower-level [`SerializedFileWriter`] and [`ArrowRowGroupWriterFactory`].
537    ///
538    /// Flushes any outstanding data before returning.
539    ///
540    /// This can be useful to provide more control over how files are written, for example
541    /// to write columns in parallel. See the example on [`ArrowColumnWriter`].
542    pub fn into_serialized_writer(
543        mut self,
544    ) -> Result<(SerializedFileWriter<W>, ArrowRowGroupWriterFactory)> {
545        self.flush()?;
546        Ok((self.writer, self.row_group_writer_factory))
547    }
548}
549
550impl<W: Write + Send> RecordBatchWriter for ArrowWriter<W> {
551    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
552        self.write(batch).map_err(|e| e.into())
553    }
554
555    fn close(self) -> std::result::Result<(), ArrowError> {
556        self.close()?;
557        Ok(())
558    }
559}
560
561/// Arrow-specific configuration settings for writing parquet files.
562///
563/// See [`ArrowWriter`] for how to configure the writer.
564#[derive(Debug, Clone, Default)]
565pub struct ArrowWriterOptions {
566    properties: WriterProperties,
567    skip_arrow_metadata: bool,
568    schema_root: Option<String>,
569    schema_descr: Option<SchemaDescriptor>,
570    page_store_factory: Option<Arc<dyn PageStoreFactory>>,
571}
572
573impl ArrowWriterOptions {
574    /// Creates a new [`ArrowWriterOptions`] with the default settings.
575    pub fn new() -> Self {
576        Self::default()
577    }
578
579    /// Sets the [`WriterProperties`] for writing parquet files.
580    pub fn with_properties(self, properties: WriterProperties) -> Self {
581        Self { properties, ..self }
582    }
583
584    /// Sets the [`PageStoreFactory`] used to buffer completed pages while a row
585    /// group is being written.
586    ///
587    /// The default implementation ([`InMemoryPageStore`]) buffers all completed
588    /// pages on the heap until the row group is flushed, so peak write memory
589    /// grows with the row group size. Using this API, pages can be spilled to a
590    /// file or object storage instead, reducing peak write memory substantially
591    /// at the expense of an extra write to and read from secondary storage.
592    ///
593    /// # Example: spilling pages to a temp file
594    ///
595    /// A simple spilling backend uses one temp file per column chunk; `put`
596    /// appends the page and `take` reads it back.
597    ///
598    /// ```
599    /// # use std::fs::File;
600    /// # use std::io::{Read, Seek, SeekFrom, Write};
601    /// # use std::sync::Arc;
602    /// # use bytes::Bytes;
603    /// # use arrow_array::{ArrayRef, Int64Array, RecordBatch};
604    /// # use parquet::arrow::arrow_writer::{
605    /// #     ArrowWriter, ArrowWriterOptions, PageKey, PageStore, PageStoreArgs, PageStoreFactory,
606    /// # };
607    /// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
608    /// # use parquet::errors::Result;
609    /// struct TempFilePageStore {
610    ///     file: File,
611    ///     /// Total size of the file
612    ///     end: u64,
613    ///     /// Location of pages: (offset, len)
614    ///     locs: Vec<(u64, usize)>,
615    /// }
616    ///
617    /// impl PageStore for TempFilePageStore {
618    ///     fn put(&mut self, value: Bytes) -> Result<PageKey> {
619    ///         // Append to the end of the file
620    ///         self.file.seek(SeekFrom::Start(self.end))?;
621    ///         self.file.write_all(&value)?;
622    ///         let key = PageKey::new(self.locs.len() as u64);
623    ///         self.locs.push((self.end, value.len()));
624    ///         self.end += value.len() as u64;
625    ///         Ok(key)
626    ///     }
627    ///
628    ///     fn take(&mut self, key: PageKey) -> Result<Bytes> {
629    ///         let (offset, len) = self.locs[key.get() as usize];
630    ///         let mut buf = vec![0u8; len];
631    ///         self.file.seek(SeekFrom::Start(offset))?;
632    ///         self.file.read_exact(&mut buf)?;
633    ///         Ok(Bytes::from(buf))
634    ///     }
635    /// }
636    ///
637    /// /// Factory for creating [`TempFilePageStore`]
638    /// #[derive(Debug)]
639    /// struct TempFilePageStoreFactory;
640    ///
641    /// impl PageStoreFactory for TempFilePageStoreFactory {
642    ///     fn create(&self, args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
643    ///         // `args` exposes the column index and descriptor (physical/logical
644    ///         // type, path), so a real backend might choose to spill only large columns.
645    ///         let _ = (args.column_index(), args.column_descriptor());
646    ///         Ok(Box::new(TempFilePageStore {
647    ///             file: tempfile::tempfile()?, // temp file is cleaned on drop
648    ///             end: 0,
649    ///             locs: Vec::new(),
650    ///         }))
651    ///     }
652    /// }
653    /// // write 1000 integers
654    /// let col = Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef;
655    /// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
656    ///
657    /// let options =
658    ///     ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory));
659    /// let mut buffer = Vec::new();
660    /// let mut writer =
661    ///     ArrowWriter::try_new_with_options(&mut buffer, to_write.schema(), options).unwrap();
662    /// writer.write(&to_write).unwrap();
663    /// writer.close().unwrap();
664    ///
665    /// // buffer now holds valid Parquet data, which can be read as normal:
666    /// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap();
667    /// assert_eq!(to_write, reader.next().unwrap().unwrap());
668    /// ```
669    pub fn with_page_store_factory(self, page_store_factory: Arc<dyn PageStoreFactory>) -> Self {
670        Self {
671            page_store_factory: Some(page_store_factory),
672            ..self
673        }
674    }
675
676    /// Skip encoding the embedded arrow metadata (defaults to `false`)
677    ///
678    /// Parquet files generated by the [`ArrowWriter`] contain embedded arrow schema
679    /// by default.
680    ///
681    /// Set `skip_arrow_metadata` to true, to skip encoding the embedded metadata.
682    pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
683        Self {
684            skip_arrow_metadata,
685            ..self
686        }
687    }
688
689    /// Set the name of the root parquet schema element (defaults to `"arrow_schema"`)
690    pub fn with_schema_root(self, schema_root: String) -> Self {
691        Self {
692            schema_root: Some(schema_root),
693            ..self
694        }
695    }
696
697    /// Explicitly specify the Parquet schema to be used
698    ///
699    /// If omitted (the default), the [`ArrowSchemaConverter`] is used to compute the
700    /// Parquet [`SchemaDescriptor`]. This may be used When the [`SchemaDescriptor`] is
701    /// already known or must be calculated using custom logic.
702    pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self {
703        Self {
704            schema_descr: Some(schema_descr),
705            ..self
706        }
707    }
708}
709
710/// A single column chunk produced by [`ArrowColumnWriter`].
711///
712/// Holds the serialized page blobs (each page's header ‖ compressed data, in
713/// write order) in a [`PageStore`], plus the handles needed to read them back,
714/// in order, when the chunk is spliced into the output file.
715struct ArrowColumnChunkData {
716    length: usize,
717    store: Box<dyn PageStore>,
718    keys: Vec<PageKey>,
719    /// Handles to the dictionary page's blobs (header then data) in the store.
720    ///
721    /// A dictionary page is produced at most once and bounded by
722    /// `dict_page_size_limit`, but it must be written *first* in the chunk even
723    /// though the data pages reach the writer before it (see
724    /// [`PageWriter::defers_dictionary_ordering`]). Its header and data are `put`
725    /// into the store like any other page — which keeps the store uniform, and
726    /// lets an oversized dictionary page spill — and their handles are held apart
727    /// so they can be emitted ahead of the data pages at splice.
728    /// Empty for non-dictionary columns.
729    dictionary_keys: Vec<PageKey>,
730    /// Serialized length of the dictionary page (0 if there is none), recorded
731    /// so the data pages can be shifted past it when offsets are rewritten to a
732    /// dictionary-first layout at splice.
733    dictionary_len: usize,
734}
735
736impl ArrowColumnChunkData {
737    fn new(store: Box<dyn PageStore>) -> Self {
738        Self {
739            length: 0,
740            store,
741            keys: Vec::new(),
742            dictionary_keys: Vec::new(),
743            dictionary_len: 0,
744        }
745    }
746
747    /// Append a data-page blob to the store, recording its handle in write
748    /// order.
749    fn push(&mut self, value: Bytes) -> Result<()> {
750        let key = self.store.put(value)?;
751        self.keys.push(key);
752        Ok(())
753    }
754
755    /// Store a dictionary-page blob (header or data) in the page store,
756    /// recording its handle (emitted first at splice) and accumulating its
757    /// serialized length.
758    fn push_dictionary(&mut self, value: Bytes) -> Result<()> {
759        self.dictionary_len += value.len();
760        let key = self.store.put(value)?;
761        self.dictionary_keys.push(key);
762        Ok(())
763    }
764
765    /// Bytes this chunk currently holds on the heap: whatever the store keeps
766    /// resident (zero for a spilling backend).
767    fn memory_size(&self) -> usize {
768        self.store.memory_size()
769    }
770}
771
772/// A streaming iterator over one column chunk's buffered page blobs, in final
773/// file order: the dictionary page (if any) first, then the data pages.
774///
775/// Each blob is taken back out of the [`PageStore`] *as it is
776/// consumed* and released immediately afterwards, so splicing a chunk into the
777/// output file never materializes more than a single page in memory at a time.
778/// This is what keeps the splice phase within the memory bound for a spilling
779/// backend (an in-memory store already holds the bytes, so it is unaffected).
780struct StreamingColumnChunkPages {
781    store: Box<dyn PageStore>,
782    /// Page handles in final file order: the dictionary page first (if any),
783    /// then the data pages.
784    keys: IntoIter<PageKey>,
785}
786
787impl StreamingColumnChunkPages {
788    fn new(data: ArrowColumnChunkData) -> Self {
789        // The dictionary page must be emitted first, ahead of the data pages,
790        // even though it was the last page produced.
791        let keys = if data.dictionary_keys.is_empty() {
792            data.keys
793        } else {
794            let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len());
795            keys.extend(data.dictionary_keys);
796            keys.extend(data.keys);
797            keys
798        };
799        Self {
800            store: data.store,
801            keys: keys.into_iter(),
802        }
803    }
804}
805
806impl Iterator for StreamingColumnChunkPages {
807    type Item = Result<Bytes>;
808
809    fn next(&mut self) -> Option<Self::Item> {
810        let key = self.keys.next()?;
811        Some(self.store.take(key))
812    }
813}
814
815/// A shared [`ArrowColumnChunkData`]
816///
817/// This allows it to be owned by [`ArrowPageWriter`] whilst allowing access via
818/// [`ArrowRowGroupWriter`] on flush, without requiring self-referential borrows
819type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
820
821struct ArrowPageWriter {
822    buffer: SharedColumnChunk,
823    #[cfg(feature = "encryption")]
824    page_encryptor: Option<PageEncryptor>,
825}
826
827impl ArrowPageWriter {
828    /// Create a page writer that buffers completed pages in `store`.
829    fn new(store: Box<dyn PageStore>) -> Self {
830        Self {
831            buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))),
832            #[cfg(feature = "encryption")]
833            page_encryptor: None,
834        }
835    }
836
837    #[cfg(feature = "encryption")]
838    pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
839        self.page_encryptor = page_encryptor;
840        self
841    }
842
843    #[cfg(feature = "encryption")]
844    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
845        self.page_encryptor.as_mut()
846    }
847
848    #[cfg(not(feature = "encryption"))]
849    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
850        None
851    }
852}
853
854impl PageWriter for ArrowPageWriter {
855    fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
856        let page = match self.page_encryptor_mut() {
857            Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
858            None => page,
859        };
860
861        let page_header = page.to_thrift_header()?;
862        let header = {
863            let mut header = Vec::with_capacity(1024);
864
865            match self.page_encryptor_mut() {
866                Some(page_encryptor) => {
867                    page_encryptor.encrypt_page_header(&page_header, &mut header)?;
868                    if page.compressed_page().is_data_page() {
869                        page_encryptor.increment_page();
870                    }
871                }
872                None => {
873                    let mut protocol = ThriftCompactOutputProtocol::new(&mut header);
874                    page_header.write_thrift(&mut protocol)?;
875                }
876            };
877
878            Bytes::from(header)
879        };
880
881        let mut buf = self.buffer.try_lock().unwrap();
882
883        let data = page.compressed_page().buffer().clone();
884        let compressed_size = data.len() + header.len();
885
886        let mut spec = PageWriteSpec::new();
887        spec.page_type = page.page_type();
888        spec.num_values = page.num_values();
889        spec.uncompressed_size = page.uncompressed_size() + header.len();
890        spec.offset = buf.length as u64;
891        spec.compressed_size = compressed_size;
892        spec.bytes_written = compressed_size as u64;
893
894        buf.length += compressed_size;
895        if spec.page_type == PageType::DICTIONARY_PAGE {
896            // Recorded apart from the data pages so it is emitted first at
897            // splice — see `ArrowColumnChunkData::dictionary_keys`.
898            buf.push_dictionary(header)?;
899            buf.push_dictionary(data)?;
900        } else {
901            buf.push(header)?;
902            buf.push(data)?;
903        }
904
905        Ok(spec)
906    }
907
908    fn defers_dictionary_ordering(&self) -> bool {
909        // The Arrow chunk is buffered in full and spliced at row-group flush, so
910        // data pages may be accepted before the dictionary page and reordered
911        // then. This lets `GenericColumnWriter` stream dictionary-column data
912        // pages straight through instead of buffering them in memory.
913        true
914    }
915
916    fn buffered_memory_size(&self) -> usize {
917        // Only what is actually resident: a spilling store reports ~0 here even
918        // though the chunk's bytes have all passed through it.
919        self.buffer.try_lock().unwrap().memory_size()
920    }
921
922    fn close(&mut self) -> Result<()> {
923        Ok(())
924    }
925}
926
927/// A leaf column that can be encoded by [`ArrowColumnWriter`]
928#[derive(Debug)]
929pub struct ArrowLeafColumn(ArrayLevels);
930
931/// Computes the [`ArrowLeafColumn`] for a potentially nested [`ArrayRef`]
932///
933/// This function can be used along with [`get_column_writers`] to encode
934/// individual columns in parallel. See example on [`ArrowColumnWriter`]
935pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result<Vec<ArrowLeafColumn>> {
936    let levels = calculate_array_levels(array, field)?;
937    Ok(levels.into_iter().map(ArrowLeafColumn).collect())
938}
939
940/// The data for a single column chunk, see [`ArrowColumnWriter`]
941pub struct ArrowColumnChunk {
942    data: ArrowColumnChunkData,
943    close: ColumnCloseResult,
944}
945
946impl std::fmt::Debug for ArrowColumnChunk {
947    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
948        f.debug_struct("ArrowColumnChunk")
949            .field("length", &self.data.length)
950            .finish_non_exhaustive()
951    }
952}
953
954impl ArrowColumnChunk {
955    /// Returns the [`ColumnCloseResult`] produced when the chunk was closed.
956    ///
957    /// Exposes encoding information, collected statistics, and the optional
958    /// [`ColumnIndexMetaData`](crate::file::page_index::column_index::ColumnIndexMetaData)
959    /// / [`OffsetIndexMetaData`](crate::file::page_index::offset_index::OffsetIndexMetaData)
960    /// gathered for the column chunk.
961    pub fn close(&self) -> &ColumnCloseResult {
962        &self.close
963    }
964
965    /// Returns a mutable reference to the [`ColumnCloseResult`].
966    ///
967    /// This allows callers to mutate the close result before the chunk is
968    /// appended to a row group — for example, clearing `column_index` or
969    /// `bloom_filter` based on a dynamic rule that inspects the encodings and
970    /// collected page statistics.
971    pub fn close_mut(&mut self) -> &mut ColumnCloseResult {
972        &mut self.close
973    }
974
975    /// Splices this column's buffered pages into the row group, streaming them
976    /// back out of the [`PageStore`] one page at a time.
977    pub fn append_to_row_group<W: Write + Send>(
978        self,
979        writer: &mut SerializedRowGroupWriter<'_, W>,
980    ) -> Result<()> {
981        let ArrowColumnChunk { data, close } = self;
982
983        // The dictionary page is produced *after* the data pages on this path (so
984        // they can stream straight through) but must be written *first*, so move
985        // it ahead of the data pages in the recorded offsets before the splice.
986        let close = close.update_dictionary_location(data.dictionary_len)?;
987
988        let pages = StreamingColumnChunkPages::new(data);
989        writer.append_column_from_pages(pages, close)
990    }
991}
992
993/// Encodes [`ArrowLeafColumn`] to [`ArrowColumnChunk`]
994///
995/// `ArrowColumnWriter` instances can be created using an [`ArrowRowGroupWriterFactory`];
996///
997/// Note: This is a low-level interface for applications that require
998/// fine-grained control of encoding (e.g. encoding using multiple threads),
999/// see [`ArrowWriter`] for a higher-level interface
1000///
1001/// # Example: Encoding two Arrow Array's in Parallel
1002/// ```
1003/// // The arrow schema
1004/// # use std::sync::Arc;
1005/// # use arrow_array::*;
1006/// # use arrow_schema::*;
1007/// # use parquet::arrow::ArrowSchemaConverter;
1008/// # use parquet::arrow::arrow_writer::{compute_leaves, ArrowColumnChunk, ArrowLeafColumn, ArrowRowGroupWriterFactory};
1009/// # use parquet::file::properties::WriterProperties;
1010/// # use parquet::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
1011/// #
1012/// let schema = Arc::new(Schema::new(vec![
1013///     Field::new("i32", DataType::Int32, false),
1014///     Field::new("f32", DataType::Float32, false),
1015/// ]));
1016///
1017/// // Compute the parquet schema
1018/// let props = Arc::new(WriterProperties::default());
1019/// let parquet_schema = ArrowSchemaConverter::new()
1020///   .with_coerce_types(props.coerce_types())
1021///   .convert(&schema)
1022///   .unwrap();
1023///
1024/// // Create parquet writer
1025/// let root_schema = parquet_schema.root_schema_ptr();
1026/// // write to memory in the example, but this could be a File
1027/// let mut out = Vec::with_capacity(1024);
1028/// let mut writer = SerializedFileWriter::new(&mut out, root_schema, props.clone())
1029///   .unwrap();
1030///
1031/// // Create a factory for building Arrow column writers
1032/// let row_group_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
1033/// // Create column writers for the 0th row group
1034/// let col_writers = row_group_factory.create_column_writers(0).unwrap();
1035///
1036/// // Spawn a worker thread for each column
1037/// //
1038/// // Note: This is for demonstration purposes, a thread-pool e.g. rayon or tokio, would be better.
1039/// // The `map` produces an iterator of type `tuple of (thread handle, send channel)`.
1040/// let mut workers: Vec<_> = col_writers
1041///     .into_iter()
1042///     .map(|mut col_writer| {
1043///         let (send, recv) = std::sync::mpsc::channel::<ArrowLeafColumn>();
1044///         let handle = std::thread::spawn(move || {
1045///             // receive Arrays to encode via the channel
1046///             for col in recv {
1047///                 col_writer.write(&col)?;
1048///             }
1049///             // once the input is complete, close the writer
1050///             // to return the newly created ArrowColumnChunk
1051///             col_writer.close()
1052///         });
1053///         (handle, send)
1054///     })
1055///     .collect();
1056///
1057/// // Start row group
1058/// let mut row_group_writer: SerializedRowGroupWriter<'_, _> = writer
1059///   .next_row_group()
1060///   .unwrap();
1061///
1062/// // Create some example input columns to encode
1063/// let to_write = vec![
1064///     Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _,
1065///     Arc::new(Float32Array::from_iter_values([1., 45., -1.])) as _,
1066/// ];
1067///
1068/// // Send the input columns to the workers
1069/// let mut worker_iter = workers.iter_mut();
1070/// for (arr, field) in to_write.iter().zip(&schema.fields) {
1071///     for leaves in compute_leaves(field, arr).unwrap() {
1072///         worker_iter.next().unwrap().1.send(leaves).unwrap();
1073///     }
1074/// }
1075///
1076/// // Wait for the workers to complete encoding, and append
1077/// // the resulting column chunks to the row group (and the file)
1078/// for (handle, send) in workers {
1079///     drop(send); // Drop send side to signal termination
1080///     // wait for the worker to send the completed chunk
1081///     let chunk: ArrowColumnChunk = handle.join().unwrap().unwrap();
1082///     chunk.append_to_row_group(&mut row_group_writer).unwrap();
1083/// }
1084/// // Close the row group which writes to the underlying file
1085/// row_group_writer.close().unwrap();
1086///
1087/// let metadata = writer.close().unwrap();
1088/// assert_eq!(metadata.file_metadata().num_rows(), 3);
1089/// ```
1090pub struct ArrowColumnWriter {
1091    writer: ArrowColumnWriterImpl,
1092    chunk: SharedColumnChunk,
1093}
1094
1095impl std::fmt::Debug for ArrowColumnWriter {
1096    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1097        f.debug_struct("ArrowColumnWriter").finish_non_exhaustive()
1098    }
1099}
1100
1101enum ArrowColumnWriterImpl {
1102    ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
1103    Column(ColumnWriter<'static>),
1104}
1105
1106impl ArrowColumnWriter {
1107    /// Write an [`ArrowLeafColumn`]
1108    pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> {
1109        self.write_internal(&col.0)
1110    }
1111
1112    /// Write with content-defined chunking, inserting page flushes at chunk boundaries.
1113    fn write_with_chunker(
1114        &mut self,
1115        col: &ArrowLeafColumn,
1116        chunker: &mut ContentDefinedChunker,
1117    ) -> Result<()> {
1118        let levels = &col.0;
1119        let chunks = chunker.get_arrow_chunks(
1120            levels.def_level_data().as_ref(),
1121            levels.rep_level_data().as_ref(),
1122            levels.array(),
1123        )?;
1124
1125        let num_chunks = chunks.len();
1126        for (i, chunk) in chunks.iter().enumerate() {
1127            let chunk_levels = levels.slice_for_chunk(chunk);
1128            self.write_internal(&chunk_levels)?;
1129
1130            // Add a page break after each chunk except the last
1131            if i + 1 < num_chunks {
1132                match &mut self.writer {
1133                    ArrowColumnWriterImpl::Column(c) => c.add_data_page()?,
1134                    ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?,
1135                }
1136            }
1137        }
1138        Ok(())
1139    }
1140
1141    fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
1142        match &mut self.writer {
1143            ArrowColumnWriterImpl::Column(c) => {
1144                let leaf = levels.array();
1145                match leaf.as_any_dictionary_opt() {
1146                    Some(dictionary) => {
1147                        let materialized =
1148                            arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?;
1149                        write_leaf(c, &materialized, levels)?
1150                    }
1151                    None => write_leaf(c, leaf, levels)?,
1152                };
1153            }
1154            ArrowColumnWriterImpl::ByteArray(c) => {
1155                write_primitive(c, levels.array().as_ref(), levels)?;
1156            }
1157        }
1158        Ok(())
1159    }
1160
1161    /// Close this column returning the written [`ArrowColumnChunk`]
1162    pub fn close(self) -> Result<ArrowColumnChunk> {
1163        let close = match self.writer {
1164            ArrowColumnWriterImpl::ByteArray(c) => c.close()?,
1165            ArrowColumnWriterImpl::Column(c) => c.close()?,
1166        };
1167        let chunk = Arc::try_unwrap(self.chunk).ok().unwrap();
1168        let data = chunk.into_inner().unwrap();
1169        Ok(ArrowColumnChunk { data, close })
1170    }
1171
1172    /// Returns the estimated total memory usage by the writer.
1173    ///
1174    /// This  [`Self::get_estimated_total_bytes`] this is an estimate
1175    /// of the current memory usage and not it's anticipated encoded size.
1176    ///
1177    /// This includes:
1178    /// 1. Data buffered in encoded form
1179    /// 2. Data buffered in un-encoded form (e.g. `usize` dictionary keys)
1180    ///
1181    /// This value should be greater than or equal to [`Self::get_estimated_total_bytes`]
1182    pub fn memory_size(&self) -> usize {
1183        match &self.writer {
1184            ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1185            ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1186        }
1187    }
1188
1189    /// Returns the estimated total encoded bytes for this column writer.
1190    ///
1191    /// This includes:
1192    /// 1. Data buffered in encoded form
1193    /// 2. An estimate of how large the data buffered in un-encoded form would be once encoded
1194    ///
1195    /// This value should be less than or equal to [`Self::memory_size`]
1196    pub fn get_estimated_total_bytes(&self) -> usize {
1197        match &self.writer {
1198            ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1199            ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1200        }
1201    }
1202}
1203
1204/// Encodes [`RecordBatch`] to a parquet row group
1205///
1206/// Note: this structure is created by [`ArrowRowGroupWriterFactory`] internally used to
1207/// create [`ArrowRowGroupWriter`]s, but it is not exposed publicly.
1208///
1209/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1210#[derive(Debug)]
1211struct ArrowRowGroupWriter {
1212    writers: Vec<ArrowColumnWriter>,
1213    schema: SchemaRef,
1214    buffered_rows: usize,
1215}
1216
1217impl ArrowRowGroupWriter {
1218    fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self {
1219        Self {
1220            writers,
1221            schema: arrow.clone(),
1222            buffered_rows: 0,
1223        }
1224    }
1225
1226    fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1227        self.buffered_rows += batch.num_rows();
1228        let mut writers = self.writers.iter_mut();
1229        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1230            for leaf in compute_leaves(field.as_ref(), column)? {
1231                writers.next().unwrap().write(&leaf)?;
1232            }
1233        }
1234        Ok(())
1235    }
1236
1237    fn write_with_chunkers(
1238        &mut self,
1239        batch: &RecordBatch,
1240        chunkers: &mut [ContentDefinedChunker],
1241    ) -> Result<()> {
1242        self.buffered_rows += batch.num_rows();
1243        let mut writers = self.writers.iter_mut();
1244        let mut chunkers = chunkers.iter_mut();
1245        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1246            for leaf in compute_leaves(field.as_ref(), column)? {
1247                writers
1248                    .next()
1249                    .unwrap()
1250                    .write_with_chunker(&leaf, chunkers.next().unwrap())?;
1251            }
1252        }
1253        Ok(())
1254    }
1255
1256    /// Returns the estimated total encoded bytes for this row group
1257    fn get_estimated_total_bytes(&self) -> usize {
1258        self.writers
1259            .iter()
1260            .map(|x| x.get_estimated_total_bytes())
1261            .sum()
1262    }
1263
1264    fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1265        self.writers
1266            .into_iter()
1267            .map(|writer| writer.close())
1268            .collect()
1269    }
1270}
1271
1272/// Factory that creates new column writers for each row group in the Parquet file.
1273///
1274/// You can create this structure via an [`ArrowWriter::into_serialized_writer`].
1275/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1276#[derive(Debug)]
1277pub struct ArrowRowGroupWriterFactory {
1278    schema: SchemaDescPtr,
1279    arrow_schema: SchemaRef,
1280    props: WriterPropertiesPtr,
1281    page_store_factory: Arc<dyn PageStoreFactory>,
1282    #[cfg(feature = "encryption")]
1283    file_encryptor: Option<Arc<FileEncryptor>>,
1284}
1285
1286impl ArrowRowGroupWriterFactory {
1287    /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema
1288    pub fn new<W: Write + Send>(
1289        file_writer: &SerializedFileWriter<W>,
1290        arrow_schema: SchemaRef,
1291    ) -> Self {
1292        let schema = Arc::clone(file_writer.schema_descr_ptr());
1293        let props = Arc::clone(file_writer.properties());
1294        Self {
1295            schema,
1296            arrow_schema,
1297            props,
1298            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1299            #[cfg(feature = "encryption")]
1300            file_encryptor: file_writer.file_encryptor(),
1301        }
1302    }
1303
1304    /// Set the [`PageStoreFactory`] used to allocate the buffer for each column
1305    /// chunk, e.g. to spill completed pages to a temp file or object storage
1306    /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`].
1307    pub fn with_page_store_factory(
1308        mut self,
1309        page_store_factory: Arc<dyn PageStoreFactory>,
1310    ) -> Self {
1311        self.page_store_factory = page_store_factory;
1312        self
1313    }
1314
1315    fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1316        let writers = self.create_column_writers(row_group_index)?;
1317        Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
1318    }
1319
1320    /// Create column writers for a new row group, with the given row group index
1321    pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1322        let mut writers = Vec::with_capacity(self.arrow_schema.fields.len());
1323        let mut leaves = self.schema.columns().iter();
1324        let column_factory = self.column_writer_factory(row_group_index);
1325        for field in &self.arrow_schema.fields {
1326            column_factory.get_arrow_column_writer(
1327                field.data_type(),
1328                &self.props,
1329                &mut leaves,
1330                &mut writers,
1331            )?;
1332        }
1333        Ok(writers)
1334    }
1335
1336    #[cfg(feature = "encryption")]
1337    fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1338        ArrowColumnWriterFactory::new()
1339            .with_page_store_factory(self.page_store_factory.clone())
1340            .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1341    }
1342
1343    #[cfg(not(feature = "encryption"))]
1344    fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1345        ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1346    }
1347}
1348
1349/// Returns [`ArrowColumnWriter`]s for each column in a given schema
1350#[deprecated(since = "57.0.0", note = "Use `ArrowRowGroupWriterFactory` instead")]
1351pub fn get_column_writers(
1352    parquet: &SchemaDescriptor,
1353    props: &WriterPropertiesPtr,
1354    arrow: &SchemaRef,
1355) -> Result<Vec<ArrowColumnWriter>> {
1356    let mut writers = Vec::with_capacity(arrow.fields.len());
1357    let mut leaves = parquet.columns().iter();
1358    let column_factory = ArrowColumnWriterFactory::new();
1359    for field in &arrow.fields {
1360        column_factory.get_arrow_column_writer(
1361            field.data_type(),
1362            props,
1363            &mut leaves,
1364            &mut writers,
1365        )?;
1366    }
1367    Ok(writers)
1368}
1369
1370/// Creates [`ArrowColumnWriter`] instances
1371struct ArrowColumnWriterFactory {
1372    /// Allocates the per-column-chunk [`PageStore`] backing each page writer.
1373    page_store_factory: Arc<dyn PageStoreFactory>,
1374    #[cfg(feature = "encryption")]
1375    row_group_index: usize,
1376    #[cfg(feature = "encryption")]
1377    file_encryptor: Option<Arc<FileEncryptor>>,
1378}
1379
1380impl ArrowColumnWriterFactory {
1381    pub fn new() -> Self {
1382        Self {
1383            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1384            #[cfg(feature = "encryption")]
1385            row_group_index: 0,
1386            #[cfg(feature = "encryption")]
1387            file_encryptor: None,
1388        }
1389    }
1390
1391    /// Use `page_store_factory` to allocate the buffer for each column chunk.
1392    pub fn with_page_store_factory(
1393        mut self,
1394        page_store_factory: Arc<dyn PageStoreFactory>,
1395    ) -> Self {
1396        self.page_store_factory = page_store_factory;
1397        self
1398    }
1399
1400    #[cfg(feature = "encryption")]
1401    pub fn with_file_encryptor(
1402        mut self,
1403        row_group_index: usize,
1404        file_encryptor: Option<Arc<FileEncryptor>>,
1405    ) -> Self {
1406        self.row_group_index = row_group_index;
1407        self.file_encryptor = file_encryptor;
1408        self
1409    }
1410
1411    #[cfg(feature = "encryption")]
1412    fn create_page_writer(
1413        &self,
1414        column_descriptor: &ColumnDescPtr,
1415        column_index: usize,
1416    ) -> Result<Box<ArrowPageWriter>> {
1417        let column_path = column_descriptor.path().string();
1418        let page_encryptor = PageEncryptor::create_if_column_encrypted(
1419            &self.file_encryptor,
1420            self.row_group_index,
1421            column_index,
1422            &column_path,
1423        )?;
1424        let args = PageStoreArgs::new(column_index, column_descriptor);
1425        let store = self.page_store_factory.create(&args)?;
1426        Ok(Box::new(
1427            ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1428        ))
1429    }
1430
1431    #[cfg(not(feature = "encryption"))]
1432    fn create_page_writer(
1433        &self,
1434        column_descriptor: &ColumnDescPtr,
1435        column_index: usize,
1436    ) -> Result<Box<ArrowPageWriter>> {
1437        let args = PageStoreArgs::new(column_index, column_descriptor);
1438        let store = self.page_store_factory.create(&args)?;
1439        Ok(Box::new(ArrowPageWriter::new(store)))
1440    }
1441
1442    /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the
1443    /// output ColumnDesc to `leaves` and the column writers to `out`
1444    fn get_arrow_column_writer(
1445        &self,
1446        data_type: &ArrowDataType,
1447        props: &WriterPropertiesPtr,
1448        leaves: &mut Iter<'_, ColumnDescPtr>,
1449        out: &mut Vec<ArrowColumnWriter>,
1450    ) -> Result<()> {
1451        // Instantiate writers for normal columns
1452        let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1453            let page_writer = self.create_page_writer(desc, out.len())?;
1454            let chunk = page_writer.buffer.clone();
1455            let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1456            Ok(ArrowColumnWriter {
1457                chunk,
1458                writer: ArrowColumnWriterImpl::Column(writer),
1459            })
1460        };
1461
1462        // Instantiate writers for byte arrays (e.g. Utf8,  Binary, etc)
1463        let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1464            let page_writer = self.create_page_writer(desc, out.len())?;
1465            let chunk = page_writer.buffer.clone();
1466            let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1467            Ok(ArrowColumnWriter {
1468                chunk,
1469                writer: ArrowColumnWriterImpl::ByteArray(writer),
1470            })
1471        };
1472
1473        match data_type {
1474            _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1475            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1476                out.push(col(leaves.next().unwrap())?)
1477            }
1478            ArrowDataType::LargeBinary
1479            | ArrowDataType::Binary
1480            | ArrowDataType::Utf8
1481            | ArrowDataType::LargeUtf8
1482            | ArrowDataType::BinaryView
1483            | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1484            ArrowDataType::List(f)
1485            | ArrowDataType::LargeList(f)
1486            | ArrowDataType::FixedSizeList(f, _)
1487            | ArrowDataType::ListView(f)
1488            | ArrowDataType::LargeListView(f) => {
1489                self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1490            }
1491            ArrowDataType::Struct(fields) => {
1492                for field in fields {
1493                    self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1494                }
1495            }
1496            ArrowDataType::Map(f, _) => match f.data_type() {
1497                ArrowDataType::Struct(f) => {
1498                    self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1499                    self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1500                }
1501                _ => unreachable!("invalid map type"),
1502            },
1503            ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1504                ArrowDataType::Utf8
1505                | ArrowDataType::LargeUtf8
1506                | ArrowDataType::Binary
1507                | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1508                ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1509                    out.push(bytes(leaves.next().unwrap())?)
1510                }
1511                ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
1512                _ => out.push(col(leaves.next().unwrap())?),
1513            },
1514            ArrowDataType::RunEndEncoded(_, value_field) => {
1515                self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1516            }
1517            _ => {
1518                return Err(ParquetError::NYI(format!(
1519                    "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1520                )));
1521            }
1522        }
1523        Ok(())
1524    }
1525}
1526
1527fn write_leaf(
1528    writer: &mut ColumnWriter<'_>,
1529    column: &dyn arrow_array::Array,
1530    levels: &ArrayLevels,
1531) -> Result<usize> {
1532    let indices = levels.non_null_indices();
1533
1534    match writer {
1535        // Note: this should match the contents of arrow_to_parquet_type
1536        ColumnWriter::Int32ColumnWriter(typed) => {
1537            match column.data_type() {
1538                ArrowDataType::Null => {
1539                    let array = Int32Array::new_null(column.len());
1540                    write_primitive(typed, array.values(), levels)
1541                }
1542                ArrowDataType::Int8 => {
1543                    let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32);
1544                    write_primitive(typed, array.values(), levels)
1545                }
1546                ArrowDataType::Int16 => {
1547                    let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32);
1548                    write_primitive(typed, array.values(), levels)
1549                }
1550                ArrowDataType::Int32 => {
1551                    write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels)
1552                }
1553                ArrowDataType::UInt8 => {
1554                    let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32);
1555                    write_primitive(typed, array.values(), levels)
1556                }
1557                ArrowDataType::UInt16 => {
1558                    let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32);
1559                    write_primitive(typed, array.values(), levels)
1560                }
1561                ArrowDataType::UInt32 => {
1562                    // follow C++ implementation and use overflow/reinterpret cast from  u32 to i32 which will map
1563                    // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0`
1564                    let array = column.as_primitive::<UInt32Type>();
1565                    write_primitive(typed, array.values().inner().typed_data(), levels)
1566                }
1567                ArrowDataType::Date32 => {
1568                    let array = column.as_primitive::<Date32Type>();
1569                    write_primitive(typed, array.values(), levels)
1570                }
1571                ArrowDataType::Time32(TimeUnit::Second) => {
1572                    let array = column.as_primitive::<Time32SecondType>();
1573                    write_primitive(typed, array.values(), levels)
1574                }
1575                ArrowDataType::Time32(TimeUnit::Millisecond) => {
1576                    let array = column.as_primitive::<Time32MillisecondType>();
1577                    write_primitive(typed, array.values(), levels)
1578                }
1579                ArrowDataType::Date64 => {
1580                    // If the column is a Date64, we truncate it
1581                    let array: Int32Array = column
1582                        .as_primitive::<Date64Type>()
1583                        .unary(|x| (x / 86_400_000) as _);
1584
1585                    write_primitive(typed, array.values(), levels)
1586                }
1587                ArrowDataType::Decimal32(_, _) => {
1588                    let array = column
1589                        .as_primitive::<Decimal32Type>()
1590                        .unary::<_, Int32Type>(|v| v);
1591                    write_primitive(typed, array.values(), levels)
1592                }
1593                ArrowDataType::Decimal64(_, _) => {
1594                    // use the int32 to represent the decimal with low precision
1595                    let array = column
1596                        .as_primitive::<Decimal64Type>()
1597                        .unary::<_, Int32Type>(|v| v as i32);
1598                    write_primitive(typed, array.values(), levels)
1599                }
1600                ArrowDataType::Decimal128(_, _) => {
1601                    // use the int32 to represent the decimal with low precision
1602                    let array = column
1603                        .as_primitive::<Decimal128Type>()
1604                        .unary::<_, Int32Type>(|v| v as i32);
1605                    write_primitive(typed, array.values(), levels)
1606                }
1607                ArrowDataType::Decimal256(_, _) => {
1608                    // use the int32 to represent the decimal with low precision
1609                    let array = column
1610                        .as_primitive::<Decimal256Type>()
1611                        .unary::<_, Int32Type>(|v| v.as_i128() as i32);
1612                    write_primitive(typed, array.values(), levels)
1613                }
1614                d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1615            }
1616        }
1617        ColumnWriter::BoolColumnWriter(typed) => {
1618            let array = column.as_boolean();
1619            let values = get_bool_array_slice(array, indices.iter().copied());
1620            typed.write_batch_internal(
1621                values.as_slice(),
1622                None,
1623                levels.def_level_data().as_ref(),
1624                levels.rep_level_data().as_ref(),
1625                None,
1626                None,
1627                None,
1628            )
1629        }
1630        ColumnWriter::Int64ColumnWriter(typed) => {
1631            match column.data_type() {
1632                ArrowDataType::Date64 => {
1633                    let array = column
1634                        .as_primitive::<Date64Type>()
1635                        .reinterpret_cast::<Int64Type>();
1636
1637                    write_primitive(typed, array.values(), levels)
1638                }
1639                ArrowDataType::Int64 => {
1640                    let array = column.as_primitive::<Int64Type>();
1641                    write_primitive(typed, array.values(), levels)
1642                }
1643                ArrowDataType::UInt64 => {
1644                    let values = column.as_primitive::<UInt64Type>().values();
1645                    // follow C++ implementation and use overflow/reinterpret cast from  u64 to i64 which will map
1646                    // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0`
1647                    let array = values.inner().typed_data::<i64>();
1648                    write_primitive(typed, array, levels)
1649                }
1650                ArrowDataType::Time64(TimeUnit::Microsecond) => {
1651                    let array = column.as_primitive::<Time64MicrosecondType>();
1652                    write_primitive(typed, array.values(), levels)
1653                }
1654                ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1655                    let array = column.as_primitive::<Time64NanosecondType>();
1656                    write_primitive(typed, array.values(), levels)
1657                }
1658                ArrowDataType::Timestamp(unit, _) => match unit {
1659                    TimeUnit::Second => {
1660                        let array = column.as_primitive::<TimestampSecondType>();
1661                        write_primitive(typed, array.values(), levels)
1662                    }
1663                    TimeUnit::Millisecond => {
1664                        let array = column.as_primitive::<TimestampMillisecondType>();
1665                        write_primitive(typed, array.values(), levels)
1666                    }
1667                    TimeUnit::Microsecond => {
1668                        let array = column.as_primitive::<TimestampMicrosecondType>();
1669                        write_primitive(typed, array.values(), levels)
1670                    }
1671                    TimeUnit::Nanosecond => {
1672                        let array = column.as_primitive::<TimestampNanosecondType>();
1673                        write_primitive(typed, array.values(), levels)
1674                    }
1675                },
1676                ArrowDataType::Duration(unit) => match unit {
1677                    TimeUnit::Second => {
1678                        let array = column.as_primitive::<DurationSecondType>();
1679                        write_primitive(typed, array.values(), levels)
1680                    }
1681                    TimeUnit::Millisecond => {
1682                        let array = column.as_primitive::<DurationMillisecondType>();
1683                        write_primitive(typed, array.values(), levels)
1684                    }
1685                    TimeUnit::Microsecond => {
1686                        let array = column.as_primitive::<DurationMicrosecondType>();
1687                        write_primitive(typed, array.values(), levels)
1688                    }
1689                    TimeUnit::Nanosecond => {
1690                        let array = column.as_primitive::<DurationNanosecondType>();
1691                        write_primitive(typed, array.values(), levels)
1692                    }
1693                },
1694                ArrowDataType::Decimal64(_, _) => {
1695                    let array = column
1696                        .as_primitive::<Decimal64Type>()
1697                        .reinterpret_cast::<Int64Type>();
1698                    write_primitive(typed, array.values(), levels)
1699                }
1700                ArrowDataType::Decimal128(_, _) => {
1701                    // use the int64 to represent the decimal with low precision
1702                    let array = column
1703                        .as_primitive::<Decimal128Type>()
1704                        .unary::<_, Int64Type>(|v| v as i64);
1705                    write_primitive(typed, array.values(), levels)
1706                }
1707                ArrowDataType::Decimal256(_, _) => {
1708                    // use the int64 to represent the decimal with low precision
1709                    let array = column
1710                        .as_primitive::<Decimal256Type>()
1711                        .unary::<_, Int64Type>(|v| v.as_i128() as i64);
1712                    write_primitive(typed, array.values(), levels)
1713                }
1714                d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1715            }
1716        }
1717        ColumnWriter::Int96ColumnWriter(_typed) => {
1718            unreachable!("Currently unreachable because data type not supported")
1719        }
1720        ColumnWriter::FloatColumnWriter(typed) => {
1721            let array = column.as_primitive::<Float32Type>();
1722            write_primitive(typed, array.values(), levels)
1723        }
1724        ColumnWriter::DoubleColumnWriter(typed) => {
1725            let array = column.as_primitive::<Float64Type>();
1726            write_primitive(typed, array.values(), levels)
1727        }
1728        ColumnWriter::ByteArrayColumnWriter(_) => {
1729            unreachable!("should use ByteArrayWriter")
1730        }
1731        ColumnWriter::FixedLenByteArrayColumnWriter(typed) => {
1732            let bytes = match column.data_type() {
1733                ArrowDataType::Interval(interval_unit) => match interval_unit {
1734                    IntervalUnit::YearMonth => {
1735                        let array = column.as_primitive::<IntervalYearMonthType>();
1736                        get_interval_ym_array_slice(array, indices.iter().copied())
1737                    }
1738                    IntervalUnit::DayTime => {
1739                        let array = column.as_primitive::<IntervalDayTimeType>();
1740                        get_interval_dt_array_slice(array, indices.iter().copied())
1741                    }
1742                    _ => {
1743                        return Err(ParquetError::NYI(format!(
1744                            "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
1745                        )));
1746                    }
1747                },
1748                ArrowDataType::FixedSizeBinary(_) => {
1749                    let array = column.as_fixed_size_binary();
1750                    get_fsb_array_slice(array, indices.iter().copied())
1751                }
1752                ArrowDataType::Decimal32(_, _) => {
1753                    let array = column.as_primitive::<Decimal32Type>();
1754                    get_decimal_array_slice(array, indices.iter().copied())
1755                }
1756                ArrowDataType::Decimal64(_, _) => {
1757                    let array = column.as_primitive::<Decimal64Type>();
1758                    get_decimal_array_slice(array, indices.iter().copied())
1759                }
1760                ArrowDataType::Decimal128(_, _) => {
1761                    let array = column.as_primitive::<Decimal128Type>();
1762                    get_decimal_array_slice(array, indices.iter().copied())
1763                }
1764                ArrowDataType::Decimal256(_, _) => {
1765                    let array = column.as_primitive::<Decimal256Type>();
1766                    get_decimal_array_slice(array, indices.iter().copied())
1767                }
1768                ArrowDataType::Float16 => {
1769                    let array = column.as_primitive::<Float16Type>();
1770                    get_float_16_array_slice(array, indices.iter().copied())
1771                }
1772                _ => {
1773                    return Err(ParquetError::NYI(
1774                        "Attempting to write an Arrow type that is not yet implemented".to_string(),
1775                    ));
1776                }
1777            };
1778            typed.write_batch_internal(
1779                bytes.as_slice(),
1780                None,
1781                levels.def_level_data().as_ref(),
1782                levels.rep_level_data().as_ref(),
1783                None,
1784                None,
1785                None,
1786            )
1787        }
1788    }
1789}
1790
1791fn write_primitive<E: ColumnValueEncoder>(
1792    writer: &mut GenericColumnWriter<E>,
1793    values: &E::Values,
1794    levels: &ArrayLevels,
1795) -> Result<usize> {
1796    writer.write_batch_internal(
1797        values,
1798        Some(levels.non_null_indices()),
1799        levels.def_level_data().as_ref(),
1800        levels.rep_level_data().as_ref(),
1801        None,
1802        None,
1803        None,
1804    )
1805}
1806
1807fn get_bool_array_slice(
1808    array: &arrow_array::BooleanArray,
1809    indices: impl ExactSizeIterator<Item = usize>,
1810) -> Vec<bool> {
1811    let mut values = Vec::with_capacity(indices.len());
1812    for i in indices {
1813        values.push(array.value(i))
1814    }
1815    values
1816}
1817
1818/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1819/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated.
1820fn get_interval_ym_array_slice(
1821    array: &arrow_array::IntervalYearMonthArray,
1822    indices: impl ExactSizeIterator<Item = usize>,
1823) -> Vec<FixedLenByteArray> {
1824    chunk_array_slice(12, indices, move |i, chunk| {
1825        let value = array.value(i);
1826        chunk[0..4].copy_from_slice(&value.to_le_bytes());
1827    })
1828}
1829
1830/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1831/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated.
1832fn get_interval_dt_array_slice(
1833    array: &arrow_array::IntervalDayTimeArray,
1834    indices: impl ExactSizeIterator<Item = usize>,
1835) -> Vec<FixedLenByteArray> {
1836    chunk_array_slice(12, indices, move |i, chunk| {
1837        let value = array.value(i);
1838        chunk[4..8].copy_from_slice(&value.days.to_le_bytes());
1839        chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1840    })
1841}
1842
1843trait NativeDecimalType: DecimalType {
1844    type NativeBytes: AsRef<[u8]>;
1845
1846    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
1847}
1848impl NativeDecimalType for Decimal32Type {
1849    type NativeBytes = [u8; Self::BYTE_LENGTH];
1850
1851    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1852        value.to_be_bytes()
1853    }
1854}
1855impl NativeDecimalType for Decimal64Type {
1856    type NativeBytes = [u8; Self::BYTE_LENGTH];
1857
1858    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1859        value.to_be_bytes()
1860    }
1861}
1862impl NativeDecimalType for Decimal128Type {
1863    type NativeBytes = [u8; Self::BYTE_LENGTH];
1864
1865    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1866        value.to_be_bytes()
1867    }
1868}
1869impl NativeDecimalType for Decimal256Type {
1870    type NativeBytes = [u8; Self::BYTE_LENGTH];
1871
1872    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1873        value.to_be_bytes()
1874    }
1875}
1876
1877fn get_decimal_array_slice<T: NativeDecimalType>(
1878    array: &PrimitiveArray<T>,
1879    indices: impl ExactSizeIterator<Item = usize>,
1880) -> Vec<FixedLenByteArray> {
1881    let chunk_size = decimal_length_from_precision(array.precision());
1882    assert!(chunk_size <= T::BYTE_LENGTH);
1883
1884    if chunk_size == T::BYTE_LENGTH {
1885        // Special-case that allows inlining memcpy.
1886        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1887            let as_be_bytes = T::to_be_bytes(array.value(i));
1888            chunk.copy_from_slice(as_be_bytes.as_ref());
1889        })
1890    } else {
1891        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1892            let as_be_bytes = T::to_be_bytes(array.value(i));
1893            let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..];
1894            chunk.copy_from_slice(resized_value);
1895        })
1896    }
1897}
1898
1899fn get_float_16_array_slice(
1900    array: &arrow_array::Float16Array,
1901    indices: impl ExactSizeIterator<Item = usize>,
1902) -> Vec<FixedLenByteArray> {
1903    chunk_array_slice(2, indices, move |i, chunk| {
1904        let value = array.value(i).to_le_bytes();
1905        chunk.copy_from_slice(&value);
1906    })
1907}
1908
1909fn get_fsb_array_slice(
1910    array: &arrow_array::FixedSizeBinaryArray,
1911    indices: impl ExactSizeIterator<Item = usize>,
1912) -> Vec<FixedLenByteArray> {
1913    chunk_array_slice(array.value_size(), indices, move |i, chunk| {
1914        let value = array.value(i);
1915        chunk.copy_from_slice(value);
1916    })
1917}
1918
1919#[inline]
1920fn chunk_array_slice(
1921    chunk_size: usize,
1922    indices: impl ExactSizeIterator<Item = usize>,
1923    writer: impl Fn(usize, &mut [u8]),
1924) -> Vec<FixedLenByteArray> {
1925    let capacity = indices.len() * chunk_size;
1926    // TODO: This could be done with Vec::spare_capacity_mut,
1927    //       but [MaybeUninit]::write_copy_of_slice is gated behind MSRV 1.93
1928    let mut arena = vec![0; capacity];
1929    for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
1930        writer(i, chunk);
1931    }
1932    chunk_contiguous_vec(arena, chunk_size)
1933}
1934
1935fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
1936    let mut values = Vec::with_capacity(arena.len() / chunk_size);
1937    let mut arena = Bytes::from(arena);
1938    while arena.len() >= chunk_size {
1939        let slice = arena.split_to(chunk_size);
1940        values.push(FixedLenByteArray::from(ByteArray::from(slice)));
1941    }
1942    values
1943}
1944
1945#[cfg(test)]
1946mod tests {
1947    use super::*;
1948    use std::cmp::Ordering;
1949    use std::collections::HashMap;
1950
1951    use std::fs::File;
1952
1953    use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1954    use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
1955    use crate::column::page::{Page, PageReader};
1956    use crate::file::metadata::thrift::PageHeader;
1957    use crate::file::page_index::column_index::ColumnIndexMetaData;
1958    use crate::file::reader::SerializedPageReader;
1959    use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
1960    use crate::schema::types::ColumnPath;
1961    use arrow::datatypes::ToByteSlice;
1962    use arrow::datatypes::{DataType, Schema};
1963    use arrow::error::Result as ArrowResult;
1964    use arrow::util::data_gen::create_random_array;
1965    use arrow::util::pretty::pretty_format_batches;
1966    use arrow::{array::*, buffer::Buffer};
1967    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
1968    use arrow_schema::Fields;
1969    use half::f16;
1970    use num_traits::{FromPrimitive, ToPrimitive};
1971    use tempfile::tempfile;
1972
1973    use crate::basic::Encoding;
1974    use crate::data_type::AsBytes;
1975    use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
1976    use crate::file::properties::{
1977        BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
1978    };
1979    use crate::file::serialized_reader::ReadOptionsBuilder;
1980    use crate::file::{
1981        reader::{FileReader, SerializedFileReader},
1982        statistics::Statistics,
1983    };
1984
1985    /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps
1986    /// blobs in a `HashMap` — nothing like the default `Vec<Bytes>`. Used to
1987    /// prove the writer relies only on the opaque-handle contract and never on
1988    /// handles being dense `Vec` indices. Records how many blobs were stored.
1989    #[derive(Debug, Default)]
1990    struct RecordingPageStore {
1991        next: u64,
1992        blobs: HashMap<u64, Bytes>,
1993        puts: Arc<std::sync::atomic::AtomicUsize>,
1994    }
1995
1996    impl PageStore for RecordingPageStore {
1997        fn put(&mut self, value: Bytes) -> Result<PageKey> {
1998            // Deliberately non-sequential, never-zero handles.
1999            let id = 100 + self.next * 7;
2000            self.next += 1;
2001            self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2002            self.blobs.insert(id, value);
2003            Ok(PageKey::new(id))
2004        }
2005
2006        fn take(&mut self, key: PageKey) -> Result<Bytes> {
2007            self.blobs
2008                .remove(&key.get())
2009                .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2010        }
2011    }
2012
2013    #[derive(Debug)]
2014    struct RecordingPageStoreFactory {
2015        puts: Arc<std::sync::atomic::AtomicUsize>,
2016    }
2017
2018    impl PageStoreFactory for RecordingPageStoreFactory {
2019        fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2020            Ok(Box::new(RecordingPageStore {
2021                puts: self.puts.clone(),
2022                ..Default::default()
2023            }))
2024        }
2025    }
2026
2027    /// A custom [`PageStore`] must produce byte-identical files to the in-memory
2028    /// default, across dictionary and non-dictionary columns and multiple row
2029    /// groups (so multiple store instances are exercised).
2030    #[test]
2031    fn custom_page_store_is_byte_identical_to_default() {
2032        let schema = Arc::new(Schema::new(vec![
2033            Field::new("i", DataType::Int32, true),
2034            // A low-cardinality string column to exercise the dictionary path.
2035            Field::new("s", DataType::Utf8, true),
2036        ]));
2037        let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2038        let s = StringArray::from(vec![
2039            Some("a"),
2040            Some("bb"),
2041            Some("a"),
2042            None,
2043            Some("bb"),
2044            Some("ccc"),
2045        ]);
2046        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2047
2048        // Small row groups so multiple column chunks (hence multiple store
2049        // instances) are produced.
2050        let props = WriterProperties::builder()
2051            .set_max_row_group_row_count(Some(3))
2052            .build();
2053
2054        let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2055            let mut buffer = Vec::new();
2056            let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2057            if let Some(factory) = factory {
2058                opts = opts.with_page_store_factory(factory);
2059            }
2060            let mut writer =
2061                ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2062            writer.write(&batch).unwrap();
2063            writer.close().unwrap();
2064            buffer
2065        };
2066
2067        let default_bytes = write(None);
2068
2069        let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2070        let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2071            puts: puts.clone(),
2072        })));
2073
2074        assert!(
2075            puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2076            "custom PageStore was never written to"
2077        );
2078        assert_eq!(
2079            default_bytes, custom_bytes,
2080            "a custom PageStore must produce byte-identical output to the default"
2081        );
2082    }
2083
2084    /// A dictionary-encoded column written through the deferred-ordering Arrow
2085    /// path must round-trip correctly even with the offset index disabled, when
2086    /// only the chunk-level dictionary/data page offsets are rewritten (there is
2087    /// no offset index to rebuild). Spans multiple data pages so the
2088    /// dictionary-first reordering is exercised.
2089    #[test]
2090    fn dictionary_column_round_trips_with_offset_index_disabled() {
2091        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2092
2093        // Low cardinality so the column stays dictionary-encoded; enough rows to
2094        // span several data pages within a single row group.
2095        let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2096        let array = Int32Array::from(values.clone());
2097        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2098
2099        let props = WriterProperties::builder()
2100            .set_offset_index_disabled(true)
2101            .set_data_page_row_count_limit(4096)
2102            .build();
2103        let opts = ArrowWriterOptions::new().with_properties(props);
2104
2105        let mut buffer = Vec::new();
2106        let mut writer =
2107            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2108        writer.write(&batch).unwrap();
2109        writer.close().unwrap();
2110
2111        let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2112        let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2113        let read_values: Vec<Option<i32>> = read
2114            .iter()
2115            .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2116            .collect();
2117        assert_eq!(read_values, values);
2118    }
2119
2120    /// The dictionary page is routed through the [`PageStore`] like any other
2121    /// page rather than held resident in memory, so a dictionary column chunk's
2122    /// *entire* serialized size — dictionary page included — passes through the
2123    /// store.
2124    #[test]
2125    fn dictionary_page_is_routed_through_the_store() {
2126        /// A store that sums the bytes handed to `put`.
2127        #[derive(Debug, Default)]
2128        struct SizeRecordingPageStore {
2129            blobs: Vec<Bytes>,
2130            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2131        }
2132        impl PageStore for SizeRecordingPageStore {
2133            fn put(&mut self, value: Bytes) -> Result<PageKey> {
2134                self.bytes_put
2135                    .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2136                let key = PageKey::new(self.blobs.len() as u64);
2137                self.blobs.push(value);
2138                Ok(key)
2139            }
2140            fn take(&mut self, key: PageKey) -> Result<Bytes> {
2141                Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2142            }
2143        }
2144        #[derive(Debug)]
2145        struct Factory {
2146            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2147        }
2148        impl PageStoreFactory for Factory {
2149            fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2150                Ok(Box::new(SizeRecordingPageStore {
2151                    bytes_put: self.bytes_put.clone(),
2152                    ..Default::default()
2153                }))
2154            }
2155        }
2156
2157        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2158        // Low cardinality keeps the column dictionary-encoded with a real,
2159        // non-empty dictionary page.
2160        let values: Vec<&str> = (0..2048)
2161            .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2162            .collect();
2163        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2164            .unwrap();
2165
2166        let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2167        let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2168            bytes_put: bytes_put.clone(),
2169        }));
2170
2171        // A single batch / single column means exactly one row group and one
2172        // store instance, so the bytes it saw map to one column chunk.
2173        let mut buffer = Vec::new();
2174        let mut writer =
2175            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2176        writer.write(&batch).unwrap();
2177        writer.close().unwrap();
2178
2179        let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2180        let column = reader.metadata().row_group(0).column(0);
2181        assert!(
2182            column.dictionary_page_offset().is_some(),
2183            "expected the column to be dictionary-encoded"
2184        );
2185
2186        // The bytes the store was handed must account for the whole chunk,
2187        // dictionary page included. Holding the dictionary page apart from the
2188        // store would make this fall short by the dictionary page's size.
2189        assert_eq!(
2190            bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2191            column.compressed_size(),
2192            "the dictionary page must pass through the store like any other page"
2193        );
2194    }
2195
2196    #[test]
2197    fn arrow_writer() {
2198        // define schema
2199        let schema = Schema::new(vec![
2200            Field::new("a", DataType::Int32, false),
2201            Field::new("b", DataType::Int32, true),
2202        ]);
2203
2204        // create some data
2205        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2206        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2207
2208        // build a record batch
2209        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2210
2211        roundtrip(batch, Some(SMALL_SIZE / 2));
2212    }
2213
2214    fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2215        let mut buffer = vec![];
2216
2217        let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2218        writer.write(expected_batch).unwrap();
2219        writer.close().unwrap();
2220
2221        buffer
2222    }
2223
2224    fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2225        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2226        writer.write(expected_batch).unwrap();
2227        writer.into_inner().unwrap()
2228    }
2229
2230    #[test]
2231    fn roundtrip_bytes() {
2232        // define schema
2233        let schema = Arc::new(Schema::new(vec![
2234            Field::new("a", DataType::Int32, false),
2235            Field::new("b", DataType::Int32, true),
2236        ]));
2237
2238        // create some data
2239        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2240        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2241
2242        // build a record batch
2243        let expected_batch =
2244            RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2245
2246        for buffer in [
2247            get_bytes_after_close(schema.clone(), &expected_batch),
2248            get_bytes_by_into_inner(schema, &expected_batch),
2249        ] {
2250            let cursor = Bytes::from(buffer);
2251            let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2252
2253            let actual_batch = record_batch_reader
2254                .next()
2255                .expect("No batch found")
2256                .expect("Unable to get batch");
2257
2258            assert_eq!(expected_batch.schema(), actual_batch.schema());
2259            assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2260            assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2261            for i in 0..expected_batch.num_columns() {
2262                let expected_data = expected_batch.column(i).to_data();
2263                let actual_data = actual_batch.column(i).to_data();
2264
2265                assert_eq!(expected_data, actual_data);
2266            }
2267        }
2268    }
2269
2270    #[test]
2271    fn arrow_writer_non_null() {
2272        // define schema
2273        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2274
2275        // create some data
2276        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2277
2278        // build a record batch
2279        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2280
2281        roundtrip(batch, Some(SMALL_SIZE / 2));
2282    }
2283
2284    #[test]
2285    fn arrow_writer_list() {
2286        // define schema
2287        let schema = Schema::new(vec![Field::new(
2288            "a",
2289            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2290            true,
2291        )]);
2292
2293        // create some data
2294        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2295
2296        // Construct a buffer for value offsets, for the nested array:
2297        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2298        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2299
2300        // Construct a list array from the above two
2301        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2302            DataType::Int32,
2303            false,
2304        ))))
2305        .len(5)
2306        .add_buffer(a_value_offsets)
2307        .add_child_data(a_values.into_data())
2308        .null_bit_buffer(Some(Buffer::from([0b00011011])))
2309        .build()
2310        .unwrap();
2311        let a = ListArray::from(a_list_data);
2312
2313        // build a record batch
2314        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2315
2316        assert_eq!(batch.column(0).null_count(), 1);
2317
2318        // This test fails if the max row group size is less than the batch's length
2319        // see https://github.com/apache/arrow-rs/issues/518
2320        roundtrip(batch, None);
2321    }
2322
2323    #[test]
2324    fn arrow_writer_list_non_null() {
2325        // define schema
2326        let schema = Schema::new(vec![Field::new(
2327            "a",
2328            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2329            false,
2330        )]);
2331
2332        // create some data
2333        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2334
2335        // Construct a buffer for value offsets, for the nested array:
2336        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2337        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2338
2339        // Construct a list array from the above two
2340        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2341            DataType::Int32,
2342            false,
2343        ))))
2344        .len(5)
2345        .add_buffer(a_value_offsets)
2346        .add_child_data(a_values.into_data())
2347        .build()
2348        .unwrap();
2349        let a = ListArray::from(a_list_data);
2350
2351        // build a record batch
2352        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2353
2354        // This test fails if the max row group size is less than the batch's length
2355        // see https://github.com/apache/arrow-rs/issues/518
2356        assert_eq!(batch.column(0).null_count(), 0);
2357
2358        roundtrip(batch, None);
2359    }
2360
2361    #[test]
2362    fn arrow_writer_list_view() {
2363        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2364        let schema = Schema::new(vec![Field::new(
2365            "a",
2366            DataType::ListView(list_field.clone()),
2367            true,
2368        )]);
2369
2370        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2371        let a = ListViewArray::new(
2372            list_field,
2373            vec![0, 1, 0, 3, 6].into(),
2374            vec![1, 2, 0, 3, 4].into(),
2375            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2376            Some(vec![true, true, false, true, true].into()),
2377        );
2378
2379        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2380
2381        assert_eq!(batch.column(0).null_count(), 1);
2382
2383        roundtrip(batch, None);
2384    }
2385
2386    #[test]
2387    fn arrow_writer_list_view_non_null() {
2388        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2389        let schema = Schema::new(vec![Field::new(
2390            "a",
2391            DataType::ListView(list_field.clone()),
2392            false,
2393        )]);
2394
2395        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2396        let a = ListViewArray::new(
2397            list_field,
2398            vec![0, 1, 0, 3, 6].into(),
2399            vec![1, 2, 0, 3, 4].into(),
2400            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2401            None,
2402        );
2403
2404        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2405
2406        assert_eq!(batch.column(0).null_count(), 0);
2407
2408        roundtrip(batch, None);
2409    }
2410
2411    #[test]
2412    fn arrow_writer_list_view_out_of_order() {
2413        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2414        let schema = Schema::new(vec![Field::new(
2415            "a",
2416            DataType::ListView(list_field.clone()),
2417            false,
2418        )]);
2419
2420        // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets
2421        let a = ListViewArray::new(
2422            list_field,
2423            vec![0, 1, 0, 6, 3].into(),
2424            vec![1, 2, 0, 4, 3].into(),
2425            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2426            None,
2427        );
2428
2429        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2430
2431        roundtrip(batch, None);
2432    }
2433
2434    #[test]
2435    fn arrow_writer_large_list_view() {
2436        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2437        let schema = Schema::new(vec![Field::new(
2438            "a",
2439            DataType::LargeListView(list_field.clone()),
2440            true,
2441        )]);
2442
2443        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2444        let a = LargeListViewArray::new(
2445            list_field,
2446            vec![0i64, 1, 0, 3, 6].into(),
2447            vec![1i64, 2, 0, 3, 4].into(),
2448            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2449            Some(vec![true, true, false, true, true].into()),
2450        );
2451
2452        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2453
2454        assert_eq!(batch.column(0).null_count(), 1);
2455
2456        roundtrip(batch, None);
2457    }
2458
2459    #[test]
2460    fn arrow_writer_list_view_with_struct() {
2461        // Test ListView containing Struct: ListView<Struct<Int32, Utf8>>
2462        let struct_fields = Fields::from(vec![
2463            Field::new("id", DataType::Int32, false),
2464            Field::new("name", DataType::Utf8, false),
2465        ]);
2466        let struct_type = DataType::Struct(struct_fields.clone());
2467        let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2468
2469        let schema = Schema::new(vec![Field::new(
2470            "a",
2471            DataType::ListView(list_field.clone()),
2472            true,
2473        )]);
2474
2475        // Create struct values
2476        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2477        let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2478        let struct_array = StructArray::new(
2479            struct_fields,
2480            vec![Arc::new(id_array), Arc::new(name_array)],
2481            None,
2482        );
2483
2484        // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"}, {5, "e"}]
2485        let list_view = ListViewArray::new(
2486            list_field,
2487            vec![0, 2, 2].into(), // offsets
2488            vec![2, 0, 3].into(), // sizes
2489            Arc::new(struct_array),
2490            Some(vec![true, false, true].into()),
2491        );
2492
2493        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap();
2494
2495        roundtrip(batch, None);
2496    }
2497
2498    #[test]
2499    fn arrow_writer_binary() {
2500        let string_field = Field::new("a", DataType::Utf8, false);
2501        let binary_field = Field::new("b", DataType::Binary, false);
2502        let schema = Schema::new(vec![string_field, binary_field]);
2503
2504        let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2505        let raw_binary_values = [
2506            b"foo".to_vec(),
2507            b"bar".to_vec(),
2508            b"baz".to_vec(),
2509            b"quux".to_vec(),
2510        ];
2511        let raw_binary_value_refs = raw_binary_values
2512            .iter()
2513            .map(|x| x.as_slice())
2514            .collect::<Vec<_>>();
2515
2516        let string_values = StringArray::from(raw_string_values.clone());
2517        let binary_values = BinaryArray::from(raw_binary_value_refs);
2518        let batch = RecordBatch::try_new(
2519            Arc::new(schema),
2520            vec![Arc::new(string_values), Arc::new(binary_values)],
2521        )
2522        .unwrap();
2523
2524        roundtrip(batch, Some(SMALL_SIZE / 2));
2525    }
2526
2527    #[test]
2528    fn arrow_writer_binary_view() {
2529        let string_field = Field::new("a", DataType::Utf8View, false);
2530        let binary_field = Field::new("b", DataType::BinaryView, false);
2531        let nullable_string_field = Field::new("a", DataType::Utf8View, true);
2532        let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]);
2533
2534        let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2535        let raw_binary_values = vec![
2536            b"foo".to_vec(),
2537            b"bar".to_vec(),
2538            b"large payload over 12 bytes".to_vec(),
2539            b"lulu".to_vec(),
2540        ];
2541        let nullable_string_values =
2542            vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2543
2544        let string_view_values = StringViewArray::from(raw_string_values);
2545        let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2546        let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2547        let batch = RecordBatch::try_new(
2548            Arc::new(schema),
2549            vec![
2550                Arc::new(string_view_values),
2551                Arc::new(binary_view_values),
2552                Arc::new(nullable_string_view_values),
2553            ],
2554        )
2555        .unwrap();
2556
2557        roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2558        roundtrip(batch, None);
2559    }
2560
2561    #[test]
2562    fn arrow_writer_binary_view_long_value() {
2563        let string_field = Field::new("a", DataType::Utf8View, false);
2564        let binary_field = Field::new("b", DataType::BinaryView, false);
2565        let schema = Schema::new(vec![string_field, binary_field]);
2566
2567        // There is special case validation for long values (greater than 128)
2568        // 128 encodes as 0x80 0x00 0x00 0x00 in little endian, which should
2569        // trigger the long-string UTF-8 validation branch in the plain decoder.
2570        let long = "a".repeat(128);
2571        let raw_string_values = vec!["foo", long.as_str(), "bar"];
2572        let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2573
2574        let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2575        let binary_view_values: ArrayRef =
2576            Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2577
2578        one_column_roundtrip(Arc::clone(&string_view_values), false);
2579        one_column_roundtrip(Arc::clone(&binary_view_values), false);
2580
2581        let batch = RecordBatch::try_new(
2582            Arc::new(schema),
2583            vec![string_view_values, binary_view_values],
2584        )
2585        .unwrap();
2586
2587        // Disable dictionary to exercise plain encoding paths in the reader.
2588        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
2589            let props = WriterProperties::builder()
2590                .set_writer_version(version)
2591                .set_dictionary_enabled(false)
2592                .build();
2593            roundtrip_opts(&batch, props);
2594        }
2595    }
2596
2597    fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2598        let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2599        let schema = Schema::new(vec![decimal_field]);
2600
2601        let decimal_values = vec![10_000, 50_000, 0, -100]
2602            .into_iter()
2603            .map(Some)
2604            .collect::<Decimal128Array>()
2605            .with_precision_and_scale(precision, scale)
2606            .unwrap();
2607
2608        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2609    }
2610
2611    #[test]
2612    fn arrow_writer_decimal() {
2613        // int32 to store the decimal value
2614        let batch_int32_decimal = get_decimal_batch(5, 2);
2615        roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2616        // int64 to store the decimal value
2617        let batch_int64_decimal = get_decimal_batch(12, 2);
2618        roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2619        // fixed_length_byte_array to store the decimal value
2620        let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2621        roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2622    }
2623
2624    #[test]
2625    fn arrow_writer_complex() {
2626        // define schema
2627        let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2628        let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2629        let struct_field_g = Arc::new(Field::new_list(
2630            "g",
2631            Field::new_list_field(DataType::Int16, true),
2632            false,
2633        ));
2634        let struct_field_h = Arc::new(Field::new_list(
2635            "h",
2636            Field::new_list_field(DataType::Int16, false),
2637            true,
2638        ));
2639        let struct_field_e = Arc::new(Field::new_struct(
2640            "e",
2641            vec![
2642                struct_field_f.clone(),
2643                struct_field_g.clone(),
2644                struct_field_h.clone(),
2645            ],
2646            false,
2647        ));
2648        let schema = Schema::new(vec![
2649            Field::new("a", DataType::Int32, false),
2650            Field::new("b", DataType::Int32, true),
2651            Field::new_struct(
2652                "c",
2653                vec![struct_field_d.clone(), struct_field_e.clone()],
2654                false,
2655            ),
2656        ]);
2657
2658        // create some data
2659        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2660        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2661        let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2662        let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2663
2664        let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2665
2666        // Construct a buffer for value offsets, for the nested array:
2667        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2668        let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2669
2670        // Construct a list array from the above two
2671        let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2672            .len(5)
2673            .add_buffer(g_value_offsets.clone())
2674            .add_child_data(g_value.to_data())
2675            .build()
2676            .unwrap();
2677        let g = ListArray::from(g_list_data);
2678        // The difference between g and h is that h has a null bitmap
2679        let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2680            .len(5)
2681            .add_buffer(g_value_offsets)
2682            .add_child_data(g_value.to_data())
2683            .null_bit_buffer(Some(Buffer::from([0b00011011])))
2684            .build()
2685            .unwrap();
2686        let h = ListArray::from(h_list_data);
2687
2688        let e = StructArray::from(vec![
2689            (struct_field_f, Arc::new(f) as ArrayRef),
2690            (struct_field_g, Arc::new(g) as ArrayRef),
2691            (struct_field_h, Arc::new(h) as ArrayRef),
2692        ]);
2693
2694        let c = StructArray::from(vec![
2695            (struct_field_d, Arc::new(d) as ArrayRef),
2696            (struct_field_e, Arc::new(e) as ArrayRef),
2697        ]);
2698
2699        // build a record batch
2700        let batch = RecordBatch::try_new(
2701            Arc::new(schema),
2702            vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2703        )
2704        .unwrap();
2705
2706        roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2707        roundtrip(batch, Some(SMALL_SIZE / 3));
2708    }
2709
2710    #[test]
2711    fn arrow_writer_complex_mixed() {
2712        // This test was added while investigating https://github.com/apache/arrow-rs/issues/244.
2713        // It was subsequently fixed while investigating https://github.com/apache/arrow-rs/issues/245.
2714
2715        // define schema
2716        let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2717        let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2718        let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2719        let schema = Schema::new(vec![Field::new(
2720            "some_nested_object",
2721            DataType::Struct(Fields::from(vec![
2722                offset_field.clone(),
2723                partition_field.clone(),
2724                topic_field.clone(),
2725            ])),
2726            false,
2727        )]);
2728
2729        // create some data
2730        let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2731        let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2732        let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2733
2734        let some_nested_object = StructArray::from(vec![
2735            (offset_field, Arc::new(offset) as ArrayRef),
2736            (partition_field, Arc::new(partition) as ArrayRef),
2737            (topic_field, Arc::new(topic) as ArrayRef),
2738        ]);
2739
2740        // build a record batch
2741        let batch =
2742            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2743
2744        roundtrip(batch, Some(SMALL_SIZE / 2));
2745    }
2746
2747    #[test]
2748    fn arrow_writer_map() {
2749        // Note: we are using the JSON Arrow reader for brevity
2750        let json_content = r#"
2751        {"stocks":{"long": "$AAA", "short": "$BBB"}}
2752        {"stocks":{"long": null, "long": "$CCC", "short": null}}
2753        {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2754        "#;
2755        let entries_struct_type = DataType::Struct(Fields::from(vec![
2756            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2757            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2758        ]));
2759        let stocks_field = Field::new(
2760            "stocks",
2761            DataType::Map(
2762                Arc::new(Field::new(
2763                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2764                    entries_struct_type,
2765                    false,
2766                )),
2767                false,
2768            ),
2769            true,
2770        );
2771        let schema = Arc::new(Schema::new(vec![stocks_field]));
2772        let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2773        let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2774
2775        let batch = reader.next().unwrap().unwrap();
2776        roundtrip(batch, None);
2777    }
2778
2779    #[test]
2780    fn arrow_writer_2_level_struct() {
2781        // tests writing <struct<struct<primitive>>
2782        let field_c = Field::new("c", DataType::Int32, true);
2783        let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2784        let type_a = DataType::Struct(vec![field_b.clone()].into());
2785        let field_a = Field::new("a", type_a, true);
2786        let schema = Schema::new(vec![field_a.clone()]);
2787
2788        // create data
2789        let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2790        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2791            .len(6)
2792            .null_bit_buffer(Some(Buffer::from([0b00100111])))
2793            .add_child_data(c.into_data())
2794            .build()
2795            .unwrap();
2796        let b = StructArray::from(b_data);
2797        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2798            .len(6)
2799            .null_bit_buffer(Some(Buffer::from([0b00101111])))
2800            .add_child_data(b.into_data())
2801            .build()
2802            .unwrap();
2803        let a = StructArray::from(a_data);
2804
2805        assert_eq!(a.null_count(), 1);
2806        assert_eq!(a.column(0).null_count(), 2);
2807
2808        // build a racord batch
2809        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2810
2811        roundtrip(batch, Some(SMALL_SIZE / 2));
2812    }
2813
2814    #[test]
2815    fn arrow_writer_2_level_struct_non_null() {
2816        // tests writing <struct<struct<primitive>>
2817        let field_c = Field::new("c", DataType::Int32, false);
2818        let type_b = DataType::Struct(vec![field_c].into());
2819        let field_b = Field::new("b", type_b.clone(), false);
2820        let type_a = DataType::Struct(vec![field_b].into());
2821        let field_a = Field::new("a", type_a.clone(), false);
2822        let schema = Schema::new(vec![field_a]);
2823
2824        // create data
2825        let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2826        let b_data = ArrayDataBuilder::new(type_b)
2827            .len(6)
2828            .add_child_data(c.into_data())
2829            .build()
2830            .unwrap();
2831        let b = StructArray::from(b_data);
2832        let a_data = ArrayDataBuilder::new(type_a)
2833            .len(6)
2834            .add_child_data(b.into_data())
2835            .build()
2836            .unwrap();
2837        let a = StructArray::from(a_data);
2838
2839        assert_eq!(a.null_count(), 0);
2840        assert_eq!(a.column(0).null_count(), 0);
2841
2842        // build a racord batch
2843        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2844
2845        roundtrip(batch, Some(SMALL_SIZE / 2));
2846    }
2847
2848    #[test]
2849    fn arrow_writer_2_level_struct_mixed_null() {
2850        // tests writing <struct<struct<primitive>>
2851        let field_c = Field::new("c", DataType::Int32, false);
2852        let type_b = DataType::Struct(vec![field_c].into());
2853        let field_b = Field::new("b", type_b.clone(), true);
2854        let type_a = DataType::Struct(vec![field_b].into());
2855        let field_a = Field::new("a", type_a.clone(), false);
2856        let schema = Schema::new(vec![field_a]);
2857
2858        // create data
2859        let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2860        let b_data = ArrayDataBuilder::new(type_b)
2861            .len(6)
2862            .null_bit_buffer(Some(Buffer::from([0b00100111])))
2863            .add_child_data(c.into_data())
2864            .build()
2865            .unwrap();
2866        let b = StructArray::from(b_data);
2867        // a intentionally has no null buffer, to test that this is handled correctly
2868        let a_data = ArrayDataBuilder::new(type_a)
2869            .len(6)
2870            .add_child_data(b.into_data())
2871            .build()
2872            .unwrap();
2873        let a = StructArray::from(a_data);
2874
2875        assert_eq!(a.null_count(), 0);
2876        assert_eq!(a.column(0).null_count(), 2);
2877
2878        // build a racord batch
2879        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2880
2881        roundtrip(batch, Some(SMALL_SIZE / 2));
2882    }
2883
2884    #[test]
2885    fn arrow_writer_2_level_struct_mixed_null_2() {
2886        // tests writing <struct<struct<primitive>>, where the primitive columns are non-null.
2887        let field_c = Field::new("c", DataType::Int32, false);
2888        let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
2889        let field_e = Field::new(
2890            "e",
2891            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2892            false,
2893        );
2894
2895        let field_b = Field::new(
2896            "b",
2897            DataType::Struct(vec![field_c, field_d, field_e].into()),
2898            false,
2899        );
2900        let type_a = DataType::Struct(vec![field_b.clone()].into());
2901        let field_a = Field::new("a", type_a, true);
2902        let schema = Schema::new(vec![field_a.clone()]);
2903
2904        // create data
2905        let c = Int32Array::from_iter_values(0..6);
2906        let d = FixedSizeBinaryArray::try_from_iter(
2907            ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
2908        )
2909        .expect("four byte values");
2910        let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
2911        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2912            .len(6)
2913            .add_child_data(c.into_data())
2914            .add_child_data(d.into_data())
2915            .add_child_data(e.into_data())
2916            .build()
2917            .unwrap();
2918        let b = StructArray::from(b_data);
2919        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2920            .len(6)
2921            .null_bit_buffer(Some(Buffer::from([0b00100101])))
2922            .add_child_data(b.into_data())
2923            .build()
2924            .unwrap();
2925        let a = StructArray::from(a_data);
2926
2927        assert_eq!(a.null_count(), 3);
2928        assert_eq!(a.column(0).null_count(), 0);
2929
2930        // build a record batch
2931        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2932
2933        roundtrip(batch, Some(SMALL_SIZE / 2));
2934    }
2935
2936    #[test]
2937    fn test_fixed_size_binary_in_dict() {
2938        fn test_fixed_size_binary_in_dict_inner<K>()
2939        where
2940            K: ArrowDictionaryKeyType,
2941            K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
2942            <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
2943        {
2944            let field = Field::new(
2945                "a",
2946                DataType::Dictionary(
2947                    Box::new(K::DATA_TYPE),
2948                    Box::new(DataType::FixedSizeBinary(4)),
2949                ),
2950                false,
2951            );
2952            let schema = Schema::new(vec![field]);
2953
2954            let keys: Vec<K::Native> = vec![
2955                K::Native::try_from(0u8).unwrap(),
2956                K::Native::try_from(0u8).unwrap(),
2957                K::Native::try_from(1u8).unwrap(),
2958            ];
2959            let keys = PrimitiveArray::<K>::from_iter_values(keys);
2960            let values = FixedSizeBinaryArray::try_from_iter(
2961                vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
2962            )
2963            .unwrap();
2964
2965            let data = DictionaryArray::<K>::new(keys, Arc::new(values));
2966            let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
2967            roundtrip(batch, None);
2968        }
2969
2970        test_fixed_size_binary_in_dict_inner::<UInt8Type>();
2971        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2972        test_fixed_size_binary_in_dict_inner::<UInt32Type>();
2973        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2974        test_fixed_size_binary_in_dict_inner::<Int8Type>();
2975        test_fixed_size_binary_in_dict_inner::<Int16Type>();
2976        test_fixed_size_binary_in_dict_inner::<Int32Type>();
2977        test_fixed_size_binary_in_dict_inner::<Int64Type>();
2978    }
2979
2980    #[test]
2981    fn test_empty_dict() {
2982        let struct_fields = Fields::from(vec![Field::new(
2983            "dict",
2984            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2985            false,
2986        )]);
2987
2988        let schema = Schema::new(vec![Field::new_struct(
2989            "struct",
2990            struct_fields.clone(),
2991            true,
2992        )]);
2993        let dictionary = Arc::new(DictionaryArray::new(
2994            Int32Array::new_null(5),
2995            Arc::new(StringArray::new_null(0)),
2996        ));
2997
2998        let s = StructArray::new(
2999            struct_fields,
3000            vec![dictionary],
3001            Some(NullBuffer::new_null(5)),
3002        );
3003
3004        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
3005        roundtrip(batch, None);
3006    }
3007    #[test]
3008    fn arrow_writer_page_size() {
3009        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3010
3011        let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3012
3013        // Generate an array of 10 unique 10 character string
3014        for i in 0..10 {
3015            let value = i
3016                .to_string()
3017                .repeat(10)
3018                .chars()
3019                .take(10)
3020                .collect::<String>();
3021
3022            builder.append_value(value);
3023        }
3024
3025        let array = Arc::new(builder.finish());
3026
3027        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3028
3029        let file = tempfile::tempfile().unwrap();
3030
3031        // Set everything very low so we fallback to PLAIN encoding after the first row
3032        let props = WriterProperties::builder()
3033            .set_data_page_size_limit(1)
3034            .set_dictionary_page_size_limit(1)
3035            .set_write_batch_size(1)
3036            .build();
3037
3038        let mut writer =
3039            ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3040                .expect("Unable to write file");
3041        writer.write(&batch).unwrap();
3042        writer.close().unwrap();
3043
3044        let options = ReadOptionsBuilder::new().with_page_index().build();
3045        let reader =
3046            SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3047
3048        let column = reader.metadata().row_group(0).columns();
3049
3050        assert_eq!(column.len(), 1);
3051
3052        // We should write one row before falling back to PLAIN encoding so there should still be a
3053        // dictionary page.
3054        assert!(
3055            column[0].dictionary_page_offset().is_some(),
3056            "Expected a dictionary page"
3057        );
3058
3059        assert!(reader.metadata().offset_index().is_some());
3060        let offset_indexes = &reader.metadata().offset_index().unwrap()[0];
3061
3062        let page_locations = offset_indexes[0].page_locations.clone();
3063
3064        // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
3065        // so we expect one dictionary encoded page and then a page per row thereafter.
3066        assert_eq!(
3067            page_locations.len(),
3068            10,
3069            "Expected 10 pages but got {page_locations:#?}"
3070        );
3071    }
3072
3073    #[test]
3074    fn arrow_writer_float_nans() {
3075        let f16_field = Field::new("a", DataType::Float16, false);
3076        let f32_field = Field::new("b", DataType::Float32, false);
3077        let f64_field = Field::new("c", DataType::Float64, false);
3078        let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3079
3080        let f16_values = (0..MEDIUM_SIZE)
3081            .map(|i| {
3082                Some(if i % 2 == 0 {
3083                    f16::NAN
3084                } else {
3085                    f16::from_f32(i as f32)
3086                })
3087            })
3088            .collect::<Float16Array>();
3089
3090        let f32_values = (0..MEDIUM_SIZE)
3091            .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3092            .collect::<Float32Array>();
3093
3094        let f64_values = (0..MEDIUM_SIZE)
3095            .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3096            .collect::<Float64Array>();
3097
3098        let batch = RecordBatch::try_new(
3099            Arc::new(schema),
3100            vec![
3101                Arc::new(f16_values),
3102                Arc::new(f32_values),
3103                Arc::new(f64_values),
3104            ],
3105        )
3106        .unwrap();
3107
3108        roundtrip(batch, None);
3109    }
3110
3111    const SMALL_SIZE: usize = 7;
3112    const MEDIUM_SIZE: usize = 63;
3113
3114    // Write the batch to parquet and read it back out, ensuring
3115    // that what comes out is the same as what was written in
3116    fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3117        let mut files = vec![];
3118        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3119            let mut props = WriterProperties::builder().set_writer_version(version);
3120
3121            if let Some(size) = max_row_group_size {
3122                props = props.set_max_row_group_row_count(Some(size))
3123            }
3124
3125            let props = props.build();
3126            files.push(roundtrip_opts(&expected_batch, props))
3127        }
3128        files
3129    }
3130
3131    // Round trip the specified record batch with the specified writer properties,
3132    // to an in-memory file, and validate the arrays using the specified function.
3133    // Returns the in-memory file.
3134    fn roundtrip_opts_with_array_validation<F>(
3135        expected_batch: &RecordBatch,
3136        props: WriterProperties,
3137        validate: F,
3138    ) -> Bytes
3139    where
3140        F: Fn(&ArrayData, &ArrayData),
3141    {
3142        let mut file = vec![];
3143
3144        let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3145            .expect("Unable to write file");
3146        writer.write(expected_batch).unwrap();
3147        writer.close().unwrap();
3148
3149        let file = Bytes::from(file);
3150        let mut record_batch_reader =
3151            ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3152
3153        let actual_batch = record_batch_reader
3154            .next()
3155            .expect("No batch found")
3156            .expect("Unable to get batch");
3157
3158        assert_eq!(expected_batch.schema(), actual_batch.schema());
3159        assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3160        assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3161        for i in 0..expected_batch.num_columns() {
3162            let expected_data = expected_batch.column(i).to_data();
3163            let actual_data = actual_batch.column(i).to_data();
3164            validate(&expected_data, &actual_data);
3165        }
3166
3167        file
3168    }
3169
3170    fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3171        roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3172            a.validate_full().expect("valid expected data");
3173            b.validate_full().expect("valid actual data");
3174            assert_eq!(a, b)
3175        })
3176    }
3177
3178    struct RoundTripOptions {
3179        values: ArrayRef,
3180        schema: SchemaRef,
3181        bloom_filter: bool,
3182        bloom_filter_ndv: Option<u64>,
3183        bloom_filter_position: BloomFilterPosition,
3184    }
3185
3186    impl RoundTripOptions {
3187        fn new(values: ArrayRef, nullable: bool) -> Self {
3188            let data_type = values.data_type().clone();
3189            let schema = Schema::new(vec![Field::new("col", data_type, nullable)]);
3190            Self {
3191                values,
3192                schema: Arc::new(schema),
3193                bloom_filter: false,
3194                bloom_filter_ndv: None,
3195                bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3196            }
3197        }
3198    }
3199
3200    fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec<Bytes> {
3201        one_column_roundtrip_with_options(RoundTripOptions::new(values, nullable))
3202    }
3203
3204    fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef) -> Vec<Bytes> {
3205        let mut options = RoundTripOptions::new(values, false);
3206        options.schema = schema;
3207        one_column_roundtrip_with_options(options)
3208    }
3209
3210    fn one_column_roundtrip_with_options(options: RoundTripOptions) -> Vec<Bytes> {
3211        let RoundTripOptions {
3212            values,
3213            schema,
3214            bloom_filter,
3215            bloom_filter_ndv,
3216            bloom_filter_position,
3217        } = options;
3218
3219        let encodings = match values.data_type() {
3220            DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3221                vec![
3222                    Encoding::PLAIN,
3223                    Encoding::DELTA_BYTE_ARRAY,
3224                    Encoding::DELTA_LENGTH_BYTE_ARRAY,
3225                ]
3226            }
3227            DataType::Int64
3228            | DataType::Int32
3229            | DataType::Int16
3230            | DataType::Int8
3231            | DataType::UInt64
3232            | DataType::UInt32
3233            | DataType::UInt16
3234            | DataType::UInt8 => vec![
3235                Encoding::PLAIN,
3236                Encoding::DELTA_BINARY_PACKED,
3237                Encoding::BYTE_STREAM_SPLIT,
3238            ],
3239            DataType::Float32 | DataType::Float64 => {
3240                vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
3241            }
3242            _ => vec![Encoding::PLAIN],
3243        };
3244
3245        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3246
3247        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3248
3249        let mut files = vec![];
3250        for dictionary_size in [0, 1, 1024] {
3251            for encoding in &encodings {
3252                for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3253                    for row_group_size in row_group_sizes {
3254                        let mut builder = WriterProperties::builder()
3255                            .set_writer_version(version)
3256                            .set_max_row_group_row_count(Some(row_group_size))
3257                            .set_dictionary_enabled(dictionary_size != 0)
3258                            .set_dictionary_page_size_limit(dictionary_size.max(1))
3259                            .set_encoding(*encoding)
3260                            .set_bloom_filter_enabled(bloom_filter)
3261                            .set_bloom_filter_position(bloom_filter_position);
3262                        if let Some(ndv) = bloom_filter_ndv {
3263                            builder = builder.set_bloom_filter_max_ndv(ndv);
3264                        }
3265                        let props = builder.build();
3266
3267                        files.push(roundtrip_opts(&expected_batch, props))
3268                    }
3269                }
3270            }
3271        }
3272        files
3273    }
3274
3275    fn values_required<A, I>(iter: I) -> Vec<Bytes>
3276    where
3277        A: From<Vec<I::Item>> + Array + 'static,
3278        I: IntoIterator,
3279    {
3280        let raw_values: Vec<_> = iter.into_iter().collect();
3281        let values = Arc::new(A::from(raw_values));
3282        one_column_roundtrip(values, false)
3283    }
3284
3285    fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3286    where
3287        A: From<Vec<Option<I::Item>>> + Array + 'static,
3288        I: IntoIterator,
3289    {
3290        let optional_raw_values: Vec<_> = iter
3291            .into_iter()
3292            .enumerate()
3293            .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3294            .collect();
3295        let optional_values = Arc::new(A::from(optional_raw_values));
3296        one_column_roundtrip(optional_values, true)
3297    }
3298
3299    fn required_and_optional<A, I>(iter: I)
3300    where
3301        A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3302        I: IntoIterator + Clone,
3303    {
3304        values_required::<A, I>(iter.clone());
3305        values_optional::<A, I>(iter);
3306    }
3307
3308    fn check_bloom_filter<T: AsBytes>(
3309        files: Vec<Bytes>,
3310        file_column: String,
3311        positive_values: Vec<T>,
3312        negative_values: Vec<T>,
3313    ) {
3314        files.into_iter().take(1).for_each(|file| {
3315            let file_reader = SerializedFileReader::new_with_options(
3316                file,
3317                ReadOptionsBuilder::new()
3318                    .with_reader_properties(
3319                        ReaderProperties::builder()
3320                            .set_read_bloom_filter(true)
3321                            .build(),
3322                    )
3323                    .build(),
3324            )
3325            .expect("Unable to open file as Parquet");
3326            let metadata = file_reader.metadata();
3327
3328            // Gets bloom filters from all row groups.
3329            let mut bloom_filters: Vec<_> = vec![];
3330            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3331                if let Some((column_index, _)) = row_group
3332                    .columns()
3333                    .iter()
3334                    .enumerate()
3335                    .find(|(_, column)| column.column_path().string() == file_column)
3336                {
3337                    let row_group_reader = file_reader
3338                        .get_row_group(ri)
3339                        .expect("Unable to read row group");
3340                    if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3341                        bloom_filters.push(sbbf.clone());
3342                    } else {
3343                        panic!("No bloom filter for column named {file_column} found");
3344                    }
3345                } else {
3346                    panic!("No column named {file_column} found");
3347                }
3348            }
3349
3350            positive_values.iter().for_each(|value| {
3351                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3352                assert!(
3353                    found.is_some(),
3354                    "{}",
3355                    format!("Value {:?} should be in bloom filter", value.as_bytes())
3356                );
3357            });
3358
3359            negative_values.iter().for_each(|value| {
3360                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3361                assert!(
3362                    found.is_none(),
3363                    "{}",
3364                    format!("Value {:?} should not be in bloom filter", value.as_bytes())
3365                );
3366            });
3367        });
3368    }
3369
3370    #[test]
3371    fn all_null_primitive_single_column() {
3372        let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3373        one_column_roundtrip(values, true);
3374    }
3375    #[test]
3376    fn null_single_column() {
3377        let values = Arc::new(NullArray::new(SMALL_SIZE));
3378        one_column_roundtrip(values, true);
3379        // null arrays are always nullable, a test with non-nullable nulls fails
3380    }
3381
3382    #[test]
3383    fn bool_single_column() {
3384        required_and_optional::<BooleanArray, _>(
3385            [true, false].iter().cycle().copied().take(SMALL_SIZE),
3386        );
3387    }
3388
3389    #[test]
3390    fn bool_large_single_column() {
3391        let values = Arc::new(
3392            [None, Some(true), Some(false)]
3393                .iter()
3394                .cycle()
3395                .copied()
3396                .take(200_000)
3397                .collect::<BooleanArray>(),
3398        );
3399        let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3400        let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3401        let file = tempfile::tempfile().unwrap();
3402
3403        let mut writer =
3404            ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3405                .expect("Unable to write file");
3406        writer.write(&expected_batch).unwrap();
3407        writer.close().unwrap();
3408    }
3409
3410    #[test]
3411    fn check_page_offset_index_with_nan() {
3412        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3413        let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3414        let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3415
3416        let mut out = Vec::with_capacity(1024);
3417        let mut writer =
3418            ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3419        writer.write(&batch).unwrap();
3420        let file_meta_data = writer.close().unwrap();
3421        for row_group in file_meta_data.row_groups() {
3422            for column in row_group.columns() {
3423                assert!(column.offset_index_offset().is_some());
3424                assert!(column.offset_index_length().is_some());
3425                assert!(column.column_index_offset().is_some());
3426                assert!(column.column_index_length().is_some());
3427            }
3428        }
3429        assert!(file_meta_data.column_index().is_some());
3430        if let Some(col_indexes) = file_meta_data.column_index() {
3431            for rg_idx in col_indexes {
3432                for idx in rg_idx {
3433                    assert!(idx.nan_counts().is_some());
3434                    let float_idx = match idx {
3435                        ColumnIndexMetaData::DOUBLE(idx) => idx,
3436                        _ => panic!("expected double statistics"),
3437                    };
3438                    for i in 0..idx.num_pages() as usize {
3439                        assert_eq!(float_idx.nan_count(i), Some(10));
3440                        assert_eq!(
3441                            f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3442                            Ordering::Equal
3443                        );
3444                        assert_eq!(
3445                            f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3446                            Ordering::Equal
3447                        );
3448                    }
3449                }
3450            }
3451        }
3452    }
3453
3454    #[test]
3455    fn check_page_offset_index_with_mixed_nan() {
3456        let schema = Arc::new(Schema::new(vec![Field::new(
3457            "col",
3458            DataType::Float64,
3459            true,
3460        )]));
3461
3462        let mut out = Vec::with_capacity(1024);
3463        let props = WriterProperties::builder()
3464            .set_data_page_row_count_limit(10)
3465            .build();
3466        let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3467            .expect("Unable to write file");
3468
3469        // write a page of all NaN (since batch min and max are NaN, global min/max are NaN)
3470        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3471        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3472        writer.write(&batch).unwrap();
3473
3474        // write a page of all -NaN (batch min/max is -NaN, should update global min to -NaN)
3475        let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3476        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3477        writer.write(&batch).unwrap();
3478
3479        // write a page of all 0 (non-NaN should override global min/max, now 0/0)
3480        let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3481        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3482        writer.write(&batch).unwrap();
3483
3484        // write a mixed page (should now have min -1, max 1)
3485        let values = Arc::new(Float64Array::from(vec![
3486            -1.0,
3487            0.0,
3488            f64::NAN,
3489            -f64::NAN,
3490            1.0,
3491        ]));
3492        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3493        writer.write(&batch).unwrap();
3494
3495        let file_meta_data = writer.close().unwrap();
3496
3497        // check the column chunk stats are correct
3498        let col_stats = file_meta_data
3499            .row_group(0)
3500            .column(0)
3501            .statistics()
3502            .expect("missing column chunk statistics");
3503
3504        assert_eq!(col_stats.nan_count_opt(), Some(22));
3505        assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3506        assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3507
3508        assert!(file_meta_data.column_index().is_some());
3509        let col_idx = &file_meta_data.column_index().as_ref().unwrap()[0][0];
3510        assert_eq!(col_idx.num_pages(), 4);
3511
3512        // test each page
3513        let float_idx = match col_idx {
3514            ColumnIndexMetaData::DOUBLE(idx) => idx,
3515            _ => panic!("expected double statistics"),
3516        };
3517
3518        assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3519        assert_eq!(
3520            f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3521            Ordering::Equal
3522        );
3523        assert_eq!(
3524            f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3525            Ordering::Equal
3526        );
3527        assert_eq!(
3528            (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3529            Ordering::Equal
3530        );
3531        assert_eq!(
3532            (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3533            Ordering::Equal
3534        );
3535        assert_eq!(float_idx.min_value(2), Some(&0.0));
3536        assert_eq!(float_idx.max_value(2), Some(&0.0));
3537        assert_eq!(float_idx.min_value(3), Some(&-1.0));
3538        assert_eq!(float_idx.max_value(3), Some(&1.0));
3539    }
3540
3541    #[test]
3542    fn i8_single_column() {
3543        required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3544    }
3545
3546    #[test]
3547    fn i16_single_column() {
3548        required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3549    }
3550
3551    #[test]
3552    fn i32_single_column() {
3553        required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3554    }
3555
3556    #[test]
3557    fn i64_single_column() {
3558        required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3559    }
3560
3561    #[test]
3562    fn u8_single_column() {
3563        required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3564    }
3565
3566    #[test]
3567    fn u16_single_column() {
3568        required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3569    }
3570
3571    #[test]
3572    fn u32_single_column() {
3573        required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3574    }
3575
3576    #[test]
3577    fn u64_single_column() {
3578        required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3579    }
3580
3581    #[test]
3582    fn f32_single_column() {
3583        required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3584    }
3585
3586    #[test]
3587    fn f64_single_column() {
3588        required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3589    }
3590
3591    // The timestamp array types don't implement From<Vec<T>> because they need the timezone
3592    // argument, and they also doesn't support building from a Vec<Option<T>>, so call
3593    // one_column_roundtrip manually instead of calling required_and_optional for these tests.
3594
3595    #[test]
3596    fn timestamp_second_single_column() {
3597        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3598        let values = Arc::new(TimestampSecondArray::from(raw_values));
3599
3600        one_column_roundtrip(values, false);
3601    }
3602
3603    #[test]
3604    fn timestamp_millisecond_single_column() {
3605        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3606        let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3607
3608        one_column_roundtrip(values, false);
3609    }
3610
3611    #[test]
3612    fn timestamp_microsecond_single_column() {
3613        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3614        let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3615
3616        one_column_roundtrip(values, false);
3617    }
3618
3619    #[test]
3620    fn timestamp_nanosecond_single_column() {
3621        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3622        let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3623
3624        one_column_roundtrip(values, false);
3625    }
3626
3627    #[test]
3628    fn date32_single_column() {
3629        required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3630    }
3631
3632    #[test]
3633    fn date64_single_column() {
3634        // Date64 must be a multiple of 86400000, see ARROW-10925
3635        required_and_optional::<Date64Array, _>(
3636            (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3637        );
3638    }
3639
3640    #[test]
3641    fn time32_second_single_column() {
3642        required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3643    }
3644
3645    #[test]
3646    fn time32_millisecond_single_column() {
3647        required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3648    }
3649
3650    #[test]
3651    fn time64_microsecond_single_column() {
3652        required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3653    }
3654
3655    #[test]
3656    fn time64_nanosecond_single_column() {
3657        required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3658    }
3659
3660    #[test]
3661    fn duration_second_single_column() {
3662        required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3663    }
3664
3665    #[test]
3666    fn duration_millisecond_single_column() {
3667        required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3668    }
3669
3670    #[test]
3671    fn duration_microsecond_single_column() {
3672        required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3673    }
3674
3675    #[test]
3676    fn duration_nanosecond_single_column() {
3677        required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3678    }
3679
3680    #[test]
3681    fn interval_year_month_single_column() {
3682        required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3683    }
3684
3685    #[test]
3686    fn interval_day_time_single_column() {
3687        required_and_optional::<IntervalDayTimeArray, _>(vec![
3688            IntervalDayTime::new(0, 1),
3689            IntervalDayTime::new(0, 3),
3690            IntervalDayTime::new(3, -2),
3691            IntervalDayTime::new(-200, 4),
3692        ]);
3693    }
3694
3695    #[test]
3696    #[should_panic(
3697        expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3698    )]
3699    fn interval_month_day_nano_single_column() {
3700        required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3701            IntervalMonthDayNano::new(0, 1, 5),
3702            IntervalMonthDayNano::new(0, 3, 2),
3703            IntervalMonthDayNano::new(3, -2, -5),
3704            IntervalMonthDayNano::new(-200, 4, -1),
3705        ]);
3706    }
3707
3708    #[test]
3709    fn binary_single_column() {
3710        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3711        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3712        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3713
3714        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3715        values_required::<BinaryArray, _>(many_vecs_iter);
3716    }
3717
3718    #[test]
3719    fn binary_view_single_column() {
3720        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3721        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3722        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3723
3724        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3725        values_required::<BinaryViewArray, _>(many_vecs_iter);
3726    }
3727
3728    #[test]
3729    fn i32_column_bloom_filter_at_end() {
3730        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3731        let mut options = RoundTripOptions::new(array, false);
3732        options.bloom_filter = true;
3733        options.bloom_filter_position = BloomFilterPosition::End;
3734
3735        let files = one_column_roundtrip_with_options(options);
3736        check_bloom_filter(
3737            files,
3738            "col".to_string(),
3739            (0..SMALL_SIZE as i32).collect(),
3740            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3741        );
3742    }
3743
3744    #[test]
3745    fn i32_column_bloom_filter() {
3746        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3747        let mut options = RoundTripOptions::new(array, false);
3748        options.bloom_filter = true;
3749
3750        let files = one_column_roundtrip_with_options(options);
3751        check_bloom_filter(
3752            files,
3753            "col".to_string(),
3754            (0..SMALL_SIZE as i32).collect(),
3755            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3756        );
3757    }
3758
3759    /// Test that bloom filter folding produces correct results even when
3760    /// the configured NDV differs significantly from actual NDV.
3761    /// A large NDV means a larger initial filter that gets folded down;
3762    /// a small NDV means a smaller initial filter.
3763    #[test]
3764    fn i32_column_bloom_filter_fixed_ndv() {
3765        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3766
3767        // NDV much larger than actual distinct values — tests folding a large filter down
3768        let mut options = RoundTripOptions::new(array.clone(), false);
3769        options.bloom_filter = true;
3770        options.bloom_filter_ndv = Some(1_000_000);
3771
3772        let files = one_column_roundtrip_with_options(options);
3773        check_bloom_filter(
3774            files,
3775            "col".to_string(),
3776            (0..SMALL_SIZE as i32).collect(),
3777            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3778        );
3779
3780        // NDV smaller than actual distinct values — tests the underestimate path
3781        let mut options = RoundTripOptions::new(array, false);
3782        options.bloom_filter = true;
3783        options.bloom_filter_ndv = Some(3);
3784
3785        let files = one_column_roundtrip_with_options(options);
3786        check_bloom_filter(
3787            files,
3788            "col".to_string(),
3789            (0..SMALL_SIZE as i32).collect(),
3790            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3791        );
3792    }
3793
3794    #[test]
3795    fn binary_column_bloom_filter() {
3796        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3797        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3798        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3799
3800        let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
3801        let mut options = RoundTripOptions::new(array, false);
3802        options.bloom_filter = true;
3803
3804        let files = one_column_roundtrip_with_options(options);
3805        check_bloom_filter(
3806            files,
3807            "col".to_string(),
3808            many_vecs,
3809            vec![vec![(SMALL_SIZE + 1) as u8]],
3810        );
3811    }
3812
3813    #[test]
3814    fn empty_string_null_column_bloom_filter() {
3815        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3816        let raw_strs = raw_values.iter().map(|s| s.as_str());
3817
3818        let array = Arc::new(StringArray::from_iter_values(raw_strs));
3819        let mut options = RoundTripOptions::new(array, false);
3820        options.bloom_filter = true;
3821
3822        let files = one_column_roundtrip_with_options(options);
3823
3824        let optional_raw_values: Vec<_> = raw_values
3825            .iter()
3826            .enumerate()
3827            .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
3828            .collect();
3829        // For null slots, empty string should not be in bloom filter.
3830        check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
3831    }
3832
3833    #[test]
3834    fn large_binary_single_column() {
3835        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3836        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3837        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3838
3839        // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3840        values_required::<LargeBinaryArray, _>(many_vecs_iter);
3841    }
3842
3843    #[test]
3844    fn fixed_size_binary_single_column() {
3845        let mut builder = FixedSizeBinaryBuilder::new(4);
3846        builder.append_value(b"0123").unwrap();
3847        builder.append_null();
3848        builder.append_value(b"8910").unwrap();
3849        builder.append_value(b"1112").unwrap();
3850        let array = Arc::new(builder.finish());
3851
3852        one_column_roundtrip(array, true);
3853    }
3854
3855    #[test]
3856    fn string_single_column() {
3857        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3858        let raw_strs = raw_values.iter().map(|s| s.as_str());
3859
3860        required_and_optional::<StringArray, _>(raw_strs);
3861    }
3862
3863    #[test]
3864    fn large_string_single_column() {
3865        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3866        let raw_strs = raw_values.iter().map(|s| s.as_str());
3867
3868        required_and_optional::<LargeStringArray, _>(raw_strs);
3869    }
3870
3871    #[test]
3872    fn string_view_single_column() {
3873        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3874        let raw_strs = raw_values.iter().map(|s| s.as_str());
3875
3876        required_and_optional::<StringViewArray, _>(raw_strs);
3877    }
3878
3879    #[test]
3880    fn null_list_single_column() {
3881        let null_field = Field::new_list_field(DataType::Null, true);
3882        let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
3883
3884        let schema = Schema::new(vec![list_field]);
3885
3886        // Build [[], null, [null, null]]
3887        let a_values = NullArray::new(2);
3888        let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
3889        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3890            DataType::Null,
3891            true,
3892        ))))
3893        .len(3)
3894        .add_buffer(a_value_offsets)
3895        .null_bit_buffer(Some(Buffer::from([0b00000101])))
3896        .add_child_data(a_values.into_data())
3897        .build()
3898        .unwrap();
3899
3900        let a = ListArray::from(a_list_data);
3901
3902        assert!(a.is_valid(0));
3903        assert!(!a.is_valid(1));
3904        assert!(a.is_valid(2));
3905
3906        assert_eq!(a.value(0).len(), 0);
3907        assert_eq!(a.value(2).len(), 2);
3908        assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
3909
3910        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3911        roundtrip(batch, None);
3912    }
3913
3914    #[test]
3915    fn list_single_column() {
3916        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3917        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
3918        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3919            DataType::Int32,
3920            false,
3921        ))))
3922        .len(5)
3923        .add_buffer(a_value_offsets)
3924        .null_bit_buffer(Some(Buffer::from([0b00011011])))
3925        .add_child_data(a_values.into_data())
3926        .build()
3927        .unwrap();
3928
3929        assert_eq!(a_list_data.null_count(), 1);
3930
3931        let a = ListArray::from(a_list_data);
3932        let values = Arc::new(a);
3933
3934        one_column_roundtrip(values, true);
3935    }
3936
3937    #[test]
3938    fn large_list_single_column() {
3939        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3940        let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
3941        let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
3942            "large_item",
3943            DataType::Int32,
3944            true,
3945        ))))
3946        .len(5)
3947        .add_buffer(a_value_offsets)
3948        .add_child_data(a_values.into_data())
3949        .null_bit_buffer(Some(Buffer::from([0b00011011])))
3950        .build()
3951        .unwrap();
3952
3953        // I think this setup is incorrect because this should pass
3954        assert_eq!(a_list_data.null_count(), 1);
3955
3956        let a = LargeListArray::from(a_list_data);
3957        let values = Arc::new(a);
3958
3959        one_column_roundtrip(values, true);
3960    }
3961
3962    #[test]
3963    fn list_nested_nulls() {
3964        use arrow::datatypes::Int32Type;
3965        let data = vec![
3966            Some(vec![Some(1)]),
3967            Some(vec![Some(2), Some(3)]),
3968            None,
3969            Some(vec![Some(4), Some(5), None]),
3970            Some(vec![None]),
3971            Some(vec![Some(6), Some(7)]),
3972        ];
3973
3974        let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
3975        one_column_roundtrip(Arc::new(list), true);
3976
3977        let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
3978        one_column_roundtrip(Arc::new(list), true);
3979    }
3980
3981    #[test]
3982    fn list_utf8_view_selective_padding_roundtrip() {
3983        let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
3984        let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
3985        builder.values().append_value("a");
3986        builder.values().append_null();
3987        builder.append(true);
3988        // The null parent list covers selective padding dropping values below
3989        // the list definition level while preserving the preceding item null.
3990        builder.append(false);
3991        // The long string covers the non-inlined Utf8View buffer path.
3992        builder.values().append_value("large payload over 12 bytes");
3993        builder.append(true);
3994
3995        one_column_roundtrip(Arc::new(builder.finish()), true);
3996    }
3997
3998    #[test]
3999    fn struct_single_column() {
4000        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4001        let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4002        let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4003
4004        let values = Arc::new(s);
4005        one_column_roundtrip(values, false);
4006    }
4007
4008    #[test]
4009    fn list_and_map_coerced_names() {
4010        // Create map and list with non-Parquet naming
4011        let list_field =
4012            Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4013        let map_field = Field::new_map(
4014            "my_map",
4015            "my_entries",
4016            Field::new("my_keys", DataType::Int32, false),
4017            Field::new("my_values", DataType::Int32, true),
4018            false,
4019            true,
4020        );
4021
4022        let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4023        let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4024
4025        let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4026
4027        // Write data to Parquet but coerce names to match spec
4028        let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4029        let file = tempfile::tempfile().unwrap();
4030        let mut writer =
4031            ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4032
4033        let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4034        writer.write(&batch).unwrap();
4035        let file_metadata = writer.close().unwrap();
4036
4037        let schema = file_metadata.file_metadata().schema();
4038        // Coerced name of "item" should be "element"
4039        let list_field = &schema.get_fields()[0].get_fields()[0];
4040        assert_eq!(list_field.get_fields()[0].name(), "element");
4041
4042        let map_field = &schema.get_fields()[1].get_fields()[0];
4043        // Coerced name of "entries" should be "key_value"
4044        assert_eq!(map_field.name(), "key_value");
4045        // Coerced name of "my_keys" should be "key"
4046        assert_eq!(map_field.get_fields()[0].name(), "key");
4047        // Coerced name of "my_values" should be "value"
4048        assert_eq!(map_field.get_fields()[1].name(), "value");
4049
4050        // Double check schema after reading from the file
4051        let reader = SerializedFileReader::new(file).unwrap();
4052        let file_schema = reader.metadata().file_metadata().schema();
4053        let fields = file_schema.get_fields();
4054        let list_field = &fields[0].get_fields()[0];
4055        assert_eq!(list_field.get_fields()[0].name(), "element");
4056        let map_field = &fields[1].get_fields()[0];
4057        assert_eq!(map_field.name(), "key_value");
4058        assert_eq!(map_field.get_fields()[0].name(), "key");
4059        assert_eq!(map_field.get_fields()[1].name(), "value");
4060    }
4061
4062    #[test]
4063    fn fallback_flush_data_page() {
4064        //tests if the Fallback::flush_data_page clears all buffers correctly
4065        let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4066        let values = Arc::new(StringArray::from(raw_values));
4067        let encodings = vec![
4068            Encoding::DELTA_BYTE_ARRAY,
4069            Encoding::DELTA_LENGTH_BYTE_ARRAY,
4070        ];
4071        let data_type = values.data_type().clone();
4072        let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4073        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4074
4075        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4076        let data_page_size_limit: usize = 32;
4077        let write_batch_size: usize = 16;
4078
4079        for encoding in &encodings {
4080            for row_group_size in row_group_sizes {
4081                let props = WriterProperties::builder()
4082                    .set_writer_version(WriterVersion::PARQUET_2_0)
4083                    .set_max_row_group_row_count(Some(row_group_size))
4084                    .set_dictionary_enabled(false)
4085                    .set_encoding(*encoding)
4086                    .set_data_page_size_limit(data_page_size_limit)
4087                    .set_write_batch_size(write_batch_size)
4088                    .build();
4089
4090                roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4091                    let string_array_a = StringArray::from(a.clone());
4092                    let string_array_b = StringArray::from(b.clone());
4093                    let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4094                    let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4095                    assert_eq!(
4096                        vec_a, vec_b,
4097                        "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4098                    );
4099                });
4100            }
4101        }
4102    }
4103
4104    #[test]
4105    fn arrow_writer_string_dictionary() {
4106        // define schema
4107        #[allow(deprecated)]
4108        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4109            "dictionary",
4110            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4111            true,
4112            42,
4113            true,
4114        )]));
4115
4116        // create some data
4117        let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4118            .iter()
4119            .copied()
4120            .collect();
4121
4122        // build a record batch
4123        one_column_roundtrip_with_schema(Arc::new(d), schema);
4124    }
4125
4126    #[test]
4127    fn arrow_writer_test_type_compatibility() {
4128        fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4129        where
4130            T1: Array + 'static,
4131            T2: Array + 'static,
4132        {
4133            let schema1 = Arc::new(Schema::new(vec![Field::new(
4134                "a",
4135                array1.data_type().clone(),
4136                false,
4137            )]));
4138
4139            let file = tempfile().unwrap();
4140            let mut writer =
4141                ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4142
4143            let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4144            writer.write(&rb1).unwrap();
4145
4146            let schema2 = Arc::new(Schema::new(vec![Field::new(
4147                "a",
4148                array2.data_type().clone(),
4149                false,
4150            )]));
4151            let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4152            writer.write(&rb2).unwrap();
4153
4154            writer.close().unwrap();
4155
4156            let mut record_batch_reader =
4157                ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4158            let actual_batch = record_batch_reader.next().unwrap().unwrap();
4159
4160            let expected_batch =
4161                RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4162            assert_eq!(actual_batch, expected_batch);
4163        }
4164
4165        // check compatibility between native and dictionaries
4166
4167        ensure_compatible_write(
4168            DictionaryArray::new(
4169                UInt8Array::from_iter_values(vec![0]),
4170                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4171            ),
4172            StringArray::from_iter_values(vec!["barquet"]),
4173            DictionaryArray::new(
4174                UInt8Array::from_iter_values(vec![0, 1]),
4175                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4176            ),
4177        );
4178
4179        ensure_compatible_write(
4180            StringArray::from_iter_values(vec!["parquet"]),
4181            DictionaryArray::new(
4182                UInt8Array::from_iter_values(vec![0]),
4183                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4184            ),
4185            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4186        );
4187
4188        // check compatibility between dictionaries with different key types
4189
4190        ensure_compatible_write(
4191            DictionaryArray::new(
4192                UInt8Array::from_iter_values(vec![0]),
4193                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4194            ),
4195            DictionaryArray::new(
4196                UInt16Array::from_iter_values(vec![0]),
4197                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4198            ),
4199            DictionaryArray::new(
4200                UInt8Array::from_iter_values(vec![0, 1]),
4201                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4202            ),
4203        );
4204
4205        // check compatibility between dictionaries with different value types
4206        ensure_compatible_write(
4207            DictionaryArray::new(
4208                UInt8Array::from_iter_values(vec![0]),
4209                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4210            ),
4211            DictionaryArray::new(
4212                UInt8Array::from_iter_values(vec![0]),
4213                Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4214            ),
4215            DictionaryArray::new(
4216                UInt8Array::from_iter_values(vec![0, 1]),
4217                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4218            ),
4219        );
4220
4221        // check compatibility between a dictionary and a native array with a different type
4222        ensure_compatible_write(
4223            DictionaryArray::new(
4224                UInt8Array::from_iter_values(vec![0]),
4225                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4226            ),
4227            LargeStringArray::from_iter_values(vec!["barquet"]),
4228            DictionaryArray::new(
4229                UInt8Array::from_iter_values(vec![0, 1]),
4230                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4231            ),
4232        );
4233
4234        // check compatibility for string types
4235
4236        ensure_compatible_write(
4237            StringArray::from_iter_values(vec!["parquet"]),
4238            LargeStringArray::from_iter_values(vec!["barquet"]),
4239            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4240        );
4241
4242        ensure_compatible_write(
4243            LargeStringArray::from_iter_values(vec!["parquet"]),
4244            StringArray::from_iter_values(vec!["barquet"]),
4245            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4246        );
4247
4248        ensure_compatible_write(
4249            StringArray::from_iter_values(vec!["parquet"]),
4250            StringViewArray::from_iter_values(vec!["barquet"]),
4251            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4252        );
4253
4254        ensure_compatible_write(
4255            StringViewArray::from_iter_values(vec!["parquet"]),
4256            StringArray::from_iter_values(vec!["barquet"]),
4257            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4258        );
4259
4260        ensure_compatible_write(
4261            LargeStringArray::from_iter_values(vec!["parquet"]),
4262            StringViewArray::from_iter_values(vec!["barquet"]),
4263            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4264        );
4265
4266        ensure_compatible_write(
4267            StringViewArray::from_iter_values(vec!["parquet"]),
4268            LargeStringArray::from_iter_values(vec!["barquet"]),
4269            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4270        );
4271
4272        // check compatibility for binary types
4273
4274        ensure_compatible_write(
4275            BinaryArray::from_iter_values(vec![b"parquet"]),
4276            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4277            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4278        );
4279
4280        ensure_compatible_write(
4281            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4282            BinaryArray::from_iter_values(vec![b"barquet"]),
4283            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4284        );
4285
4286        ensure_compatible_write(
4287            BinaryArray::from_iter_values(vec![b"parquet"]),
4288            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4289            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4290        );
4291
4292        ensure_compatible_write(
4293            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4294            BinaryArray::from_iter_values(vec![b"barquet"]),
4295            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4296        );
4297
4298        ensure_compatible_write(
4299            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4300            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4301            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4302        );
4303
4304        ensure_compatible_write(
4305            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4306            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4307            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4308        );
4309
4310        // check compatibility for list types
4311
4312        let list_field_metadata = HashMap::from_iter(vec![(
4313            PARQUET_FIELD_ID_META_KEY.to_string(),
4314            "1".to_string(),
4315        )]);
4316        let list_field = Field::new_list_field(DataType::Int32, false);
4317
4318        let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4319        let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4320
4321        let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4322        let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4323
4324        let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4325        let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4326
4327        ensure_compatible_write(
4328            // when the initial schema has the metadata ...
4329            ListArray::try_new(
4330                Arc::new(
4331                    list_field
4332                        .clone()
4333                        .with_metadata(list_field_metadata.clone()),
4334                ),
4335                offsets1,
4336                values1,
4337                None,
4338            )
4339            .unwrap(),
4340            // ... and some intermediate schema doesn't have the metadata
4341            ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4342            // ... the write will still go through, and the resulting schema will inherit the initial metadata
4343            ListArray::try_new(
4344                Arc::new(
4345                    list_field
4346                        .clone()
4347                        .with_metadata(list_field_metadata.clone()),
4348                ),
4349                offsets_expected,
4350                values_expected,
4351                None,
4352            )
4353            .unwrap(),
4354        );
4355    }
4356
4357    #[test]
4358    fn arrow_writer_primitive_dictionary() {
4359        // define schema
4360        #[allow(deprecated)]
4361        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4362            "dictionary",
4363            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4364            true,
4365            42,
4366            true,
4367        )]));
4368
4369        // create some data
4370        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4371        builder.append(12345678).unwrap();
4372        builder.append_null();
4373        builder.append(22345678).unwrap();
4374        builder.append(12345678).unwrap();
4375        let d = builder.finish();
4376
4377        one_column_roundtrip_with_schema(Arc::new(d), schema);
4378    }
4379
4380    #[test]
4381    fn arrow_writer_decimal32_dictionary() {
4382        let integers = vec![12345, 56789, 34567];
4383
4384        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4385
4386        let values = Decimal32Array::from(integers.clone())
4387            .with_precision_and_scale(5, 2)
4388            .unwrap();
4389
4390        let array = DictionaryArray::new(keys, Arc::new(values));
4391        one_column_roundtrip(Arc::new(array.clone()), true);
4392
4393        let values = Decimal32Array::from(integers)
4394            .with_precision_and_scale(9, 2)
4395            .unwrap();
4396
4397        let array = array.with_values(Arc::new(values));
4398        one_column_roundtrip(Arc::new(array), true);
4399    }
4400
4401    #[test]
4402    fn arrow_writer_decimal64_dictionary() {
4403        let integers = vec![12345, 56789, 34567];
4404
4405        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4406
4407        let values = Decimal64Array::from(integers.clone())
4408            .with_precision_and_scale(5, 2)
4409            .unwrap();
4410
4411        let array = DictionaryArray::new(keys, Arc::new(values));
4412        one_column_roundtrip(Arc::new(array.clone()), true);
4413
4414        let values = Decimal64Array::from(integers)
4415            .with_precision_and_scale(12, 2)
4416            .unwrap();
4417
4418        let array = array.with_values(Arc::new(values));
4419        one_column_roundtrip(Arc::new(array), true);
4420    }
4421
4422    #[test]
4423    fn arrow_writer_decimal128_dictionary() {
4424        let integers = vec![12345, 56789, 34567];
4425
4426        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4427
4428        let values = Decimal128Array::from(integers.clone())
4429            .with_precision_and_scale(5, 2)
4430            .unwrap();
4431
4432        let array = DictionaryArray::new(keys, Arc::new(values));
4433        one_column_roundtrip(Arc::new(array.clone()), true);
4434
4435        let values = Decimal128Array::from(integers)
4436            .with_precision_and_scale(12, 2)
4437            .unwrap();
4438
4439        let array = array.with_values(Arc::new(values));
4440        one_column_roundtrip(Arc::new(array), true);
4441    }
4442
4443    #[test]
4444    fn arrow_writer_decimal256_dictionary() {
4445        let integers = vec![
4446            i256::from_i128(12345),
4447            i256::from_i128(56789),
4448            i256::from_i128(34567),
4449        ];
4450
4451        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4452
4453        let values = Decimal256Array::from(integers.clone())
4454            .with_precision_and_scale(5, 2)
4455            .unwrap();
4456
4457        let array = DictionaryArray::new(keys, Arc::new(values));
4458        one_column_roundtrip(Arc::new(array.clone()), true);
4459
4460        let values = Decimal256Array::from(integers)
4461            .with_precision_and_scale(12, 2)
4462            .unwrap();
4463
4464        let array = array.with_values(Arc::new(values));
4465        one_column_roundtrip(Arc::new(array), true);
4466    }
4467
4468    #[test]
4469    fn arrow_writer_string_dictionary_unsigned_index() {
4470        // define schema
4471        #[allow(deprecated)]
4472        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4473            "dictionary",
4474            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4475            true,
4476            42,
4477            true,
4478        )]));
4479
4480        // create some data
4481        let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4482            .iter()
4483            .copied()
4484            .collect();
4485
4486        one_column_roundtrip_with_schema(Arc::new(d), schema);
4487    }
4488
4489    #[test]
4490    fn u32_min_max() {
4491        // check values roundtrip through parquet
4492        let src = [
4493            u32::MIN,
4494            1,
4495            (i32::MAX as u32) - 1,
4496            i32::MAX as u32,
4497            (i32::MAX as u32) + 1,
4498            u32::MAX - 1,
4499            u32::MAX,
4500        ];
4501        let values = Arc::new(UInt32Array::from_iter_values(src.iter().cloned()));
4502        let files = one_column_roundtrip(values, false);
4503
4504        for file in files {
4505            // check statistics are valid
4506            let reader = SerializedFileReader::new(file).unwrap();
4507            let metadata = reader.metadata();
4508
4509            let mut row_offset = 0;
4510            for row_group in metadata.row_groups() {
4511                assert_eq!(row_group.num_columns(), 1);
4512                let column = row_group.column(0);
4513
4514                let num_values = column.num_values() as usize;
4515                let src_slice = &src[row_offset..row_offset + num_values];
4516                row_offset += column.num_values() as usize;
4517
4518                let stats = column.statistics().unwrap();
4519                if let Statistics::Int32(stats) = stats {
4520                    assert_eq!(
4521                        *stats.min_opt().unwrap() as u32,
4522                        *src_slice.iter().min().unwrap()
4523                    );
4524                    assert_eq!(
4525                        *stats.max_opt().unwrap() as u32,
4526                        *src_slice.iter().max().unwrap()
4527                    );
4528                } else {
4529                    panic!("Statistics::Int32 missing")
4530                }
4531            }
4532        }
4533    }
4534
4535    #[test]
4536    fn u64_min_max() {
4537        // check values roundtrip through parquet
4538        let src = [
4539            u64::MIN,
4540            1,
4541            (i64::MAX as u64) - 1,
4542            i64::MAX as u64,
4543            (i64::MAX as u64) + 1,
4544            u64::MAX - 1,
4545            u64::MAX,
4546        ];
4547        let values = Arc::new(UInt64Array::from_iter_values(src.iter().cloned()));
4548        let files = one_column_roundtrip(values, false);
4549
4550        for file in files {
4551            // check statistics are valid
4552            let reader = SerializedFileReader::new(file).unwrap();
4553            let metadata = reader.metadata();
4554
4555            let mut row_offset = 0;
4556            for row_group in metadata.row_groups() {
4557                assert_eq!(row_group.num_columns(), 1);
4558                let column = row_group.column(0);
4559
4560                let num_values = column.num_values() as usize;
4561                let src_slice = &src[row_offset..row_offset + num_values];
4562                row_offset += column.num_values() as usize;
4563
4564                let stats = column.statistics().unwrap();
4565                if let Statistics::Int64(stats) = stats {
4566                    assert_eq!(
4567                        *stats.min_opt().unwrap() as u64,
4568                        *src_slice.iter().min().unwrap()
4569                    );
4570                    assert_eq!(
4571                        *stats.max_opt().unwrap() as u64,
4572                        *src_slice.iter().max().unwrap()
4573                    );
4574                } else {
4575                    panic!("Statistics::Int64 missing")
4576                }
4577            }
4578        }
4579    }
4580
4581    #[test]
4582    fn statistics_null_counts_only_nulls() {
4583        // check that null-count statistics for "only NULL"-columns are correct
4584        let values = Arc::new(UInt64Array::from(vec![None, None]));
4585        let files = one_column_roundtrip(values, true);
4586
4587        for file in files {
4588            // check statistics are valid
4589            let reader = SerializedFileReader::new(file).unwrap();
4590            let metadata = reader.metadata();
4591            assert_eq!(metadata.num_row_groups(), 1);
4592            let row_group = metadata.row_group(0);
4593            assert_eq!(row_group.num_columns(), 1);
4594            let column = row_group.column(0);
4595            let stats = column.statistics().unwrap();
4596            assert_eq!(stats.null_count_opt(), Some(2));
4597        }
4598    }
4599
4600    #[test]
4601    fn test_list_of_struct_roundtrip() {
4602        // define schema
4603        let int_field = Field::new("a", DataType::Int32, true);
4604        let int_field2 = Field::new("b", DataType::Int32, true);
4605
4606        let int_builder = Int32Builder::with_capacity(10);
4607        let int_builder2 = Int32Builder::with_capacity(10);
4608
4609        let struct_builder = StructBuilder::new(
4610            vec![int_field, int_field2],
4611            vec![Box::new(int_builder), Box::new(int_builder2)],
4612        );
4613        let mut list_builder = ListBuilder::new(struct_builder);
4614
4615        // Construct the following array
4616        // [{a: 1, b: 2}], [], null, [null, null], [{a: null, b: 3}], [{a: 2, b: null}]
4617
4618        // [{a: 1, b: 2}]
4619        let values = list_builder.values();
4620        values
4621            .field_builder::<Int32Builder>(0)
4622            .unwrap()
4623            .append_value(1);
4624        values
4625            .field_builder::<Int32Builder>(1)
4626            .unwrap()
4627            .append_value(2);
4628        values.append(true);
4629        list_builder.append(true);
4630
4631        // []
4632        list_builder.append(true);
4633
4634        // null
4635        list_builder.append(false);
4636
4637        // [null, null]
4638        let values = list_builder.values();
4639        values
4640            .field_builder::<Int32Builder>(0)
4641            .unwrap()
4642            .append_null();
4643        values
4644            .field_builder::<Int32Builder>(1)
4645            .unwrap()
4646            .append_null();
4647        values.append(false);
4648        values
4649            .field_builder::<Int32Builder>(0)
4650            .unwrap()
4651            .append_null();
4652        values
4653            .field_builder::<Int32Builder>(1)
4654            .unwrap()
4655            .append_null();
4656        values.append(false);
4657        list_builder.append(true);
4658
4659        // [{a: null, b: 3}]
4660        let values = list_builder.values();
4661        values
4662            .field_builder::<Int32Builder>(0)
4663            .unwrap()
4664            .append_null();
4665        values
4666            .field_builder::<Int32Builder>(1)
4667            .unwrap()
4668            .append_value(3);
4669        values.append(true);
4670        list_builder.append(true);
4671
4672        // [{a: 2, b: null}]
4673        let values = list_builder.values();
4674        values
4675            .field_builder::<Int32Builder>(0)
4676            .unwrap()
4677            .append_value(2);
4678        values
4679            .field_builder::<Int32Builder>(1)
4680            .unwrap()
4681            .append_null();
4682        values.append(true);
4683        list_builder.append(true);
4684
4685        let array = Arc::new(list_builder.finish());
4686
4687        one_column_roundtrip(array, true);
4688    }
4689
4690    fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4691        metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4692    }
4693
4694    #[test]
4695    fn test_aggregates_records() {
4696        let arrays = [
4697            Int32Array::from((0..100).collect::<Vec<_>>()),
4698            Int32Array::from((0..50).collect::<Vec<_>>()),
4699            Int32Array::from((200..500).collect::<Vec<_>>()),
4700        ];
4701
4702        let schema = Arc::new(Schema::new(vec![Field::new(
4703            "int",
4704            ArrowDataType::Int32,
4705            false,
4706        )]));
4707
4708        let file = tempfile::tempfile().unwrap();
4709
4710        let props = WriterProperties::builder()
4711            .set_max_row_group_row_count(Some(200))
4712            .build();
4713
4714        let mut writer =
4715            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4716
4717        for array in arrays {
4718            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4719            writer.write(&batch).unwrap();
4720        }
4721
4722        writer.close().unwrap();
4723
4724        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4725        assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4726
4727        let batches = builder
4728            .with_batch_size(100)
4729            .build()
4730            .unwrap()
4731            .collect::<ArrowResult<Vec<_>>>()
4732            .unwrap();
4733
4734        assert_eq!(batches.len(), 5);
4735        assert!(batches.iter().all(|x| x.num_columns() == 1));
4736
4737        let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4738
4739        assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4740
4741        let values: Vec<_> = batches
4742            .iter()
4743            .flat_map(|x| {
4744                x.column(0)
4745                    .as_any()
4746                    .downcast_ref::<Int32Array>()
4747                    .unwrap()
4748                    .values()
4749                    .iter()
4750                    .cloned()
4751            })
4752            .collect();
4753
4754        let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4755        assert_eq!(&values, &expected_values)
4756    }
4757
4758    #[test]
4759    fn complex_aggregate() {
4760        // Tests aggregating nested data
4761        let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4762        let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4763        let struct_a = Arc::new(Field::new(
4764            "struct_a",
4765            DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4766            true,
4767        ));
4768
4769        let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4770        let struct_b = Arc::new(Field::new(
4771            "struct_b",
4772            DataType::Struct(vec![list_a.clone()].into()),
4773            false,
4774        ));
4775
4776        let schema = Arc::new(Schema::new(vec![struct_b]));
4777
4778        // create nested data
4779        let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4780        let field_b_array =
4781            Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4782
4783        let struct_a_array = StructArray::from(vec![
4784            (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4785            (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4786        ]);
4787
4788        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4789            .len(5)
4790            .add_buffer(Buffer::from_iter(vec![
4791                0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4792            ]))
4793            .null_bit_buffer(Some(Buffer::from_iter(vec![
4794                true, false, true, false, true,
4795            ])))
4796            .child_data(vec![struct_a_array.into_data()])
4797            .build()
4798            .unwrap();
4799
4800        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4801        let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4802
4803        let batch1 =
4804            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4805                .unwrap();
4806
4807        let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4808        let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4809
4810        let struct_a_array = StructArray::from(vec![
4811            (field_a, Arc::new(field_a_array) as ArrayRef),
4812            (field_b, Arc::new(field_b_array) as ArrayRef),
4813        ]);
4814
4815        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4816            .len(2)
4817            .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4818            .child_data(vec![struct_a_array.into_data()])
4819            .build()
4820            .unwrap();
4821
4822        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4823        let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4824
4825        let batch2 =
4826            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4827                .unwrap();
4828
4829        let batches = &[batch1, batch2];
4830
4831        // Verify data is as expected
4832
4833        let expected = r#"
4834            +-------------------------------------------------------------------------------------------------------+
4835            | struct_b                                                                                              |
4836            +-------------------------------------------------------------------------------------------------------+
4837            | {list: [{leaf_a: 1, leaf_b: 1}]}                                                                      |
4838            | {list: }                                                                                              |
4839            | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]}                                               |
4840            | {list: }                                                                                              |
4841            | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]}                                                |
4842            | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4843            | {list: [{leaf_a: 10, leaf_b: }]}                                                                      |
4844            +-------------------------------------------------------------------------------------------------------+
4845        "#.trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4846
4847        let actual = pretty_format_batches(batches).unwrap().to_string();
4848        assert_eq!(actual, expected);
4849
4850        // Write data
4851        let file = tempfile::tempfile().unwrap();
4852        let props = WriterProperties::builder()
4853            .set_max_row_group_row_count(Some(6))
4854            .build();
4855
4856        let mut writer =
4857            ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4858
4859        for batch in batches {
4860            writer.write(batch).unwrap();
4861        }
4862        writer.close().unwrap();
4863
4864        // Read Data
4865        // Should have written entire first batch and first row of second to the first row group
4866        // leaving a single row in the second row group
4867
4868        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4869        assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4870
4871        let batches = builder
4872            .with_batch_size(2)
4873            .build()
4874            .unwrap()
4875            .collect::<ArrowResult<Vec<_>>>()
4876            .unwrap();
4877
4878        assert_eq!(batches.len(), 4);
4879        let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4880        assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4881
4882        let actual = pretty_format_batches(&batches).unwrap().to_string();
4883        assert_eq!(actual, expected);
4884    }
4885
4886    #[test]
4887    fn test_arrow_writer_metadata() {
4888        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4889        let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4890
4891        let batch = RecordBatch::try_new(
4892            Arc::new(batch_schema),
4893            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4894        )
4895        .unwrap();
4896
4897        let mut buf = Vec::with_capacity(1024);
4898        let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4899        writer.write(&batch).unwrap();
4900        writer.close().unwrap();
4901    }
4902
4903    #[test]
4904    fn test_arrow_writer_nullable() {
4905        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4906        let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
4907        let file_schema = Arc::new(file_schema);
4908
4909        let batch = RecordBatch::try_new(
4910            Arc::new(batch_schema),
4911            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4912        )
4913        .unwrap();
4914
4915        let mut buf = Vec::with_capacity(1024);
4916        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
4917        writer.write(&batch).unwrap();
4918        writer.close().unwrap();
4919
4920        let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
4921        let back = read.next().unwrap().unwrap();
4922        assert_eq!(back.schema(), file_schema);
4923        assert_ne!(back.schema(), batch.schema());
4924        assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
4925    }
4926
4927    #[test]
4928    fn in_progress_accounting() {
4929        // define schema
4930        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4931
4932        // create some data
4933        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4934
4935        // build a record batch
4936        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4937
4938        let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4939
4940        // starts empty
4941        assert_eq!(writer.in_progress_size(), 0);
4942        assert_eq!(writer.in_progress_rows(), 0);
4943        assert_eq!(writer.memory_size(), 0);
4944        assert_eq!(writer.bytes_written(), 4); // Initial header
4945        writer.write(&batch).unwrap();
4946
4947        // updated on write
4948        let initial_size = writer.in_progress_size();
4949        assert!(initial_size > 0);
4950        assert_eq!(writer.in_progress_rows(), 5);
4951        let initial_memory = writer.memory_size();
4952        assert!(initial_memory > 0);
4953        // memory estimate is larger than estimated encoded size
4954        assert!(
4955            initial_size <= initial_memory,
4956            "{initial_size} <= {initial_memory}"
4957        );
4958
4959        // updated on second write
4960        writer.write(&batch).unwrap();
4961        assert!(writer.in_progress_size() > initial_size);
4962        assert_eq!(writer.in_progress_rows(), 10);
4963        assert!(writer.memory_size() > initial_memory);
4964        assert!(
4965            writer.in_progress_size() <= writer.memory_size(),
4966            "in_progress_size {} <= memory_size {}",
4967            writer.in_progress_size(),
4968            writer.memory_size()
4969        );
4970
4971        // in progress tracking is cleared, but the overall data written is updated
4972        let pre_flush_bytes_written = writer.bytes_written();
4973        writer.flush().unwrap();
4974        assert_eq!(writer.in_progress_size(), 0);
4975        assert_eq!(writer.memory_size(), 0);
4976        assert!(writer.bytes_written() > pre_flush_bytes_written);
4977
4978        writer.close().unwrap();
4979    }
4980
4981    #[test]
4982    fn test_writer_all_null() {
4983        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4984        let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4985        let batch = RecordBatch::try_from_iter(vec![
4986            ("a", Arc::new(a) as ArrayRef),
4987            ("b", Arc::new(b) as ArrayRef),
4988        ])
4989        .unwrap();
4990
4991        let mut buf = Vec::with_capacity(1024);
4992        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4993        writer.write(&batch).unwrap();
4994        writer.close().unwrap();
4995
4996        let bytes = Bytes::from(buf);
4997        let options = ReadOptionsBuilder::new().with_page_index().build();
4998        let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4999        let index = reader.metadata().offset_index().unwrap();
5000
5001        assert_eq!(index.len(), 1);
5002        assert_eq!(index[0].len(), 2); // 2 columns
5003        assert_eq!(index[0][0].page_locations().len(), 1); // 1 page
5004        assert_eq!(index[0][1].page_locations().len(), 1); // 1 page
5005    }
5006
5007    #[test]
5008    fn test_disabled_statistics_with_page() {
5009        let file_schema = Schema::new(vec![
5010            Field::new("a", DataType::Utf8, true),
5011            Field::new("b", DataType::Utf8, true),
5012        ]);
5013        let file_schema = Arc::new(file_schema);
5014
5015        let batch = RecordBatch::try_new(
5016            file_schema.clone(),
5017            vec![
5018                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5019                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5020            ],
5021        )
5022        .unwrap();
5023
5024        let props = WriterProperties::builder()
5025            .set_statistics_enabled(EnabledStatistics::None)
5026            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5027            .build();
5028
5029        let mut buf = Vec::with_capacity(1024);
5030        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5031        writer.write(&batch).unwrap();
5032
5033        let metadata = writer.close().unwrap();
5034        assert_eq!(metadata.num_row_groups(), 1);
5035        let row_group = metadata.row_group(0);
5036        assert_eq!(row_group.num_columns(), 2);
5037        // Column "a" has both offset and column index, as requested
5038        assert!(row_group.column(0).offset_index_offset().is_some());
5039        assert!(row_group.column(0).column_index_offset().is_some());
5040        // Column "b" should only have offset index
5041        assert!(row_group.column(1).offset_index_offset().is_some());
5042        assert!(row_group.column(1).column_index_offset().is_none());
5043
5044        let options = ReadOptionsBuilder::new().with_page_index().build();
5045        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5046
5047        let row_group = reader.get_row_group(0).unwrap();
5048        let a_col = row_group.metadata().column(0);
5049        let b_col = row_group.metadata().column(1);
5050
5051        // Column chunk of column "a" should have chunk level statistics
5052        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5053            let min = byte_array_stats.min_opt().unwrap();
5054            let max = byte_array_stats.max_opt().unwrap();
5055
5056            assert_eq!(min.as_bytes(), b"a");
5057            assert_eq!(max.as_bytes(), b"d");
5058        } else {
5059            panic!("expecting Statistics::ByteArray");
5060        }
5061
5062        // The column chunk for column "b" shouldn't have statistics
5063        assert!(b_col.statistics().is_none());
5064
5065        let offset_index = reader.metadata().offset_index().unwrap();
5066        assert_eq!(offset_index.len(), 1); // 1 row group
5067        assert_eq!(offset_index[0].len(), 2); // 2 columns
5068
5069        let column_index = reader.metadata().column_index().unwrap();
5070        assert_eq!(column_index.len(), 1); // 1 row group
5071        assert_eq!(column_index[0].len(), 2); // 2 columns
5072
5073        let a_idx = &column_index[0][0];
5074        assert!(
5075            matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)),
5076            "{a_idx:?}"
5077        );
5078        let b_idx = &column_index[0][1];
5079        assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5080    }
5081
5082    #[test]
5083    fn test_disabled_statistics_with_chunk() {
5084        let file_schema = Schema::new(vec![
5085            Field::new("a", DataType::Utf8, true),
5086            Field::new("b", DataType::Utf8, true),
5087        ]);
5088        let file_schema = Arc::new(file_schema);
5089
5090        let batch = RecordBatch::try_new(
5091            file_schema.clone(),
5092            vec![
5093                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5094                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5095            ],
5096        )
5097        .unwrap();
5098
5099        let props = WriterProperties::builder()
5100            .set_statistics_enabled(EnabledStatistics::None)
5101            .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5102            .build();
5103
5104        let mut buf = Vec::with_capacity(1024);
5105        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5106        writer.write(&batch).unwrap();
5107
5108        let metadata = writer.close().unwrap();
5109        assert_eq!(metadata.num_row_groups(), 1);
5110        let row_group = metadata.row_group(0);
5111        assert_eq!(row_group.num_columns(), 2);
5112        // Column "a" should only have offset index
5113        assert!(row_group.column(0).offset_index_offset().is_some());
5114        assert!(row_group.column(0).column_index_offset().is_none());
5115        // Column "b" should only have offset index
5116        assert!(row_group.column(1).offset_index_offset().is_some());
5117        assert!(row_group.column(1).column_index_offset().is_none());
5118
5119        let options = ReadOptionsBuilder::new().with_page_index().build();
5120        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5121
5122        let row_group = reader.get_row_group(0).unwrap();
5123        let a_col = row_group.metadata().column(0);
5124        let b_col = row_group.metadata().column(1);
5125
5126        // Column chunk of column "a" should have chunk level statistics
5127        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5128            let min = byte_array_stats.min_opt().unwrap();
5129            let max = byte_array_stats.max_opt().unwrap();
5130
5131            assert_eq!(min.as_bytes(), b"a");
5132            assert_eq!(max.as_bytes(), b"d");
5133        } else {
5134            panic!("expecting Statistics::ByteArray");
5135        }
5136
5137        // The column chunk for column "b"  shouldn't have statistics
5138        assert!(b_col.statistics().is_none());
5139
5140        let column_index = reader.metadata().column_index().unwrap();
5141        assert_eq!(column_index.len(), 1); // 1 row group
5142        assert_eq!(column_index[0].len(), 2); // 2 columns
5143
5144        let a_idx = &column_index[0][0];
5145        assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}");
5146        let b_idx = &column_index[0][1];
5147        assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5148    }
5149
5150    #[test]
5151    fn test_arrow_writer_skip_metadata() {
5152        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5153        let file_schema = Arc::new(batch_schema.clone());
5154
5155        let batch = RecordBatch::try_new(
5156            Arc::new(batch_schema),
5157            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5158        )
5159        .unwrap();
5160        let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5161
5162        let mut buf = Vec::with_capacity(1024);
5163        let mut writer =
5164            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5165        writer.write(&batch).unwrap();
5166        writer.close().unwrap();
5167
5168        let bytes = Bytes::from(buf);
5169        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5170        assert_eq!(file_schema, *reader_builder.schema());
5171        if let Some(key_value_metadata) = reader_builder
5172            .metadata()
5173            .file_metadata()
5174            .key_value_metadata()
5175        {
5176            assert!(
5177                !key_value_metadata
5178                    .iter()
5179                    .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5180            );
5181        }
5182    }
5183
5184    #[test]
5185    fn test_arrow_writer_skip_path_in_schema() {
5186        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5187        let file_schema = Arc::new(batch_schema.clone());
5188
5189        let batch = RecordBatch::try_new(
5190            Arc::new(batch_schema),
5191            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5192        )
5193        .unwrap();
5194
5195        // default options should still write path_in_schema
5196        let skip_options = ArrowWriterOptions::new();
5197
5198        let mut buf = Vec::with_capacity(1024);
5199        let mut writer =
5200            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5201        writer.write(&batch).unwrap();
5202        writer.close().unwrap();
5203
5204        // override to not write path_in_schema
5205        let skip_options = ArrowWriterOptions::new().with_properties(
5206            WriterProperties::builder()
5207                .set_write_path_in_schema(false)
5208                .build(),
5209        );
5210
5211        let mut buf2 = Vec::with_capacity(1024);
5212        let mut writer =
5213            ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5214                .unwrap();
5215        writer.write(&batch).unwrap();
5216        writer.close().unwrap();
5217
5218        // buf2 should be a bit smaller due to lack of path_in_schema
5219        assert!(buf.len() > buf2.len());
5220    }
5221
5222    #[test]
5223    fn mismatched_schemas() {
5224        let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5225        let file_schema = Arc::new(Schema::new(vec![Field::new(
5226            "temperature",
5227            DataType::Float64,
5228            false,
5229        )]));
5230
5231        let batch = RecordBatch::try_new(
5232            Arc::new(batch_schema),
5233            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5234        )
5235        .unwrap();
5236
5237        let mut buf = Vec::with_capacity(1024);
5238        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5239
5240        let err = writer.write(&batch).unwrap_err().to_string();
5241        assert_eq!(
5242            err,
5243            "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5244        );
5245    }
5246
5247    #[test]
5248    // https://github.com/apache/arrow-rs/issues/6988
5249    fn test_roundtrip_empty_schema() {
5250        // create empty record batch with empty schema
5251        let empty_batch = RecordBatch::try_new_with_options(
5252            Arc::new(Schema::empty()),
5253            vec![],
5254            &RecordBatchOptions::default().with_row_count(Some(0)),
5255        )
5256        .unwrap();
5257
5258        // write to parquet
5259        let mut parquet_bytes: Vec<u8> = Vec::new();
5260        let mut writer =
5261            ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5262        writer.write(&empty_batch).unwrap();
5263        writer.close().unwrap();
5264
5265        // read from parquet
5266        let bytes = Bytes::from(parquet_bytes);
5267        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5268        assert_eq!(reader.schema(), &empty_batch.schema());
5269        let batches: Vec<_> = reader
5270            .build()
5271            .unwrap()
5272            .collect::<ArrowResult<Vec<_>>>()
5273            .unwrap();
5274        assert_eq!(batches.len(), 0);
5275    }
5276
5277    #[test]
5278    fn test_page_stats_not_written_by_default() {
5279        let string_field = Field::new("a", DataType::Utf8, false);
5280        let schema = Schema::new(vec![string_field]);
5281        let raw_string_values = vec!["Blart Versenwald III"];
5282        let string_values = StringArray::from(raw_string_values.clone());
5283        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5284
5285        let props = WriterProperties::builder()
5286            .set_statistics_enabled(EnabledStatistics::Page)
5287            .set_dictionary_enabled(false)
5288            .set_encoding(Encoding::PLAIN)
5289            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5290            .build();
5291
5292        let file = roundtrip_opts(&batch, props);
5293
5294        // read file and decode page headers
5295        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5296
5297        // decode first page header
5298        let first_page = &file[4..];
5299        let mut prot = ThriftSliceInputProtocol::new(first_page);
5300        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5301        let stats = hdr.data_page_header.unwrap().statistics;
5302
5303        assert!(stats.is_none());
5304    }
5305
5306    #[test]
5307    fn test_page_stats_when_enabled() {
5308        let string_field = Field::new("a", DataType::Utf8, false);
5309        let schema = Schema::new(vec![string_field]);
5310        let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5311        let string_values = StringArray::from(raw_string_values.clone());
5312        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5313
5314        let props = WriterProperties::builder()
5315            .set_statistics_enabled(EnabledStatistics::Page)
5316            .set_dictionary_enabled(false)
5317            .set_encoding(Encoding::PLAIN)
5318            .set_write_page_header_statistics(true)
5319            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5320            .build();
5321
5322        let file = roundtrip_opts(&batch, props);
5323
5324        // read file and decode page headers
5325        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5326
5327        // decode first page header
5328        let first_page = &file[4..];
5329        let mut prot = ThriftSliceInputProtocol::new(first_page);
5330        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5331        let stats = hdr.data_page_header.unwrap().statistics;
5332
5333        let stats = stats.unwrap();
5334        // check that min/max were actually written to the page
5335        assert!(stats.is_max_value_exact.unwrap());
5336        assert!(stats.is_min_value_exact.unwrap());
5337        assert_eq!(stats.max_value.unwrap(), "Blart Versenwald III".as_bytes());
5338        assert_eq!(stats.min_value.unwrap(), "Andrew Lamb".as_bytes());
5339    }
5340
5341    #[test]
5342    fn test_page_stats_truncation() {
5343        let string_field = Field::new("a", DataType::Utf8, false);
5344        let binary_field = Field::new("b", DataType::Binary, false);
5345        let schema = Schema::new(vec![string_field, binary_field]);
5346
5347        let raw_string_values = vec!["Blart Versenwald III"];
5348        let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5349        let raw_binary_value_refs = raw_binary_values
5350            .iter()
5351            .map(|x| x.as_slice())
5352            .collect::<Vec<_>>();
5353
5354        let string_values = StringArray::from(raw_string_values.clone());
5355        let binary_values = BinaryArray::from(raw_binary_value_refs);
5356        let batch = RecordBatch::try_new(
5357            Arc::new(schema),
5358            vec![Arc::new(string_values), Arc::new(binary_values)],
5359        )
5360        .unwrap();
5361
5362        let props = WriterProperties::builder()
5363            .set_statistics_truncate_length(Some(2))
5364            .set_dictionary_enabled(false)
5365            .set_encoding(Encoding::PLAIN)
5366            .set_write_page_header_statistics(true)
5367            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5368            .build();
5369
5370        let file = roundtrip_opts(&batch, props);
5371
5372        // read file and decode page headers
5373        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5374
5375        // decode first page header
5376        let first_page = &file[4..];
5377        let mut prot = ThriftSliceInputProtocol::new(first_page);
5378        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5379        let stats = hdr.data_page_header.unwrap().statistics;
5380        assert!(stats.is_some());
5381        let stats = stats.unwrap();
5382        // check that min/max were properly truncated
5383        assert!(!stats.is_max_value_exact.unwrap());
5384        assert!(!stats.is_min_value_exact.unwrap());
5385        assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5386        assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5387
5388        // check second page now
5389        let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5390        let mut prot = ThriftSliceInputProtocol::new(second_page);
5391        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5392        let stats = hdr.data_page_header.unwrap().statistics;
5393        assert!(stats.is_some());
5394        let stats = stats.unwrap();
5395        // check that min/max were properly truncated
5396        assert!(!stats.is_max_value_exact.unwrap());
5397        assert!(!stats.is_min_value_exact.unwrap());
5398        assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5399        assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5400    }
5401
5402    #[test]
5403    fn test_page_encoding_statistics_roundtrip() {
5404        let batch_schema = Schema::new(vec![Field::new(
5405            "int32",
5406            arrow_schema::DataType::Int32,
5407            false,
5408        )]);
5409
5410        let batch = RecordBatch::try_new(
5411            Arc::new(batch_schema.clone()),
5412            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5413        )
5414        .unwrap();
5415
5416        let mut file: File = tempfile::tempfile().unwrap();
5417        let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5418        writer.write(&batch).unwrap();
5419        let file_metadata = writer.close().unwrap();
5420
5421        assert_eq!(file_metadata.num_row_groups(), 1);
5422        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5423        assert!(
5424            file_metadata
5425                .row_group(0)
5426                .column(0)
5427                .page_encoding_stats()
5428                .is_some()
5429        );
5430        let chunk_page_stats = file_metadata
5431            .row_group(0)
5432            .column(0)
5433            .page_encoding_stats()
5434            .unwrap();
5435
5436        // check that the read metadata is also correct
5437        let options = ReadOptionsBuilder::new()
5438            .with_page_index()
5439            .with_encoding_stats_as_mask(false)
5440            .build();
5441        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5442
5443        let rowgroup = reader.get_row_group(0).expect("row group missing");
5444        assert_eq!(rowgroup.num_columns(), 1);
5445        let column = rowgroup.metadata().column(0);
5446        assert!(column.page_encoding_stats().is_some());
5447        let file_page_stats = column.page_encoding_stats().unwrap();
5448        assert_eq!(chunk_page_stats, file_page_stats);
5449    }
5450
5451    #[test]
5452    fn test_different_dict_page_size_limit() {
5453        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5454        let schema = Arc::new(Schema::new(vec![
5455            Field::new("col0", arrow_schema::DataType::Int64, false),
5456            Field::new("col1", arrow_schema::DataType::Int64, false),
5457        ]));
5458        let batch =
5459            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5460
5461        let props = WriterProperties::builder()
5462            .set_dictionary_page_size_limit(1024 * 1024)
5463            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5464            .build();
5465        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5466        writer.write(&batch).unwrap();
5467        let data = Bytes::from(writer.into_inner().unwrap());
5468
5469        let mut metadata = ParquetMetaDataReader::new();
5470        metadata.try_parse(&data).unwrap();
5471        let metadata = metadata.finish().unwrap();
5472        let col0_meta = metadata.row_group(0).column(0);
5473        let col1_meta = metadata.row_group(0).column(1);
5474
5475        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5476            let mut reader =
5477                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5478            let page = reader.get_next_page().unwrap().unwrap();
5479            match page {
5480                Page::DictionaryPage { buf, .. } => buf.len(),
5481                _ => panic!("expected DictionaryPage"),
5482            }
5483        };
5484
5485        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5486        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5487    }
5488
5489    #[test]
5490    fn test_arrow_writer_granular_mode_roundtrip() {
5491        // Granular mode subdivides chunks and writes more pages than the
5492        // default batched path. Make sure the data we write back is
5493        // bit-identical to what went in — page-count assertions elsewhere
5494        // only prove pages were cut, not that the encoded data is correct.
5495        //
5496        // Mix value sizes so that the cumulative-byte-budget cutoff
5497        // lands mid-chunk, exercising both batched and granular paths
5498        // within the same `write_batch_internal` call.
5499        let small = "tiny".to_string();
5500        let big = "x".repeat(64 * 1024);
5501        let strings: Vec<String> = (0..256)
5502            .map(|i| {
5503                if i % 16 == 0 {
5504                    big.clone()
5505                } else {
5506                    small.clone()
5507                }
5508            })
5509            .collect();
5510
5511        let schema = Arc::new(Schema::new(vec![Field::new(
5512            "col",
5513            ArrowDataType::Utf8,
5514            false,
5515        )]));
5516        let batch = RecordBatch::try_new(
5517            schema.clone(),
5518            vec![Arc::new(StringArray::from(strings.clone())) as _],
5519        )
5520        .unwrap();
5521
5522        let props = WriterProperties::builder()
5523            .set_dictionary_enabled(false)
5524            .set_data_page_size_limit(16 * 1024)
5525            .build();
5526        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5527        writer.write(&batch).unwrap();
5528        let data = Bytes::from(writer.into_inner().unwrap());
5529
5530        let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5531        let read = reader.next().unwrap().unwrap();
5532        assert!(reader.next().is_none(), "expected one batch");
5533        let col = read
5534            .column(0)
5535            .as_any()
5536            .downcast_ref::<StringArray>()
5537            .unwrap();
5538        assert_eq!(col.len(), strings.len());
5539        for (i, expected) in strings.iter().enumerate() {
5540            assert_eq!(
5541                col.value(i),
5542                expected.as_str(),
5543                "value mismatch at index {i}"
5544            );
5545        }
5546    }
5547
5548    #[test]
5549    fn test_arrow_writer_all_null_string_column() {
5550        // The `LevelDataRef::value_count` Uniform branch with
5551        // `value != max_def` (entirely-null chunk) must return 0 so the
5552        // sub-batch sizer short-circuits to batch mode without trying
5553        // to estimate byte budgets for non-existent values.
5554        let num_rows = 1024;
5555        let schema = Arc::new(Schema::new(vec![Field::new(
5556            "col",
5557            ArrowDataType::Utf8,
5558            true,
5559        )]));
5560        let nulls: Vec<Option<&str>> = vec![None; num_rows];
5561        let batch = RecordBatch::try_new(
5562            schema.clone(),
5563            vec![Arc::new(StringArray::from(nulls)) as _],
5564        )
5565        .unwrap();
5566
5567        let props = WriterProperties::builder()
5568            .set_dictionary_enabled(false)
5569            .set_data_page_size_limit(16 * 1024)
5570            .build();
5571        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5572        writer.write(&batch).unwrap();
5573        let data = Bytes::from(writer.into_inner().unwrap());
5574
5575        // Re-parse the file: row group has one column, every row is
5576        // null, all data pages report `num_rows / page_count` rows.
5577        let mut metadata = ParquetMetaDataReader::new();
5578        metadata.try_parse(&data).unwrap();
5579        let metadata = metadata.finish().unwrap();
5580        let row_group = metadata.row_group(0);
5581        let col_meta = row_group.column(0);
5582        assert_eq!(row_group.num_rows() as usize, num_rows);
5583        // Statistics record `null_count = num_rows` — proves every value
5584        // was written as null.
5585        if let Some(stats) = col_meta.statistics() {
5586            assert_eq!(
5587                stats.null_count_opt().unwrap_or(0) as usize,
5588                num_rows,
5589                "expected all-null column to report null_count = num_rows"
5590            );
5591        }
5592
5593        let mut reader =
5594            SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5595        let mut total_values = 0u32;
5596        while let Some(page) = reader.get_next_page().unwrap() {
5597            if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5598                total_values += page.num_values();
5599            }
5600        }
5601        assert_eq!(
5602            total_values as usize, num_rows,
5603            "expected every level position to be represented in some page"
5604        );
5605    }
5606
5607    struct WriteBatchesShape {
5608        num_batches: usize,
5609        rows_per_batch: usize,
5610        row_size: usize,
5611    }
5612
5613    /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter`
5614    fn write_batches(
5615        WriteBatchesShape {
5616            num_batches,
5617            rows_per_batch,
5618            row_size,
5619        }: WriteBatchesShape,
5620        props: WriterProperties,
5621    ) -> ParquetRecordBatchReaderBuilder<File> {
5622        let schema = Arc::new(Schema::new(vec![Field::new(
5623            "str",
5624            ArrowDataType::Utf8,
5625            false,
5626        )]));
5627        let file = tempfile::tempfile().unwrap();
5628        let mut writer =
5629            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5630
5631        for batch_idx in 0..num_batches {
5632            let strings: Vec<String> = (0..rows_per_batch)
5633                .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5634                .collect();
5635            let array = StringArray::from(strings);
5636            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5637            writer.write(&batch).unwrap();
5638        }
5639        writer.close().unwrap();
5640        ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5641    }
5642
5643    #[test]
5644    // When both limits are None, all data should go into a single row group
5645    fn test_row_group_limit_none_writes_single_row_group() {
5646        let props = WriterProperties::builder()
5647            .set_max_row_group_row_count(None)
5648            .set_max_row_group_bytes(None)
5649            .build();
5650
5651        let builder = write_batches(
5652            WriteBatchesShape {
5653                num_batches: 1,
5654                rows_per_batch: 1000,
5655                row_size: 4,
5656            },
5657            props,
5658        );
5659
5660        assert_eq!(
5661            &row_group_sizes(builder.metadata()),
5662            &[1000],
5663            "With no limits, all rows should be in a single row group"
5664        );
5665    }
5666
5667    #[test]
5668    // When only max_row_group_size is set, respect the row limit
5669    fn test_row_group_limit_rows_only() {
5670        let props = WriterProperties::builder()
5671            .set_max_row_group_row_count(Some(300))
5672            .set_max_row_group_bytes(None)
5673            .build();
5674
5675        let builder = write_batches(
5676            WriteBatchesShape {
5677                num_batches: 1,
5678                rows_per_batch: 1000,
5679                row_size: 4,
5680            },
5681            props,
5682        );
5683
5684        assert_eq!(
5685            &row_group_sizes(builder.metadata()),
5686            &[300, 300, 300, 100],
5687            "Row groups should be split by row count"
5688        );
5689    }
5690
5691    #[test]
5692    // When only max_row_group_bytes is set, respect the byte limit
5693    fn test_row_group_limit_bytes_only() {
5694        let props = WriterProperties::builder()
5695            .set_max_row_group_row_count(None)
5696            // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
5697            .set_max_row_group_bytes(Some(3500))
5698            .build();
5699
5700        let builder = write_batches(
5701            WriteBatchesShape {
5702                num_batches: 10,
5703                rows_per_batch: 10,
5704                row_size: 100,
5705            },
5706            props,
5707        );
5708
5709        let sizes = row_group_sizes(builder.metadata());
5710
5711        assert!(
5712            sizes.len() > 1,
5713            "Should have multiple row groups due to byte limit, got {sizes:?}",
5714        );
5715
5716        let total_rows: i64 = sizes.iter().sum();
5717        assert_eq!(total_rows, 100, "Total rows should be preserved");
5718    }
5719
5720    #[test]
5721    // If an in-progress row group is already oversized, it should be flushed before writing more.
5722    fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5723        let schema = Arc::new(Schema::new(vec![Field::new(
5724            "str",
5725            ArrowDataType::Utf8,
5726            false,
5727        )]));
5728        let file = tempfile::tempfile().unwrap();
5729
5730        // Start with no byte limit so we can intentionally build an oversized in-progress row group.
5731        let props = WriterProperties::builder()
5732            .set_max_row_group_row_count(None)
5733            .set_max_row_group_bytes(None)
5734            .build();
5735        let mut writer =
5736            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5737
5738        let first_array = StringArray::from(
5739            (0..10)
5740                .map(|i| format!("{:0>100}", i))
5741                .collect::<Vec<String>>(),
5742        );
5743        let first_batch =
5744            RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5745        writer.write(&first_batch).unwrap();
5746        assert_eq!(writer.in_progress_rows(), 10);
5747
5748        // Tighten the limit below the current in-progress bytes to exercise:
5749        // `if current_bytes >= max_bytes { self.flush()?; ... }`
5750        writer.max_row_group_bytes = Some(1);
5751
5752        let second_array = StringArray::from(vec!["x".to_string()]);
5753        let second_batch =
5754            RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5755        writer.write(&second_batch).unwrap();
5756        writer.close().unwrap();
5757        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5758
5759        assert_eq!(
5760            &row_group_sizes(builder.metadata()),
5761            &[10, 1],
5762            "The second write should flush an oversized in-progress row group first",
5763        );
5764    }
5765
5766    #[test]
5767    // When both limits are set, the row limit triggers first
5768    fn test_row_group_limit_both_row_wins_single_batch() {
5769        let props = WriterProperties::builder()
5770            .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows
5771            .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
5772            .build();
5773
5774        let builder = write_batches(
5775            WriteBatchesShape {
5776                num_batches: 1,
5777                row_size: 4,
5778                rows_per_batch: 1000,
5779            },
5780            props,
5781        );
5782
5783        assert_eq!(
5784            &row_group_sizes(builder.metadata()),
5785            &[200, 200, 200, 200, 200],
5786            "Row limit should trigger before byte limit"
5787        );
5788    }
5789
5790    #[test]
5791    // When both limits are set, the row limit triggers first
5792    fn test_row_group_limit_both_row_wins_multiple_batches() {
5793        let props = WriterProperties::builder()
5794            .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows
5795            .set_max_row_group_bytes(Some(9999)) // Won't trigger
5796            .build();
5797
5798        let builder = write_batches(
5799            WriteBatchesShape {
5800                num_batches: 10,
5801                rows_per_batch: 10,
5802                row_size: 100,
5803            },
5804            props,
5805        );
5806
5807        assert_eq!(
5808            &row_group_sizes(builder.metadata()),
5809            &[5; 20],
5810            "Row limit should trigger before byte limit"
5811        );
5812    }
5813
5814    #[test]
5815    // When both limits are set, the byte limit triggers first
5816    fn test_row_group_limit_both_bytes_wins() {
5817        let props = WriterProperties::builder()
5818            .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows
5819            .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows
5820            .build();
5821
5822        let builder = write_batches(
5823            WriteBatchesShape {
5824                num_batches: 10,
5825                rows_per_batch: 10,
5826                row_size: 100,
5827            },
5828            props,
5829        );
5830
5831        let sizes = row_group_sizes(builder.metadata());
5832
5833        assert!(
5834            sizes.len() > 1,
5835            "Byte limit should trigger before row limit, got {sizes:?}",
5836        );
5837
5838        assert!(
5839            sizes.iter().all(|&s| s < 1000),
5840            "No row group should hit the row limit"
5841        );
5842
5843        let total_rows: i64 = sizes.iter().sum();
5844        assert_eq!(total_rows, 100, "Total rows should be preserved");
5845    }
5846
5847    #[test]
5848    fn arrow_column_chunk_close_mut_drops_column_index() {
5849        use crate::arrow::ArrowSchemaConverter;
5850        use crate::file::writer::SerializedFileWriter;
5851
5852        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5853        let props = Arc::new(
5854            WriterProperties::builder()
5855                .set_statistics_enabled(EnabledStatistics::Page)
5856                .build(),
5857        );
5858        let parquet_schema = ArrowSchemaConverter::new()
5859            .with_coerce_types(props.coerce_types())
5860            .convert(&schema)
5861            .unwrap();
5862
5863        let mut buf = Vec::with_capacity(1024);
5864        let mut writer =
5865            SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5866                .unwrap();
5867
5868        let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
5869        let mut col_writers = factory.create_column_writers(0).unwrap();
5870        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5871        for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5872            col_writers[0].write(&leaves).unwrap();
5873        }
5874        let mut chunk = col_writers.pop().unwrap().close().unwrap();
5875
5876        // Immutable accessor exposes the close result produced at close time.
5877        assert!(
5878            chunk.close().column_index.is_some(),
5879            "EnabledStatistics::Page should produce a column_index"
5880        );
5881
5882        // Mutable accessor lets callers drop the page-level index before append.
5883        chunk.close_mut().column_index = None;
5884        assert!(chunk.close().column_index.is_none());
5885
5886        let mut rg = writer.next_row_group().unwrap();
5887        chunk.append_to_row_group(&mut rg).unwrap();
5888        rg.close().unwrap();
5889        let file_meta = writer.close().unwrap();
5890
5891        // After dropping column_index, the resulting file records no column
5892        // index offset/length for this chunk.
5893        let cc = file_meta.row_group(0).column(0);
5894        assert!(cc.column_index_range().is_none());
5895    }
5896
5897    /// Writes a single-column RecordBatch to an in-memory Parquet buffer.
5898    fn write_column_to_bytes(array: ArrayRef) -> Bytes {
5899        let schema = Arc::new(Schema::new(vec![Field::new(
5900            "col",
5901            array.data_type().clone(),
5902            true,
5903        )]));
5904        let buf = get_bytes_after_close(
5905            schema.clone(),
5906            &RecordBatch::try_new(schema, vec![array]).unwrap(),
5907        );
5908        Bytes::from(buf)
5909    }
5910
5911    /// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema.
5912    /// Passing a flat schema when the buffer was written from a REE array lets callers decode
5913    /// the physical values without the run-end encoding wrapper.
5914    fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
5915        let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
5916        ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
5917            .unwrap()
5918            .build()
5919            .unwrap()
5920            .next()
5921            .unwrap()
5922            .unwrap()
5923            .column(0)
5924            .clone()
5925    }
5926
5927    fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
5928        let flat_schema = Arc::new(Schema::new(vec![Field::new(
5929            "col",
5930            flat.data_type().clone(),
5931            true,
5932        )]));
5933        let ree_bytes = write_column_to_bytes(ree);
5934        let flat_bytes = write_column_to_bytes(flat.clone());
5935        assert_eq!(
5936            ree_bytes, flat_bytes,
5937            "REE and flat bytes should be identical"
5938        );
5939
5940        let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
5941        let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
5942
5943        assert_eq!(decoded_ree.as_ref(), flat.as_ref());
5944        assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
5945    }
5946
5947    #[test]
5948    fn ree_string() {
5949        let ree: ArrayRef = Arc::new(
5950            [Some("a"), Some("a"), None, Some("b"), Some("b")]
5951                .into_iter()
5952                .collect::<Int32RunArray>(),
5953        );
5954        let flat: ArrayRef = Arc::new(StringArray::from(vec![
5955            Some("a"),
5956            Some("a"),
5957            None,
5958            Some("b"),
5959            Some("b"),
5960        ]));
5961        ree_write_read_roundtrip(ree, flat);
5962    }
5963
5964    #[test]
5965    fn ree_int32() {
5966        let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
5967        for v in [Some(1), Some(1), None, Some(2), Some(2)] {
5968            b.append_option(v);
5969        }
5970        let ree: ArrayRef = Arc::new(b.finish());
5971        let flat: ArrayRef = Arc::new(Int32Array::from(vec![
5972            Some(1),
5973            Some(1),
5974            None,
5975            Some(2),
5976            Some(2),
5977        ]));
5978        ree_write_read_roundtrip(ree, flat);
5979    }
5980
5981    #[test]
5982    fn ree_bool() {
5983        // run_ends [3, 5, 7] → [T,T,T, null,null, F,F]
5984        let ree: ArrayRef = Arc::new(
5985            RunArray::try_new(
5986                &Int32Array::from(vec![3, 5, 7]),
5987                &BooleanArray::from(vec![Some(true), None, Some(false)]),
5988            )
5989            .unwrap(),
5990        );
5991        let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
5992            Some(true),
5993            Some(true),
5994            Some(true),
5995            None,
5996            None,
5997            Some(false),
5998            Some(false),
5999        ]));
6000        ree_write_read_roundtrip(ree, flat);
6001    }
6002
6003    #[test]
6004    fn ree_fixed_size_binary() {
6005        let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6006            let mut b = FixedSizeBinaryBuilder::new(2);
6007            for v in vals {
6008                match v {
6009                    Some(x) => b.append_value(x).unwrap(),
6010                    None => b.append_null(),
6011                }
6012            }
6013            b.finish()
6014        };
6015        // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb]
6016        let ree: ArrayRef = Arc::new(
6017            RunArray::try_new(
6018                &Int32Array::from(vec![2, 4, 6]),
6019                &mk(&[Some(b"aa"), None, Some(b"bb")]),
6020            )
6021            .unwrap(),
6022        );
6023        let flat: ArrayRef = Arc::new(mk(&[
6024            Some(b"aa"),
6025            Some(b"aa"),
6026            None,
6027            None,
6028            Some(b"bb"),
6029            Some(b"bb"),
6030        ]));
6031        ree_write_read_roundtrip(ree, flat);
6032    }
6033
6034    #[test]
6035    fn ree_single_run() {
6036        let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6037        let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6038        ree_write_read_roundtrip(ree, flat);
6039    }
6040
6041    #[test]
6042    fn ree_float32() {
6043        // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5]
6044        let ree: ArrayRef = Arc::new(
6045            RunArray::try_new(
6046                &Int32Array::from(vec![2, 4, 5]),
6047                &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6048            )
6049            .unwrap(),
6050        );
6051        let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6052            Some(1.0_f32),
6053            Some(1.0_f32),
6054            None,
6055            None,
6056            Some(2.5_f32),
6057        ]));
6058        ree_write_read_roundtrip(ree, flat);
6059    }
6060
6061    #[test]
6062    fn ree_sliced() {
6063        // A sliced (non-zero offset) REE array: verify that get_physical_index
6064        // correctly accounts for the logical offset when expanding.
6065        // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c]
6066        // After slice(2, 5) the logical view is [a, b, b, c, c].
6067        let full: ArrayRef = Arc::new(
6068            RunArray::try_new(
6069                &Int32Array::from(vec![3, 5, 7]),
6070                &StringArray::from(vec!["a", "b", "c"]),
6071            )
6072            .unwrap(),
6073        );
6074        let sliced = full.slice(2, 5);
6075        let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6076        ree_write_read_roundtrip(sliced, flat);
6077    }
6078
6079    #[test]
6080    fn ree_struct_with_ree_child() {
6081        // Struct with a REE string field and a REE int field — confirms
6082        // recursion visits every child and each collapses to the right leaf type.
6083        let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6084
6085        let col_a: ArrayRef = Arc::new(
6086            RunArray::try_new(
6087                &run_ends,
6088                &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6089            )
6090            .unwrap(),
6091        );
6092        let col_b: ArrayRef = Arc::new(
6093            RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6094        );
6095
6096        let struct_array: ArrayRef = Arc::new(StructArray::new(
6097            Fields::from(vec![
6098                Field::new("a", col_a.data_type().clone(), true),
6099                Field::new("b", col_b.data_type().clone(), true),
6100            ]),
6101            vec![col_a, col_b],
6102            None,
6103        ));
6104
6105        let schema = Arc::new(Schema::new(vec![Field::new(
6106            "row",
6107            struct_array.data_type().clone(),
6108            true,
6109        )]));
6110        let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6111
6112        let mut buf = Vec::new();
6113        let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6114        writer.write(&batch).unwrap();
6115        let metadata = writer.close().unwrap();
6116
6117        let parquet_schema = metadata.file_metadata().schema_descr();
6118        assert_eq!(parquet_schema.num_columns(), 2);
6119        assert_eq!(
6120            parquet_schema.column(0).physical_type(),
6121            crate::basic::Type::BYTE_ARRAY
6122        );
6123        assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6124        assert_eq!(
6125            parquet_schema.column(1).physical_type(),
6126            crate::basic::Type::INT32
6127        );
6128        assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6129    }
6130}