1use 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};
47use std::collections::HashSet;
48type DistinctValuesSet = HashSet<u64>;
49#[cfg(feature = "encryption")]
50use crate::encryption::encrypt::FileEncryptor;
51use crate::errors::{ParquetError, Result};
52use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData};
53use crate::file::properties::{WriterProperties, WriterPropertiesPtr};
54use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
55use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
56use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor};
57use levels::{ArrayLevels, calculate_array_levels};
58
59mod byte_array;
60mod levels;
61
62#[doc(inline)]
63pub use crate::column::page_store::{
64 InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs,
65 PageStoreFactory,
66};
67
68pub struct ArrowWriter<W: Write> {
185 writer: SerializedFileWriter<W>,
187
188 in_progress: Option<ArrowRowGroupWriter>,
190
191 arrow_schema: SchemaRef,
195
196 row_group_writer_factory: ArrowRowGroupWriterFactory,
198
199 max_row_group_row_count: Option<usize>,
201
202 max_row_group_bytes: Option<usize>,
204
205 cdc_chunkers: Option<Vec<ContentDefinedChunker>>,
207}
208
209impl<W: Write + Send> std::fmt::Debug for ArrowWriter<W> {
210 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
211 let buffered_memory = self.in_progress_size();
212 f.debug_struct("ArrowWriter")
213 .field("writer", &self.writer)
214 .field("in_progress_size", &format_args!("{buffered_memory} bytes"))
215 .field("in_progress_rows", &self.in_progress_rows())
216 .field("arrow_schema", &self.arrow_schema)
217 .field("max_row_group_row_count", &self.max_row_group_row_count)
218 .field("max_row_group_bytes", &self.max_row_group_bytes)
219 .finish()
220 }
221}
222
223impl<W: Write + Send> ArrowWriter<W> {
224 pub fn try_new(
230 writer: W,
231 arrow_schema: SchemaRef,
232 props: Option<WriterProperties>,
233 ) -> Result<Self> {
234 let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
235 Self::try_new_with_options(writer, arrow_schema, options)
236 }
237
238 pub fn try_new_with_options(
244 writer: W,
245 arrow_schema: SchemaRef,
246 options: ArrowWriterOptions,
247 ) -> Result<Self> {
248 let mut props = options.properties;
249
250 let schema = if let Some(parquet_schema) = options.schema_descr {
251 parquet_schema.clone()
252 } else {
253 let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types());
254 if let Some(schema_root) = &options.schema_root {
255 converter = converter.schema_root(schema_root);
256 }
257
258 converter.convert(&arrow_schema)?
259 };
260
261 if !options.skip_arrow_metadata {
262 add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
264 }
265
266 let max_row_group_row_count = props.max_row_group_row_count();
267 let max_row_group_bytes = props.max_row_group_bytes();
268
269 let props_ptr = Arc::new(props);
270 let file_writer =
271 SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
272
273 let mut row_group_writer_factory =
274 ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone());
275 if let Some(page_store_factory) = options.page_store_factory {
276 row_group_writer_factory =
277 row_group_writer_factory.with_page_store_factory(page_store_factory);
278 }
279
280 let cdc_chunkers = props_ptr
281 .content_defined_chunking()
282 .map(|opts| {
283 file_writer
284 .schema_descr()
285 .columns()
286 .iter()
287 .map(|desc| ContentDefinedChunker::new(desc, opts))
288 .collect::<Result<Vec<_>>>()
289 })
290 .transpose()?;
291
292 Ok(Self {
293 writer: file_writer,
294 in_progress: None,
295 arrow_schema,
296 row_group_writer_factory,
297 max_row_group_row_count,
298 max_row_group_bytes,
299 cdc_chunkers,
300 })
301 }
302
303 pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
305 self.writer.flushed_row_groups()
306 }
307
308 pub fn memory_size(&self) -> usize {
313 match &self.in_progress {
314 Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(),
315 None => 0,
316 }
317 }
318
319 pub fn in_progress_size(&self) -> usize {
326 match &self.in_progress {
327 Some(in_progress) => in_progress
328 .writers
329 .iter()
330 .map(|x| x.get_estimated_total_bytes())
331 .sum(),
332 None => 0,
333 }
334 }
335
336 pub fn in_progress_rows(&self) -> usize {
338 self.in_progress
339 .as_ref()
340 .map(|x| x.buffered_rows)
341 .unwrap_or_default()
342 }
343
344 pub fn bytes_written(&self) -> usize {
346 self.writer.bytes_written()
347 }
348
349 pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
361 if batch.num_rows() == 0 {
362 return Ok(());
363 }
364
365 let mut remaining = batch.clone();
368
369 loop {
370 let in_progress = match &mut self.in_progress {
371 Some(in_progress) => in_progress,
372 x => x.insert(
373 self.row_group_writer_factory
374 .create_row_group_writer(self.writer.flushed_row_groups().len())?,
375 ),
376 };
377 let buffered_rows = in_progress.buffered_rows;
378
379 let mut split_at = match self.max_row_group_row_count {
382 Some(max_rows) if buffered_rows + remaining.num_rows() > max_rows => {
383 Some(max_rows - buffered_rows)
384 }
385 _ => None,
386 };
387
388 let candidate_rows = split_at.unwrap_or_else(|| remaining.num_rows());
393
394 if let Some(max_bytes) = self.max_row_group_bytes
395 && buffered_rows > 0
396 {
397 let current_bytes = in_progress.get_estimated_total_bytes();
398
399 if current_bytes >= max_bytes {
400 self.flush()?;
401 continue;
402 }
403
404 if let Some(avg_row_bytes) = current_bytes
405 .checked_div(buffered_rows)
406 .filter(|avg_row_bytes| *avg_row_bytes > 0)
407 {
408 let remaining_bytes = max_bytes - current_bytes;
410 let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
411
412 if candidate_rows > rows_that_fit {
413 if rows_that_fit > 0 {
414 split_at = Some(rows_that_fit);
415 } else {
416 self.flush()?;
417 continue;
418 }
419 }
420 }
421 }
422
423 let rest = split_at.map(|to_write| {
424 let rest = remaining.slice(to_write, remaining.num_rows() - to_write);
425 remaining = remaining.slice(0, to_write);
426 rest
427 });
428
429 let in_progress = self.in_progress.as_mut().unwrap();
430 match self.cdc_chunkers.as_mut() {
431 Some(chunkers) => in_progress.write_with_chunkers(&remaining, chunkers)?,
432 None => in_progress.write(&remaining)?,
433 }
434
435 let should_flush = self
436 .max_row_group_row_count
437 .is_some_and(|max| in_progress.buffered_rows >= max)
438 || self
439 .max_row_group_bytes
440 .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
441
442 if should_flush {
443 self.flush()?
444 }
445
446 match rest {
447 Some(rest) => remaining = rest,
448 None => return Ok(()),
449 }
450 }
451 }
452
453 pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
458 self.writer.write_all(buf)
459 }
460
461 pub fn sync(&mut self) -> std::io::Result<()> {
463 self.writer.flush()
464 }
465
466 pub fn flush(&mut self) -> Result<()> {
471 let Some(in_progress) = self.in_progress.take() else {
472 return Ok(());
473 };
474
475 let mut row_group_writer = self.writer.next_row_group()?;
476 for chunk in in_progress.close()? {
477 chunk.append_to_row_group(&mut row_group_writer)?;
478 }
479 row_group_writer.close()?;
480 Ok(())
481 }
482
483 pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
487 self.writer.append_key_value_metadata(kv_metadata)
488 }
489
490 pub fn inner(&self) -> &W {
492 self.writer.inner()
493 }
494
495 pub fn inner_mut(&mut self) -> &mut W {
504 self.writer.inner_mut()
505 }
506
507 pub fn into_inner(mut self) -> Result<W> {
509 self.flush()?;
510 self.writer.into_inner()
511 }
512
513 pub fn finish(&mut self) -> Result<ParquetMetaData> {
519 self.flush()?;
520 self.writer.finish()
521 }
522
523 pub fn close(mut self) -> Result<ParquetMetaData> {
525 self.finish()
526 }
527
528 pub fn into_serialized_writer(
535 mut self,
536 ) -> Result<(SerializedFileWriter<W>, ArrowRowGroupWriterFactory)> {
537 self.flush()?;
538 Ok((self.writer, self.row_group_writer_factory))
539 }
540}
541
542impl<W: Write + Send> RecordBatchWriter for ArrowWriter<W> {
543 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
544 self.write(batch).map_err(|e| e.into())
545 }
546
547 fn close(self) -> std::result::Result<(), ArrowError> {
548 self.close()?;
549 Ok(())
550 }
551}
552
553#[derive(Debug, Clone, Default)]
557pub struct ArrowWriterOptions {
558 properties: WriterProperties,
559 skip_arrow_metadata: bool,
560 schema_root: Option<String>,
561 schema_descr: Option<SchemaDescriptor>,
562 page_store_factory: Option<Arc<dyn PageStoreFactory>>,
563}
564
565impl ArrowWriterOptions {
566 pub fn new() -> Self {
568 Self::default()
569 }
570
571 pub fn with_properties(self, properties: WriterProperties) -> Self {
573 Self { properties, ..self }
574 }
575
576 pub fn with_page_store_factory(self, page_store_factory: Arc<dyn PageStoreFactory>) -> Self {
662 Self {
663 page_store_factory: Some(page_store_factory),
664 ..self
665 }
666 }
667
668 pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
675 Self {
676 skip_arrow_metadata,
677 ..self
678 }
679 }
680
681 pub fn with_schema_root(self, schema_root: String) -> Self {
683 Self {
684 schema_root: Some(schema_root),
685 ..self
686 }
687 }
688
689 pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self {
695 Self {
696 schema_descr: Some(schema_descr),
697 ..self
698 }
699 }
700}
701
702struct ArrowColumnChunkData {
708 length: usize,
709 store: Box<dyn PageStore>,
710 keys: Vec<PageKey>,
711 dictionary_keys: Vec<PageKey>,
722 dictionary_len: usize,
726}
727
728impl ArrowColumnChunkData {
729 fn new(store: Box<dyn PageStore>) -> Self {
730 Self {
731 length: 0,
732 store,
733 keys: Vec::new(),
734 dictionary_keys: Vec::new(),
735 dictionary_len: 0,
736 }
737 }
738
739 fn push(&mut self, value: Bytes) -> Result<()> {
742 let key = self.store.put(value)?;
743 self.keys.push(key);
744 Ok(())
745 }
746
747 fn push_dictionary(&mut self, value: Bytes) -> Result<()> {
751 self.dictionary_len += value.len();
752 let key = self.store.put(value)?;
753 self.dictionary_keys.push(key);
754 Ok(())
755 }
756
757 fn memory_size(&self) -> usize {
760 self.store.memory_size()
761 }
762}
763
764struct StreamingColumnChunkPages {
773 store: Box<dyn PageStore>,
774 keys: IntoIter<PageKey>,
777}
778
779impl StreamingColumnChunkPages {
780 fn new(data: ArrowColumnChunkData) -> Self {
781 let keys = if data.dictionary_keys.is_empty() {
784 data.keys
785 } else {
786 let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len());
787 keys.extend(data.dictionary_keys);
788 keys.extend(data.keys);
789 keys
790 };
791 Self {
792 store: data.store,
793 keys: keys.into_iter(),
794 }
795 }
796}
797
798impl Iterator for StreamingColumnChunkPages {
799 type Item = Result<Bytes>;
800
801 fn next(&mut self) -> Option<Self::Item> {
802 let key = self.keys.next()?;
803 Some(self.store.take(key))
804 }
805}
806
807type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
812
813struct ArrowPageWriter {
814 buffer: SharedColumnChunk,
815 #[cfg(feature = "encryption")]
816 page_encryptor: Option<PageEncryptor>,
817}
818
819impl ArrowPageWriter {
820 fn new(store: Box<dyn PageStore>) -> Self {
822 Self {
823 buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))),
824 #[cfg(feature = "encryption")]
825 page_encryptor: None,
826 }
827 }
828
829 #[cfg(feature = "encryption")]
830 pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
831 self.page_encryptor = page_encryptor;
832 self
833 }
834
835 #[cfg(feature = "encryption")]
836 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
837 self.page_encryptor.as_mut()
838 }
839
840 #[cfg(not(feature = "encryption"))]
843 #[expect(
844 clippy::needless_pass_by_ref_mut,
845 reason = "mirrors the encryption-enabled signature"
846 )]
847 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
848 None
849 }
850}
851
852impl PageWriter for ArrowPageWriter {
853 fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
854 let page = match self.page_encryptor_mut() {
855 Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
856 None => page,
857 };
858
859 let page_header = page.to_thrift_header()?;
860 let header = {
861 let mut header = Vec::with_capacity(1024);
862
863 match self.page_encryptor_mut() {
864 Some(page_encryptor) => {
865 page_encryptor.encrypt_page_header(&page_header, &mut header)?;
866 if page.compressed_page().is_data_page() {
867 page_encryptor.increment_page();
868 }
869 }
870 None => {
871 let mut protocol = ThriftCompactOutputProtocol::new(&mut header);
872 page_header.write_thrift(&mut protocol)?;
873 }
874 }
875
876 Bytes::from(header)
877 };
878
879 let mut buf = self.buffer.try_lock().unwrap();
880
881 let data = page.compressed_page().buffer().clone();
882 let compressed_size = data.len() + header.len();
883
884 let mut spec = PageWriteSpec::new();
885 spec.page_type = page.page_type();
886 spec.num_values = page.num_values();
887 spec.uncompressed_size = page.uncompressed_size() + header.len();
888 spec.offset = buf.length as u64;
889 spec.compressed_size = compressed_size;
890 spec.bytes_written = compressed_size as u64;
891
892 buf.length += compressed_size;
893 if spec.page_type == PageType::DICTIONARY_PAGE {
894 buf.push_dictionary(header)?;
897 buf.push_dictionary(data)?;
898 } else {
899 buf.push(header)?;
900 buf.push(data)?;
901 }
902
903 Ok(spec)
904 }
905
906 fn defers_dictionary_ordering(&self) -> bool {
907 true
912 }
913
914 fn buffered_memory_size(&self) -> usize {
915 self.buffer.try_lock().unwrap().memory_size()
918 }
919
920 fn close(&mut self) -> Result<()> {
921 Ok(())
922 }
923}
924
925#[derive(Debug)]
927pub struct ArrowLeafColumn(ArrayLevels);
928
929pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result<Vec<ArrowLeafColumn>> {
934 let levels = calculate_array_levels(array, field)?;
935 Ok(levels.into_iter().map(ArrowLeafColumn).collect())
936}
937
938pub struct ArrowColumnChunk {
940 data: ArrowColumnChunkData,
941 close: ColumnCloseResult,
942}
943
944impl std::fmt::Debug for ArrowColumnChunk {
945 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
946 f.debug_struct("ArrowColumnChunk")
947 .field("length", &self.data.length)
948 .finish_non_exhaustive()
949 }
950}
951
952impl ArrowColumnChunk {
953 pub fn close(&self) -> &ColumnCloseResult {
960 &self.close
961 }
962
963 pub fn close_mut(&mut self) -> &mut ColumnCloseResult {
970 &mut self.close
971 }
972
973 pub fn append_to_row_group<W: Write + Send>(
976 self,
977 writer: &mut SerializedRowGroupWriter<'_, W>,
978 ) -> Result<()> {
979 let ArrowColumnChunk { data, close } = self;
980
981 let close = close.update_dictionary_location(data.dictionary_len)?;
985
986 let pages = StreamingColumnChunkPages::new(data);
987 writer.append_column_from_pages(pages, close)
988 }
989}
990
991pub struct ArrowColumnWriter {
1089 writer: ArrowColumnWriterImpl,
1090 chunk: SharedColumnChunk,
1091 distinct_values_seen: Option<DistinctValuesSet>,
1094}
1095
1096impl std::fmt::Debug for ArrowColumnWriter {
1097 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1098 f.debug_struct("ArrowColumnWriter").finish_non_exhaustive()
1099 }
1100}
1101
1102enum ArrowColumnWriterImpl {
1103 ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
1104 Column(ColumnWriter<'static>),
1105}
1106
1107impl ArrowColumnWriter {
1108 pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> {
1110 self.write_internal(&col.0)
1111 }
1112
1113 fn write_with_chunker(
1115 &mut self,
1116 col: &ArrowLeafColumn,
1117 chunker: &mut ContentDefinedChunker,
1118 ) -> Result<()> {
1119 let levels = &col.0;
1120 let chunks = chunker.get_arrow_chunks(
1121 levels.def_level_data().as_ref(),
1122 levels.rep_level_data().as_ref(),
1123 levels.array(),
1124 )?;
1125
1126 let num_chunks = chunks.len();
1127 for (i, chunk) in chunks.iter().enumerate() {
1128 let chunk_levels = levels.slice_for_chunk(chunk);
1129 self.write_internal(&chunk_levels)?;
1130
1131 if i + 1 < num_chunks {
1133 match &mut self.writer {
1134 ArrowColumnWriterImpl::Column(c) => c.add_data_page()?,
1135 ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?,
1136 }
1137 }
1138 }
1139 Ok(())
1140 }
1141
1142 fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
1143 if let Some(seen) = &mut self.distinct_values_seen {
1144 let array = levels.array();
1145 let non_null = levels.non_null_indices();
1146 match array.as_any_dictionary_opt() {
1147 Some(dict) => {
1148 let keys = dict.keys();
1152 let key_data = keys.to_data();
1153 let offset = key_data.offset();
1154 let width = arrow_key_byte_width(keys.data_type());
1155 if width > 0 {
1156 let buffer = key_data.buffers()[0].as_slice();
1157 for &row in non_null {
1159 let pos = (offset + row) * width;
1160 seen.insert(hash_bytes(&buffer[pos..pos + width]));
1161 }
1162 }
1163 }
1164 None => update_distinct_values_seen(array.as_ref(), non_null, seen),
1166 }
1167 }
1168
1169 match &mut self.writer {
1170 ArrowColumnWriterImpl::Column(c) => {
1171 let leaf = levels.array();
1172 match leaf.as_any_dictionary_opt() {
1173 Some(dictionary) => {
1174 let materialized =
1175 arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?;
1176 write_leaf(c, &materialized, levels)?
1177 }
1178 None => write_leaf(c, leaf, levels)?,
1179 };
1180 }
1181 ArrowColumnWriterImpl::ByteArray(c) => {
1182 write_primitive(c, levels.array().as_ref(), levels)?;
1183 }
1184 }
1185 Ok(())
1186 }
1187
1188 pub fn close(self) -> Result<ArrowColumnChunk> {
1195 let distinct_count = self
1196 .distinct_values_seen
1197 .as_ref()
1198 .filter(|s| !s.is_empty())
1199 .map(|s| s.len() as u64);
1200 let close = match self.writer {
1201 ArrowColumnWriterImpl::ByteArray(mut c) => {
1202 if let Some(count) = distinct_count {
1203 c.set_distinct_count_override(count);
1204 }
1205 c.close()?
1206 }
1207 ArrowColumnWriterImpl::Column(mut c) => {
1208 if let Some(count) = distinct_count {
1209 c.set_distinct_count_override(count);
1210 }
1211 c.close()?
1212 }
1213 };
1214 let chunk = Arc::try_unwrap(self.chunk)
1216 .map_err(|_| general_err!("Internal Error: the column chunk is still shared"))?;
1217 let data = chunk
1218 .into_inner()
1219 .map_err(|_| general_err!("The column chunk lock is poisoned"))?;
1220 Ok(ArrowColumnChunk { data, close })
1221 }
1222
1223 pub fn memory_size(&self) -> usize {
1234 match &self.writer {
1235 ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1236 ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1237 }
1238 }
1239
1240 pub fn get_estimated_total_bytes(&self) -> usize {
1248 match &self.writer {
1249 ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1250 ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1251 }
1252 }
1253}
1254
1255#[derive(Debug)]
1262struct ArrowRowGroupWriter {
1263 writers: Vec<ArrowColumnWriter>,
1264 schema: SchemaRef,
1265 buffered_rows: usize,
1266}
1267
1268impl ArrowRowGroupWriter {
1269 fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self {
1270 Self {
1271 writers,
1272 schema: arrow.clone(),
1273 buffered_rows: 0,
1274 }
1275 }
1276
1277 fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1278 self.buffered_rows += batch.num_rows();
1279 let mut writers = self.writers.iter_mut();
1280 for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1281 for leaf in compute_leaves(field.as_ref(), column)? {
1282 writers.next().unwrap().write(&leaf)?;
1283 }
1284 }
1285 Ok(())
1286 }
1287
1288 fn write_with_chunkers(
1289 &mut self,
1290 batch: &RecordBatch,
1291 chunkers: &mut [ContentDefinedChunker],
1292 ) -> Result<()> {
1293 self.buffered_rows += batch.num_rows();
1294 let mut writers = self.writers.iter_mut();
1295 let mut chunkers = chunkers.iter_mut();
1296 for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1297 for leaf in compute_leaves(field.as_ref(), column)? {
1298 writers
1299 .next()
1300 .unwrap()
1301 .write_with_chunker(&leaf, chunkers.next().unwrap())?;
1302 }
1303 }
1304 Ok(())
1305 }
1306
1307 fn get_estimated_total_bytes(&self) -> usize {
1309 self.writers
1310 .iter()
1311 .map(|x| x.get_estimated_total_bytes())
1312 .sum()
1313 }
1314
1315 fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1316 self.writers
1317 .into_iter()
1318 .map(|writer| writer.close())
1319 .collect()
1320 }
1321}
1322
1323#[derive(Debug)]
1328pub struct ArrowRowGroupWriterFactory {
1329 schema: SchemaDescPtr,
1330 arrow_schema: SchemaRef,
1331 props: WriterPropertiesPtr,
1332 page_store_factory: Arc<dyn PageStoreFactory>,
1333 #[cfg(feature = "encryption")]
1334 file_encryptor: Option<Arc<FileEncryptor>>,
1335}
1336
1337impl ArrowRowGroupWriterFactory {
1338 pub fn new<W: Write + Send>(
1340 file_writer: &SerializedFileWriter<W>,
1341 arrow_schema: SchemaRef,
1342 ) -> Self {
1343 let schema = Arc::clone(file_writer.schema_descr_ptr());
1344 let props = Arc::clone(file_writer.properties());
1345 Self {
1346 schema,
1347 arrow_schema,
1348 props,
1349 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1350 #[cfg(feature = "encryption")]
1351 file_encryptor: file_writer.file_encryptor(),
1352 }
1353 }
1354
1355 pub fn with_page_store_factory(
1359 mut self,
1360 page_store_factory: Arc<dyn PageStoreFactory>,
1361 ) -> Self {
1362 self.page_store_factory = page_store_factory;
1363 self
1364 }
1365
1366 fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1367 let writers = self.create_column_writers(row_group_index)?;
1368 Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
1369 }
1370
1371 pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1373 let mut writers = Vec::with_capacity(self.arrow_schema.fields.len());
1374 let mut leaves = self.schema.columns().iter();
1375 let column_factory = self.column_writer_factory(row_group_index);
1376 for field in &self.arrow_schema.fields {
1377 column_factory.get_arrow_column_writer(
1378 field.data_type(),
1379 &self.props,
1380 &mut leaves,
1381 &mut writers,
1382 )?;
1383 }
1384 Ok(writers)
1385 }
1386
1387 #[cfg(feature = "encryption")]
1388 fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1389 ArrowColumnWriterFactory::new()
1390 .with_page_store_factory(self.page_store_factory.clone())
1391 .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1392 }
1393
1394 #[cfg(not(feature = "encryption"))]
1395 fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1396 ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1397 }
1398}
1399
1400struct ArrowColumnWriterFactory {
1402 page_store_factory: Arc<dyn PageStoreFactory>,
1404 #[cfg(feature = "encryption")]
1405 row_group_index: usize,
1406 #[cfg(feature = "encryption")]
1407 file_encryptor: Option<Arc<FileEncryptor>>,
1408}
1409
1410impl ArrowColumnWriterFactory {
1411 pub fn new() -> Self {
1412 Self {
1413 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1414 #[cfg(feature = "encryption")]
1415 row_group_index: 0,
1416 #[cfg(feature = "encryption")]
1417 file_encryptor: None,
1418 }
1419 }
1420
1421 pub fn with_page_store_factory(
1423 mut self,
1424 page_store_factory: Arc<dyn PageStoreFactory>,
1425 ) -> Self {
1426 self.page_store_factory = page_store_factory;
1427 self
1428 }
1429
1430 #[cfg(feature = "encryption")]
1431 pub fn with_file_encryptor(
1432 mut self,
1433 row_group_index: usize,
1434 file_encryptor: Option<Arc<FileEncryptor>>,
1435 ) -> Self {
1436 self.row_group_index = row_group_index;
1437 self.file_encryptor = file_encryptor;
1438 self
1439 }
1440
1441 #[cfg(feature = "encryption")]
1442 fn create_page_writer(
1443 &self,
1444 column_descriptor: &ColumnDescPtr,
1445 column_index: usize,
1446 ) -> Result<Box<ArrowPageWriter>> {
1447 let column_path = column_descriptor.path().string();
1448 let page_encryptor = PageEncryptor::create_if_column_encrypted(
1449 self.file_encryptor.as_ref(),
1450 self.row_group_index,
1451 column_index,
1452 &column_path,
1453 )?;
1454 let args = PageStoreArgs::new(column_index, column_descriptor);
1455 let store = self.page_store_factory.create(&args)?;
1456 Ok(Box::new(
1457 ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1458 ))
1459 }
1460
1461 #[cfg(not(feature = "encryption"))]
1462 fn create_page_writer(
1463 &self,
1464 column_descriptor: &ColumnDescPtr,
1465 column_index: usize,
1466 ) -> Result<Box<ArrowPageWriter>> {
1467 let args = PageStoreArgs::new(column_index, column_descriptor);
1468 let store = self.page_store_factory.create(&args)?;
1469 Ok(Box::new(ArrowPageWriter::new(store)))
1470 }
1471
1472 fn get_arrow_column_writer(
1475 &self,
1476 data_type: &ArrowDataType,
1477 props: &WriterPropertiesPtr,
1478 leaves: &mut Iter<'_, ColumnDescPtr>,
1479 out: &mut Vec<ArrowColumnWriter>,
1480 ) -> Result<()> {
1481 let write_distinct_values = props.write_row_group_number_distinct_values();
1482
1483 let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1485 let page_writer = self.create_page_writer(desc, out.len())?;
1486 let chunk = page_writer.buffer.clone();
1487 let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1488 Ok(ArrowColumnWriter {
1489 chunk,
1490 writer: ArrowColumnWriterImpl::Column(writer),
1491 distinct_values_seen: write_distinct_values.then(HashSet::new),
1492 })
1493 };
1494
1495 let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1497 let page_writer = self.create_page_writer(desc, out.len())?;
1498 let chunk = page_writer.buffer.clone();
1499 let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1500 Ok(ArrowColumnWriter {
1501 chunk,
1502 writer: ArrowColumnWriterImpl::ByteArray(writer),
1503 distinct_values_seen: write_distinct_values.then(HashSet::new),
1504 })
1505 };
1506
1507 match data_type {
1508 _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1509 ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1510 out.push(col(leaves.next().unwrap())?)
1511 }
1512 ArrowDataType::LargeBinary
1513 | ArrowDataType::Binary
1514 | ArrowDataType::Utf8
1515 | ArrowDataType::LargeUtf8
1516 | ArrowDataType::BinaryView
1517 | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1518 ArrowDataType::List(f)
1519 | ArrowDataType::LargeList(f)
1520 | ArrowDataType::FixedSizeList(f, _)
1521 | ArrowDataType::ListView(f)
1522 | ArrowDataType::LargeListView(f) => {
1523 self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1524 }
1525 ArrowDataType::Struct(fields) => {
1526 for field in fields {
1527 self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1528 }
1529 }
1530 ArrowDataType::Map(f, _) => match f.data_type() {
1531 ArrowDataType::Struct(f) => {
1532 self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1533 self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1534 }
1535 _ => unreachable!("invalid map type"),
1536 },
1537 ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1538 ArrowDataType::Utf8
1539 | ArrowDataType::LargeUtf8
1540 | ArrowDataType::Binary
1541 | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1542 ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1543 out.push(bytes(leaves.next().unwrap())?)
1544 }
1545 ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
1546 _ => out.push(col(leaves.next().unwrap())?),
1547 },
1548 ArrowDataType::RunEndEncoded(_, value_field) => {
1549 self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1550 }
1551 _ => {
1552 return Err(ParquetError::NYI(format!(
1553 "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1554 )));
1555 }
1556 }
1557 Ok(())
1558 }
1559}
1560
1561fn write_leaf(
1562 writer: &mut ColumnWriter<'_>,
1563 column: &dyn arrow_array::Array,
1564 levels: &ArrayLevels,
1565) -> Result<usize> {
1566 let indices = levels.non_null_indices();
1567
1568 match writer {
1569 ColumnWriter::Int32ColumnWriter(typed) => {
1571 match column.data_type() {
1572 ArrowDataType::Null => {
1573 let array = Int32Array::new_null(column.len());
1574 write_primitive(typed, array.values(), levels)
1575 }
1576 ArrowDataType::Int8 => {
1577 let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32);
1578 write_primitive(typed, array.values(), levels)
1579 }
1580 ArrowDataType::Int16 => {
1581 let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32);
1582 write_primitive(typed, array.values(), levels)
1583 }
1584 ArrowDataType::Int32 => {
1585 write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels)
1586 }
1587 ArrowDataType::UInt8 => {
1588 let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32);
1589 write_primitive(typed, array.values(), levels)
1590 }
1591 ArrowDataType::UInt16 => {
1592 let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32);
1593 write_primitive(typed, array.values(), levels)
1594 }
1595 ArrowDataType::UInt32 => {
1596 let array = column.as_primitive::<UInt32Type>();
1599 write_primitive(typed, array.values().inner().typed_data(), levels)
1600 }
1601 ArrowDataType::Date32 => {
1602 let array = column.as_primitive::<Date32Type>();
1603 write_primitive(typed, array.values(), levels)
1604 }
1605 ArrowDataType::Time32(TimeUnit::Second) => {
1606 let array = column.as_primitive::<Time32SecondType>();
1607 write_primitive(typed, array.values(), levels)
1608 }
1609 ArrowDataType::Time32(TimeUnit::Millisecond) => {
1610 let array = column.as_primitive::<Time32MillisecondType>();
1611 write_primitive(typed, array.values(), levels)
1612 }
1613 ArrowDataType::Date64 => {
1614 let array: Int32Array = column
1616 .as_primitive::<Date64Type>()
1617 .unary(|x| (x / 86_400_000) as _);
1618
1619 write_primitive(typed, array.values(), levels)
1620 }
1621 ArrowDataType::Decimal32(_, _) => {
1622 let array = column
1623 .as_primitive::<Decimal32Type>()
1624 .unary::<_, Int32Type>(|v| v);
1625 write_primitive(typed, array.values(), levels)
1626 }
1627 ArrowDataType::Decimal64(_, _) => {
1628 let array = column
1630 .as_primitive::<Decimal64Type>()
1631 .unary::<_, Int32Type>(|v| v as i32);
1632 write_primitive(typed, array.values(), levels)
1633 }
1634 ArrowDataType::Decimal128(_, _) => {
1635 let array = column
1637 .as_primitive::<Decimal128Type>()
1638 .unary::<_, Int32Type>(|v| v as i32);
1639 write_primitive(typed, array.values(), levels)
1640 }
1641 ArrowDataType::Decimal256(_, _) => {
1642 let array = column
1644 .as_primitive::<Decimal256Type>()
1645 .unary::<_, Int32Type>(|v| v.as_i128() as i32);
1646 write_primitive(typed, array.values(), levels)
1647 }
1648 d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1649 }
1650 }
1651 ColumnWriter::BoolColumnWriter(typed) => {
1652 let array = column.as_boolean();
1653 let values = get_bool_array_slice(array, indices.iter().copied());
1654 typed.write_batch_internal(
1655 values.as_slice(),
1656 None,
1657 levels.def_level_data().as_ref(),
1658 levels.rep_level_data().as_ref(),
1659 None,
1660 None,
1661 None,
1662 )
1663 }
1664 ColumnWriter::Int64ColumnWriter(typed) => {
1665 match column.data_type() {
1666 ArrowDataType::Date64 => {
1667 let array = column
1668 .as_primitive::<Date64Type>()
1669 .reinterpret_cast::<Int64Type>();
1670
1671 write_primitive(typed, array.values(), levels)
1672 }
1673 ArrowDataType::Int64 => {
1674 let array = column.as_primitive::<Int64Type>();
1675 write_primitive(typed, array.values(), levels)
1676 }
1677 ArrowDataType::UInt64 => {
1678 let values = column.as_primitive::<UInt64Type>().values();
1679 let array = values.inner().typed_data::<i64>();
1682 write_primitive(typed, array, levels)
1683 }
1684 ArrowDataType::Time64(TimeUnit::Microsecond) => {
1685 let array = column.as_primitive::<Time64MicrosecondType>();
1686 write_primitive(typed, array.values(), levels)
1687 }
1688 ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1689 let array = column.as_primitive::<Time64NanosecondType>();
1690 write_primitive(typed, array.values(), levels)
1691 }
1692 ArrowDataType::Timestamp(unit, _) => match unit {
1693 TimeUnit::Second => {
1694 let array = column.as_primitive::<TimestampSecondType>();
1695 write_primitive(typed, array.values(), levels)
1696 }
1697 TimeUnit::Millisecond => {
1698 let array = column.as_primitive::<TimestampMillisecondType>();
1699 write_primitive(typed, array.values(), levels)
1700 }
1701 TimeUnit::Microsecond => {
1702 let array = column.as_primitive::<TimestampMicrosecondType>();
1703 write_primitive(typed, array.values(), levels)
1704 }
1705 TimeUnit::Nanosecond => {
1706 let array = column.as_primitive::<TimestampNanosecondType>();
1707 write_primitive(typed, array.values(), levels)
1708 }
1709 },
1710 ArrowDataType::Duration(unit) => match unit {
1711 TimeUnit::Second => {
1712 let array = column.as_primitive::<DurationSecondType>();
1713 write_primitive(typed, array.values(), levels)
1714 }
1715 TimeUnit::Millisecond => {
1716 let array = column.as_primitive::<DurationMillisecondType>();
1717 write_primitive(typed, array.values(), levels)
1718 }
1719 TimeUnit::Microsecond => {
1720 let array = column.as_primitive::<DurationMicrosecondType>();
1721 write_primitive(typed, array.values(), levels)
1722 }
1723 TimeUnit::Nanosecond => {
1724 let array = column.as_primitive::<DurationNanosecondType>();
1725 write_primitive(typed, array.values(), levels)
1726 }
1727 },
1728 ArrowDataType::Decimal64(_, _) => {
1729 let array = column
1730 .as_primitive::<Decimal64Type>()
1731 .reinterpret_cast::<Int64Type>();
1732 write_primitive(typed, array.values(), levels)
1733 }
1734 ArrowDataType::Decimal128(_, _) => {
1735 let array = column
1737 .as_primitive::<Decimal128Type>()
1738 .unary::<_, Int64Type>(|v| v as i64);
1739 write_primitive(typed, array.values(), levels)
1740 }
1741 ArrowDataType::Decimal256(_, _) => {
1742 let array = column
1744 .as_primitive::<Decimal256Type>()
1745 .unary::<_, Int64Type>(|v| v.as_i128() as i64);
1746 write_primitive(typed, array.values(), levels)
1747 }
1748 d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1749 }
1750 }
1751 ColumnWriter::Int96ColumnWriter(_typed) => {
1752 unreachable!("Currently unreachable because data type not supported")
1753 }
1754 ColumnWriter::FloatColumnWriter(typed) => {
1755 let array = column.as_primitive::<Float32Type>();
1756 write_primitive(typed, array.values(), levels)
1757 }
1758 ColumnWriter::DoubleColumnWriter(typed) => {
1759 let array = column.as_primitive::<Float64Type>();
1760 write_primitive(typed, array.values(), levels)
1761 }
1762 ColumnWriter::ByteArrayColumnWriter(_) => {
1763 unreachable!("should use ByteArrayWriter")
1764 }
1765 ColumnWriter::FixedLenByteArrayColumnWriter(typed) => {
1766 let bytes = match column.data_type() {
1767 ArrowDataType::Interval(interval_unit) => match interval_unit {
1768 IntervalUnit::YearMonth => {
1769 let array = column.as_primitive::<IntervalYearMonthType>();
1770 get_interval_ym_array_slice(array, indices.iter().copied())
1771 }
1772 IntervalUnit::DayTime => {
1773 let array = column.as_primitive::<IntervalDayTimeType>();
1774 get_interval_dt_array_slice(array, indices.iter().copied())
1775 }
1776 IntervalUnit::MonthDayNano => {
1777 return Err(ParquetError::NYI(format!(
1778 "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
1779 )));
1780 }
1781 },
1782 ArrowDataType::FixedSizeBinary(_) => {
1783 let array = column.as_fixed_size_binary();
1784 get_fsb_array_slice(array, indices.iter().copied())
1785 }
1786 ArrowDataType::Decimal32(_, _) => {
1787 let array = column.as_primitive::<Decimal32Type>();
1788 get_decimal_array_slice(array, indices.iter().copied())
1789 }
1790 ArrowDataType::Decimal64(_, _) => {
1791 let array = column.as_primitive::<Decimal64Type>();
1792 get_decimal_array_slice(array, indices.iter().copied())
1793 }
1794 ArrowDataType::Decimal128(_, _) => {
1795 let array = column.as_primitive::<Decimal128Type>();
1796 get_decimal_array_slice(array, indices.iter().copied())
1797 }
1798 ArrowDataType::Decimal256(_, _) => {
1799 let array = column.as_primitive::<Decimal256Type>();
1800 get_decimal_array_slice(array, indices.iter().copied())
1801 }
1802 ArrowDataType::Float16 => {
1803 let array = column.as_primitive::<Float16Type>();
1804 get_float_16_array_slice(array, indices.iter().copied())
1805 }
1806 _ => {
1807 return Err(ParquetError::NYI(
1808 "Attempting to write an Arrow type that is not yet implemented".to_string(),
1809 ));
1810 }
1811 };
1812 typed.write_batch_internal(
1813 bytes.as_slice(),
1814 None,
1815 levels.def_level_data().as_ref(),
1816 levels.rep_level_data().as_ref(),
1817 None,
1818 None,
1819 None,
1820 )
1821 }
1822 }
1823}
1824
1825fn write_primitive<E: ColumnValueEncoder>(
1826 writer: &mut GenericColumnWriter<E>,
1827 values: &E::Values,
1828 levels: &ArrayLevels,
1829) -> Result<usize> {
1830 writer.write_batch_internal(
1831 values,
1832 Some(levels.non_null_indices()),
1833 levels.def_level_data().as_ref(),
1834 levels.rep_level_data().as_ref(),
1835 None,
1836 None,
1837 None,
1838 )
1839}
1840
1841fn get_bool_array_slice(
1842 array: &arrow_array::BooleanArray,
1843 indices: impl ExactSizeIterator<Item = usize>,
1844) -> Vec<bool> {
1845 let mut values = Vec::with_capacity(indices.len());
1846 for i in indices {
1847 values.push(array.value(i))
1848 }
1849 values
1850}
1851
1852fn get_interval_ym_array_slice(
1855 array: &arrow_array::IntervalYearMonthArray,
1856 indices: impl ExactSizeIterator<Item = usize>,
1857) -> Vec<FixedLenByteArray> {
1858 chunk_array_slice(12, indices, move |i, chunk| {
1859 let value = array.value(i);
1860 chunk[0..4].copy_from_slice(&value.to_le_bytes());
1861 })
1862}
1863
1864fn get_interval_dt_array_slice(
1867 array: &arrow_array::IntervalDayTimeArray,
1868 indices: impl ExactSizeIterator<Item = usize>,
1869) -> Vec<FixedLenByteArray> {
1870 chunk_array_slice(12, indices, move |i, chunk| {
1871 let value = array.value(i);
1872 chunk[4..8].copy_from_slice(&value.days.to_le_bytes());
1873 chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1874 })
1875}
1876
1877trait NativeDecimalType: DecimalType {
1878 type NativeBytes: AsRef<[u8]>;
1879
1880 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
1881}
1882impl NativeDecimalType for Decimal32Type {
1883 type NativeBytes = [u8; Self::BYTE_LENGTH];
1884
1885 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1886 value.to_be_bytes()
1887 }
1888}
1889impl NativeDecimalType for Decimal64Type {
1890 type NativeBytes = [u8; Self::BYTE_LENGTH];
1891
1892 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1893 value.to_be_bytes()
1894 }
1895}
1896impl NativeDecimalType for Decimal128Type {
1897 type NativeBytes = [u8; Self::BYTE_LENGTH];
1898
1899 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1900 value.to_be_bytes()
1901 }
1902}
1903impl NativeDecimalType for Decimal256Type {
1904 type NativeBytes = [u8; Self::BYTE_LENGTH];
1905
1906 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1907 value.to_be_bytes()
1908 }
1909}
1910
1911fn get_decimal_array_slice<T: NativeDecimalType>(
1912 array: &PrimitiveArray<T>,
1913 indices: impl ExactSizeIterator<Item = usize>,
1914) -> Vec<FixedLenByteArray> {
1915 let chunk_size = decimal_length_from_precision(array.precision());
1916 assert!(chunk_size <= T::BYTE_LENGTH);
1917
1918 if chunk_size == T::BYTE_LENGTH {
1919 chunk_array_slice(chunk_size, indices, move |i, chunk| {
1921 let as_be_bytes = T::to_be_bytes(array.value(i));
1922 chunk.copy_from_slice(as_be_bytes.as_ref());
1923 })
1924 } else {
1925 chunk_array_slice(chunk_size, indices, move |i, chunk| {
1926 let as_be_bytes = T::to_be_bytes(array.value(i));
1927 let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..];
1928 chunk.copy_from_slice(resized_value);
1929 })
1930 }
1931}
1932
1933fn get_float_16_array_slice(
1934 array: &arrow_array::Float16Array,
1935 indices: impl ExactSizeIterator<Item = usize>,
1936) -> Vec<FixedLenByteArray> {
1937 chunk_array_slice(2, indices, move |i, chunk| {
1938 let value = array.value(i).to_le_bytes();
1939 chunk.copy_from_slice(&value);
1940 })
1941}
1942
1943fn get_fsb_array_slice(
1944 array: &arrow_array::FixedSizeBinaryArray,
1945 indices: impl ExactSizeIterator<Item = usize>,
1946) -> Vec<FixedLenByteArray> {
1947 chunk_array_slice(array.value_size(), indices, move |i, chunk| {
1948 let value = array.value(i);
1949 chunk.copy_from_slice(value);
1950 })
1951}
1952
1953#[inline]
1954fn chunk_array_slice(
1955 chunk_size: usize,
1956 indices: impl ExactSizeIterator<Item = usize>,
1957 writer: impl Fn(usize, &mut [u8]),
1958) -> Vec<FixedLenByteArray> {
1959 let capacity = indices.len() * chunk_size;
1960 let mut arena = vec![0; capacity];
1963 for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
1964 writer(i, chunk);
1965 }
1966 chunk_contiguous_vec(arena, chunk_size)
1967}
1968
1969fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
1970 let mut values = Vec::with_capacity(arena.len() / chunk_size);
1971 let mut arena = Bytes::from(arena);
1972 while arena.len() >= chunk_size {
1973 let slice = arena.split_to(chunk_size);
1974 values.push(FixedLenByteArray::from(ByteArray::from(slice)));
1975 }
1976 values
1977}
1978
1979#[inline]
1981fn hash_bytes(bytes: &[u8]) -> u64 {
1982 twox_hash::XxHash64::oneshot(0, bytes)
1983}
1984
1985fn arrow_key_byte_width(dt: &ArrowDataType) -> usize {
1987 match dt {
1988 ArrowDataType::Int8 | ArrowDataType::UInt8 => 1,
1989 ArrowDataType::Int16 | ArrowDataType::UInt16 => 2,
1990 ArrowDataType::Int32 | ArrowDataType::UInt32 => 4,
1991 ArrowDataType::Int64 | ArrowDataType::UInt64 => 8,
1992 _ => 0,
1993 }
1994}
1995
1996fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
1998 use ArrowDataType::*;
1999 match dt {
2000 Int8 | UInt8 => Some(1),
2001 Int16 | UInt16 | Float16 => Some(2),
2002 Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) => Some(4),
2003 Int64
2004 | UInt64
2005 | Float64
2006 | Date64
2007 | Time64(_)
2008 | Timestamp(_, _)
2009 | Duration(_)
2010 | Decimal64(_, _) => Some(8),
2011 Interval(IntervalUnit::YearMonth) => Some(4),
2012 Interval(IntervalUnit::DayTime) => Some(8),
2013 Interval(IntervalUnit::MonthDayNano) => Some(16),
2014 Decimal128(_, _) => Some(16),
2015 Decimal256(_, _) => Some(32),
2016 _ => None,
2017 }
2018}
2019
2020fn update_distinct_values_seen(
2026 array: &dyn arrow_array::Array,
2027 non_null_indices: &[usize],
2028 seen: &mut DistinctValuesSet,
2029) {
2030 let data = array.to_data();
2031 let offset = data.offset();
2032
2033 match array.data_type() {
2034 ArrowDataType::Boolean => {
2035 let arr = array
2036 .as_any()
2037 .downcast_ref::<arrow_array::BooleanArray>()
2038 .unwrap();
2039 for &row in non_null_indices {
2040 seen.insert(arr.value(row) as u64);
2041 }
2042 }
2043 ArrowDataType::Utf8 | ArrowDataType::Binary => {
2044 let offsets = data.buffers()[0].typed_data::<i32>();
2045 let values = data.buffers()[1].as_slice();
2046 for &row in non_null_indices {
2047 let start = offsets[offset + row] as usize;
2048 let end = offsets[offset + row + 1] as usize;
2049 seen.insert(hash_bytes(&values[start..end]));
2050 }
2051 }
2052 ArrowDataType::LargeUtf8 | ArrowDataType::LargeBinary => {
2053 let offsets = data.buffers()[0].typed_data::<i64>();
2054 let values = data.buffers()[1].as_slice();
2055 for &row in non_null_indices {
2056 let start = offsets[offset + row] as usize;
2057 let end = offsets[offset + row + 1] as usize;
2058 seen.insert(hash_bytes(&values[start..end]));
2059 }
2060 }
2061 ArrowDataType::FixedSizeBinary(byte_width) => {
2062 let byte_width = *byte_width as usize;
2063 let buffer = data.buffers()[0].as_slice();
2064 for &row in non_null_indices {
2065 let start = (offset + row) * byte_width;
2066 seen.insert(hash_bytes(&buffer[start..start + byte_width]));
2067 }
2068 }
2069 ArrowDataType::Utf8View => {
2070 let string_view_array = array.as_string_view();
2071 for &row in non_null_indices {
2072 seen.insert(hash_bytes(string_view_array.value(row).as_bytes()));
2073 }
2074 }
2075 ArrowDataType::BinaryView => {
2076 let binary_view_array = array.as_binary_view();
2077 for &row in non_null_indices {
2078 seen.insert(hash_bytes(binary_view_array.value(row)));
2079 }
2080 }
2081 data_type => {
2082 if let Some(width) = fixed_byte_width(data_type) {
2083 let buffer = data.buffers()[0].as_slice();
2084 for &row in non_null_indices {
2085 let pos = (offset + row) * width;
2086 seen.insert(hash_bytes(&buffer[pos..pos + width]));
2087 }
2088 }
2089 }
2091 }
2092}
2093
2094#[cfg(test)]
2095mod tests {
2096 use super::*;
2097 use std::cmp::Ordering;
2098 use std::collections::HashMap;
2099
2100 use std::fs::File;
2101
2102 use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
2103 use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
2104 use crate::column::page::{Page, PageReader};
2105 use crate::file::metadata::thrift::PageHeader;
2106 use crate::file::page_index::column_index::ColumnIndexMetaData;
2107 use crate::file::reader::SerializedPageReader;
2108 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
2109 use crate::schema::types::ColumnPath;
2110 use arrow::datatypes::ToByteSlice;
2111 use arrow::datatypes::{DataType, Schema};
2112 use arrow::error::Result as ArrowResult;
2113 use arrow::util::data_gen::create_random_array;
2114 use arrow::util::pretty::pretty_format_batches;
2115 use arrow::{array::*, buffer::Buffer};
2116 use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
2117 use arrow_schema::Fields;
2118 use half::f16;
2119 use num_traits::{FromPrimitive, ToPrimitive};
2120 use tempfile::tempfile;
2121
2122 use crate::basic::{Encoding, EncodingMask};
2123 use crate::data_type::AsBytes;
2124 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
2125 use crate::file::properties::{
2126 BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
2127 };
2128 use crate::file::serialized_reader::ReadOptionsBuilder;
2129 use crate::file::{
2130 reader::{FileReader, SerializedFileReader},
2131 statistics::Statistics,
2132 };
2133
2134 #[derive(Debug, Default)]
2139 struct RecordingPageStore {
2140 next: u64,
2141 blobs: HashMap<u64, Bytes>,
2142 puts: Arc<std::sync::atomic::AtomicUsize>,
2143 }
2144
2145 impl PageStore for RecordingPageStore {
2146 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2147 let id = 100 + self.next * 7;
2149 self.next += 1;
2150 self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2151 self.blobs.insert(id, value);
2152 Ok(PageKey::new(id))
2153 }
2154
2155 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2156 self.blobs
2157 .remove(&key.get())
2158 .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2159 }
2160 }
2161
2162 #[derive(Debug)]
2163 struct RecordingPageStoreFactory {
2164 puts: Arc<std::sync::atomic::AtomicUsize>,
2165 }
2166
2167 impl PageStoreFactory for RecordingPageStoreFactory {
2168 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2169 Ok(Box::new(RecordingPageStore {
2170 puts: self.puts.clone(),
2171 ..Default::default()
2172 }))
2173 }
2174 }
2175
2176 #[test]
2180 fn custom_page_store_is_byte_identical_to_default() {
2181 let schema = Arc::new(Schema::new(vec![
2182 Field::new("i", DataType::Int32, true),
2183 Field::new("s", DataType::Utf8, true),
2185 ]));
2186 let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2187 let s = StringArray::from(vec![
2188 Some("a"),
2189 Some("bb"),
2190 Some("a"),
2191 None,
2192 Some("bb"),
2193 Some("ccc"),
2194 ]);
2195 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2196
2197 let props = WriterProperties::builder()
2200 .set_max_row_group_row_count(Some(3))
2201 .build();
2202
2203 let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2204 let mut buffer = Vec::new();
2205 let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2206 if let Some(factory) = factory {
2207 opts = opts.with_page_store_factory(factory);
2208 }
2209 let mut writer =
2210 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2211 writer.write(&batch).unwrap();
2212 writer.close().unwrap();
2213 buffer
2214 };
2215
2216 let default_bytes = write(None);
2217
2218 let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2219 let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2220 puts: puts.clone(),
2221 })));
2222
2223 assert!(
2224 puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2225 "custom PageStore was never written to"
2226 );
2227 assert_eq!(
2228 default_bytes, custom_bytes,
2229 "a custom PageStore must produce byte-identical output to the default"
2230 );
2231 }
2232
2233 #[test]
2239 #[cfg_attr(miri, ignore)] fn dictionary_column_round_trips_with_offset_index_disabled() {
2241 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2242
2243 let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2246 let array = Int32Array::from(values.clone());
2247 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2248
2249 let props = WriterProperties::builder()
2250 .set_offset_index_disabled(true)
2251 .set_data_page_row_count_limit(4096)
2252 .build();
2253 let opts = ArrowWriterOptions::new().with_properties(props);
2254
2255 let mut buffer = Vec::new();
2256 let mut writer =
2257 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2258 writer.write(&batch).unwrap();
2259 writer.close().unwrap();
2260
2261 let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2262 let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2263 let read_values: Vec<Option<i32>> = read
2264 .iter()
2265 .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2266 .collect();
2267 assert_eq!(read_values, values);
2268 }
2269
2270 #[test]
2275 fn dictionary_page_is_routed_through_the_store() {
2276 #[derive(Debug, Default)]
2278 struct SizeRecordingPageStore {
2279 blobs: Vec<Bytes>,
2280 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2281 }
2282 impl PageStore for SizeRecordingPageStore {
2283 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2284 self.bytes_put
2285 .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2286 let key = PageKey::new(self.blobs.len() as u64);
2287 self.blobs.push(value);
2288 Ok(key)
2289 }
2290 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2291 Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2292 }
2293 }
2294 #[derive(Debug)]
2295 struct Factory {
2296 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2297 }
2298 impl PageStoreFactory for Factory {
2299 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2300 Ok(Box::new(SizeRecordingPageStore {
2301 bytes_put: self.bytes_put.clone(),
2302 ..Default::default()
2303 }))
2304 }
2305 }
2306
2307 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2308 let values: Vec<&str> = (0..2048)
2311 .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2312 .collect();
2313 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2314 .unwrap();
2315
2316 let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2317 let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2318 bytes_put: bytes_put.clone(),
2319 }));
2320
2321 let mut buffer = Vec::new();
2324 let mut writer =
2325 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2326 writer.write(&batch).unwrap();
2327 writer.close().unwrap();
2328
2329 let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2330 let column = reader.metadata().row_group(0).column(0);
2331 assert!(
2332 column.dictionary_page_offset().is_some(),
2333 "expected the column to be dictionary-encoded"
2334 );
2335
2336 assert_eq!(
2340 bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2341 column.compressed_size(),
2342 "the dictionary page must pass through the store like any other page"
2343 );
2344 }
2345
2346 #[test]
2347 fn arrow_writer() {
2348 let schema = Schema::new(vec![
2350 Field::new("a", DataType::Int32, false),
2351 Field::new("b", DataType::Int32, true),
2352 ]);
2353
2354 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2356 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2357
2358 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2360
2361 roundtrip(batch, Some(SMALL_SIZE / 2));
2362 }
2363
2364 fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2365 let mut buffer = vec![];
2366
2367 let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2368 writer.write(expected_batch).unwrap();
2369 writer.close().unwrap();
2370
2371 buffer
2372 }
2373
2374 fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2375 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2376 writer.write(expected_batch).unwrap();
2377 writer.into_inner().unwrap()
2378 }
2379
2380 #[test]
2381 fn roundtrip_bytes() {
2382 let schema = Arc::new(Schema::new(vec![
2384 Field::new("a", DataType::Int32, false),
2385 Field::new("b", DataType::Int32, true),
2386 ]));
2387
2388 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2390 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2391
2392 let expected_batch =
2394 RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2395
2396 for buffer in [
2397 get_bytes_after_close(schema.clone(), &expected_batch),
2398 get_bytes_by_into_inner(schema, &expected_batch),
2399 ] {
2400 let cursor = Bytes::from(buffer);
2401 let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2402
2403 let actual_batch = record_batch_reader
2404 .next()
2405 .expect("No batch found")
2406 .expect("Unable to get batch");
2407
2408 assert_eq!(expected_batch.schema(), actual_batch.schema());
2409 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2410 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2411 for i in 0..expected_batch.num_columns() {
2412 let expected_data = expected_batch.column(i).to_data();
2413 let actual_data = actual_batch.column(i).to_data();
2414
2415 assert_eq!(expected_data, actual_data);
2416 }
2417 }
2418 }
2419
2420 #[test]
2421 #[cfg_attr(miri, ignore)] fn arrow_writer_non_null() {
2423 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2424 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2425
2426 RoundTripTest::new(Arc::new(a))
2427 .with_schema(Arc::new(schema))
2428 .run();
2429 }
2430
2431 #[test]
2432 #[cfg_attr(miri, ignore)] fn arrow_writer_list() {
2434 let schema = Schema::new(vec![Field::new(
2436 "a",
2437 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2438 true,
2439 )]);
2440
2441 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2443
2444 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2447
2448 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2450 DataType::Int32,
2451 false,
2452 ))))
2453 .len(5)
2454 .add_buffer(a_value_offsets)
2455 .add_child_data(a_values.into_data())
2456 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2457 .build()
2458 .unwrap();
2459 let a = ListArray::from(a_list_data);
2460 assert_eq!(a.null_count(), 1);
2461
2462 RoundTripTest::new(Arc::new(a))
2463 .with_schema(Arc::new(schema))
2464 .run();
2465 }
2466
2467 #[test]
2468 #[cfg_attr(miri, ignore)] fn arrow_writer_list_non_null() {
2470 let schema = Schema::new(vec![Field::new(
2472 "a",
2473 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2474 false,
2475 )]);
2476
2477 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2479
2480 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2483
2484 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2486 DataType::Int32,
2487 false,
2488 ))))
2489 .len(5)
2490 .add_buffer(a_value_offsets)
2491 .add_child_data(a_values.into_data())
2492 .build()
2493 .unwrap();
2494 let a = ListArray::from(a_list_data);
2495 assert_eq!(a.null_count(), 0);
2496
2497 RoundTripTest::new(Arc::new(a))
2498 .with_schema(Arc::new(schema))
2499 .run();
2500 }
2501
2502 #[test]
2503 #[cfg_attr(miri, ignore)] fn arrow_writer_list_view() {
2505 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2506 let schema = Schema::new(vec![Field::new(
2507 "a",
2508 DataType::ListView(list_field.clone()),
2509 true,
2510 )]);
2511
2512 let a = ListViewArray::new(
2514 list_field,
2515 vec![0, 1, 0, 3, 6].into(),
2516 vec![1, 2, 0, 3, 4].into(),
2517 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2518 Some(vec![true, true, false, true, true].into()),
2519 );
2520 assert_eq!(a.null_count(), 1);
2521
2522 RoundTripTest::new(Arc::new(a))
2523 .with_schema(Arc::new(schema))
2524 .run();
2525 }
2526
2527 #[test]
2528 #[cfg_attr(miri, ignore)] fn arrow_writer_list_view_non_null() {
2530 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2531 let schema = Schema::new(vec![Field::new(
2532 "a",
2533 DataType::ListView(list_field.clone()),
2534 false,
2535 )]);
2536
2537 let a = ListViewArray::new(
2539 list_field,
2540 vec![0, 1, 0, 3, 6].into(),
2541 vec![1, 2, 0, 3, 4].into(),
2542 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2543 None,
2544 );
2545 assert_eq!(a.null_count(), 0);
2546
2547 RoundTripTest::new(Arc::new(a))
2548 .with_schema(Arc::new(schema))
2549 .run();
2550 }
2551
2552 #[test]
2553 #[cfg_attr(miri, ignore)] fn arrow_writer_list_view_out_of_order() {
2555 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2556 let schema = Schema::new(vec![Field::new(
2557 "a",
2558 DataType::ListView(list_field.clone()),
2559 false,
2560 )]);
2561
2562 let a = ListViewArray::new(
2564 list_field,
2565 vec![0, 1, 0, 6, 3].into(),
2566 vec![1, 2, 0, 4, 3].into(),
2567 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2568 None,
2569 );
2570 assert_eq!(a.null_count(), 0);
2571
2572 RoundTripTest::new(Arc::new(a))
2573 .with_schema(Arc::new(schema))
2574 .run();
2575 }
2576
2577 #[test]
2578 #[cfg_attr(miri, ignore)] fn arrow_writer_large_list_view() {
2580 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2581 let schema = Schema::new(vec![Field::new(
2582 "a",
2583 DataType::LargeListView(list_field.clone()),
2584 true,
2585 )]);
2586
2587 let a = LargeListViewArray::new(
2589 list_field,
2590 vec![0i64, 1, 0, 3, 6].into(),
2591 vec![1i64, 2, 0, 3, 4].into(),
2592 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2593 Some(vec![true, true, false, true, true].into()),
2594 );
2595 assert_eq!(a.null_count(), 1);
2596
2597 RoundTripTest::new(Arc::new(a))
2598 .with_schema(Arc::new(schema))
2599 .run();
2600 }
2601
2602 #[test]
2603 #[cfg_attr(miri, ignore)] fn arrow_writer_list_view_with_struct() {
2605 let struct_fields = Fields::from(vec![
2607 Field::new("id", DataType::Int32, false),
2608 Field::new("name", DataType::Utf8, false),
2609 ]);
2610 let struct_type = DataType::Struct(struct_fields.clone());
2611 let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2612
2613 let schema = Schema::new(vec![Field::new(
2614 "a",
2615 DataType::ListView(list_field.clone()),
2616 true,
2617 )]);
2618
2619 let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2621 let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2622 let struct_array = StructArray::new(
2623 struct_fields,
2624 vec![Arc::new(id_array), Arc::new(name_array)],
2625 None,
2626 );
2627
2628 let list_view = ListViewArray::new(
2630 list_field,
2631 vec![0, 2, 2].into(), vec![2, 0, 3].into(), Arc::new(struct_array),
2634 Some(vec![true, false, true].into()),
2635 );
2636 assert_eq!(list_view.null_count(), 1);
2637
2638 RoundTripTest::new(Arc::new(list_view))
2639 .with_schema(Arc::new(schema))
2640 .run();
2641 }
2642
2643 #[test]
2644 #[cfg_attr(miri, ignore)] fn arrow_writer_binary() {
2646 let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2647 let raw_binary_values = [
2648 b"foo".to_vec(),
2649 b"bar".to_vec(),
2650 b"baz".to_vec(),
2651 b"quux".to_vec(),
2652 ];
2653 let raw_binary_value_refs = raw_binary_values
2654 .iter()
2655 .map(|x| x.as_slice())
2656 .collect::<Vec<_>>();
2657
2658 let string_values = StringArray::from(raw_string_values.clone());
2659 let binary_values = BinaryArray::from(raw_binary_value_refs);
2660 assert_eq!(string_values.null_count(), 0);
2661 assert_eq!(binary_values.null_count(), 0);
2662
2663 RoundTripTest::new(Arc::new(string_values)).run();
2664 RoundTripTest::new(Arc::new(binary_values)).run();
2665 }
2666
2667 #[test]
2668 #[cfg_attr(miri, ignore)] fn arrow_writer_binary_view() {
2670 let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2671 let raw_binary_values = vec![
2672 b"foo".to_vec(),
2673 b"bar".to_vec(),
2674 b"large payload over 12 bytes".to_vec(),
2675 b"lulu".to_vec(),
2676 ];
2677 let nullable_string_values =
2678 vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2679
2680 let string_view_values = StringViewArray::from(raw_string_values);
2681 let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2682 let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2683
2684 RoundTripTest::new(Arc::new(string_view_values)).run();
2685 RoundTripTest::new(Arc::new(binary_view_values)).run();
2686 RoundTripTest::new(Arc::new(nullable_string_view_values)).run();
2687 }
2688
2689 #[test]
2690 #[cfg_attr(miri, ignore)] fn arrow_writer_binary_view_long_value() {
2692 let long = "a".repeat(128);
2696 let raw_string_values = vec!["foo", long.as_str(), "bar"];
2697 let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2698
2699 let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2700 let binary_view_values: ArrayRef =
2701 Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2702
2703 RoundTripTest::new(Arc::clone(&string_view_values))
2704 .with_nullable(false)
2705 .run();
2706 RoundTripTest::new(Arc::clone(&binary_view_values))
2707 .with_nullable(false)
2708 .run();
2709 }
2710
2711 #[test]
2714 fn arrow_writer_string_view_dictionary() {
2715 let raw_string_values = vec!["a", "b", "large payload over 12 bytes"];
2716 let raw_binary_values = vec![
2717 b"a".to_vec(),
2718 b"b".to_vec(),
2719 b"large payload over 12 bytes".to_vec(),
2720 ];
2721
2722 let keys = UInt32Array::from(vec![Some(0), None, Some(2), Some(1), None]);
2723
2724 let string_view_values = Arc::new(StringViewArray::from(raw_string_values));
2725 let string_dict: ArrayRef = Arc::new(
2726 DictionaryArray::<UInt32Type>::try_new(keys.clone(), string_view_values).unwrap(),
2727 );
2728
2729 let binary_view_values = Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2730 let binary_dict: ArrayRef =
2731 Arc::new(DictionaryArray::<UInt32Type>::try_new(keys, binary_view_values).unwrap());
2732
2733 RoundTripTest::new(string_dict).run();
2734 RoundTripTest::new(binary_dict).run();
2735 }
2736
2737 fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2738 let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2739 let schema = Schema::new(vec![decimal_field]);
2740
2741 let decimal_values = vec![10_000, 50_000, 0, -100]
2742 .into_iter()
2743 .map(Some)
2744 .collect::<Decimal128Array>()
2745 .with_precision_and_scale(precision, scale)
2746 .unwrap();
2747
2748 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2749 }
2750
2751 #[test]
2752 fn arrow_writer_decimal() {
2753 let batch_int32_decimal = get_decimal_batch(5, 2);
2755 roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2756 let batch_int64_decimal = get_decimal_batch(12, 2);
2758 roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2759 let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2761 roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2762 }
2763
2764 #[test]
2765 #[cfg_attr(miri, ignore)] fn arrow_writer_complex() {
2767 let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2769 let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2770 let struct_field_g = Arc::new(Field::new_list(
2771 "g",
2772 Field::new_list_field(DataType::Int16, true),
2773 false,
2774 ));
2775 let struct_field_h = Arc::new(Field::new_list(
2776 "h",
2777 Field::new_list_field(DataType::Int16, false),
2778 true,
2779 ));
2780 let struct_field_e = Arc::new(Field::new_struct(
2781 "e",
2782 vec![
2783 struct_field_f.clone(),
2784 struct_field_g.clone(),
2785 struct_field_h.clone(),
2786 ],
2787 false,
2788 ));
2789 let schema = Schema::new(vec![
2790 Field::new("a", DataType::Int32, false),
2791 Field::new("b", DataType::Int32, true),
2792 Field::new_struct(
2793 "c",
2794 vec![struct_field_d.clone(), struct_field_e.clone()],
2795 false,
2796 ),
2797 ]);
2798
2799 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2801 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2802 let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2803 let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2804
2805 let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2806
2807 let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2810
2811 let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2813 .len(5)
2814 .add_buffer(g_value_offsets.clone())
2815 .add_child_data(g_value.to_data())
2816 .build()
2817 .unwrap();
2818 let g = ListArray::from(g_list_data);
2819 let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2821 .len(5)
2822 .add_buffer(g_value_offsets)
2823 .add_child_data(g_value.to_data())
2824 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2825 .build()
2826 .unwrap();
2827 let h = ListArray::from(h_list_data);
2828
2829 let e = StructArray::from(vec![
2830 (struct_field_f, Arc::new(f) as ArrayRef),
2831 (struct_field_g, Arc::new(g) as ArrayRef),
2832 (struct_field_h, Arc::new(h) as ArrayRef),
2833 ]);
2834
2835 let c = StructArray::from(vec![
2836 (struct_field_d, Arc::new(d) as ArrayRef),
2837 (struct_field_e, Arc::new(e) as ArrayRef),
2838 ]);
2839
2840 let batch = RecordBatch::try_new(
2842 Arc::new(schema),
2843 vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2844 )
2845 .unwrap();
2846
2847 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2848 roundtrip(batch, Some(SMALL_SIZE / 3));
2849 }
2850
2851 #[test]
2852 fn arrow_writer_complex_mixed() {
2853 let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2858 let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2859 let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2860 let schema = Schema::new(vec![Field::new(
2861 "some_nested_object",
2862 DataType::Struct(Fields::from(vec![
2863 offset_field.clone(),
2864 partition_field.clone(),
2865 topic_field.clone(),
2866 ])),
2867 false,
2868 )]);
2869
2870 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2872 let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2873 let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2874
2875 let some_nested_object = StructArray::from(vec![
2876 (offset_field, Arc::new(offset) as ArrayRef),
2877 (partition_field, Arc::new(partition) as ArrayRef),
2878 (topic_field, Arc::new(topic) as ArrayRef),
2879 ]);
2880
2881 let batch =
2883 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2884
2885 roundtrip(batch, Some(SMALL_SIZE / 2));
2886 }
2887
2888 #[test]
2889 fn arrow_writer_map() {
2890 let json_content = r#"
2892 {"stocks":{"long": "$AAA", "short": "$BBB"}}
2893 {"stocks":{"long": null, "long": "$CCC", "short": null}}
2894 {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2895 "#;
2896 let entries_struct_type = DataType::Struct(Fields::from(vec![
2897 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2898 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2899 ]));
2900 let stocks_field = Field::new(
2901 "stocks",
2902 DataType::Map(
2903 Arc::new(Field::new(
2904 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2905 entries_struct_type,
2906 false,
2907 )),
2908 false,
2909 ),
2910 true,
2911 );
2912 let schema = Arc::new(Schema::new(vec![stocks_field]));
2913 let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2914 let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2915
2916 let batch = reader.next().unwrap().unwrap();
2917 roundtrip(batch, None);
2918 }
2919
2920 #[test]
2921 fn arrow_writer_2_level_struct() {
2922 let field_c = Field::new("c", DataType::Int32, true);
2924 let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2925 let type_a = DataType::Struct(vec![field_b.clone()].into());
2926 let field_a = Field::new("a", type_a, true);
2927 let schema = Schema::new(vec![field_a.clone()]);
2928
2929 let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2931 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2932 .len(6)
2933 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2934 .add_child_data(c.into_data())
2935 .build()
2936 .unwrap();
2937 let b = StructArray::from(b_data);
2938 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2939 .len(6)
2940 .null_bit_buffer(Some(Buffer::from([0b00101111])))
2941 .add_child_data(b.into_data())
2942 .build()
2943 .unwrap();
2944 let a = StructArray::from(a_data);
2945
2946 assert_eq!(a.null_count(), 1);
2947 assert_eq!(a.column(0).null_count(), 2);
2948
2949 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2951
2952 roundtrip(batch, Some(SMALL_SIZE / 2));
2953 }
2954
2955 #[test]
2956 fn arrow_writer_2_level_struct_non_null() {
2957 let field_c = Field::new("c", DataType::Int32, false);
2959 let type_b = DataType::Struct(vec![field_c].into());
2960 let field_b = Field::new("b", type_b.clone(), false);
2961 let type_a = DataType::Struct(vec![field_b].into());
2962 let field_a = Field::new("a", type_a.clone(), false);
2963 let schema = Schema::new(vec![field_a]);
2964
2965 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2967 let b_data = ArrayDataBuilder::new(type_b)
2968 .len(6)
2969 .add_child_data(c.into_data())
2970 .build()
2971 .unwrap();
2972 let b = StructArray::from(b_data);
2973 let a_data = ArrayDataBuilder::new(type_a)
2974 .len(6)
2975 .add_child_data(b.into_data())
2976 .build()
2977 .unwrap();
2978 let a = StructArray::from(a_data);
2979
2980 assert_eq!(a.null_count(), 0);
2981 assert_eq!(a.column(0).null_count(), 0);
2982
2983 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2985
2986 roundtrip(batch, Some(SMALL_SIZE / 2));
2987 }
2988
2989 #[test]
2990 fn arrow_writer_2_level_struct_mixed_null() {
2991 let field_c = Field::new("c", DataType::Int32, false);
2993 let type_b = DataType::Struct(vec![field_c].into());
2994 let field_b = Field::new("b", type_b.clone(), true);
2995 let type_a = DataType::Struct(vec![field_b].into());
2996 let field_a = Field::new("a", type_a.clone(), false);
2997 let schema = Schema::new(vec![field_a]);
2998
2999 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
3001 let b_data = ArrayDataBuilder::new(type_b)
3002 .len(6)
3003 .null_bit_buffer(Some(Buffer::from([0b00100111])))
3004 .add_child_data(c.into_data())
3005 .build()
3006 .unwrap();
3007 let b = StructArray::from(b_data);
3008 let a_data = ArrayDataBuilder::new(type_a)
3010 .len(6)
3011 .add_child_data(b.into_data())
3012 .build()
3013 .unwrap();
3014 let a = StructArray::from(a_data);
3015
3016 assert_eq!(a.null_count(), 0);
3017 assert_eq!(a.column(0).null_count(), 2);
3018
3019 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3021
3022 roundtrip(batch, Some(SMALL_SIZE / 2));
3023 }
3024
3025 #[test]
3026 fn arrow_writer_2_level_struct_mixed_null_2() {
3027 let field_c = Field::new("c", DataType::Int32, false);
3029 let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
3030 let field_e = Field::new(
3031 "e",
3032 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3033 false,
3034 );
3035
3036 let field_b = Field::new(
3037 "b",
3038 DataType::Struct(vec![field_c, field_d, field_e].into()),
3039 false,
3040 );
3041 let type_a = DataType::Struct(vec![field_b.clone()].into());
3042 let field_a = Field::new("a", type_a, true);
3043 let schema = Schema::new(vec![field_a.clone()]);
3044
3045 let c = Int32Array::from_iter_values(0..6);
3047 let d = FixedSizeBinaryArray::try_from_iter(
3048 ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
3049 )
3050 .expect("four byte values");
3051 let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
3052 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
3053 .len(6)
3054 .add_child_data(c.into_data())
3055 .add_child_data(d.into_data())
3056 .add_child_data(e.into_data())
3057 .build()
3058 .unwrap();
3059 let b = StructArray::from(b_data);
3060 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
3061 .len(6)
3062 .null_bit_buffer(Some(Buffer::from([0b00100101])))
3063 .add_child_data(b.into_data())
3064 .build()
3065 .unwrap();
3066 let a = StructArray::from(a_data);
3067
3068 assert_eq!(a.null_count(), 3);
3069 assert_eq!(a.column(0).null_count(), 0);
3070
3071 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3073
3074 roundtrip(batch, Some(SMALL_SIZE / 2));
3075 }
3076
3077 #[test]
3078 fn test_fixed_size_binary_in_dict() {
3079 fn test_fixed_size_binary_in_dict_inner<K>()
3080 where
3081 K: ArrowDictionaryKeyType,
3082 K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
3083 <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
3084 {
3085 let field = Field::new(
3086 "a",
3087 DataType::Dictionary(
3088 Box::new(K::DATA_TYPE),
3089 Box::new(DataType::FixedSizeBinary(4)),
3090 ),
3091 false,
3092 );
3093 let schema = Schema::new(vec![field]);
3094
3095 let keys: Vec<K::Native> = vec![
3096 K::Native::try_from(0u8).unwrap(),
3097 K::Native::try_from(0u8).unwrap(),
3098 K::Native::try_from(1u8).unwrap(),
3099 ];
3100 let keys = PrimitiveArray::<K>::from_iter_values(keys);
3101 let values = FixedSizeBinaryArray::try_from_iter(
3102 vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
3103 )
3104 .unwrap();
3105
3106 let data = DictionaryArray::<K>::new(keys, Arc::new(values));
3107 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
3108 roundtrip(batch, None);
3109 }
3110
3111 test_fixed_size_binary_in_dict_inner::<UInt8Type>();
3112 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3113 test_fixed_size_binary_in_dict_inner::<UInt32Type>();
3114 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3115 test_fixed_size_binary_in_dict_inner::<Int8Type>();
3116 test_fixed_size_binary_in_dict_inner::<Int16Type>();
3117 test_fixed_size_binary_in_dict_inner::<Int32Type>();
3118 test_fixed_size_binary_in_dict_inner::<Int64Type>();
3119 }
3120
3121 #[test]
3122 fn test_empty_dict() {
3123 let struct_fields = Fields::from(vec![Field::new(
3124 "dict",
3125 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3126 false,
3127 )]);
3128
3129 let schema = Schema::new(vec![Field::new_struct(
3130 "struct",
3131 struct_fields.clone(),
3132 true,
3133 )]);
3134 let dictionary = Arc::new(DictionaryArray::new(
3135 Int32Array::new_null(5),
3136 Arc::new(StringArray::new_null(0)),
3137 ));
3138
3139 let s = StructArray::new(
3140 struct_fields,
3141 vec![dictionary],
3142 Some(NullBuffer::new_null(5)),
3143 );
3144
3145 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
3146 roundtrip(batch, None);
3147 }
3148 #[test]
3149 fn arrow_writer_page_size() {
3150 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3151
3152 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3153
3154 for i in 0..10 {
3156 let value = i
3157 .to_string()
3158 .repeat(10)
3159 .chars()
3160 .take(10)
3161 .collect::<String>();
3162
3163 builder.append_value(value);
3164 }
3165
3166 let array = Arc::new(builder.finish());
3167
3168 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3169
3170 let file = tempfile::tempfile().unwrap();
3171
3172 let props = WriterProperties::builder()
3174 .set_data_page_size_limit(1)
3175 .set_dictionary_page_size_limit(1)
3176 .set_write_batch_size(1)
3177 .build();
3178
3179 let mut writer =
3180 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3181 .expect("Unable to write file");
3182 writer.write(&batch).unwrap();
3183 writer.close().unwrap();
3184
3185 let options = ReadOptionsBuilder::new().with_page_index().build();
3186 let reader =
3187 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3188
3189 let column = reader.metadata().row_group(0).columns();
3190
3191 assert_eq!(column.len(), 1);
3192
3193 assert!(
3196 column[0].dictionary_page_offset().is_some(),
3197 "Expected a dictionary page"
3198 );
3199
3200 let page_index = reader
3201 .metadata()
3202 .page_index()
3203 .expect("page index should be present");
3204 let page_locations = page_index
3205 .page_locations(0, 0)
3206 .expect("page locations should exist");
3207
3208 assert_eq!(
3211 page_locations.len(),
3212 10,
3213 "Expected 10 pages but got {page_locations:#?}"
3214 );
3215 }
3216
3217 #[test]
3218 #[cfg_attr(miri, ignore)] fn arrow_writer_float_nans() {
3220 let f16_field = Field::new("a", DataType::Float16, false);
3221 let f32_field = Field::new("b", DataType::Float32, false);
3222 let f64_field = Field::new("c", DataType::Float64, false);
3223 let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3224
3225 let f16_values = (0..MEDIUM_SIZE)
3226 .map(|i| {
3227 Some(if i % 2 == 0 {
3228 f16::NAN
3229 } else {
3230 f16::from_f32(i as f32)
3231 })
3232 })
3233 .collect::<Float16Array>();
3234
3235 let f32_values = (0..MEDIUM_SIZE)
3236 .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3237 .collect::<Float32Array>();
3238
3239 let f64_values = (0..MEDIUM_SIZE)
3240 .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3241 .collect::<Float64Array>();
3242
3243 let batch = RecordBatch::try_new(
3244 Arc::new(schema),
3245 vec![
3246 Arc::new(f16_values),
3247 Arc::new(f32_values),
3248 Arc::new(f64_values),
3249 ],
3250 )
3251 .unwrap();
3252
3253 roundtrip(batch, None);
3254 }
3255
3256 const SMALL_SIZE: usize = 7;
3257 const MEDIUM_SIZE: usize = 63;
3258
3259 fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3262 let mut files = vec![];
3263 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3264 let mut props = WriterProperties::builder().set_writer_version(version);
3265
3266 if let Some(size) = max_row_group_size {
3267 props = props.set_max_row_group_row_count(Some(size))
3268 }
3269
3270 let props = props.build();
3271 files.push(roundtrip_opts(&expected_batch, props))
3272 }
3273 files
3274 }
3275
3276 fn roundtrip_opts_with_array_validation<F>(
3280 expected_batch: &RecordBatch,
3281 props: WriterProperties,
3282 validate: F,
3283 ) -> Bytes
3284 where
3285 F: Fn(&ArrayData, &ArrayData),
3286 {
3287 let mut file = vec![];
3288
3289 let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3290 .expect("Unable to write file");
3291 writer.write(expected_batch).unwrap();
3292 writer.close().unwrap();
3293
3294 let file = Bytes::from(file);
3295 let mut record_batch_reader =
3296 ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3297
3298 let actual_batch = record_batch_reader
3299 .next()
3300 .expect("No batch found")
3301 .expect("Unable to get batch");
3302
3303 assert_eq!(expected_batch.schema(), actual_batch.schema());
3304 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3305 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3306 for i in 0..expected_batch.num_columns() {
3307 let expected_data = expected_batch.column(i).to_data();
3308 let actual_data = actual_batch.column(i).to_data();
3309 validate(&expected_data, &actual_data);
3310 }
3311
3312 file
3313 }
3314
3315 fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3316 roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3317 a.validate_full().expect("valid expected data");
3318 b.validate_full().expect("valid actual data");
3319 assert_eq!(a, b)
3320 })
3321 }
3322
3323 struct RoundTripTest {
3327 values: ArrayRef,
3328 schema: Option<SchemaRef>,
3330 nullable: bool,
3333 bloom_filter: bool,
3334 bloom_filter_ndv: Option<u64>,
3335 bloom_filter_position: BloomFilterPosition,
3336 }
3337
3338 impl RoundTripTest {
3339 fn new(values: ArrayRef) -> Self {
3341 Self {
3342 values,
3343 schema: None,
3344 nullable: true,
3345 bloom_filter: false,
3346 bloom_filter_ndv: None,
3347 bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3348 }
3349 }
3350
3351 fn with_schema(mut self, schema: SchemaRef) -> Self {
3353 self.schema = Some(schema);
3354 self
3355 }
3356
3357 fn with_nullable(mut self, nullable: bool) -> Self {
3359 self.nullable = nullable;
3360 self
3361 }
3362
3363 fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
3365 self.bloom_filter = bloom_filter;
3366 self
3367 }
3368
3369 fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) -> Self {
3371 self.bloom_filter_ndv = Some(bloom_filter_ndv);
3372 self
3373 }
3374
3375 fn with_bloom_filter_position(
3377 mut self,
3378 bloom_filter_position: BloomFilterPosition,
3379 ) -> Self {
3380 self.bloom_filter_position = bloom_filter_position;
3381 self
3382 }
3383
3384 fn run(self) -> Vec<Bytes> {
3386 let RoundTripTest {
3387 values,
3388 schema,
3389 nullable,
3390 bloom_filter,
3391 bloom_filter_ndv,
3392 bloom_filter_position,
3393 } = self;
3394
3395 let schema = schema.unwrap_or_else(|| {
3396 let data_type = values.data_type().clone();
3397 Arc::new(Schema::new(vec![Field::new("col", data_type, nullable)]))
3398 });
3399
3400 let encodings = match values.data_type() {
3401 DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3402 vec![
3403 Encoding::PLAIN,
3404 Encoding::DELTA_BYTE_ARRAY,
3405 Encoding::DELTA_LENGTH_BYTE_ARRAY,
3406 ]
3407 }
3408 DataType::Int64
3409 | DataType::Int32
3410 | DataType::Int16
3411 | DataType::Int8
3412 | DataType::UInt64
3413 | DataType::UInt32
3414 | DataType::UInt16
3415 | DataType::UInt8 => vec![
3416 Encoding::PLAIN,
3417 Encoding::DELTA_BINARY_PACKED,
3418 Encoding::BYTE_STREAM_SPLIT,
3419 ],
3420 DataType::Float32 | DataType::Float64 => {
3421 vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT, Encoding::ALP]
3422 }
3423 _ => vec![Encoding::PLAIN],
3424 };
3425
3426 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3427
3428 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3429
3430 let mut files = vec![];
3431 for dictionary_size in [0, 1, 1024] {
3432 for encoding in &encodings {
3433 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3434 for row_group_size in row_group_sizes {
3435 let mut builder = WriterProperties::builder()
3436 .set_writer_version(version)
3437 .set_max_row_group_row_count(Some(row_group_size))
3438 .set_dictionary_enabled(dictionary_size != 0)
3439 .set_dictionary_page_size_limit(dictionary_size.max(1))
3440 .set_encoding(*encoding)
3441 .set_bloom_filter_enabled(bloom_filter)
3442 .set_bloom_filter_position(bloom_filter_position);
3443 if let Some(ndv) = bloom_filter_ndv {
3444 builder = builder.set_bloom_filter_max_ndv(ndv);
3445 }
3446 let props = builder.build();
3447
3448 files.push(roundtrip_opts(&expected_batch, props))
3449 }
3450 }
3451 }
3452 }
3453 files
3454 }
3455 }
3456
3457 fn values_required<A, I>(iter: I) -> Vec<Bytes>
3458 where
3459 A: From<Vec<I::Item>> + Array + 'static,
3460 I: IntoIterator,
3461 {
3462 let raw_values: Vec<_> = iter.into_iter().collect();
3463 let values = Arc::new(A::from(raw_values));
3464 RoundTripTest::new(values).with_nullable(false).run()
3465 }
3466
3467 fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3468 where
3469 A: From<Vec<Option<I::Item>>> + Array + 'static,
3470 I: IntoIterator,
3471 {
3472 let optional_raw_values: Vec<_> = iter
3473 .into_iter()
3474 .enumerate()
3475 .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3476 .collect();
3477 let optional_values = Arc::new(A::from(optional_raw_values));
3478 RoundTripTest::new(optional_values).run()
3479 }
3480
3481 fn required_and_optional<A, I>(iter: I)
3482 where
3483 A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3484 I: IntoIterator + Clone,
3485 {
3486 values_required::<A, I>(iter.clone());
3487 values_optional::<A, I>(iter);
3488 }
3489
3490 fn check_bloom_filter<T: AsBytes>(
3491 files: Vec<Bytes>,
3492 file_column: String,
3493 positive_values: Vec<T>,
3494 negative_values: Vec<T>,
3495 ) {
3496 files.into_iter().take(1).for_each(|file| {
3497 let file_reader = SerializedFileReader::new_with_options(
3498 file,
3499 ReadOptionsBuilder::new()
3500 .with_reader_properties(
3501 ReaderProperties::builder()
3502 .set_read_bloom_filter(true)
3503 .build(),
3504 )
3505 .build(),
3506 )
3507 .expect("Unable to open file as Parquet");
3508 let metadata = file_reader.metadata();
3509
3510 let mut bloom_filters: Vec<_> = vec![];
3512 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3513 if let Some((column_index, _)) = row_group
3514 .columns()
3515 .iter()
3516 .enumerate()
3517 .find(|(_, column)| column.column_path().string() == file_column)
3518 {
3519 let row_group_reader = file_reader
3520 .get_row_group(ri)
3521 .expect("Unable to read row group");
3522 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3523 bloom_filters.push(sbbf.clone());
3524 } else {
3525 panic!("No bloom filter for column named {file_column} found");
3526 }
3527 } else {
3528 panic!("No column named {file_column} found");
3529 }
3530 }
3531
3532 positive_values.iter().for_each(|value| {
3533 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3534 assert!(
3535 found.is_some(),
3536 "{}",
3537 format!("Value {:?} should be in bloom filter", value.as_bytes())
3538 );
3539 });
3540
3541 negative_values.iter().for_each(|value| {
3542 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3543 assert!(
3544 found.is_none(),
3545 "{}",
3546 format!("Value {:?} should not be in bloom filter", value.as_bytes())
3547 );
3548 });
3549 });
3550 }
3551
3552 #[test]
3553 #[cfg_attr(miri, ignore)] fn all_null_primitive_single_column() {
3555 let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3556 RoundTripTest::new(values).run();
3557 }
3558 #[test]
3559 #[cfg_attr(miri, ignore)] fn null_single_column() {
3561 let values = Arc::new(NullArray::new(SMALL_SIZE));
3562 RoundTripTest::new(values).run();
3563 }
3565
3566 #[test]
3567 #[cfg_attr(miri, ignore)] fn bool_single_column() {
3569 required_and_optional::<BooleanArray, _>(
3570 [true, false].iter().cycle().copied().take(SMALL_SIZE),
3571 );
3572 }
3573
3574 #[test]
3575 #[cfg_attr(miri, ignore)] fn bool_large_single_column() {
3577 let values = Arc::new(
3578 [None, Some(true), Some(false)]
3579 .iter()
3580 .cycle()
3581 .copied()
3582 .take(200_000)
3583 .collect::<BooleanArray>(),
3584 );
3585 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3586 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3587 let file = tempfile::tempfile().unwrap();
3588
3589 let mut writer =
3590 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3591 .expect("Unable to write file");
3592 writer.write(&expected_batch).unwrap();
3593 writer.close().unwrap();
3594 }
3595
3596 #[test]
3597 fn check_page_offset_index_with_nan() {
3598 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3599 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3600 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3601
3602 let mut out = Vec::with_capacity(1024);
3603 let mut writer =
3604 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3605 writer.write(&batch).unwrap();
3606 let file_meta_data = writer.close().unwrap();
3607 for row_group in file_meta_data.row_groups() {
3608 for column in row_group.columns() {
3609 assert!(column.offset_index_offset().is_some());
3610 assert!(column.offset_index_length().is_some());
3611 assert!(column.column_index_offset().is_some());
3612 assert!(column.column_index_length().is_some());
3613 }
3614 }
3615 if let Some(page_index) = file_meta_data.page_index() {
3616 for rg in 0..file_meta_data.num_row_groups() {
3617 for col in 0..file_meta_data.row_group(rg).num_columns() {
3618 let idx = page_index
3619 .column_index(rg, col)
3620 .expect("column index should exist");
3621 assert!(idx.nan_counts().is_some());
3622 let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
3623 panic!("expected double statistics")
3624 };
3625 for i in 0..idx.num_pages() as usize {
3626 assert_eq!(float_idx.nan_count(i), Some(10));
3627 assert_eq!(
3628 f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3629 Ordering::Equal
3630 );
3631 assert_eq!(
3632 f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3633 Ordering::Equal
3634 );
3635 }
3636 }
3637 }
3638 } else {
3639 panic!("page index should be present");
3640 }
3641 }
3642
3643 #[test]
3644 fn check_page_offset_index_with_mixed_nan() {
3645 let schema = Arc::new(Schema::new(vec![Field::new(
3646 "col",
3647 DataType::Float64,
3648 true,
3649 )]));
3650
3651 let mut out = Vec::with_capacity(1024);
3652 let props = WriterProperties::builder()
3653 .set_data_page_row_count_limit(10)
3654 .build();
3655 let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3656 .expect("Unable to write file");
3657
3658 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3660 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3661 writer.write(&batch).unwrap();
3662
3663 let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3665 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3666 writer.write(&batch).unwrap();
3667
3668 let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3670 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3671 writer.write(&batch).unwrap();
3672
3673 let values = Arc::new(Float64Array::from(vec![
3675 -1.0,
3676 0.0,
3677 f64::NAN,
3678 -f64::NAN,
3679 1.0,
3680 ]));
3681 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3682 writer.write(&batch).unwrap();
3683
3684 let file_meta_data = writer.close().unwrap();
3685
3686 let col_stats = file_meta_data
3688 .row_group(0)
3689 .column(0)
3690 .statistics()
3691 .expect("missing column chunk statistics");
3692
3693 assert_eq!(col_stats.nan_count_opt(), Some(22));
3694 assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3695 assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3696
3697 assert!(file_meta_data.page_index().is_some());
3698 let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
3699 assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
3700
3701 let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
3703 panic!("expected double statistics")
3704 };
3705
3706 assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3707 assert_eq!(
3708 f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3709 Ordering::Equal
3710 );
3711 assert_eq!(
3712 f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3713 Ordering::Equal
3714 );
3715 assert_eq!(
3716 (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3717 Ordering::Equal
3718 );
3719 assert_eq!(
3720 (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3721 Ordering::Equal
3722 );
3723 assert_eq!(float_idx.min_value(2), Some(&0.0));
3724 assert_eq!(float_idx.max_value(2), Some(&0.0));
3725 assert_eq!(float_idx.min_value(3), Some(&-1.0));
3726 assert_eq!(float_idx.max_value(3), Some(&1.0));
3727 }
3728
3729 #[test]
3730 #[cfg_attr(miri, ignore)] fn i8_single_column() {
3732 required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3733 }
3734
3735 #[test]
3736 #[cfg_attr(miri, ignore)] fn i16_single_column() {
3738 required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3739 }
3740
3741 #[test]
3742 #[cfg_attr(miri, ignore)] fn i32_single_column() {
3744 required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3745 }
3746
3747 #[test]
3748 #[cfg_attr(miri, ignore)] fn i64_single_column() {
3750 required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3751 }
3752
3753 #[test]
3754 #[cfg_attr(miri, ignore)] fn u8_single_column() {
3756 required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3757 }
3758
3759 #[test]
3760 #[cfg_attr(miri, ignore)] fn u16_single_column() {
3762 required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3763 }
3764
3765 #[test]
3766 #[cfg_attr(miri, ignore)] fn u32_single_column() {
3768 required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3769 }
3770
3771 #[test]
3772 #[cfg_attr(miri, ignore)] fn u64_single_column() {
3774 required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3775 }
3776
3777 #[test]
3778 #[cfg_attr(miri, ignore)] fn f32_single_column() {
3780 required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3781 }
3782
3783 #[test]
3784 #[cfg_attr(miri, ignore)] fn f64_single_column() {
3786 required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3787 }
3788
3789 #[test]
3794 #[cfg_attr(miri, ignore)] fn timestamp_second_single_column() {
3796 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3797 let values = Arc::new(TimestampSecondArray::from(raw_values));
3798
3799 RoundTripTest::new(values).with_nullable(false).run();
3800 }
3801
3802 #[test]
3803 #[cfg_attr(miri, ignore)] fn timestamp_millisecond_single_column() {
3805 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3806 let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3807
3808 RoundTripTest::new(values).with_nullable(false).run();
3809 }
3810
3811 #[test]
3812 #[cfg_attr(miri, ignore)] fn timestamp_microsecond_single_column() {
3814 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3815 let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3816
3817 RoundTripTest::new(values).with_nullable(false).run();
3818 }
3819
3820 #[test]
3821 #[cfg_attr(miri, ignore)] fn timestamp_nanosecond_single_column() {
3823 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3824 let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3825
3826 RoundTripTest::new(values).with_nullable(false).run();
3827 }
3828
3829 #[test]
3830 #[cfg_attr(miri, ignore)] fn date32_single_column() {
3832 required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3833 }
3834
3835 #[test]
3836 #[cfg_attr(miri, ignore)] fn date64_single_column() {
3838 required_and_optional::<Date64Array, _>(
3840 (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3841 );
3842 }
3843
3844 #[test]
3845 #[cfg_attr(miri, ignore)] fn time32_second_single_column() {
3847 required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3848 }
3849
3850 #[test]
3851 #[cfg_attr(miri, ignore)] fn time32_millisecond_single_column() {
3853 required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3854 }
3855
3856 #[test]
3857 #[cfg_attr(miri, ignore)] fn time64_microsecond_single_column() {
3859 required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3860 }
3861
3862 #[test]
3863 #[cfg_attr(miri, ignore)] fn time64_nanosecond_single_column() {
3865 required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3866 }
3867
3868 #[test]
3869 #[cfg_attr(miri, ignore)] fn duration_second_single_column() {
3871 required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3872 }
3873
3874 #[test]
3875 #[cfg_attr(miri, ignore)] fn duration_millisecond_single_column() {
3877 required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3878 }
3879
3880 #[test]
3881 #[cfg_attr(miri, ignore)] fn duration_microsecond_single_column() {
3883 required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3884 }
3885
3886 #[test]
3887 #[cfg_attr(miri, ignore)] fn duration_nanosecond_single_column() {
3889 required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3890 }
3891
3892 #[test]
3893 #[cfg_attr(miri, ignore)] fn interval_year_month_single_column() {
3895 required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3896 }
3897
3898 #[test]
3899 #[cfg_attr(miri, ignore)] fn interval_day_time_single_column() {
3901 required_and_optional::<IntervalDayTimeArray, _>(vec![
3902 IntervalDayTime::new(0, 1),
3903 IntervalDayTime::new(0, 3),
3904 IntervalDayTime::new(3, -2),
3905 IntervalDayTime::new(-200, 4),
3906 ]);
3907 }
3908
3909 #[test]
3910 #[should_panic(
3911 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3912 )]
3913 fn interval_month_day_nano_single_column() {
3914 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3915 IntervalMonthDayNano::new(0, 1, 5),
3916 IntervalMonthDayNano::new(0, 3, 2),
3917 IntervalMonthDayNano::new(3, -2, -5),
3918 IntervalMonthDayNano::new(-200, 4, -1),
3919 ]);
3920 }
3921
3922 #[test]
3923 #[cfg_attr(miri, ignore)] fn binary_single_column() {
3925 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3926 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3927 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3928
3929 values_required::<BinaryArray, _>(many_vecs_iter);
3931 }
3932
3933 #[test]
3934 #[cfg_attr(miri, ignore)] fn binary_view_single_column() {
3936 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3937 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3938 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3939
3940 values_required::<BinaryViewArray, _>(many_vecs_iter);
3942 }
3943
3944 #[test]
3945 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_at_end() {
3947 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3948 let files = RoundTripTest::new(array)
3949 .with_nullable(false)
3950 .with_bloom_filter(true)
3951 .with_bloom_filter_position(BloomFilterPosition::End)
3952 .run();
3953
3954 check_bloom_filter(
3955 files,
3956 "col".to_string(),
3957 (0..SMALL_SIZE as i32).collect(),
3958 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3959 );
3960 }
3961
3962 #[test]
3963 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter() {
3965 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3966 let files = RoundTripTest::new(array)
3967 .with_nullable(false)
3968 .with_bloom_filter(true)
3969 .run();
3970
3971 check_bloom_filter(
3972 files,
3973 "col".to_string(),
3974 (0..SMALL_SIZE as i32).collect(),
3975 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3976 );
3977 }
3978
3979 fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
3980 let schema = Arc::new(Schema::new(vec![Field::new(
3981 "col",
3982 array.data_type().clone(),
3983 false,
3984 )]));
3985 let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
3986 let props = WriterProperties::builder()
3987 .set_dictionary_enabled(true)
3988 .set_dictionary_page_size_limit(dictionary_page_size_limit)
3989 .set_write_batch_size(256)
3990 .set_bloom_filter_enabled(true)
3991 .build();
3992 let mut buf = Vec::new();
3993 let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
3994 writer.write(&batch).unwrap();
3995 writer.close().unwrap();
3996 Bytes::from(buf)
3997 }
3998
3999 fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4000 let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4001 *metadata
4002 .row_group(0)
4003 .column(0)
4004 .page_encoding_stats_mask()
4005 .unwrap()
4006 }
4007
4008 #[test]
4011 fn string_column_bloom_filter_populated_from_dictionary() {
4012 let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4013 let array = Arc::new(StringArray::from_iter_values(&values));
4014 let file = write_with_bloom_filter(array, 1024 * 1024);
4015 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4016
4017 check_bloom_filter(
4018 vec![file],
4019 "col".to_string(),
4020 (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4021 (10..20)
4022 .map(|i| format!("value-{i}").into_bytes())
4023 .collect(),
4024 );
4025 }
4026
4027 #[test]
4030 fn string_column_bloom_filter_across_dictionary_fallback() {
4031 let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4032 let array = Arc::new(StringArray::from_iter_values(&values));
4033 let file = write_with_bloom_filter(array, 1024);
4034 let encodings = data_page_encoding_mask(&file);
4035 assert!(
4036 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4037 "expected dictionary and plain data pages, got {encodings:?}"
4038 );
4039
4040 check_bloom_filter(
4041 vec![file],
4042 "col".to_string(),
4043 values.into_iter().map(String::into_bytes).collect(),
4044 (2000..2010)
4045 .map(|i| format!("value-{i}").into_bytes())
4046 .collect(),
4047 );
4048 }
4049
4050 #[test]
4051 fn i64_column_bloom_filter_populated_from_dictionary() {
4052 let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4053 let file = write_with_bloom_filter(array, 1024 * 1024);
4054 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4055
4056 check_bloom_filter(
4057 vec![file],
4058 "col".to_string(),
4059 (0..10i64).collect(),
4060 (10..20i64).collect(),
4061 );
4062 }
4063
4064 #[test]
4065 fn i64_column_bloom_filter_across_dictionary_fallback() {
4066 let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4067 let file = write_with_bloom_filter(array, 1024);
4068 let encodings = data_page_encoding_mask(&file);
4069 assert!(
4070 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4071 "expected dictionary and plain data pages, got {encodings:?}"
4072 );
4073
4074 check_bloom_filter(
4075 vec![file],
4076 "col".to_string(),
4077 (0..2000i64).collect(),
4078 (2000..2010i64).collect(),
4079 );
4080 }
4081
4082 #[test]
4087 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_fixed_ndv() {
4089 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4090
4091 let files = RoundTripTest::new(array.clone())
4093 .with_nullable(false)
4094 .with_bloom_filter(true)
4095 .with_bloom_filter_ndv(1_000_000)
4096 .run();
4097
4098 check_bloom_filter(
4099 files,
4100 "col".to_string(),
4101 (0..SMALL_SIZE as i32).collect(),
4102 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4103 );
4104
4105 let files = RoundTripTest::new(array)
4107 .with_nullable(false)
4108 .with_bloom_filter(true)
4109 .with_bloom_filter_ndv(3)
4110 .run();
4111
4112 check_bloom_filter(
4113 files,
4114 "col".to_string(),
4115 (0..SMALL_SIZE as i32).collect(),
4116 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4117 );
4118 }
4119
4120 #[test]
4121 #[cfg_attr(miri, ignore)] fn binary_column_bloom_filter() {
4123 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4124 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4125 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4126
4127 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4128 let files = RoundTripTest::new(array)
4129 .with_nullable(false)
4130 .with_bloom_filter(true)
4131 .run();
4132
4133 check_bloom_filter(
4134 files,
4135 "col".to_string(),
4136 many_vecs,
4137 vec![vec![(SMALL_SIZE + 1) as u8]],
4138 );
4139 }
4140
4141 #[test]
4142 #[cfg_attr(miri, ignore)] fn empty_string_null_column_bloom_filter() {
4144 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4145 let raw_strs = raw_values.iter().map(|s| s.as_str());
4146
4147 let array = Arc::new(StringArray::from_iter_values(raw_strs));
4148 let files = RoundTripTest::new(array)
4149 .with_nullable(false)
4150 .with_bloom_filter(true)
4151 .run();
4152
4153 let optional_raw_values: Vec<_> = raw_values
4154 .iter()
4155 .enumerate()
4156 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4157 .collect();
4158 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4160 }
4161
4162 #[test]
4163 #[cfg_attr(miri, ignore)] fn large_binary_single_column() {
4165 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4166 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4167 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4168
4169 values_required::<LargeBinaryArray, _>(many_vecs_iter);
4171 }
4172
4173 #[test]
4174 #[cfg_attr(miri, ignore)] fn fixed_size_binary_single_column() {
4176 let mut builder = FixedSizeBinaryBuilder::new(4);
4177 builder.append_value(b"0123").unwrap();
4178 builder.append_null();
4179 builder.append_value(b"8910").unwrap();
4180 builder.append_value(b"1112").unwrap();
4181 let array = Arc::new(builder.finish());
4182
4183 RoundTripTest::new(array).run();
4184 }
4185
4186 #[test]
4187 #[cfg_attr(miri, ignore)] fn string_single_column() {
4189 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4190 let raw_strs = raw_values.iter().map(|s| s.as_str());
4191
4192 required_and_optional::<StringArray, _>(raw_strs);
4193 }
4194
4195 #[test]
4196 #[cfg_attr(miri, ignore)] fn large_string_single_column() {
4198 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4199 let raw_strs = raw_values.iter().map(|s| s.as_str());
4200
4201 required_and_optional::<LargeStringArray, _>(raw_strs);
4202 }
4203
4204 #[test]
4205 #[cfg_attr(miri, ignore)] fn string_view_single_column() {
4207 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4208 let raw_strs = raw_values.iter().map(|s| s.as_str());
4209
4210 required_and_optional::<StringViewArray, _>(raw_strs);
4211 }
4212
4213 #[test]
4214 fn null_list_single_column() {
4215 let null_field = Field::new_list_field(DataType::Null, true);
4216 let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
4217
4218 let schema = Schema::new(vec![list_field]);
4219
4220 let a_values = NullArray::new(2);
4222 let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
4223 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4224 DataType::Null,
4225 true,
4226 ))))
4227 .len(3)
4228 .add_buffer(a_value_offsets)
4229 .null_bit_buffer(Some(Buffer::from([0b00000101])))
4230 .add_child_data(a_values.into_data())
4231 .build()
4232 .unwrap();
4233
4234 let a = ListArray::from(a_list_data);
4235
4236 assert!(a.is_valid(0));
4237 assert!(!a.is_valid(1));
4238 assert!(a.is_valid(2));
4239
4240 assert_eq!(a.value(0).len(), 0);
4241 assert_eq!(a.value(2).len(), 2);
4242 assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
4243
4244 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4245 roundtrip(batch, None);
4246 }
4247
4248 #[test]
4249 #[cfg_attr(miri, ignore)] fn list_single_column() {
4251 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4252 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
4253 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4254 DataType::Int32,
4255 false,
4256 ))))
4257 .len(5)
4258 .add_buffer(a_value_offsets)
4259 .null_bit_buffer(Some(Buffer::from([0b00011011])))
4260 .add_child_data(a_values.into_data())
4261 .build()
4262 .unwrap();
4263
4264 assert_eq!(a_list_data.null_count(), 1);
4265
4266 let a = ListArray::from(a_list_data);
4267 let values = Arc::new(a);
4268
4269 RoundTripTest::new(values).run();
4270 }
4271
4272 #[test]
4273 #[cfg_attr(miri, ignore)] fn large_list_single_column() {
4275 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4276 let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
4277 let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
4278 "large_item",
4279 DataType::Int32,
4280 true,
4281 ))))
4282 .len(5)
4283 .add_buffer(a_value_offsets)
4284 .add_child_data(a_values.into_data())
4285 .null_bit_buffer(Some(Buffer::from([0b00011011])))
4286 .build()
4287 .unwrap();
4288
4289 assert_eq!(a_list_data.null_count(), 1);
4291
4292 let a = LargeListArray::from(a_list_data);
4293 let values = Arc::new(a);
4294
4295 RoundTripTest::new(values).run();
4296 }
4297
4298 #[test]
4299 #[cfg_attr(miri, ignore)] fn list_nested_nulls() {
4301 use arrow::datatypes::Int32Type;
4302 let data = vec![
4303 Some(vec![Some(1)]),
4304 Some(vec![Some(2), Some(3)]),
4305 None,
4306 Some(vec![Some(4), Some(5), None]),
4307 Some(vec![None]),
4308 Some(vec![Some(6), Some(7)]),
4309 ];
4310
4311 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
4312 RoundTripTest::new(Arc::new(list)).run();
4313
4314 let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
4315 RoundTripTest::new(Arc::new(list)).run();
4316 }
4317
4318 #[test]
4319 #[cfg_attr(miri, ignore)] fn list_utf8_view_selective_padding_roundtrip() {
4321 let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
4322 let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
4323 builder.values().append_value("a");
4324 builder.values().append_null();
4325 builder.append(true);
4326 builder.append(false);
4329 builder.values().append_value("large payload over 12 bytes");
4331 builder.append(true);
4332
4333 RoundTripTest::new(Arc::new(builder.finish())).run();
4334 }
4335
4336 #[test]
4337 #[cfg_attr(miri, ignore)] fn struct_single_column() {
4339 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4340 let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4341 let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4342
4343 let values = Arc::new(s);
4344 RoundTripTest::new(values).with_nullable(false).run();
4345 }
4346
4347 #[test]
4348 fn list_and_map_coerced_names() {
4349 let list_field =
4351 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4352 let map_field = Field::new_map(
4353 "my_map",
4354 "my_entries",
4355 Field::new("my_keys", DataType::Int32, false),
4356 Field::new("my_values", DataType::Int32, true),
4357 false,
4358 true,
4359 );
4360
4361 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4362 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4363
4364 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4365
4366 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4368 let file = tempfile::tempfile().unwrap();
4369 let mut writer =
4370 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4371
4372 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4373 writer.write(&batch).unwrap();
4374 let file_metadata = writer.close().unwrap();
4375
4376 let schema = file_metadata.file_metadata().schema();
4377 let list_field = &schema.get_fields()[0].get_fields()[0];
4379 assert_eq!(list_field.get_fields()[0].name(), "element");
4380
4381 let map_field = &schema.get_fields()[1].get_fields()[0];
4382 assert_eq!(map_field.name(), "key_value");
4384 assert_eq!(map_field.get_fields()[0].name(), "key");
4386 assert_eq!(map_field.get_fields()[1].name(), "value");
4388
4389 let reader = SerializedFileReader::new(file).unwrap();
4391 let file_schema = reader.metadata().file_metadata().schema();
4392 let fields = file_schema.get_fields();
4393 let list_field = &fields[0].get_fields()[0];
4394 assert_eq!(list_field.get_fields()[0].name(), "element");
4395 let map_field = &fields[1].get_fields()[0];
4396 assert_eq!(map_field.name(), "key_value");
4397 assert_eq!(map_field.get_fields()[0].name(), "key");
4398 assert_eq!(map_field.get_fields()[1].name(), "value");
4399 }
4400
4401 #[test]
4402 #[cfg_attr(miri, ignore)] fn fallback_flush_data_page() {
4404 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4406 let values = Arc::new(StringArray::from(raw_values));
4407 let encodings = vec![
4408 Encoding::DELTA_BYTE_ARRAY,
4409 Encoding::DELTA_LENGTH_BYTE_ARRAY,
4410 ];
4411 let data_type = values.data_type().clone();
4412 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4413 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4414
4415 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4416 let data_page_size_limit: usize = 32;
4417 let write_batch_size: usize = 16;
4418
4419 for encoding in &encodings {
4420 for row_group_size in row_group_sizes {
4421 let props = WriterProperties::builder()
4422 .set_writer_version(WriterVersion::PARQUET_2_0)
4423 .set_max_row_group_row_count(Some(row_group_size))
4424 .set_dictionary_enabled(false)
4425 .set_encoding(*encoding)
4426 .set_data_page_size_limit(data_page_size_limit)
4427 .set_write_batch_size(write_batch_size)
4428 .build();
4429
4430 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4431 let string_array_a = StringArray::from(a.clone());
4432 let string_array_b = StringArray::from(b.clone());
4433 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4434 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4435 assert_eq!(
4436 vec_a, vec_b,
4437 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4438 );
4439 });
4440 }
4441 }
4442 }
4443
4444 #[test]
4445 #[cfg_attr(miri, ignore)] fn arrow_writer_string_dictionary() {
4447 #[expect(deprecated)]
4449 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4450 "dictionary",
4451 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4452 true,
4453 42,
4454 true,
4455 )]));
4456
4457 let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4459 .iter()
4460 .copied()
4461 .collect();
4462
4463 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4465 }
4466
4467 #[test]
4468 fn arrow_writer_test_type_compatibility() {
4469 fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4470 where
4471 T1: Array + 'static,
4472 T2: Array + 'static,
4473 {
4474 let schema1 = Arc::new(Schema::new(vec![Field::new(
4475 "a",
4476 array1.data_type().clone(),
4477 false,
4478 )]));
4479
4480 let file = tempfile().unwrap();
4481 let mut writer =
4482 ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4483
4484 let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4485 writer.write(&rb1).unwrap();
4486
4487 let schema2 = Arc::new(Schema::new(vec![Field::new(
4488 "a",
4489 array2.data_type().clone(),
4490 false,
4491 )]));
4492 let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4493 writer.write(&rb2).unwrap();
4494
4495 writer.close().unwrap();
4496
4497 let mut record_batch_reader =
4498 ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4499 let actual_batch = record_batch_reader.next().unwrap().unwrap();
4500
4501 let expected_batch =
4502 RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4503 assert_eq!(actual_batch, expected_batch);
4504 }
4505
4506 ensure_compatible_write(
4509 DictionaryArray::new(
4510 UInt8Array::from_iter_values(vec![0]),
4511 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4512 ),
4513 StringArray::from_iter_values(vec!["barquet"]),
4514 DictionaryArray::new(
4515 UInt8Array::from_iter_values(vec![0, 1]),
4516 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4517 ),
4518 );
4519
4520 ensure_compatible_write(
4521 StringArray::from_iter_values(vec!["parquet"]),
4522 DictionaryArray::new(
4523 UInt8Array::from_iter_values(vec![0]),
4524 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4525 ),
4526 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4527 );
4528
4529 ensure_compatible_write(
4532 DictionaryArray::new(
4533 UInt8Array::from_iter_values(vec![0]),
4534 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4535 ),
4536 DictionaryArray::new(
4537 UInt16Array::from_iter_values(vec![0]),
4538 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4539 ),
4540 DictionaryArray::new(
4541 UInt8Array::from_iter_values(vec![0, 1]),
4542 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4543 ),
4544 );
4545
4546 ensure_compatible_write(
4548 DictionaryArray::new(
4549 UInt8Array::from_iter_values(vec![0]),
4550 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4551 ),
4552 DictionaryArray::new(
4553 UInt8Array::from_iter_values(vec![0]),
4554 Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4555 ),
4556 DictionaryArray::new(
4557 UInt8Array::from_iter_values(vec![0, 1]),
4558 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4559 ),
4560 );
4561
4562 ensure_compatible_write(
4564 DictionaryArray::new(
4565 UInt8Array::from_iter_values(vec![0]),
4566 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4567 ),
4568 LargeStringArray::from_iter_values(vec!["barquet"]),
4569 DictionaryArray::new(
4570 UInt8Array::from_iter_values(vec![0, 1]),
4571 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4572 ),
4573 );
4574
4575 ensure_compatible_write(
4578 StringArray::from_iter_values(vec!["parquet"]),
4579 LargeStringArray::from_iter_values(vec!["barquet"]),
4580 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4581 );
4582
4583 ensure_compatible_write(
4584 LargeStringArray::from_iter_values(vec!["parquet"]),
4585 StringArray::from_iter_values(vec!["barquet"]),
4586 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4587 );
4588
4589 ensure_compatible_write(
4590 StringArray::from_iter_values(vec!["parquet"]),
4591 StringViewArray::from_iter_values(vec!["barquet"]),
4592 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4593 );
4594
4595 ensure_compatible_write(
4596 StringViewArray::from_iter_values(vec!["parquet"]),
4597 StringArray::from_iter_values(vec!["barquet"]),
4598 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4599 );
4600
4601 ensure_compatible_write(
4602 LargeStringArray::from_iter_values(vec!["parquet"]),
4603 StringViewArray::from_iter_values(vec!["barquet"]),
4604 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4605 );
4606
4607 ensure_compatible_write(
4608 StringViewArray::from_iter_values(vec!["parquet"]),
4609 LargeStringArray::from_iter_values(vec!["barquet"]),
4610 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4611 );
4612
4613 ensure_compatible_write(
4616 BinaryArray::from_iter_values(vec![b"parquet"]),
4617 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4618 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4619 );
4620
4621 ensure_compatible_write(
4622 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4623 BinaryArray::from_iter_values(vec![b"barquet"]),
4624 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4625 );
4626
4627 ensure_compatible_write(
4628 BinaryArray::from_iter_values(vec![b"parquet"]),
4629 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4630 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4631 );
4632
4633 ensure_compatible_write(
4634 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4635 BinaryArray::from_iter_values(vec![b"barquet"]),
4636 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4637 );
4638
4639 ensure_compatible_write(
4640 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4641 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4642 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4643 );
4644
4645 ensure_compatible_write(
4646 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4647 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4648 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4649 );
4650
4651 let list_field_metadata = HashMap::from_iter(vec![(
4654 PARQUET_FIELD_ID_META_KEY.to_string(),
4655 "1".to_string(),
4656 )]);
4657 let list_field = Field::new_list_field(DataType::Int32, false);
4658
4659 let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4660 let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4661
4662 let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4663 let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4664
4665 let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4666 let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4667
4668 ensure_compatible_write(
4669 ListArray::try_new(
4671 Arc::new(
4672 list_field
4673 .clone()
4674 .with_metadata(list_field_metadata.clone()),
4675 ),
4676 offsets1,
4677 values1,
4678 None,
4679 )
4680 .unwrap(),
4681 ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4683 ListArray::try_new(
4685 Arc::new(
4686 list_field
4687 .clone()
4688 .with_metadata(list_field_metadata.clone()),
4689 ),
4690 offsets_expected,
4691 values_expected,
4692 None,
4693 )
4694 .unwrap(),
4695 );
4696 }
4697
4698 #[test]
4699 #[cfg_attr(miri, ignore)] fn arrow_writer_primitive_dictionary() {
4701 #[expect(deprecated)]
4703 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4704 "dictionary",
4705 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4706 true,
4707 42,
4708 true,
4709 )]));
4710
4711 let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4713 builder.append(12345678).unwrap();
4714 builder.append_null();
4715 builder.append(22345678).unwrap();
4716 builder.append(12345678).unwrap();
4717 let d = builder.finish();
4718
4719 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4720 }
4721
4722 #[test]
4723 #[cfg_attr(miri, ignore)] fn arrow_writer_decimal32_dictionary() {
4725 let integers = vec![12345, 56789, 34567];
4726
4727 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4728
4729 let values = Decimal32Array::from(integers.clone())
4730 .with_precision_and_scale(5, 2)
4731 .unwrap();
4732
4733 let array = DictionaryArray::new(keys, Arc::new(values));
4734 RoundTripTest::new(Arc::new(array.clone())).run();
4735
4736 let values = Decimal32Array::from(integers)
4737 .with_precision_and_scale(9, 2)
4738 .unwrap();
4739
4740 let array = array.with_values(Arc::new(values));
4741 RoundTripTest::new(Arc::new(array)).run();
4742 }
4743
4744 #[test]
4745 #[cfg_attr(miri, ignore)] fn arrow_writer_decimal64_dictionary() {
4747 let integers = vec![12345, 56789, 34567];
4748
4749 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4750
4751 let values = Decimal64Array::from(integers.clone())
4752 .with_precision_and_scale(5, 2)
4753 .unwrap();
4754
4755 let array = DictionaryArray::new(keys, Arc::new(values));
4756 RoundTripTest::new(Arc::new(array.clone())).run();
4757
4758 let values = Decimal64Array::from(integers)
4759 .with_precision_and_scale(12, 2)
4760 .unwrap();
4761
4762 let array = array.with_values(Arc::new(values));
4763 RoundTripTest::new(Arc::new(array)).run();
4764 }
4765
4766 #[test]
4767 #[cfg_attr(miri, ignore)] fn arrow_writer_decimal128_dictionary() {
4769 let integers = vec![12345, 56789, 34567];
4770
4771 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4772
4773 let values = Decimal128Array::from(integers.clone())
4774 .with_precision_and_scale(5, 2)
4775 .unwrap();
4776
4777 let array = DictionaryArray::new(keys, Arc::new(values));
4778 RoundTripTest::new(Arc::new(array.clone())).run();
4779
4780 let values = Decimal128Array::from(integers)
4781 .with_precision_and_scale(12, 2)
4782 .unwrap();
4783
4784 let array = array.with_values(Arc::new(values));
4785 RoundTripTest::new(Arc::new(array)).run();
4786 }
4787
4788 #[test]
4789 #[cfg_attr(miri, ignore)] fn arrow_writer_decimal256_dictionary() {
4791 let integers = vec![
4792 i256::from_i128(12345),
4793 i256::from_i128(56789),
4794 i256::from_i128(34567),
4795 ];
4796
4797 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4798
4799 let values = Decimal256Array::from(integers.clone())
4800 .with_precision_and_scale(5, 2)
4801 .unwrap();
4802
4803 let array = DictionaryArray::new(keys, Arc::new(values));
4804 RoundTripTest::new(Arc::new(array.clone())).run();
4805
4806 let values = Decimal256Array::from(integers)
4807 .with_precision_and_scale(12, 2)
4808 .unwrap();
4809
4810 let array = array.with_values(Arc::new(values));
4811 RoundTripTest::new(Arc::new(array)).run();
4812 }
4813
4814 #[test]
4815 #[cfg_attr(miri, ignore)] fn arrow_writer_string_dictionary_unsigned_index() {
4817 #[expect(deprecated)]
4819 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4820 "dictionary",
4821 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4822 true,
4823 42,
4824 true,
4825 )]));
4826
4827 let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4829 .iter()
4830 .copied()
4831 .collect();
4832
4833 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4834 }
4835
4836 #[test]
4837 #[cfg_attr(miri, ignore)] fn u32_min_max() {
4839 let src = [
4841 u32::MIN,
4842 1,
4843 (i32::MAX as u32) - 1,
4844 i32::MAX as u32,
4845 (i32::MAX as u32) + 1,
4846 u32::MAX - 1,
4847 u32::MAX,
4848 ];
4849 let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4850 let files = RoundTripTest::new(values).with_nullable(false).run();
4851
4852 for file in files {
4853 let reader = SerializedFileReader::new(file).unwrap();
4855 let metadata = reader.metadata();
4856
4857 let mut row_offset = 0;
4858 for row_group in metadata.row_groups() {
4859 assert_eq!(row_group.num_columns(), 1);
4860 let column = row_group.column(0);
4861
4862 let num_values = column.num_values() as usize;
4863 let src_slice = &src[row_offset..row_offset + num_values];
4864 row_offset += column.num_values() as usize;
4865
4866 let stats = column.statistics().unwrap();
4867 if let Statistics::Int32(stats) = stats {
4868 assert_eq!(
4869 *stats.min_opt().unwrap() as u32,
4870 *src_slice.iter().min().unwrap()
4871 );
4872 assert_eq!(
4873 *stats.max_opt().unwrap() as u32,
4874 *src_slice.iter().max().unwrap()
4875 );
4876 } else {
4877 panic!("Statistics::Int32 missing")
4878 }
4879 }
4880 }
4881 }
4882
4883 #[test]
4884 #[cfg_attr(miri, ignore)] fn u64_min_max() {
4886 let src = [
4888 u64::MIN,
4889 1,
4890 (i64::MAX as u64) - 1,
4891 i64::MAX as u64,
4892 (i64::MAX as u64) + 1,
4893 u64::MAX - 1,
4894 u64::MAX,
4895 ];
4896 let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4897 let files = RoundTripTest::new(values).with_nullable(false).run();
4898
4899 for file in files {
4900 let reader = SerializedFileReader::new(file).unwrap();
4902 let metadata = reader.metadata();
4903
4904 let mut row_offset = 0;
4905 for row_group in metadata.row_groups() {
4906 assert_eq!(row_group.num_columns(), 1);
4907 let column = row_group.column(0);
4908
4909 let num_values = column.num_values() as usize;
4910 let src_slice = &src[row_offset..row_offset + num_values];
4911 row_offset += column.num_values() as usize;
4912
4913 let stats = column.statistics().unwrap();
4914 if let Statistics::Int64(stats) = stats {
4915 assert_eq!(
4916 *stats.min_opt().unwrap() as u64,
4917 *src_slice.iter().min().unwrap()
4918 );
4919 assert_eq!(
4920 *stats.max_opt().unwrap() as u64,
4921 *src_slice.iter().max().unwrap()
4922 );
4923 } else {
4924 panic!("Statistics::Int64 missing")
4925 }
4926 }
4927 }
4928 }
4929
4930 #[test]
4931 #[cfg_attr(miri, ignore)] fn statistics_null_counts_only_nulls() {
4933 let values = Arc::new(UInt64Array::from(vec![None, None]));
4935 let files = RoundTripTest::new(values).run();
4936
4937 for file in files {
4938 let reader = SerializedFileReader::new(file).unwrap();
4940 let metadata = reader.metadata();
4941 assert_eq!(metadata.num_row_groups(), 1);
4942 let row_group = metadata.row_group(0);
4943 assert_eq!(row_group.num_columns(), 1);
4944 let column = row_group.column(0);
4945 let stats = column.statistics().unwrap();
4946 assert_eq!(stats.null_count_opt(), Some(2));
4947 }
4948 }
4949
4950 #[test]
4951 #[cfg_attr(miri, ignore)] fn test_list_of_struct_roundtrip() {
4953 let int_field = Field::new("a", DataType::Int32, true);
4955 let int_field2 = Field::new("b", DataType::Int32, true);
4956
4957 let int_builder = Int32Builder::with_capacity(10);
4958 let int_builder2 = Int32Builder::with_capacity(10);
4959
4960 let struct_builder = StructBuilder::new(
4961 vec![int_field, int_field2],
4962 vec![Box::new(int_builder), Box::new(int_builder2)],
4963 );
4964 let mut list_builder = ListBuilder::new(struct_builder);
4965
4966 let values = list_builder.values();
4971 values
4972 .field_builder::<Int32Builder>(0)
4973 .unwrap()
4974 .append_value(1);
4975 values
4976 .field_builder::<Int32Builder>(1)
4977 .unwrap()
4978 .append_value(2);
4979 values.append(true);
4980 list_builder.append(true);
4981
4982 list_builder.append(true);
4984
4985 list_builder.append(false);
4987
4988 let values = list_builder.values();
4990 values
4991 .field_builder::<Int32Builder>(0)
4992 .unwrap()
4993 .append_null();
4994 values
4995 .field_builder::<Int32Builder>(1)
4996 .unwrap()
4997 .append_null();
4998 values.append(false);
4999 values
5000 .field_builder::<Int32Builder>(0)
5001 .unwrap()
5002 .append_null();
5003 values
5004 .field_builder::<Int32Builder>(1)
5005 .unwrap()
5006 .append_null();
5007 values.append(false);
5008 list_builder.append(true);
5009
5010 let values = list_builder.values();
5012 values
5013 .field_builder::<Int32Builder>(0)
5014 .unwrap()
5015 .append_null();
5016 values
5017 .field_builder::<Int32Builder>(1)
5018 .unwrap()
5019 .append_value(3);
5020 values.append(true);
5021 list_builder.append(true);
5022
5023 let values = list_builder.values();
5025 values
5026 .field_builder::<Int32Builder>(0)
5027 .unwrap()
5028 .append_value(2);
5029 values
5030 .field_builder::<Int32Builder>(1)
5031 .unwrap()
5032 .append_null();
5033 values.append(true);
5034 list_builder.append(true);
5035
5036 let array = Arc::new(list_builder.finish());
5037
5038 RoundTripTest::new(array).run();
5039 }
5040
5041 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
5042 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
5043 }
5044
5045 #[test]
5046 fn test_aggregates_records() {
5047 let arrays = [
5048 Int32Array::from((0..100).collect::<Vec<_>>()),
5049 Int32Array::from((0..50).collect::<Vec<_>>()),
5050 Int32Array::from((200..500).collect::<Vec<_>>()),
5051 ];
5052
5053 let schema = Arc::new(Schema::new(vec![Field::new(
5054 "int",
5055 ArrowDataType::Int32,
5056 false,
5057 )]));
5058
5059 let file = tempfile::tempfile().unwrap();
5060
5061 let props = WriterProperties::builder()
5062 .set_max_row_group_row_count(Some(200))
5063 .build();
5064
5065 let mut writer =
5066 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5067
5068 for array in arrays {
5069 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5070 writer.write(&batch).unwrap();
5071 }
5072
5073 writer.close().unwrap();
5074
5075 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5076 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
5077
5078 let batches = builder
5079 .with_batch_size(100)
5080 .build()
5081 .unwrap()
5082 .collect::<ArrowResult<Vec<_>>>()
5083 .unwrap();
5084
5085 assert_eq!(batches.len(), 5);
5086 assert!(batches.iter().all(|x| x.num_columns() == 1));
5087
5088 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
5089
5090 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
5091
5092 let values: Vec<_> = batches
5093 .iter()
5094 .flat_map(|x| {
5095 x.column(0)
5096 .as_any()
5097 .downcast_ref::<Int32Array>()
5098 .unwrap()
5099 .values()
5100 .iter()
5101 .copied()
5102 })
5103 .collect();
5104
5105 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
5106 assert_eq!(&values, &expected_values)
5107 }
5108
5109 #[test]
5110 fn complex_aggregate() {
5111 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
5113 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
5114 let struct_a = Arc::new(Field::new(
5115 "struct_a",
5116 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
5117 true,
5118 ));
5119
5120 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
5121 let struct_b = Arc::new(Field::new(
5122 "struct_b",
5123 DataType::Struct(vec![list_a.clone()].into()),
5124 false,
5125 ));
5126
5127 let schema = Arc::new(Schema::new(vec![struct_b]));
5128
5129 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
5131 let field_b_array =
5132 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
5133
5134 let struct_a_array = StructArray::from(vec![
5135 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
5136 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
5137 ]);
5138
5139 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
5140 .len(5)
5141 .add_buffer(Buffer::from_iter(vec![
5142 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
5143 ]))
5144 .null_bit_buffer(Some(Buffer::from_iter(vec![
5145 true, false, true, false, true,
5146 ])))
5147 .child_data(vec![struct_a_array.into_data()])
5148 .build()
5149 .unwrap();
5150
5151 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
5152 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
5153
5154 let batch1 =
5155 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
5156 .unwrap();
5157
5158 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
5159 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
5160
5161 let struct_a_array = StructArray::from(vec![
5162 (field_a, Arc::new(field_a_array) as ArrayRef),
5163 (field_b, Arc::new(field_b_array) as ArrayRef),
5164 ]);
5165
5166 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
5167 .len(2)
5168 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
5169 .child_data(vec![struct_a_array.into_data()])
5170 .build()
5171 .unwrap();
5172
5173 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
5174 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
5175
5176 let batch2 =
5177 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
5178 .unwrap();
5179
5180 let batches = &[batch1, batch2];
5181
5182 let expected = r"
5185 +-------------------------------------------------------------------------------------------------------+
5186 | struct_b |
5187 +-------------------------------------------------------------------------------------------------------+
5188 | {list: [{leaf_a: 1, leaf_b: 1}]} |
5189 | {list: } |
5190 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
5191 | {list: } |
5192 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
5193 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
5194 | {list: [{leaf_a: 10, leaf_b: }]} |
5195 +-------------------------------------------------------------------------------------------------------+
5196 ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
5197
5198 let actual = pretty_format_batches(batches).unwrap().to_string();
5199 assert_eq!(actual, expected);
5200
5201 let file = tempfile::tempfile().unwrap();
5203 let props = WriterProperties::builder()
5204 .set_max_row_group_row_count(Some(6))
5205 .build();
5206
5207 let mut writer =
5208 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
5209
5210 for batch in batches {
5211 writer.write(batch).unwrap();
5212 }
5213 writer.close().unwrap();
5214
5215 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5220 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
5221
5222 let batches = builder
5223 .with_batch_size(2)
5224 .build()
5225 .unwrap()
5226 .collect::<ArrowResult<Vec<_>>>()
5227 .unwrap();
5228
5229 assert_eq!(batches.len(), 4);
5230 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
5231 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
5232
5233 let actual = pretty_format_batches(&batches).unwrap().to_string();
5234 assert_eq!(actual, expected);
5235 }
5236
5237 #[test]
5238 fn test_arrow_writer_metadata() {
5239 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5240 let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
5241
5242 let batch = RecordBatch::try_new(
5243 Arc::new(batch_schema),
5244 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5245 )
5246 .unwrap();
5247
5248 let mut buf = Vec::with_capacity(1024);
5249 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
5250 writer.write(&batch).unwrap();
5251 writer.close().unwrap();
5252 }
5253
5254 #[test]
5255 fn test_arrow_writer_nullable() {
5256 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5257 let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
5258 let file_schema = Arc::new(file_schema);
5259
5260 let batch = RecordBatch::try_new(
5261 Arc::new(batch_schema),
5262 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5263 )
5264 .unwrap();
5265
5266 let mut buf = Vec::with_capacity(1024);
5267 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5268 writer.write(&batch).unwrap();
5269 writer.close().unwrap();
5270
5271 let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
5272 let back = read.next().unwrap().unwrap();
5273 assert_eq!(back.schema(), file_schema);
5274 assert_ne!(back.schema(), batch.schema());
5275 assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
5276 }
5277
5278 #[test]
5279 fn in_progress_accounting() {
5280 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
5282
5283 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5285
5286 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
5288
5289 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
5290
5291 assert_eq!(writer.in_progress_size(), 0);
5293 assert_eq!(writer.in_progress_rows(), 0);
5294 assert_eq!(writer.memory_size(), 0);
5295 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
5297
5298 let initial_size = writer.in_progress_size();
5300 assert!(initial_size > 0);
5301 assert_eq!(writer.in_progress_rows(), 5);
5302 let initial_memory = writer.memory_size();
5303 assert!(initial_memory > 0);
5304 assert!(
5306 initial_size <= initial_memory,
5307 "{initial_size} <= {initial_memory}"
5308 );
5309
5310 writer.write(&batch).unwrap();
5312 assert!(writer.in_progress_size() > initial_size);
5313 assert_eq!(writer.in_progress_rows(), 10);
5314 assert!(writer.memory_size() > initial_memory);
5315 assert!(
5316 writer.in_progress_size() <= writer.memory_size(),
5317 "in_progress_size {} <= memory_size {}",
5318 writer.in_progress_size(),
5319 writer.memory_size()
5320 );
5321
5322 let pre_flush_bytes_written = writer.bytes_written();
5324 writer.flush().unwrap();
5325 assert_eq!(writer.in_progress_size(), 0);
5326 assert_eq!(writer.memory_size(), 0);
5327 assert!(writer.bytes_written() > pre_flush_bytes_written);
5328
5329 writer.close().unwrap();
5330 }
5331
5332 #[test]
5333 fn test_writer_all_null() {
5334 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5335 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
5336 let batch = RecordBatch::try_from_iter(vec![
5337 ("a", Arc::new(a) as ArrayRef),
5338 ("b", Arc::new(b) as ArrayRef),
5339 ])
5340 .unwrap();
5341
5342 let mut buf = Vec::with_capacity(1024);
5343 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5344 writer.write(&batch).unwrap();
5345 writer.close().unwrap();
5346
5347 let bytes = Bytes::from(buf);
5348 let options = ReadOptionsBuilder::new().with_page_index().build();
5349 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
5350 let index = reader.metadata().page_index().unwrap();
5351
5352 assert_eq!(index.num_data_pages(0, 0), Some(1)); assert_eq!(index.num_data_pages(0, 1), Some(1)); }
5355
5356 #[test]
5357 fn test_disabled_statistics_with_page() {
5358 let file_schema = Schema::new(vec![
5359 Field::new("a", DataType::Utf8, true),
5360 Field::new("b", DataType::Utf8, true),
5361 ]);
5362 let file_schema = Arc::new(file_schema);
5363
5364 let batch = RecordBatch::try_new(
5365 file_schema.clone(),
5366 vec![
5367 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5368 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5369 ],
5370 )
5371 .unwrap();
5372
5373 let props = WriterProperties::builder()
5374 .set_statistics_enabled(EnabledStatistics::None)
5375 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5376 .build();
5377
5378 let mut buf = Vec::with_capacity(1024);
5379 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5380 writer.write(&batch).unwrap();
5381
5382 let metadata = writer.close().unwrap();
5383 assert_eq!(metadata.num_row_groups(), 1);
5384 let row_group = metadata.row_group(0);
5385 assert_eq!(row_group.num_columns(), 2);
5386 assert!(row_group.column(0).offset_index_offset().is_some());
5388 assert!(row_group.column(0).column_index_offset().is_some());
5389 assert!(row_group.column(1).offset_index_offset().is_some());
5391 assert!(row_group.column(1).column_index_offset().is_none());
5392
5393 let options = ReadOptionsBuilder::new().with_page_index().build();
5394 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5395
5396 let row_group = reader.get_row_group(0).unwrap();
5397 let a_col = row_group.metadata().column(0);
5398 let b_col = row_group.metadata().column(1);
5399
5400 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5402 let min = byte_array_stats.min_opt().unwrap();
5403 let max = byte_array_stats.max_opt().unwrap();
5404
5405 assert_eq!(min.as_bytes(), b"a");
5406 assert_eq!(max.as_bytes(), b"d");
5407 } else {
5408 panic!("expecting Statistics::ByteArray");
5409 }
5410
5411 assert!(b_col.statistics().is_none());
5413
5414 let page_index = reader.metadata().page_index().unwrap();
5415
5416 let a_idx = page_index.column_index(0, 0);
5417 assert!(
5418 matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
5419 "{a_idx:?}"
5420 );
5421 let b_idx = page_index.column_index(0, 1);
5422 assert!(b_idx.is_none(), "{b_idx:?}");
5423 }
5424
5425 #[test]
5426 fn test_disabled_statistics_with_chunk() {
5427 let file_schema = Schema::new(vec![
5428 Field::new("a", DataType::Utf8, true),
5429 Field::new("b", DataType::Utf8, true),
5430 ]);
5431 let file_schema = Arc::new(file_schema);
5432
5433 let batch = RecordBatch::try_new(
5434 file_schema.clone(),
5435 vec![
5436 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5437 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5438 ],
5439 )
5440 .unwrap();
5441
5442 let props = WriterProperties::builder()
5443 .set_statistics_enabled(EnabledStatistics::None)
5444 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5445 .build();
5446
5447 let mut buf = Vec::with_capacity(1024);
5448 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5449 writer.write(&batch).unwrap();
5450
5451 let metadata = writer.close().unwrap();
5452 assert_eq!(metadata.num_row_groups(), 1);
5453 let row_group = metadata.row_group(0);
5454 assert_eq!(row_group.num_columns(), 2);
5455 assert!(row_group.column(0).offset_index_offset().is_some());
5457 assert!(row_group.column(0).column_index_offset().is_none());
5458 assert!(row_group.column(1).offset_index_offset().is_some());
5460 assert!(row_group.column(1).column_index_offset().is_none());
5461
5462 let options = ReadOptionsBuilder::new().with_page_index().build();
5463 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5464
5465 let row_group = reader.get_row_group(0).unwrap();
5466 let a_col = row_group.metadata().column(0);
5467 let b_col = row_group.metadata().column(1);
5468
5469 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5471 let min = byte_array_stats.min_opt().unwrap();
5472 let max = byte_array_stats.max_opt().unwrap();
5473
5474 assert_eq!(min.as_bytes(), b"a");
5475 assert_eq!(max.as_bytes(), b"d");
5476 } else {
5477 panic!("expecting Statistics::ByteArray");
5478 }
5479
5480 assert!(b_col.statistics().is_none());
5482
5483 let page_index = reader.metadata().page_index().unwrap();
5484
5485 let a_idx = page_index.column_index(0, 0);
5486 assert!(a_idx.is_none(), "{a_idx:?}");
5487 let b_idx = page_index.column_index(0, 1);
5488 assert!(b_idx.is_none(), "{b_idx:?}");
5489 }
5490
5491 #[test]
5492 fn test_arrow_writer_skip_metadata() {
5493 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5494 let file_schema = Arc::new(batch_schema.clone());
5495
5496 let batch = RecordBatch::try_new(
5497 Arc::new(batch_schema),
5498 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5499 )
5500 .unwrap();
5501 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5502
5503 let mut buf = Vec::with_capacity(1024);
5504 let mut writer =
5505 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5506 writer.write(&batch).unwrap();
5507 writer.close().unwrap();
5508
5509 let bytes = Bytes::from(buf);
5510 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5511 assert_eq!(file_schema, *reader_builder.schema());
5512 if let Some(key_value_metadata) = reader_builder
5513 .metadata()
5514 .file_metadata()
5515 .key_value_metadata()
5516 {
5517 assert!(
5518 !key_value_metadata
5519 .iter()
5520 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5521 );
5522 }
5523 }
5524
5525 #[test]
5526 fn test_arrow_writer_skip_path_in_schema() {
5527 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5528 let file_schema = Arc::new(batch_schema.clone());
5529
5530 let batch = RecordBatch::try_new(
5531 Arc::new(batch_schema),
5532 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5533 )
5534 .unwrap();
5535
5536 let skip_options = ArrowWriterOptions::new();
5538
5539 let mut buf = Vec::with_capacity(1024);
5540 let mut writer =
5541 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5542 writer.write(&batch).unwrap();
5543 writer.close().unwrap();
5544
5545 let skip_options = ArrowWriterOptions::new().with_properties(
5547 WriterProperties::builder()
5548 .set_write_path_in_schema(false)
5549 .build(),
5550 );
5551
5552 let mut buf2 = Vec::with_capacity(1024);
5553 let mut writer =
5554 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5555 .unwrap();
5556 writer.write(&batch).unwrap();
5557 writer.close().unwrap();
5558
5559 assert!(buf.len() > buf2.len());
5561 }
5562
5563 #[test]
5564 fn mismatched_schemas() {
5565 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5566 let file_schema = Arc::new(Schema::new(vec![Field::new(
5567 "temperature",
5568 DataType::Float64,
5569 false,
5570 )]));
5571
5572 let batch = RecordBatch::try_new(
5573 Arc::new(batch_schema),
5574 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5575 )
5576 .unwrap();
5577
5578 let mut buf = Vec::with_capacity(1024);
5579 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5580
5581 let err = writer.write(&batch).unwrap_err().to_string();
5582 assert_eq!(
5583 err,
5584 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5585 );
5586 }
5587
5588 #[test]
5589 fn test_roundtrip_empty_schema() {
5591 let empty_batch = RecordBatch::try_new_with_options(
5593 Arc::new(Schema::empty()),
5594 vec![],
5595 &RecordBatchOptions::default().with_row_count(Some(0)),
5596 )
5597 .unwrap();
5598
5599 let mut parquet_bytes: Vec<u8> = Vec::new();
5601 let mut writer =
5602 ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5603 writer.write(&empty_batch).unwrap();
5604 writer.close().unwrap();
5605
5606 let bytes = Bytes::from(parquet_bytes);
5608 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5609 assert_eq!(reader.schema(), &empty_batch.schema());
5610 let batches: Vec<_> = reader
5611 .build()
5612 .unwrap()
5613 .collect::<ArrowResult<Vec<_>>>()
5614 .unwrap();
5615 assert_eq!(batches.len(), 0);
5616 }
5617
5618 #[test]
5619 fn test_page_stats_not_written_by_default() {
5620 let string_field = Field::new("a", DataType::Utf8, false);
5621 let schema = Schema::new(vec![string_field]);
5622 let raw_string_values = vec!["Blart Versenwald III"];
5623 let string_values = StringArray::from(raw_string_values.clone());
5624 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5625
5626 let props = WriterProperties::builder()
5627 .set_statistics_enabled(EnabledStatistics::Page)
5628 .set_dictionary_enabled(false)
5629 .set_encoding(Encoding::PLAIN)
5630 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5631 .build();
5632
5633 let file = roundtrip_opts(&batch, props);
5634
5635 let first_page = &file[4..];
5640 let mut prot = ThriftSliceInputProtocol::new(first_page);
5641 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5642 let stats = hdr.data_page_header.unwrap().statistics;
5643
5644 assert!(stats.is_none());
5645 }
5646
5647 #[test]
5648 fn test_page_stats_when_enabled() {
5649 let string_field = Field::new("a", DataType::Utf8, false);
5650 let schema = Schema::new(vec![string_field]);
5651 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5652 let string_values = StringArray::from(raw_string_values.clone());
5653 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5654
5655 let props = WriterProperties::builder()
5656 .set_statistics_enabled(EnabledStatistics::Page)
5657 .set_dictionary_enabled(false)
5658 .set_encoding(Encoding::PLAIN)
5659 .set_write_page_header_statistics(true)
5660 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5661 .build();
5662
5663 let file = roundtrip_opts(&batch, props);
5664
5665 let first_page = &file[4..];
5670 let mut prot = ThriftSliceInputProtocol::new(first_page);
5671 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5672 let stats = hdr.data_page_header.unwrap().statistics;
5673
5674 let stats = stats.unwrap();
5675 assert!(stats.is_max_value_exact.unwrap());
5677 assert!(stats.is_min_value_exact.unwrap());
5678 assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5679 assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5680 }
5681
5682 #[test]
5683 fn test_page_stats_truncation() {
5684 let string_field = Field::new("a", DataType::Utf8, false);
5685 let binary_field = Field::new("b", DataType::Binary, false);
5686 let schema = Schema::new(vec![string_field, binary_field]);
5687
5688 let raw_string_values = vec!["Blart Versenwald III"];
5689 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5690 let raw_binary_value_refs = raw_binary_values
5691 .iter()
5692 .map(|x| x.as_slice())
5693 .collect::<Vec<_>>();
5694
5695 let string_values = StringArray::from(raw_string_values.clone());
5696 let binary_values = BinaryArray::from(raw_binary_value_refs);
5697 let batch = RecordBatch::try_new(
5698 Arc::new(schema),
5699 vec![Arc::new(string_values), Arc::new(binary_values)],
5700 )
5701 .unwrap();
5702
5703 let props = WriterProperties::builder()
5704 .set_statistics_truncate_length(Some(2))
5705 .set_dictionary_enabled(false)
5706 .set_encoding(Encoding::PLAIN)
5707 .set_write_page_header_statistics(true)
5708 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5709 .build();
5710
5711 let file = roundtrip_opts(&batch, props);
5712
5713 let first_page = &file[4..];
5718 let mut prot = ThriftSliceInputProtocol::new(first_page);
5719 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5720 let stats = hdr.data_page_header.unwrap().statistics;
5721 assert!(stats.is_some());
5722 let stats = stats.unwrap();
5723 assert!(!stats.is_max_value_exact.unwrap());
5725 assert!(!stats.is_min_value_exact.unwrap());
5726 assert_eq!(stats.max_value.unwrap(), b"Bm");
5727 assert_eq!(stats.min_value.unwrap(), b"Bl");
5728
5729 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5731 let mut prot = ThriftSliceInputProtocol::new(second_page);
5732 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5733 let stats = hdr.data_page_header.unwrap().statistics;
5734 assert!(stats.is_some());
5735 let stats = stats.unwrap();
5736 assert!(!stats.is_max_value_exact.unwrap());
5738 assert!(!stats.is_min_value_exact.unwrap());
5739 assert_eq!(stats.max_value.unwrap(), b"Bm");
5740 assert_eq!(stats.min_value.unwrap(), b"Bl");
5741 }
5742
5743 #[test]
5744 fn test_page_encoding_statistics_roundtrip() {
5745 let batch_schema = Schema::new(vec![Field::new(
5746 "int32",
5747 arrow_schema::DataType::Int32,
5748 false,
5749 )]);
5750
5751 let batch = RecordBatch::try_new(
5752 Arc::new(batch_schema.clone()),
5753 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5754 )
5755 .unwrap();
5756
5757 let mut file: File = tempfile::tempfile().unwrap();
5758 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5759 writer.write(&batch).unwrap();
5760 let file_metadata = writer.close().unwrap();
5761
5762 assert_eq!(file_metadata.num_row_groups(), 1);
5763 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5764 assert!(
5765 file_metadata
5766 .row_group(0)
5767 .column(0)
5768 .page_encoding_stats()
5769 .is_some()
5770 );
5771 let chunk_page_stats = file_metadata
5772 .row_group(0)
5773 .column(0)
5774 .page_encoding_stats()
5775 .unwrap();
5776
5777 let options = ReadOptionsBuilder::new()
5779 .with_page_index()
5780 .with_encoding_stats_as_mask(false)
5781 .build();
5782 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5783
5784 let rowgroup = reader.get_row_group(0).expect("row group missing");
5785 assert_eq!(rowgroup.num_columns(), 1);
5786 let column = rowgroup.metadata().column(0);
5787 assert!(column.page_encoding_stats().is_some());
5788 let file_page_stats = column.page_encoding_stats().unwrap();
5789 assert_eq!(chunk_page_stats, file_page_stats);
5790 }
5791
5792 #[test]
5793 #[cfg_attr(miri, ignore)] fn test_different_dict_page_size_limit() {
5795 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5796 let schema = Arc::new(Schema::new(vec![
5797 Field::new("col0", arrow_schema::DataType::Int64, false),
5798 Field::new("col1", arrow_schema::DataType::Int64, false),
5799 ]));
5800 let batch =
5801 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5802
5803 let props = WriterProperties::builder()
5804 .set_dictionary_page_size_limit(1024 * 1024)
5805 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5806 .build();
5807 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5808 writer.write(&batch).unwrap();
5809 let data = Bytes::from(writer.into_inner().unwrap());
5810
5811 let mut metadata = ParquetMetaDataReader::new();
5812 metadata.try_parse(&data).unwrap();
5813 let metadata = metadata.finish().unwrap();
5814 let col0_meta = metadata.row_group(0).column(0);
5815 let col1_meta = metadata.row_group(0).column(1);
5816
5817 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5818 let mut reader =
5819 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5820 let page = reader.get_next_page().unwrap().unwrap();
5821 match page {
5822 Page::DictionaryPage { buf, .. } => buf.len(),
5823 _ => panic!("expected DictionaryPage"),
5824 }
5825 };
5826
5827 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5828 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5829 }
5830
5831 #[test]
5832 #[cfg_attr(miri, ignore)] fn test_arrow_writer_granular_mode_roundtrip() {
5834 let small = "tiny".to_string();
5843 let big = "x".repeat(64 * 1024);
5844 let strings: Vec<String> = (0..256)
5845 .map(|i| {
5846 if i % 16 == 0 {
5847 big.clone()
5848 } else {
5849 small.clone()
5850 }
5851 })
5852 .collect();
5853
5854 let schema = Arc::new(Schema::new(vec![Field::new(
5855 "col",
5856 ArrowDataType::Utf8,
5857 false,
5858 )]));
5859 let batch = RecordBatch::try_new(
5860 schema.clone(),
5861 vec![Arc::new(StringArray::from(strings.clone())) as _],
5862 )
5863 .unwrap();
5864
5865 let props = WriterProperties::builder()
5866 .set_dictionary_enabled(false)
5867 .set_data_page_size_limit(16 * 1024)
5868 .build();
5869 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5870 writer.write(&batch).unwrap();
5871 let data = Bytes::from(writer.into_inner().unwrap());
5872
5873 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5874 let read = reader.next().unwrap().unwrap();
5875 assert!(reader.next().is_none(), "expected one batch");
5876 let col = read
5877 .column(0)
5878 .as_any()
5879 .downcast_ref::<StringArray>()
5880 .unwrap();
5881 assert_eq!(col.len(), strings.len());
5882 for (i, expected) in strings.iter().enumerate() {
5883 assert_eq!(
5884 col.value(i),
5885 expected.as_str(),
5886 "value mismatch at index {i}"
5887 );
5888 }
5889 }
5890
5891 #[test]
5892 fn test_arrow_writer_all_null_string_column() {
5893 let num_rows = 1024;
5898 let schema = Arc::new(Schema::new(vec![Field::new(
5899 "col",
5900 ArrowDataType::Utf8,
5901 true,
5902 )]));
5903 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5904 let batch = RecordBatch::try_new(
5905 schema.clone(),
5906 vec![Arc::new(StringArray::from(nulls)) as _],
5907 )
5908 .unwrap();
5909
5910 let props = WriterProperties::builder()
5911 .set_dictionary_enabled(false)
5912 .set_data_page_size_limit(16 * 1024)
5913 .build();
5914 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5915 writer.write(&batch).unwrap();
5916 let data = Bytes::from(writer.into_inner().unwrap());
5917
5918 let mut metadata = ParquetMetaDataReader::new();
5921 metadata.try_parse(&data).unwrap();
5922 let metadata = metadata.finish().unwrap();
5923 let row_group = metadata.row_group(0);
5924 let col_meta = row_group.column(0);
5925 assert_eq!(row_group.num_rows() as usize, num_rows);
5926 if let Some(stats) = col_meta.statistics() {
5929 assert_eq!(
5930 stats.null_count_opt().unwrap_or(0) as usize,
5931 num_rows,
5932 "expected all-null column to report null_count = num_rows"
5933 );
5934 }
5935
5936 let mut reader =
5937 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5938 let mut total_values = 0u32;
5939 while let Some(page) = reader.get_next_page().unwrap() {
5940 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5941 total_values += page.num_values();
5942 }
5943 }
5944 assert_eq!(
5945 total_values as usize, num_rows,
5946 "expected every level position to be represented in some page"
5947 );
5948 }
5949
5950 struct WriteBatchesShape {
5951 num_batches: usize,
5952 rows_per_batch: usize,
5953 row_size: usize,
5954 }
5955
5956 fn write_batches(
5958 WriteBatchesShape {
5959 num_batches,
5960 rows_per_batch,
5961 row_size,
5962 }: WriteBatchesShape,
5963 props: WriterProperties,
5964 ) -> ParquetRecordBatchReaderBuilder<File> {
5965 let schema = Arc::new(Schema::new(vec![Field::new(
5966 "str",
5967 ArrowDataType::Utf8,
5968 false,
5969 )]));
5970 let file = tempfile::tempfile().unwrap();
5971 let mut writer =
5972 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5973
5974 for batch_idx in 0..num_batches {
5975 let strings: Vec<String> = (0..rows_per_batch)
5976 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5977 .collect();
5978 let array = StringArray::from(strings);
5979 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5980 writer.write(&batch).unwrap();
5981 }
5982 writer.close().unwrap();
5983 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5984 }
5985
5986 #[test]
5987 fn test_row_group_limit_none_writes_single_row_group() {
5989 let props = WriterProperties::builder()
5990 .set_max_row_group_row_count(None)
5991 .set_max_row_group_bytes(None)
5992 .build();
5993
5994 let builder = write_batches(
5995 WriteBatchesShape {
5996 num_batches: 1,
5997 rows_per_batch: 1000,
5998 row_size: 4,
5999 },
6000 props,
6001 );
6002
6003 assert_eq!(
6004 &row_group_sizes(builder.metadata()),
6005 &[1000],
6006 "With no limits, all rows should be in a single row group"
6007 );
6008 }
6009
6010 #[test]
6011 fn test_row_group_limit_rows_only() {
6013 let props = WriterProperties::builder()
6014 .set_max_row_group_row_count(Some(300))
6015 .set_max_row_group_bytes(None)
6016 .build();
6017
6018 let builder = write_batches(
6019 WriteBatchesShape {
6020 num_batches: 1,
6021 rows_per_batch: 1000,
6022 row_size: 4,
6023 },
6024 props,
6025 );
6026
6027 assert_eq!(
6028 &row_group_sizes(builder.metadata()),
6029 &[300, 300, 300, 100],
6030 "Row groups should be split by row count"
6031 );
6032 }
6033
6034 #[test]
6035 #[cfg_attr(miri, ignore)] fn test_row_group_limit_rows_only_many_splits() {
6039 let props = WriterProperties::builder()
6040 .set_max_row_group_row_count(Some(1))
6041 .set_max_row_group_bytes(None)
6042 .build();
6043
6044 let rows = 50_000;
6045 let builder = write_batches(
6046 WriteBatchesShape {
6047 num_batches: 1,
6048 rows_per_batch: rows,
6049 row_size: 4,
6050 },
6051 props,
6052 );
6053
6054 let sizes = row_group_sizes(builder.metadata());
6055 assert_eq!(sizes.len(), rows, "Every row should get its own row group");
6056 assert_eq!(
6057 sizes.iter().sum::<i64>(),
6058 rows as i64,
6059 "Total rows should be preserved"
6060 );
6061 }
6062
6063 #[test]
6064 fn test_row_group_limit_bytes_only() {
6066 let props = WriterProperties::builder()
6067 .set_max_row_group_row_count(None)
6068 .set_max_row_group_bytes(Some(3500))
6070 .build();
6071
6072 let builder = write_batches(
6073 WriteBatchesShape {
6074 num_batches: 10,
6075 rows_per_batch: 10,
6076 row_size: 100,
6077 },
6078 props,
6079 );
6080
6081 let sizes = row_group_sizes(builder.metadata());
6082
6083 assert!(
6084 sizes.len() > 1,
6085 "Should have multiple row groups due to byte limit, got {sizes:?}",
6086 );
6087
6088 let total_rows: i64 = sizes.iter().sum();
6089 assert_eq!(total_rows, 100, "Total rows should be preserved");
6090 }
6091
6092 #[test]
6093 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
6095 let schema = Arc::new(Schema::new(vec![Field::new(
6096 "str",
6097 ArrowDataType::Utf8,
6098 false,
6099 )]));
6100 let file = tempfile::tempfile().unwrap();
6101
6102 let props = WriterProperties::builder()
6104 .set_max_row_group_row_count(None)
6105 .set_max_row_group_bytes(None)
6106 .build();
6107 let mut writer =
6108 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
6109
6110 let first_array = StringArray::from(
6111 (0..10)
6112 .map(|i| format!("{i:0>100}"))
6113 .collect::<Vec<String>>(),
6114 );
6115 let first_batch =
6116 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
6117 writer.write(&first_batch).unwrap();
6118 assert_eq!(writer.in_progress_rows(), 10);
6119
6120 writer.max_row_group_bytes = Some(1);
6123
6124 let second_array = StringArray::from(vec!["x".to_string()]);
6125 let second_batch =
6126 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
6127 writer.write(&second_batch).unwrap();
6128 writer.close().unwrap();
6129 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
6130
6131 assert_eq!(
6132 &row_group_sizes(builder.metadata()),
6133 &[10, 1],
6134 "The second write should flush an oversized in-progress row group first",
6135 );
6136 }
6137
6138 #[test]
6139 fn test_row_group_limit_both_row_wins_single_batch() {
6141 let props = WriterProperties::builder()
6142 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
6145
6146 let builder = write_batches(
6147 WriteBatchesShape {
6148 num_batches: 1,
6149 row_size: 4,
6150 rows_per_batch: 1000,
6151 },
6152 props,
6153 );
6154
6155 assert_eq!(
6156 &row_group_sizes(builder.metadata()),
6157 &[200, 200, 200, 200, 200],
6158 "Row limit should trigger before byte limit"
6159 );
6160 }
6161
6162 #[test]
6163 fn test_row_group_limit_both_row_wins_multiple_batches() {
6165 let props = WriterProperties::builder()
6166 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
6169
6170 let builder = write_batches(
6171 WriteBatchesShape {
6172 num_batches: 10,
6173 rows_per_batch: 10,
6174 row_size: 100,
6175 },
6176 props,
6177 );
6178
6179 assert_eq!(
6180 &row_group_sizes(builder.metadata()),
6181 &[5; 20],
6182 "Row limit should trigger before byte limit"
6183 );
6184 }
6185
6186 #[test]
6187 fn test_row_group_limit_both_bytes_wins() {
6189 let props = WriterProperties::builder()
6190 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
6193
6194 let builder = write_batches(
6195 WriteBatchesShape {
6196 num_batches: 10,
6197 rows_per_batch: 10,
6198 row_size: 100,
6199 },
6200 props,
6201 );
6202
6203 let sizes = row_group_sizes(builder.metadata());
6204
6205 assert!(
6206 sizes.len() > 1,
6207 "Byte limit should trigger before row limit, got {sizes:?}",
6208 );
6209
6210 assert!(
6211 sizes.iter().all(|&s| s < 1000),
6212 "No row group should hit the row limit"
6213 );
6214
6215 let total_rows: i64 = sizes.iter().sum();
6216 assert_eq!(total_rows, 100, "Total rows should be preserved");
6217 }
6218
6219 #[test]
6220 fn test_row_group_limit_both_apply_to_same_batch() {
6223 let props = WriterProperties::builder()
6224 .set_max_row_group_row_count(Some(15))
6225 .set_max_row_group_bytes(Some(1500))
6226 .build();
6227
6228 let builder = write_batches(
6229 WriteBatchesShape {
6230 num_batches: 2,
6231 rows_per_batch: 10,
6232 row_size: 100,
6233 },
6234 props,
6235 );
6236
6237 assert_eq!(
6238 &row_group_sizes(builder.metadata()),
6239 &[14, 6],
6240 "Byte limit should still apply to a batch the row limit already split"
6241 );
6242 }
6243
6244 #[test]
6245 fn arrow_column_chunk_close_mut_drops_column_index() {
6246 use crate::arrow::ArrowSchemaConverter;
6247 use crate::file::writer::SerializedFileWriter;
6248
6249 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
6250 let props = Arc::new(
6251 WriterProperties::builder()
6252 .set_statistics_enabled(EnabledStatistics::Page)
6253 .build(),
6254 );
6255 let parquet_schema = ArrowSchemaConverter::new()
6256 .with_coerce_types(props.coerce_types())
6257 .convert(&schema)
6258 .unwrap();
6259
6260 let mut buf = Vec::with_capacity(1024);
6261 let mut writer =
6262 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
6263 .unwrap();
6264
6265 let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
6266 let mut col_writers = factory.create_column_writers(0).unwrap();
6267 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
6268 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
6269 col_writers[0].write(&leaves).unwrap();
6270 }
6271 let mut chunk = col_writers.pop().unwrap().close().unwrap();
6272
6273 assert!(
6275 chunk.close().column_index.is_some(),
6276 "EnabledStatistics::Page should produce a column_index"
6277 );
6278
6279 chunk.close_mut().column_index = None;
6281 assert!(chunk.close().column_index.is_none());
6282
6283 let mut rg = writer.next_row_group().unwrap();
6284 chunk.append_to_row_group(&mut rg).unwrap();
6285 rg.close().unwrap();
6286 let file_meta = writer.close().unwrap();
6287
6288 let cc = file_meta.row_group(0).column(0);
6291 assert!(cc.column_index_range().is_none());
6292 }
6293
6294 fn write_column_to_bytes(array: ArrayRef) -> Bytes {
6296 let schema = Arc::new(Schema::new(vec![Field::new(
6297 "col",
6298 array.data_type().clone(),
6299 true,
6300 )]));
6301 let buf = get_bytes_after_close(
6302 schema.clone(),
6303 &RecordBatch::try_new(schema, vec![array]).unwrap(),
6304 );
6305 Bytes::from(buf)
6306 }
6307
6308 fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
6312 let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
6313 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
6314 .unwrap()
6315 .build()
6316 .unwrap()
6317 .next()
6318 .unwrap()
6319 .unwrap()
6320 .column(0)
6321 .clone()
6322 }
6323
6324 fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
6325 let flat_schema = Arc::new(Schema::new(vec![Field::new(
6326 "col",
6327 flat.data_type().clone(),
6328 true,
6329 )]));
6330 let ree_bytes = write_column_to_bytes(ree);
6331 let flat_bytes = write_column_to_bytes(flat.clone());
6332 assert_eq!(
6333 ree_bytes, flat_bytes,
6334 "REE and flat bytes should be identical"
6335 );
6336
6337 let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
6338 let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
6339
6340 assert_eq!(decoded_ree.as_ref(), flat.as_ref());
6341 assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
6342 }
6343
6344 #[test]
6345 fn ree_string() {
6346 let ree: ArrayRef = Arc::new(
6347 [Some("a"), Some("a"), None, Some("b"), Some("b")]
6348 .into_iter()
6349 .collect::<Int32RunArray>(),
6350 );
6351 let flat: ArrayRef = Arc::new(StringArray::from(vec![
6352 Some("a"),
6353 Some("a"),
6354 None,
6355 Some("b"),
6356 Some("b"),
6357 ]));
6358 ree_write_read_roundtrip(ree, flat);
6359 }
6360
6361 #[test]
6362 fn ree_int32() {
6363 let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
6364 for v in [Some(1), Some(1), None, Some(2), Some(2)] {
6365 b.append_option(v);
6366 }
6367 let ree: ArrayRef = Arc::new(b.finish());
6368 let flat: ArrayRef = Arc::new(Int32Array::from(vec![
6369 Some(1),
6370 Some(1),
6371 None,
6372 Some(2),
6373 Some(2),
6374 ]));
6375 ree_write_read_roundtrip(ree, flat);
6376 }
6377
6378 #[test]
6379 fn ree_bool() {
6380 let ree: ArrayRef = Arc::new(
6382 RunArray::try_new(
6383 &Int32Array::from(vec![3, 5, 7]),
6384 &BooleanArray::from(vec![Some(true), None, Some(false)]),
6385 )
6386 .unwrap(),
6387 );
6388 let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
6389 Some(true),
6390 Some(true),
6391 Some(true),
6392 None,
6393 None,
6394 Some(false),
6395 Some(false),
6396 ]));
6397 ree_write_read_roundtrip(ree, flat);
6398 }
6399
6400 #[test]
6401 fn ree_fixed_size_binary() {
6402 let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6403 let mut b = FixedSizeBinaryBuilder::new(2);
6404 for v in vals {
6405 match v {
6406 Some(x) => b.append_value(x).unwrap(),
6407 None => b.append_null(),
6408 }
6409 }
6410 b.finish()
6411 };
6412 let ree: ArrayRef = Arc::new(
6414 RunArray::try_new(
6415 &Int32Array::from(vec![2, 4, 6]),
6416 &mk(&[Some(b"aa"), None, Some(b"bb")]),
6417 )
6418 .unwrap(),
6419 );
6420 let flat: ArrayRef = Arc::new(mk(&[
6421 Some(b"aa"),
6422 Some(b"aa"),
6423 None,
6424 None,
6425 Some(b"bb"),
6426 Some(b"bb"),
6427 ]));
6428 ree_write_read_roundtrip(ree, flat);
6429 }
6430
6431 #[test]
6432 fn ree_single_run() {
6433 let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6434 let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6435 ree_write_read_roundtrip(ree, flat);
6436 }
6437
6438 #[test]
6439 fn ree_float32() {
6440 let ree: ArrayRef = Arc::new(
6442 RunArray::try_new(
6443 &Int32Array::from(vec![2, 4, 5]),
6444 &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6445 )
6446 .unwrap(),
6447 );
6448 let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6449 Some(1.0_f32),
6450 Some(1.0_f32),
6451 None,
6452 None,
6453 Some(2.5_f32),
6454 ]));
6455 ree_write_read_roundtrip(ree, flat);
6456 }
6457
6458 #[test]
6459 fn ree_sliced() {
6460 let full: ArrayRef = Arc::new(
6465 RunArray::try_new(
6466 &Int32Array::from(vec![3, 5, 7]),
6467 &StringArray::from(vec!["a", "b", "c"]),
6468 )
6469 .unwrap(),
6470 );
6471 let sliced = full.slice(2, 5);
6472 let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6473 ree_write_read_roundtrip(sliced, flat);
6474 }
6475
6476 #[test]
6477 #[cfg_attr(miri, ignore)] fn test_number_distinct_values_exact_count() {
6479 let cardinality = 50u32;
6482 let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
6483 if i % 7 == 0 {
6484 None
6485 } else {
6486 Some((i % cardinality) as i32)
6487 }
6488 })));
6489 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
6490 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6491
6492 let props = WriterProperties::builder()
6493 .set_write_row_group_number_distinct_values(true)
6494 .build();
6495 let mut buf = Vec::new();
6496 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6497 writer.write(&batch).unwrap();
6498 let metadata = writer.close().unwrap();
6499
6500 let count = metadata
6501 .row_group(0)
6502 .column(0)
6503 .statistics()
6504 .and_then(|s| s.distinct_count_opt())
6505 .expect("distinct_count should be set");
6506 assert_eq!(count, cardinality as u64);
6508 }
6509
6510 #[test]
6511 fn test_number_distinct_values_view_types() {
6512 let cardinality = 5u32;
6515 let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
6516
6517 let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
6518 if i % 4 == 0 {
6519 None
6520 } else {
6521 Some(distinct_strings[(i % cardinality) as usize])
6522 }
6523 })));
6524
6525 let schema = Arc::new(Schema::new(vec![Field::new(
6526 "string_view_col",
6527 DataType::Utf8View,
6528 true,
6529 )]));
6530 let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
6531
6532 let props = WriterProperties::builder()
6533 .set_write_row_group_number_distinct_values(true)
6534 .build();
6535 let mut parquet_bytes = Vec::new();
6536 let mut writer =
6537 ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
6538 writer.write(&batch).unwrap();
6539 let metadata = writer.close().unwrap();
6540
6541 let distinct_count = metadata
6542 .row_group(0)
6543 .column(0)
6544 .statistics()
6545 .and_then(|s| s.distinct_count_opt())
6546 .expect("distinct_count should be set for Utf8View column");
6547 assert_eq!(distinct_count, cardinality as u64);
6548 }
6549
6550 #[test]
6551 fn test_number_distinct_values_not_written_by_default() {
6552 let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
6553 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
6554 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6555
6556 let mut buf = Vec::new();
6557 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
6558 writer.write(&batch).unwrap();
6559 let metadata = writer.close().unwrap();
6560
6561 let count = metadata
6562 .row_group(0)
6563 .column(0)
6564 .statistics()
6565 .and_then(|s| s.distinct_count_opt());
6566 assert!(count.is_none());
6567 }
6568
6569 #[test]
6570 fn ree_struct_with_ree_child() {
6571 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6574
6575 let col_a: ArrayRef = Arc::new(
6576 RunArray::try_new(
6577 &run_ends,
6578 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6579 )
6580 .unwrap(),
6581 );
6582 let col_b: ArrayRef = Arc::new(
6583 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6584 );
6585
6586 let struct_array: ArrayRef = Arc::new(StructArray::new(
6587 Fields::from(vec![
6588 Field::new("a", col_a.data_type().clone(), true),
6589 Field::new("b", col_b.data_type().clone(), true),
6590 ]),
6591 vec![col_a, col_b],
6592 None,
6593 ));
6594
6595 let schema = Arc::new(Schema::new(vec![Field::new(
6596 "row",
6597 struct_array.data_type().clone(),
6598 true,
6599 )]));
6600 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6601
6602 let mut buf = Vec::new();
6603 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6604 writer.write(&batch).unwrap();
6605 let metadata = writer.close().unwrap();
6606
6607 let parquet_schema = metadata.file_metadata().schema_descr();
6608 assert_eq!(parquet_schema.num_columns(), 2);
6609 assert_eq!(
6610 parquet_schema.column(0).physical_type(),
6611 crate::basic::Type::BYTE_ARRAY
6612 );
6613 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6614 assert_eq!(
6615 parquet_schema.column(1).physical_type(),
6616 crate::basic::Type::INT32
6617 );
6618 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6619 }
6620}