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::types::*;
30use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter};
31use arrow_schema::{
32 ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit,
33};
34
35use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_precision};
36
37use crate::arrow::ArrowSchemaConverter;
38use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder;
39use crate::basic::PageType;
40use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter};
41use crate::column::page_encryption::PageEncryptor;
42use crate::column::writer::encoder::ColumnValueEncoder;
43use crate::column::writer::{
44 ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer,
45};
46use crate::data_type::{ByteArray, FixedLenByteArray};
47#[cfg(feature = "encryption")]
48use crate::encryption::encrypt::FileEncryptor;
49use crate::errors::{ParquetError, Result};
50use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData};
51use crate::file::properties::{WriterProperties, WriterPropertiesPtr};
52use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
53use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
54use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor};
55use levels::{ArrayLevels, calculate_array_levels};
56
57mod byte_array;
58mod levels;
59
60#[doc(inline)]
61pub use crate::column::page_store::{
62 InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs,
63 PageStoreFactory,
64};
65
66pub struct ArrowWriter<W: Write> {
183 writer: SerializedFileWriter<W>,
185
186 in_progress: Option<ArrowRowGroupWriter>,
188
189 arrow_schema: SchemaRef,
193
194 row_group_writer_factory: ArrowRowGroupWriterFactory,
196
197 max_row_group_row_count: Option<usize>,
199
200 max_row_group_bytes: Option<usize>,
202
203 cdc_chunkers: Option<Vec<ContentDefinedChunker>>,
205}
206
207impl<W: Write + Send> std::fmt::Debug for ArrowWriter<W> {
208 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209 let buffered_memory = self.in_progress_size();
210 f.debug_struct("ArrowWriter")
211 .field("writer", &self.writer)
212 .field("in_progress_size", &format_args!("{buffered_memory} bytes"))
213 .field("in_progress_rows", &self.in_progress_rows())
214 .field("arrow_schema", &self.arrow_schema)
215 .field("max_row_group_row_count", &self.max_row_group_row_count)
216 .field("max_row_group_bytes", &self.max_row_group_bytes)
217 .finish()
218 }
219}
220
221impl<W: Write + Send> ArrowWriter<W> {
222 pub fn try_new(
228 writer: W,
229 arrow_schema: SchemaRef,
230 props: Option<WriterProperties>,
231 ) -> Result<Self> {
232 let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
233 Self::try_new_with_options(writer, arrow_schema, options)
234 }
235
236 pub fn try_new_with_options(
242 writer: W,
243 arrow_schema: SchemaRef,
244 options: ArrowWriterOptions,
245 ) -> Result<Self> {
246 let mut props = options.properties;
247
248 let schema = if let Some(parquet_schema) = options.schema_descr {
249 parquet_schema.clone()
250 } else {
251 let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types());
252 if let Some(schema_root) = &options.schema_root {
253 converter = converter.schema_root(schema_root);
254 }
255
256 converter.convert(&arrow_schema)?
257 };
258
259 if !options.skip_arrow_metadata {
260 add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
262 }
263
264 let max_row_group_row_count = props.max_row_group_row_count();
265 let max_row_group_bytes = props.max_row_group_bytes();
266
267 let props_ptr = Arc::new(props);
268 let file_writer =
269 SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
270
271 let mut row_group_writer_factory =
272 ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone());
273 if let Some(page_store_factory) = options.page_store_factory {
274 row_group_writer_factory =
275 row_group_writer_factory.with_page_store_factory(page_store_factory);
276 }
277
278 let cdc_chunkers = props_ptr
279 .content_defined_chunking()
280 .map(|opts| {
281 file_writer
282 .schema_descr()
283 .columns()
284 .iter()
285 .map(|desc| ContentDefinedChunker::new(desc, opts))
286 .collect::<Result<Vec<_>>>()
287 })
288 .transpose()?;
289
290 Ok(Self {
291 writer: file_writer,
292 in_progress: None,
293 arrow_schema,
294 row_group_writer_factory,
295 max_row_group_row_count,
296 max_row_group_bytes,
297 cdc_chunkers,
298 })
299 }
300
301 pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
303 self.writer.flushed_row_groups()
304 }
305
306 pub fn memory_size(&self) -> usize {
311 match &self.in_progress {
312 Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(),
313 None => 0,
314 }
315 }
316
317 pub fn in_progress_size(&self) -> usize {
324 match &self.in_progress {
325 Some(in_progress) => in_progress
326 .writers
327 .iter()
328 .map(|x| x.get_estimated_total_bytes())
329 .sum(),
330 None => 0,
331 }
332 }
333
334 pub fn in_progress_rows(&self) -> usize {
336 self.in_progress
337 .as_ref()
338 .map(|x| x.buffered_rows)
339 .unwrap_or_default()
340 }
341
342 pub fn bytes_written(&self) -> usize {
344 self.writer.bytes_written()
345 }
346
347 pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
359 if batch.num_rows() == 0 {
360 return Ok(());
361 }
362
363 let in_progress = match &mut self.in_progress {
364 Some(in_progress) => in_progress,
365 x => x.insert(
366 self.row_group_writer_factory
367 .create_row_group_writer(self.writer.flushed_row_groups().len())?,
368 ),
369 };
370
371 if let Some(max_rows) = self.max_row_group_row_count {
372 if in_progress.buffered_rows + batch.num_rows() > max_rows {
373 let to_write = max_rows - in_progress.buffered_rows;
374 let a = batch.slice(0, to_write);
375 let b = batch.slice(to_write, batch.num_rows() - to_write);
376 self.write(&a)?;
377 return self.write(&b);
378 }
379 }
380
381 if let Some(max_bytes) = self.max_row_group_bytes {
384 if in_progress.buffered_rows > 0 {
385 let current_bytes = in_progress.get_estimated_total_bytes();
386
387 if current_bytes >= max_bytes {
388 self.flush()?;
389 return self.write(batch);
390 }
391
392 if let Some(avg_row_bytes) = current_bytes
393 .checked_div(in_progress.buffered_rows)
394 .filter(|avg_row_bytes| *avg_row_bytes > 0)
395 {
396 let remaining_bytes = max_bytes - current_bytes;
398 let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
399
400 if batch.num_rows() > rows_that_fit {
401 if rows_that_fit > 0 {
402 let a = batch.slice(0, rows_that_fit);
403 let b = batch.slice(rows_that_fit, batch.num_rows() - rows_that_fit);
404 self.write(&a)?;
405 return self.write(&b);
406 } else {
407 self.flush()?;
408 return self.write(batch);
409 }
410 }
411 }
412 }
413 }
414
415 match self.cdc_chunkers.as_mut() {
416 Some(chunkers) => in_progress.write_with_chunkers(batch, chunkers)?,
417 None => in_progress.write(batch)?,
418 }
419
420 let should_flush = self
421 .max_row_group_row_count
422 .is_some_and(|max| in_progress.buffered_rows >= max)
423 || self
424 .max_row_group_bytes
425 .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
426
427 if should_flush {
428 self.flush()?
429 }
430 Ok(())
431 }
432
433 pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
438 self.writer.write_all(buf)
439 }
440
441 pub fn sync(&mut self) -> std::io::Result<()> {
443 self.writer.flush()
444 }
445
446 pub fn flush(&mut self) -> Result<()> {
451 let in_progress = match self.in_progress.take() {
452 Some(in_progress) => in_progress,
453 None => return Ok(()),
454 };
455
456 let mut row_group_writer = self.writer.next_row_group()?;
457 for chunk in in_progress.close()? {
458 chunk.append_to_row_group(&mut row_group_writer)?;
459 }
460 row_group_writer.close()?;
461 Ok(())
462 }
463
464 pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
468 self.writer.append_key_value_metadata(kv_metadata)
469 }
470
471 pub fn inner(&self) -> &W {
473 self.writer.inner()
474 }
475
476 pub fn inner_mut(&mut self) -> &mut W {
485 self.writer.inner_mut()
486 }
487
488 pub fn into_inner(mut self) -> Result<W> {
490 self.flush()?;
491 self.writer.into_inner()
492 }
493
494 pub fn finish(&mut self) -> Result<ParquetMetaData> {
500 self.flush()?;
501 self.writer.finish()
502 }
503
504 pub fn close(mut self) -> Result<ParquetMetaData> {
506 self.finish()
507 }
508
509 #[deprecated(
511 since = "56.2.0",
512 note = "Use `ArrowRowGroupWriterFactory` instead, see `ArrowColumnWriter` for an example"
513 )]
514 pub fn get_column_writers(&mut self) -> Result<Vec<ArrowColumnWriter>> {
515 self.flush()?;
516 let in_progress = self
517 .row_group_writer_factory
518 .create_row_group_writer(self.writer.flushed_row_groups().len())?;
519 Ok(in_progress.writers)
520 }
521
522 #[deprecated(
524 since = "56.2.0",
525 note = "Use `SerializedFileWriter` directly instead, see `ArrowColumnWriter` for an example"
526 )]
527 pub fn append_row_group(&mut self, chunks: Vec<ArrowColumnChunk>) -> Result<()> {
528 let mut row_group_writer = self.writer.next_row_group()?;
529 for chunk in chunks {
530 chunk.append_to_row_group(&mut row_group_writer)?;
531 }
532 row_group_writer.close()?;
533 Ok(())
534 }
535
536 pub fn into_serialized_writer(
543 mut self,
544 ) -> Result<(SerializedFileWriter<W>, ArrowRowGroupWriterFactory)> {
545 self.flush()?;
546 Ok((self.writer, self.row_group_writer_factory))
547 }
548}
549
550impl<W: Write + Send> RecordBatchWriter for ArrowWriter<W> {
551 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
552 self.write(batch).map_err(|e| e.into())
553 }
554
555 fn close(self) -> std::result::Result<(), ArrowError> {
556 self.close()?;
557 Ok(())
558 }
559}
560
561#[derive(Debug, Clone, Default)]
565pub struct ArrowWriterOptions {
566 properties: WriterProperties,
567 skip_arrow_metadata: bool,
568 schema_root: Option<String>,
569 schema_descr: Option<SchemaDescriptor>,
570 page_store_factory: Option<Arc<dyn PageStoreFactory>>,
571}
572
573impl ArrowWriterOptions {
574 pub fn new() -> Self {
576 Self::default()
577 }
578
579 pub fn with_properties(self, properties: WriterProperties) -> Self {
581 Self { properties, ..self }
582 }
583
584 pub fn with_page_store_factory(self, page_store_factory: Arc<dyn PageStoreFactory>) -> Self {
670 Self {
671 page_store_factory: Some(page_store_factory),
672 ..self
673 }
674 }
675
676 pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
683 Self {
684 skip_arrow_metadata,
685 ..self
686 }
687 }
688
689 pub fn with_schema_root(self, schema_root: String) -> Self {
691 Self {
692 schema_root: Some(schema_root),
693 ..self
694 }
695 }
696
697 pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self {
703 Self {
704 schema_descr: Some(schema_descr),
705 ..self
706 }
707 }
708}
709
710struct ArrowColumnChunkData {
716 length: usize,
717 store: Box<dyn PageStore>,
718 keys: Vec<PageKey>,
719 dictionary_keys: Vec<PageKey>,
730 dictionary_len: usize,
734}
735
736impl ArrowColumnChunkData {
737 fn new(store: Box<dyn PageStore>) -> Self {
738 Self {
739 length: 0,
740 store,
741 keys: Vec::new(),
742 dictionary_keys: Vec::new(),
743 dictionary_len: 0,
744 }
745 }
746
747 fn push(&mut self, value: Bytes) -> Result<()> {
750 let key = self.store.put(value)?;
751 self.keys.push(key);
752 Ok(())
753 }
754
755 fn push_dictionary(&mut self, value: Bytes) -> Result<()> {
759 self.dictionary_len += value.len();
760 let key = self.store.put(value)?;
761 self.dictionary_keys.push(key);
762 Ok(())
763 }
764
765 fn memory_size(&self) -> usize {
768 self.store.memory_size()
769 }
770}
771
772struct StreamingColumnChunkPages {
781 store: Box<dyn PageStore>,
782 keys: IntoIter<PageKey>,
785}
786
787impl StreamingColumnChunkPages {
788 fn new(data: ArrowColumnChunkData) -> Self {
789 let keys = if data.dictionary_keys.is_empty() {
792 data.keys
793 } else {
794 let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len());
795 keys.extend(data.dictionary_keys);
796 keys.extend(data.keys);
797 keys
798 };
799 Self {
800 store: data.store,
801 keys: keys.into_iter(),
802 }
803 }
804}
805
806impl Iterator for StreamingColumnChunkPages {
807 type Item = Result<Bytes>;
808
809 fn next(&mut self) -> Option<Self::Item> {
810 let key = self.keys.next()?;
811 Some(self.store.take(key))
812 }
813}
814
815type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
820
821struct ArrowPageWriter {
822 buffer: SharedColumnChunk,
823 #[cfg(feature = "encryption")]
824 page_encryptor: Option<PageEncryptor>,
825}
826
827impl ArrowPageWriter {
828 fn new(store: Box<dyn PageStore>) -> Self {
830 Self {
831 buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))),
832 #[cfg(feature = "encryption")]
833 page_encryptor: None,
834 }
835 }
836
837 #[cfg(feature = "encryption")]
838 pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
839 self.page_encryptor = page_encryptor;
840 self
841 }
842
843 #[cfg(feature = "encryption")]
844 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
845 self.page_encryptor.as_mut()
846 }
847
848 #[cfg(not(feature = "encryption"))]
849 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
850 None
851 }
852}
853
854impl PageWriter for ArrowPageWriter {
855 fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
856 let page = match self.page_encryptor_mut() {
857 Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
858 None => page,
859 };
860
861 let page_header = page.to_thrift_header()?;
862 let header = {
863 let mut header = Vec::with_capacity(1024);
864
865 match self.page_encryptor_mut() {
866 Some(page_encryptor) => {
867 page_encryptor.encrypt_page_header(&page_header, &mut header)?;
868 if page.compressed_page().is_data_page() {
869 page_encryptor.increment_page();
870 }
871 }
872 None => {
873 let mut protocol = ThriftCompactOutputProtocol::new(&mut header);
874 page_header.write_thrift(&mut protocol)?;
875 }
876 };
877
878 Bytes::from(header)
879 };
880
881 let mut buf = self.buffer.try_lock().unwrap();
882
883 let data = page.compressed_page().buffer().clone();
884 let compressed_size = data.len() + header.len();
885
886 let mut spec = PageWriteSpec::new();
887 spec.page_type = page.page_type();
888 spec.num_values = page.num_values();
889 spec.uncompressed_size = page.uncompressed_size() + header.len();
890 spec.offset = buf.length as u64;
891 spec.compressed_size = compressed_size;
892 spec.bytes_written = compressed_size as u64;
893
894 buf.length += compressed_size;
895 if spec.page_type == PageType::DICTIONARY_PAGE {
896 buf.push_dictionary(header)?;
899 buf.push_dictionary(data)?;
900 } else {
901 buf.push(header)?;
902 buf.push(data)?;
903 }
904
905 Ok(spec)
906 }
907
908 fn defers_dictionary_ordering(&self) -> bool {
909 true
914 }
915
916 fn buffered_memory_size(&self) -> usize {
917 self.buffer.try_lock().unwrap().memory_size()
920 }
921
922 fn close(&mut self) -> Result<()> {
923 Ok(())
924 }
925}
926
927#[derive(Debug)]
929pub struct ArrowLeafColumn(ArrayLevels);
930
931pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result<Vec<ArrowLeafColumn>> {
936 let levels = calculate_array_levels(array, field)?;
937 Ok(levels.into_iter().map(ArrowLeafColumn).collect())
938}
939
940pub struct ArrowColumnChunk {
942 data: ArrowColumnChunkData,
943 close: ColumnCloseResult,
944}
945
946impl std::fmt::Debug for ArrowColumnChunk {
947 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
948 f.debug_struct("ArrowColumnChunk")
949 .field("length", &self.data.length)
950 .finish_non_exhaustive()
951 }
952}
953
954impl ArrowColumnChunk {
955 pub fn close(&self) -> &ColumnCloseResult {
962 &self.close
963 }
964
965 pub fn close_mut(&mut self) -> &mut ColumnCloseResult {
972 &mut self.close
973 }
974
975 pub fn append_to_row_group<W: Write + Send>(
978 self,
979 writer: &mut SerializedRowGroupWriter<'_, W>,
980 ) -> Result<()> {
981 let ArrowColumnChunk { data, close } = self;
982
983 let close = close.update_dictionary_location(data.dictionary_len)?;
987
988 let pages = StreamingColumnChunkPages::new(data);
989 writer.append_column_from_pages(pages, close)
990 }
991}
992
993pub struct ArrowColumnWriter {
1091 writer: ArrowColumnWriterImpl,
1092 chunk: SharedColumnChunk,
1093}
1094
1095impl std::fmt::Debug for ArrowColumnWriter {
1096 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1097 f.debug_struct("ArrowColumnWriter").finish_non_exhaustive()
1098 }
1099}
1100
1101enum ArrowColumnWriterImpl {
1102 ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
1103 Column(ColumnWriter<'static>),
1104}
1105
1106impl ArrowColumnWriter {
1107 pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> {
1109 self.write_internal(&col.0)
1110 }
1111
1112 fn write_with_chunker(
1114 &mut self,
1115 col: &ArrowLeafColumn,
1116 chunker: &mut ContentDefinedChunker,
1117 ) -> Result<()> {
1118 let levels = &col.0;
1119 let chunks = chunker.get_arrow_chunks(
1120 levels.def_level_data().as_ref(),
1121 levels.rep_level_data().as_ref(),
1122 levels.array(),
1123 )?;
1124
1125 let num_chunks = chunks.len();
1126 for (i, chunk) in chunks.iter().enumerate() {
1127 let chunk_levels = levels.slice_for_chunk(chunk);
1128 self.write_internal(&chunk_levels)?;
1129
1130 if i + 1 < num_chunks {
1132 match &mut self.writer {
1133 ArrowColumnWriterImpl::Column(c) => c.add_data_page()?,
1134 ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?,
1135 }
1136 }
1137 }
1138 Ok(())
1139 }
1140
1141 fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
1142 match &mut self.writer {
1143 ArrowColumnWriterImpl::Column(c) => {
1144 let leaf = levels.array();
1145 match leaf.as_any_dictionary_opt() {
1146 Some(dictionary) => {
1147 let materialized =
1148 arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?;
1149 write_leaf(c, &materialized, levels)?
1150 }
1151 None => write_leaf(c, leaf, levels)?,
1152 };
1153 }
1154 ArrowColumnWriterImpl::ByteArray(c) => {
1155 write_primitive(c, levels.array().as_ref(), levels)?;
1156 }
1157 }
1158 Ok(())
1159 }
1160
1161 pub fn close(self) -> Result<ArrowColumnChunk> {
1163 let close = match self.writer {
1164 ArrowColumnWriterImpl::ByteArray(c) => c.close()?,
1165 ArrowColumnWriterImpl::Column(c) => c.close()?,
1166 };
1167 let chunk = Arc::try_unwrap(self.chunk).ok().unwrap();
1168 let data = chunk.into_inner().unwrap();
1169 Ok(ArrowColumnChunk { data, close })
1170 }
1171
1172 pub fn memory_size(&self) -> usize {
1183 match &self.writer {
1184 ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1185 ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1186 }
1187 }
1188
1189 pub fn get_estimated_total_bytes(&self) -> usize {
1197 match &self.writer {
1198 ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1199 ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1200 }
1201 }
1202}
1203
1204#[derive(Debug)]
1211struct ArrowRowGroupWriter {
1212 writers: Vec<ArrowColumnWriter>,
1213 schema: SchemaRef,
1214 buffered_rows: usize,
1215}
1216
1217impl ArrowRowGroupWriter {
1218 fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self {
1219 Self {
1220 writers,
1221 schema: arrow.clone(),
1222 buffered_rows: 0,
1223 }
1224 }
1225
1226 fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1227 self.buffered_rows += batch.num_rows();
1228 let mut writers = self.writers.iter_mut();
1229 for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1230 for leaf in compute_leaves(field.as_ref(), column)? {
1231 writers.next().unwrap().write(&leaf)?;
1232 }
1233 }
1234 Ok(())
1235 }
1236
1237 fn write_with_chunkers(
1238 &mut self,
1239 batch: &RecordBatch,
1240 chunkers: &mut [ContentDefinedChunker],
1241 ) -> Result<()> {
1242 self.buffered_rows += batch.num_rows();
1243 let mut writers = self.writers.iter_mut();
1244 let mut chunkers = chunkers.iter_mut();
1245 for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1246 for leaf in compute_leaves(field.as_ref(), column)? {
1247 writers
1248 .next()
1249 .unwrap()
1250 .write_with_chunker(&leaf, chunkers.next().unwrap())?;
1251 }
1252 }
1253 Ok(())
1254 }
1255
1256 fn get_estimated_total_bytes(&self) -> usize {
1258 self.writers
1259 .iter()
1260 .map(|x| x.get_estimated_total_bytes())
1261 .sum()
1262 }
1263
1264 fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1265 self.writers
1266 .into_iter()
1267 .map(|writer| writer.close())
1268 .collect()
1269 }
1270}
1271
1272#[derive(Debug)]
1277pub struct ArrowRowGroupWriterFactory {
1278 schema: SchemaDescPtr,
1279 arrow_schema: SchemaRef,
1280 props: WriterPropertiesPtr,
1281 page_store_factory: Arc<dyn PageStoreFactory>,
1282 #[cfg(feature = "encryption")]
1283 file_encryptor: Option<Arc<FileEncryptor>>,
1284}
1285
1286impl ArrowRowGroupWriterFactory {
1287 pub fn new<W: Write + Send>(
1289 file_writer: &SerializedFileWriter<W>,
1290 arrow_schema: SchemaRef,
1291 ) -> Self {
1292 let schema = Arc::clone(file_writer.schema_descr_ptr());
1293 let props = Arc::clone(file_writer.properties());
1294 Self {
1295 schema,
1296 arrow_schema,
1297 props,
1298 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1299 #[cfg(feature = "encryption")]
1300 file_encryptor: file_writer.file_encryptor(),
1301 }
1302 }
1303
1304 pub fn with_page_store_factory(
1308 mut self,
1309 page_store_factory: Arc<dyn PageStoreFactory>,
1310 ) -> Self {
1311 self.page_store_factory = page_store_factory;
1312 self
1313 }
1314
1315 fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1316 let writers = self.create_column_writers(row_group_index)?;
1317 Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
1318 }
1319
1320 pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1322 let mut writers = Vec::with_capacity(self.arrow_schema.fields.len());
1323 let mut leaves = self.schema.columns().iter();
1324 let column_factory = self.column_writer_factory(row_group_index);
1325 for field in &self.arrow_schema.fields {
1326 column_factory.get_arrow_column_writer(
1327 field.data_type(),
1328 &self.props,
1329 &mut leaves,
1330 &mut writers,
1331 )?;
1332 }
1333 Ok(writers)
1334 }
1335
1336 #[cfg(feature = "encryption")]
1337 fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1338 ArrowColumnWriterFactory::new()
1339 .with_page_store_factory(self.page_store_factory.clone())
1340 .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1341 }
1342
1343 #[cfg(not(feature = "encryption"))]
1344 fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1345 ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1346 }
1347}
1348
1349#[deprecated(since = "57.0.0", note = "Use `ArrowRowGroupWriterFactory` instead")]
1351pub fn get_column_writers(
1352 parquet: &SchemaDescriptor,
1353 props: &WriterPropertiesPtr,
1354 arrow: &SchemaRef,
1355) -> Result<Vec<ArrowColumnWriter>> {
1356 let mut writers = Vec::with_capacity(arrow.fields.len());
1357 let mut leaves = parquet.columns().iter();
1358 let column_factory = ArrowColumnWriterFactory::new();
1359 for field in &arrow.fields {
1360 column_factory.get_arrow_column_writer(
1361 field.data_type(),
1362 props,
1363 &mut leaves,
1364 &mut writers,
1365 )?;
1366 }
1367 Ok(writers)
1368}
1369
1370struct ArrowColumnWriterFactory {
1372 page_store_factory: Arc<dyn PageStoreFactory>,
1374 #[cfg(feature = "encryption")]
1375 row_group_index: usize,
1376 #[cfg(feature = "encryption")]
1377 file_encryptor: Option<Arc<FileEncryptor>>,
1378}
1379
1380impl ArrowColumnWriterFactory {
1381 pub fn new() -> Self {
1382 Self {
1383 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1384 #[cfg(feature = "encryption")]
1385 row_group_index: 0,
1386 #[cfg(feature = "encryption")]
1387 file_encryptor: None,
1388 }
1389 }
1390
1391 pub fn with_page_store_factory(
1393 mut self,
1394 page_store_factory: Arc<dyn PageStoreFactory>,
1395 ) -> Self {
1396 self.page_store_factory = page_store_factory;
1397 self
1398 }
1399
1400 #[cfg(feature = "encryption")]
1401 pub fn with_file_encryptor(
1402 mut self,
1403 row_group_index: usize,
1404 file_encryptor: Option<Arc<FileEncryptor>>,
1405 ) -> Self {
1406 self.row_group_index = row_group_index;
1407 self.file_encryptor = file_encryptor;
1408 self
1409 }
1410
1411 #[cfg(feature = "encryption")]
1412 fn create_page_writer(
1413 &self,
1414 column_descriptor: &ColumnDescPtr,
1415 column_index: usize,
1416 ) -> Result<Box<ArrowPageWriter>> {
1417 let column_path = column_descriptor.path().string();
1418 let page_encryptor = PageEncryptor::create_if_column_encrypted(
1419 &self.file_encryptor,
1420 self.row_group_index,
1421 column_index,
1422 &column_path,
1423 )?;
1424 let args = PageStoreArgs::new(column_index, column_descriptor);
1425 let store = self.page_store_factory.create(&args)?;
1426 Ok(Box::new(
1427 ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1428 ))
1429 }
1430
1431 #[cfg(not(feature = "encryption"))]
1432 fn create_page_writer(
1433 &self,
1434 column_descriptor: &ColumnDescPtr,
1435 column_index: usize,
1436 ) -> Result<Box<ArrowPageWriter>> {
1437 let args = PageStoreArgs::new(column_index, column_descriptor);
1438 let store = self.page_store_factory.create(&args)?;
1439 Ok(Box::new(ArrowPageWriter::new(store)))
1440 }
1441
1442 fn get_arrow_column_writer(
1445 &self,
1446 data_type: &ArrowDataType,
1447 props: &WriterPropertiesPtr,
1448 leaves: &mut Iter<'_, ColumnDescPtr>,
1449 out: &mut Vec<ArrowColumnWriter>,
1450 ) -> Result<()> {
1451 let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1453 let page_writer = self.create_page_writer(desc, out.len())?;
1454 let chunk = page_writer.buffer.clone();
1455 let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1456 Ok(ArrowColumnWriter {
1457 chunk,
1458 writer: ArrowColumnWriterImpl::Column(writer),
1459 })
1460 };
1461
1462 let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1464 let page_writer = self.create_page_writer(desc, out.len())?;
1465 let chunk = page_writer.buffer.clone();
1466 let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1467 Ok(ArrowColumnWriter {
1468 chunk,
1469 writer: ArrowColumnWriterImpl::ByteArray(writer),
1470 })
1471 };
1472
1473 match data_type {
1474 _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1475 ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1476 out.push(col(leaves.next().unwrap())?)
1477 }
1478 ArrowDataType::LargeBinary
1479 | ArrowDataType::Binary
1480 | ArrowDataType::Utf8
1481 | ArrowDataType::LargeUtf8
1482 | ArrowDataType::BinaryView
1483 | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1484 ArrowDataType::List(f)
1485 | ArrowDataType::LargeList(f)
1486 | ArrowDataType::FixedSizeList(f, _)
1487 | ArrowDataType::ListView(f)
1488 | ArrowDataType::LargeListView(f) => {
1489 self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1490 }
1491 ArrowDataType::Struct(fields) => {
1492 for field in fields {
1493 self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1494 }
1495 }
1496 ArrowDataType::Map(f, _) => match f.data_type() {
1497 ArrowDataType::Struct(f) => {
1498 self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1499 self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1500 }
1501 _ => unreachable!("invalid map type"),
1502 },
1503 ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1504 ArrowDataType::Utf8
1505 | ArrowDataType::LargeUtf8
1506 | ArrowDataType::Binary
1507 | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1508 ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1509 out.push(bytes(leaves.next().unwrap())?)
1510 }
1511 ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
1512 _ => out.push(col(leaves.next().unwrap())?),
1513 },
1514 ArrowDataType::RunEndEncoded(_, value_field) => {
1515 self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1516 }
1517 _ => {
1518 return Err(ParquetError::NYI(format!(
1519 "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1520 )));
1521 }
1522 }
1523 Ok(())
1524 }
1525}
1526
1527fn write_leaf(
1528 writer: &mut ColumnWriter<'_>,
1529 column: &dyn arrow_array::Array,
1530 levels: &ArrayLevels,
1531) -> Result<usize> {
1532 let indices = levels.non_null_indices();
1533
1534 match writer {
1535 ColumnWriter::Int32ColumnWriter(typed) => {
1537 match column.data_type() {
1538 ArrowDataType::Null => {
1539 let array = Int32Array::new_null(column.len());
1540 write_primitive(typed, array.values(), levels)
1541 }
1542 ArrowDataType::Int8 => {
1543 let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32);
1544 write_primitive(typed, array.values(), levels)
1545 }
1546 ArrowDataType::Int16 => {
1547 let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32);
1548 write_primitive(typed, array.values(), levels)
1549 }
1550 ArrowDataType::Int32 => {
1551 write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels)
1552 }
1553 ArrowDataType::UInt8 => {
1554 let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32);
1555 write_primitive(typed, array.values(), levels)
1556 }
1557 ArrowDataType::UInt16 => {
1558 let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32);
1559 write_primitive(typed, array.values(), levels)
1560 }
1561 ArrowDataType::UInt32 => {
1562 let array = column.as_primitive::<UInt32Type>();
1565 write_primitive(typed, array.values().inner().typed_data(), levels)
1566 }
1567 ArrowDataType::Date32 => {
1568 let array = column.as_primitive::<Date32Type>();
1569 write_primitive(typed, array.values(), levels)
1570 }
1571 ArrowDataType::Time32(TimeUnit::Second) => {
1572 let array = column.as_primitive::<Time32SecondType>();
1573 write_primitive(typed, array.values(), levels)
1574 }
1575 ArrowDataType::Time32(TimeUnit::Millisecond) => {
1576 let array = column.as_primitive::<Time32MillisecondType>();
1577 write_primitive(typed, array.values(), levels)
1578 }
1579 ArrowDataType::Date64 => {
1580 let array: Int32Array = column
1582 .as_primitive::<Date64Type>()
1583 .unary(|x| (x / 86_400_000) as _);
1584
1585 write_primitive(typed, array.values(), levels)
1586 }
1587 ArrowDataType::Decimal32(_, _) => {
1588 let array = column
1589 .as_primitive::<Decimal32Type>()
1590 .unary::<_, Int32Type>(|v| v);
1591 write_primitive(typed, array.values(), levels)
1592 }
1593 ArrowDataType::Decimal64(_, _) => {
1594 let array = column
1596 .as_primitive::<Decimal64Type>()
1597 .unary::<_, Int32Type>(|v| v as i32);
1598 write_primitive(typed, array.values(), levels)
1599 }
1600 ArrowDataType::Decimal128(_, _) => {
1601 let array = column
1603 .as_primitive::<Decimal128Type>()
1604 .unary::<_, Int32Type>(|v| v as i32);
1605 write_primitive(typed, array.values(), levels)
1606 }
1607 ArrowDataType::Decimal256(_, _) => {
1608 let array = column
1610 .as_primitive::<Decimal256Type>()
1611 .unary::<_, Int32Type>(|v| v.as_i128() as i32);
1612 write_primitive(typed, array.values(), levels)
1613 }
1614 d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1615 }
1616 }
1617 ColumnWriter::BoolColumnWriter(typed) => {
1618 let array = column.as_boolean();
1619 let values = get_bool_array_slice(array, indices.iter().copied());
1620 typed.write_batch_internal(
1621 values.as_slice(),
1622 None,
1623 levels.def_level_data().as_ref(),
1624 levels.rep_level_data().as_ref(),
1625 None,
1626 None,
1627 None,
1628 )
1629 }
1630 ColumnWriter::Int64ColumnWriter(typed) => {
1631 match column.data_type() {
1632 ArrowDataType::Date64 => {
1633 let array = column
1634 .as_primitive::<Date64Type>()
1635 .reinterpret_cast::<Int64Type>();
1636
1637 write_primitive(typed, array.values(), levels)
1638 }
1639 ArrowDataType::Int64 => {
1640 let array = column.as_primitive::<Int64Type>();
1641 write_primitive(typed, array.values(), levels)
1642 }
1643 ArrowDataType::UInt64 => {
1644 let values = column.as_primitive::<UInt64Type>().values();
1645 let array = values.inner().typed_data::<i64>();
1648 write_primitive(typed, array, levels)
1649 }
1650 ArrowDataType::Time64(TimeUnit::Microsecond) => {
1651 let array = column.as_primitive::<Time64MicrosecondType>();
1652 write_primitive(typed, array.values(), levels)
1653 }
1654 ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1655 let array = column.as_primitive::<Time64NanosecondType>();
1656 write_primitive(typed, array.values(), levels)
1657 }
1658 ArrowDataType::Timestamp(unit, _) => match unit {
1659 TimeUnit::Second => {
1660 let array = column.as_primitive::<TimestampSecondType>();
1661 write_primitive(typed, array.values(), levels)
1662 }
1663 TimeUnit::Millisecond => {
1664 let array = column.as_primitive::<TimestampMillisecondType>();
1665 write_primitive(typed, array.values(), levels)
1666 }
1667 TimeUnit::Microsecond => {
1668 let array = column.as_primitive::<TimestampMicrosecondType>();
1669 write_primitive(typed, array.values(), levels)
1670 }
1671 TimeUnit::Nanosecond => {
1672 let array = column.as_primitive::<TimestampNanosecondType>();
1673 write_primitive(typed, array.values(), levels)
1674 }
1675 },
1676 ArrowDataType::Duration(unit) => match unit {
1677 TimeUnit::Second => {
1678 let array = column.as_primitive::<DurationSecondType>();
1679 write_primitive(typed, array.values(), levels)
1680 }
1681 TimeUnit::Millisecond => {
1682 let array = column.as_primitive::<DurationMillisecondType>();
1683 write_primitive(typed, array.values(), levels)
1684 }
1685 TimeUnit::Microsecond => {
1686 let array = column.as_primitive::<DurationMicrosecondType>();
1687 write_primitive(typed, array.values(), levels)
1688 }
1689 TimeUnit::Nanosecond => {
1690 let array = column.as_primitive::<DurationNanosecondType>();
1691 write_primitive(typed, array.values(), levels)
1692 }
1693 },
1694 ArrowDataType::Decimal64(_, _) => {
1695 let array = column
1696 .as_primitive::<Decimal64Type>()
1697 .reinterpret_cast::<Int64Type>();
1698 write_primitive(typed, array.values(), levels)
1699 }
1700 ArrowDataType::Decimal128(_, _) => {
1701 let array = column
1703 .as_primitive::<Decimal128Type>()
1704 .unary::<_, Int64Type>(|v| v as i64);
1705 write_primitive(typed, array.values(), levels)
1706 }
1707 ArrowDataType::Decimal256(_, _) => {
1708 let array = column
1710 .as_primitive::<Decimal256Type>()
1711 .unary::<_, Int64Type>(|v| v.as_i128() as i64);
1712 write_primitive(typed, array.values(), levels)
1713 }
1714 d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1715 }
1716 }
1717 ColumnWriter::Int96ColumnWriter(_typed) => {
1718 unreachable!("Currently unreachable because data type not supported")
1719 }
1720 ColumnWriter::FloatColumnWriter(typed) => {
1721 let array = column.as_primitive::<Float32Type>();
1722 write_primitive(typed, array.values(), levels)
1723 }
1724 ColumnWriter::DoubleColumnWriter(typed) => {
1725 let array = column.as_primitive::<Float64Type>();
1726 write_primitive(typed, array.values(), levels)
1727 }
1728 ColumnWriter::ByteArrayColumnWriter(_) => {
1729 unreachable!("should use ByteArrayWriter")
1730 }
1731 ColumnWriter::FixedLenByteArrayColumnWriter(typed) => {
1732 let bytes = match column.data_type() {
1733 ArrowDataType::Interval(interval_unit) => match interval_unit {
1734 IntervalUnit::YearMonth => {
1735 let array = column.as_primitive::<IntervalYearMonthType>();
1736 get_interval_ym_array_slice(array, indices.iter().copied())
1737 }
1738 IntervalUnit::DayTime => {
1739 let array = column.as_primitive::<IntervalDayTimeType>();
1740 get_interval_dt_array_slice(array, indices.iter().copied())
1741 }
1742 _ => {
1743 return Err(ParquetError::NYI(format!(
1744 "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
1745 )));
1746 }
1747 },
1748 ArrowDataType::FixedSizeBinary(_) => {
1749 let array = column.as_fixed_size_binary();
1750 get_fsb_array_slice(array, indices.iter().copied())
1751 }
1752 ArrowDataType::Decimal32(_, _) => {
1753 let array = column.as_primitive::<Decimal32Type>();
1754 get_decimal_32_array_slice(array, indices.iter().copied())
1755 }
1756 ArrowDataType::Decimal64(_, _) => {
1757 let array = column.as_primitive::<Decimal64Type>();
1758 get_decimal_64_array_slice(array, indices.iter().copied())
1759 }
1760 ArrowDataType::Decimal128(_, _) => {
1761 let array = column.as_primitive::<Decimal128Type>();
1762 get_decimal_128_array_slice(array, indices.iter().copied())
1763 }
1764 ArrowDataType::Decimal256(_, _) => {
1765 let array = column.as_primitive::<Decimal256Type>();
1766 get_decimal_256_array_slice(array, indices.iter().copied())
1767 }
1768 ArrowDataType::Float16 => {
1769 let array = column.as_primitive::<Float16Type>();
1770 get_float_16_array_slice(array, indices.iter().copied())
1771 }
1772 _ => {
1773 return Err(ParquetError::NYI(
1774 "Attempting to write an Arrow type that is not yet implemented".to_string(),
1775 ));
1776 }
1777 };
1778 typed.write_batch_internal(
1779 bytes.as_slice(),
1780 None,
1781 levels.def_level_data().as_ref(),
1782 levels.rep_level_data().as_ref(),
1783 None,
1784 None,
1785 None,
1786 )
1787 }
1788 }
1789}
1790
1791fn write_primitive<E: ColumnValueEncoder>(
1792 writer: &mut GenericColumnWriter<E>,
1793 values: &E::Values,
1794 levels: &ArrayLevels,
1795) -> Result<usize> {
1796 writer.write_batch_internal(
1797 values,
1798 Some(levels.non_null_indices()),
1799 levels.def_level_data().as_ref(),
1800 levels.rep_level_data().as_ref(),
1801 None,
1802 None,
1803 None,
1804 )
1805}
1806
1807fn get_bool_array_slice(
1808 array: &arrow_array::BooleanArray,
1809 indices: impl ExactSizeIterator<Item = usize>,
1810) -> Vec<bool> {
1811 let mut values = Vec::with_capacity(indices.len());
1812 for i in indices {
1813 values.push(array.value(i))
1814 }
1815 values
1816}
1817
1818fn get_interval_ym_array_slice(
1821 array: &arrow_array::IntervalYearMonthArray,
1822 indices: impl ExactSizeIterator<Item = usize>,
1823) -> Vec<FixedLenByteArray> {
1824 let mut values = Vec::with_capacity(indices.len());
1825 for i in indices {
1826 let mut value = array.value(i).to_le_bytes().to_vec();
1827 let mut suffix = vec![0; 8];
1828 value.append(&mut suffix);
1829 values.push(FixedLenByteArray::from(ByteArray::from(value)))
1830 }
1831 values
1832}
1833
1834fn get_interval_dt_array_slice(
1837 array: &arrow_array::IntervalDayTimeArray,
1838 indices: impl ExactSizeIterator<Item = usize>,
1839) -> Vec<FixedLenByteArray> {
1840 let mut values = Vec::with_capacity(indices.len());
1841 for i in indices {
1842 let mut out = [0; 12];
1843 let value = array.value(i);
1844 out[4..8].copy_from_slice(&value.days.to_le_bytes());
1845 out[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1846 values.push(FixedLenByteArray::from(ByteArray::from(out.to_vec())));
1847 }
1848 values
1849}
1850
1851fn get_decimal_32_array_slice(
1852 array: &arrow_array::Decimal32Array,
1853 indices: impl ExactSizeIterator<Item = usize>,
1854) -> Vec<FixedLenByteArray> {
1855 let mut values = Vec::with_capacity(indices.len());
1856 let size = decimal_length_from_precision(array.precision());
1857 for i in indices {
1858 let as_be_bytes = array.value(i).to_be_bytes();
1859 let resized_value = as_be_bytes[(4 - size)..].to_vec();
1860 values.push(FixedLenByteArray::from(ByteArray::from(resized_value)));
1861 }
1862 values
1863}
1864
1865fn get_decimal_64_array_slice(
1866 array: &arrow_array::Decimal64Array,
1867 indices: impl ExactSizeIterator<Item = usize>,
1868) -> Vec<FixedLenByteArray> {
1869 let mut values = Vec::with_capacity(indices.len());
1870 let size = decimal_length_from_precision(array.precision());
1871 for i in indices {
1872 let as_be_bytes = array.value(i).to_be_bytes();
1873 let resized_value = as_be_bytes[(8 - size)..].to_vec();
1874 values.push(FixedLenByteArray::from(ByteArray::from(resized_value)));
1875 }
1876 values
1877}
1878
1879fn get_decimal_128_array_slice(
1880 array: &arrow_array::Decimal128Array,
1881 indices: impl ExactSizeIterator<Item = usize>,
1882) -> Vec<FixedLenByteArray> {
1883 let mut values = Vec::with_capacity(indices.len());
1884 let size = decimal_length_from_precision(array.precision());
1885 for i in indices {
1886 let as_be_bytes = array.value(i).to_be_bytes();
1887 let resized_value = as_be_bytes[(16 - size)..].to_vec();
1888 values.push(FixedLenByteArray::from(ByteArray::from(resized_value)));
1889 }
1890 values
1891}
1892
1893fn get_decimal_256_array_slice(
1894 array: &arrow_array::Decimal256Array,
1895 indices: impl ExactSizeIterator<Item = usize>,
1896) -> Vec<FixedLenByteArray> {
1897 let mut values = Vec::with_capacity(indices.len());
1898 let size = decimal_length_from_precision(array.precision());
1899 for i in indices {
1900 let as_be_bytes = array.value(i).to_be_bytes();
1901 let resized_value = as_be_bytes[(32 - size)..].to_vec();
1902 values.push(FixedLenByteArray::from(ByteArray::from(resized_value)));
1903 }
1904 values
1905}
1906
1907fn get_float_16_array_slice(
1908 array: &arrow_array::Float16Array,
1909 indices: impl ExactSizeIterator<Item = usize>,
1910) -> Vec<FixedLenByteArray> {
1911 let mut values = Vec::with_capacity(indices.len());
1912 for i in indices {
1913 let value = array.value(i).to_le_bytes().to_vec();
1914 values.push(FixedLenByteArray::from(ByteArray::from(value)));
1915 }
1916 values
1917}
1918
1919fn get_fsb_array_slice(
1920 array: &arrow_array::FixedSizeBinaryArray,
1921 indices: impl ExactSizeIterator<Item = usize>,
1922) -> Vec<FixedLenByteArray> {
1923 let mut values = Vec::with_capacity(indices.len());
1924 for i in indices {
1925 let value = array.value(i).to_vec();
1926 values.push(FixedLenByteArray::from(ByteArray::from(value)))
1927 }
1928 values
1929}
1930
1931#[cfg(test)]
1932mod tests {
1933 use super::*;
1934 use std::collections::HashMap;
1935
1936 use std::fs::File;
1937
1938 use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1939 use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
1940 use crate::column::page::{Page, PageReader};
1941 use crate::file::metadata::thrift::PageHeader;
1942 use crate::file::page_index::column_index::ColumnIndexMetaData;
1943 use crate::file::reader::SerializedPageReader;
1944 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
1945 use crate::schema::types::ColumnPath;
1946 use arrow::datatypes::ToByteSlice;
1947 use arrow::datatypes::{DataType, Schema};
1948 use arrow::error::Result as ArrowResult;
1949 use arrow::util::data_gen::create_random_array;
1950 use arrow::util::pretty::pretty_format_batches;
1951 use arrow::{array::*, buffer::Buffer};
1952 use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
1953 use arrow_schema::Fields;
1954 use half::f16;
1955 use num_traits::{FromPrimitive, ToPrimitive};
1956 use tempfile::tempfile;
1957
1958 use crate::basic::Encoding;
1959 use crate::data_type::AsBytes;
1960 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
1961 use crate::file::properties::{
1962 BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
1963 };
1964 use crate::file::serialized_reader::ReadOptionsBuilder;
1965 use crate::file::{
1966 reader::{FileReader, SerializedFileReader},
1967 statistics::Statistics,
1968 };
1969
1970 #[derive(Debug, Default)]
1975 struct RecordingPageStore {
1976 next: u64,
1977 blobs: HashMap<u64, Bytes>,
1978 puts: Arc<std::sync::atomic::AtomicUsize>,
1979 }
1980
1981 impl PageStore for RecordingPageStore {
1982 fn put(&mut self, value: Bytes) -> Result<PageKey> {
1983 let id = 100 + self.next * 7;
1985 self.next += 1;
1986 self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1987 self.blobs.insert(id, value);
1988 Ok(PageKey::new(id))
1989 }
1990
1991 fn take(&mut self, key: PageKey) -> Result<Bytes> {
1992 self.blobs
1993 .remove(&key.get())
1994 .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
1995 }
1996 }
1997
1998 #[derive(Debug)]
1999 struct RecordingPageStoreFactory {
2000 puts: Arc<std::sync::atomic::AtomicUsize>,
2001 }
2002
2003 impl PageStoreFactory for RecordingPageStoreFactory {
2004 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2005 Ok(Box::new(RecordingPageStore {
2006 puts: self.puts.clone(),
2007 ..Default::default()
2008 }))
2009 }
2010 }
2011
2012 #[test]
2016 fn custom_page_store_is_byte_identical_to_default() {
2017 let schema = Arc::new(Schema::new(vec![
2018 Field::new("i", DataType::Int32, true),
2019 Field::new("s", DataType::Utf8, true),
2021 ]));
2022 let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2023 let s = StringArray::from(vec![
2024 Some("a"),
2025 Some("bb"),
2026 Some("a"),
2027 None,
2028 Some("bb"),
2029 Some("ccc"),
2030 ]);
2031 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2032
2033 let props = WriterProperties::builder()
2036 .set_max_row_group_row_count(Some(3))
2037 .build();
2038
2039 let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2040 let mut buffer = Vec::new();
2041 let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2042 if let Some(factory) = factory {
2043 opts = opts.with_page_store_factory(factory);
2044 }
2045 let mut writer =
2046 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2047 writer.write(&batch).unwrap();
2048 writer.close().unwrap();
2049 buffer
2050 };
2051
2052 let default_bytes = write(None);
2053
2054 let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2055 let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2056 puts: puts.clone(),
2057 })));
2058
2059 assert!(
2060 puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2061 "custom PageStore was never written to"
2062 );
2063 assert_eq!(
2064 default_bytes, custom_bytes,
2065 "a custom PageStore must produce byte-identical output to the default"
2066 );
2067 }
2068
2069 #[test]
2075 fn dictionary_column_round_trips_with_offset_index_disabled() {
2076 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2077
2078 let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2081 let array = Int32Array::from(values.clone());
2082 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2083
2084 let props = WriterProperties::builder()
2085 .set_offset_index_disabled(true)
2086 .set_data_page_row_count_limit(4096)
2087 .build();
2088 let opts = ArrowWriterOptions::new().with_properties(props);
2089
2090 let mut buffer = Vec::new();
2091 let mut writer =
2092 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2093 writer.write(&batch).unwrap();
2094 writer.close().unwrap();
2095
2096 let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2097 let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2098 let read_values: Vec<Option<i32>> = read
2099 .iter()
2100 .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2101 .collect();
2102 assert_eq!(read_values, values);
2103 }
2104
2105 #[test]
2110 fn dictionary_page_is_routed_through_the_store() {
2111 #[derive(Debug, Default)]
2113 struct SizeRecordingPageStore {
2114 blobs: Vec<Bytes>,
2115 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2116 }
2117 impl PageStore for SizeRecordingPageStore {
2118 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2119 self.bytes_put
2120 .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2121 let key = PageKey::new(self.blobs.len() as u64);
2122 self.blobs.push(value);
2123 Ok(key)
2124 }
2125 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2126 Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2127 }
2128 }
2129 #[derive(Debug)]
2130 struct Factory {
2131 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2132 }
2133 impl PageStoreFactory for Factory {
2134 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2135 Ok(Box::new(SizeRecordingPageStore {
2136 bytes_put: self.bytes_put.clone(),
2137 ..Default::default()
2138 }))
2139 }
2140 }
2141
2142 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2143 let values: Vec<&str> = (0..2048)
2146 .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2147 .collect();
2148 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2149 .unwrap();
2150
2151 let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2152 let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2153 bytes_put: bytes_put.clone(),
2154 }));
2155
2156 let mut buffer = Vec::new();
2159 let mut writer =
2160 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2161 writer.write(&batch).unwrap();
2162 writer.close().unwrap();
2163
2164 let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2165 let column = reader.metadata().row_group(0).column(0);
2166 assert!(
2167 column.dictionary_page_offset().is_some(),
2168 "expected the column to be dictionary-encoded"
2169 );
2170
2171 assert_eq!(
2175 bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2176 column.compressed_size(),
2177 "the dictionary page must pass through the store like any other page"
2178 );
2179 }
2180
2181 #[test]
2182 fn arrow_writer() {
2183 let schema = Schema::new(vec![
2185 Field::new("a", DataType::Int32, false),
2186 Field::new("b", DataType::Int32, true),
2187 ]);
2188
2189 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2191 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2192
2193 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2195
2196 roundtrip(batch, Some(SMALL_SIZE / 2));
2197 }
2198
2199 fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2200 let mut buffer = vec![];
2201
2202 let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2203 writer.write(expected_batch).unwrap();
2204 writer.close().unwrap();
2205
2206 buffer
2207 }
2208
2209 fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2210 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2211 writer.write(expected_batch).unwrap();
2212 writer.into_inner().unwrap()
2213 }
2214
2215 #[test]
2216 fn roundtrip_bytes() {
2217 let schema = Arc::new(Schema::new(vec![
2219 Field::new("a", DataType::Int32, false),
2220 Field::new("b", DataType::Int32, true),
2221 ]));
2222
2223 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2225 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2226
2227 let expected_batch =
2229 RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2230
2231 for buffer in [
2232 get_bytes_after_close(schema.clone(), &expected_batch),
2233 get_bytes_by_into_inner(schema, &expected_batch),
2234 ] {
2235 let cursor = Bytes::from(buffer);
2236 let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2237
2238 let actual_batch = record_batch_reader
2239 .next()
2240 .expect("No batch found")
2241 .expect("Unable to get batch");
2242
2243 assert_eq!(expected_batch.schema(), actual_batch.schema());
2244 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2245 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2246 for i in 0..expected_batch.num_columns() {
2247 let expected_data = expected_batch.column(i).to_data();
2248 let actual_data = actual_batch.column(i).to_data();
2249
2250 assert_eq!(expected_data, actual_data);
2251 }
2252 }
2253 }
2254
2255 #[test]
2256 fn arrow_writer_non_null() {
2257 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2259
2260 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2262
2263 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2265
2266 roundtrip(batch, Some(SMALL_SIZE / 2));
2267 }
2268
2269 #[test]
2270 fn arrow_writer_list() {
2271 let schema = Schema::new(vec![Field::new(
2273 "a",
2274 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2275 true,
2276 )]);
2277
2278 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2280
2281 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2284
2285 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2287 DataType::Int32,
2288 false,
2289 ))))
2290 .len(5)
2291 .add_buffer(a_value_offsets)
2292 .add_child_data(a_values.into_data())
2293 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2294 .build()
2295 .unwrap();
2296 let a = ListArray::from(a_list_data);
2297
2298 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2300
2301 assert_eq!(batch.column(0).null_count(), 1);
2302
2303 roundtrip(batch, None);
2306 }
2307
2308 #[test]
2309 fn arrow_writer_list_non_null() {
2310 let schema = Schema::new(vec![Field::new(
2312 "a",
2313 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2314 false,
2315 )]);
2316
2317 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2319
2320 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2323
2324 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2326 DataType::Int32,
2327 false,
2328 ))))
2329 .len(5)
2330 .add_buffer(a_value_offsets)
2331 .add_child_data(a_values.into_data())
2332 .build()
2333 .unwrap();
2334 let a = ListArray::from(a_list_data);
2335
2336 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2338
2339 assert_eq!(batch.column(0).null_count(), 0);
2342
2343 roundtrip(batch, None);
2344 }
2345
2346 #[test]
2347 fn arrow_writer_list_view() {
2348 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2349 let schema = Schema::new(vec![Field::new(
2350 "a",
2351 DataType::ListView(list_field.clone()),
2352 true,
2353 )]);
2354
2355 let a = ListViewArray::new(
2357 list_field,
2358 vec![0, 1, 0, 3, 6].into(),
2359 vec![1, 2, 0, 3, 4].into(),
2360 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2361 Some(vec![true, true, false, true, true].into()),
2362 );
2363
2364 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2365
2366 assert_eq!(batch.column(0).null_count(), 1);
2367
2368 roundtrip(batch, None);
2369 }
2370
2371 #[test]
2372 fn arrow_writer_list_view_non_null() {
2373 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2374 let schema = Schema::new(vec![Field::new(
2375 "a",
2376 DataType::ListView(list_field.clone()),
2377 false,
2378 )]);
2379
2380 let a = ListViewArray::new(
2382 list_field,
2383 vec![0, 1, 0, 3, 6].into(),
2384 vec![1, 2, 0, 3, 4].into(),
2385 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2386 None,
2387 );
2388
2389 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2390
2391 assert_eq!(batch.column(0).null_count(), 0);
2392
2393 roundtrip(batch, None);
2394 }
2395
2396 #[test]
2397 fn arrow_writer_list_view_out_of_order() {
2398 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2399 let schema = Schema::new(vec![Field::new(
2400 "a",
2401 DataType::ListView(list_field.clone()),
2402 false,
2403 )]);
2404
2405 let a = ListViewArray::new(
2407 list_field,
2408 vec![0, 1, 0, 6, 3].into(),
2409 vec![1, 2, 0, 4, 3].into(),
2410 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2411 None,
2412 );
2413
2414 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2415
2416 roundtrip(batch, None);
2417 }
2418
2419 #[test]
2420 fn arrow_writer_large_list_view() {
2421 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2422 let schema = Schema::new(vec![Field::new(
2423 "a",
2424 DataType::LargeListView(list_field.clone()),
2425 true,
2426 )]);
2427
2428 let a = LargeListViewArray::new(
2430 list_field,
2431 vec![0i64, 1, 0, 3, 6].into(),
2432 vec![1i64, 2, 0, 3, 4].into(),
2433 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2434 Some(vec![true, true, false, true, true].into()),
2435 );
2436
2437 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2438
2439 assert_eq!(batch.column(0).null_count(), 1);
2440
2441 roundtrip(batch, None);
2442 }
2443
2444 #[test]
2445 fn arrow_writer_list_view_with_struct() {
2446 let struct_fields = Fields::from(vec![
2448 Field::new("id", DataType::Int32, false),
2449 Field::new("name", DataType::Utf8, false),
2450 ]);
2451 let struct_type = DataType::Struct(struct_fields.clone());
2452 let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2453
2454 let schema = Schema::new(vec![Field::new(
2455 "a",
2456 DataType::ListView(list_field.clone()),
2457 true,
2458 )]);
2459
2460 let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2462 let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2463 let struct_array = StructArray::new(
2464 struct_fields,
2465 vec![Arc::new(id_array), Arc::new(name_array)],
2466 None,
2467 );
2468
2469 let list_view = ListViewArray::new(
2471 list_field,
2472 vec![0, 2, 2].into(), vec![2, 0, 3].into(), Arc::new(struct_array),
2475 Some(vec![true, false, true].into()),
2476 );
2477
2478 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap();
2479
2480 roundtrip(batch, None);
2481 }
2482
2483 #[test]
2484 fn arrow_writer_binary() {
2485 let string_field = Field::new("a", DataType::Utf8, false);
2486 let binary_field = Field::new("b", DataType::Binary, false);
2487 let schema = Schema::new(vec![string_field, binary_field]);
2488
2489 let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2490 let raw_binary_values = [
2491 b"foo".to_vec(),
2492 b"bar".to_vec(),
2493 b"baz".to_vec(),
2494 b"quux".to_vec(),
2495 ];
2496 let raw_binary_value_refs = raw_binary_values
2497 .iter()
2498 .map(|x| x.as_slice())
2499 .collect::<Vec<_>>();
2500
2501 let string_values = StringArray::from(raw_string_values.clone());
2502 let binary_values = BinaryArray::from(raw_binary_value_refs);
2503 let batch = RecordBatch::try_new(
2504 Arc::new(schema),
2505 vec![Arc::new(string_values), Arc::new(binary_values)],
2506 )
2507 .unwrap();
2508
2509 roundtrip(batch, Some(SMALL_SIZE / 2));
2510 }
2511
2512 #[test]
2513 fn arrow_writer_binary_view() {
2514 let string_field = Field::new("a", DataType::Utf8View, false);
2515 let binary_field = Field::new("b", DataType::BinaryView, false);
2516 let nullable_string_field = Field::new("a", DataType::Utf8View, true);
2517 let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]);
2518
2519 let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2520 let raw_binary_values = vec![
2521 b"foo".to_vec(),
2522 b"bar".to_vec(),
2523 b"large payload over 12 bytes".to_vec(),
2524 b"lulu".to_vec(),
2525 ];
2526 let nullable_string_values =
2527 vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2528
2529 let string_view_values = StringViewArray::from(raw_string_values);
2530 let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2531 let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2532 let batch = RecordBatch::try_new(
2533 Arc::new(schema),
2534 vec![
2535 Arc::new(string_view_values),
2536 Arc::new(binary_view_values),
2537 Arc::new(nullable_string_view_values),
2538 ],
2539 )
2540 .unwrap();
2541
2542 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2543 roundtrip(batch, None);
2544 }
2545
2546 #[test]
2547 fn arrow_writer_binary_view_long_value() {
2548 let string_field = Field::new("a", DataType::Utf8View, false);
2549 let binary_field = Field::new("b", DataType::BinaryView, false);
2550 let schema = Schema::new(vec![string_field, binary_field]);
2551
2552 let long = "a".repeat(128);
2556 let raw_string_values = vec!["foo", long.as_str(), "bar"];
2557 let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2558
2559 let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2560 let binary_view_values: ArrayRef =
2561 Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2562
2563 one_column_roundtrip(Arc::clone(&string_view_values), false);
2564 one_column_roundtrip(Arc::clone(&binary_view_values), false);
2565
2566 let batch = RecordBatch::try_new(
2567 Arc::new(schema),
2568 vec![string_view_values, binary_view_values],
2569 )
2570 .unwrap();
2571
2572 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
2574 let props = WriterProperties::builder()
2575 .set_writer_version(version)
2576 .set_dictionary_enabled(false)
2577 .build();
2578 roundtrip_opts(&batch, props);
2579 }
2580 }
2581
2582 fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2583 let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2584 let schema = Schema::new(vec![decimal_field]);
2585
2586 let decimal_values = vec![10_000, 50_000, 0, -100]
2587 .into_iter()
2588 .map(Some)
2589 .collect::<Decimal128Array>()
2590 .with_precision_and_scale(precision, scale)
2591 .unwrap();
2592
2593 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2594 }
2595
2596 #[test]
2597 fn arrow_writer_decimal() {
2598 let batch_int32_decimal = get_decimal_batch(5, 2);
2600 roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2601 let batch_int64_decimal = get_decimal_batch(12, 2);
2603 roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2604 let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2606 roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2607 }
2608
2609 #[test]
2610 fn arrow_writer_complex() {
2611 let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2613 let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2614 let struct_field_g = Arc::new(Field::new_list(
2615 "g",
2616 Field::new_list_field(DataType::Int16, true),
2617 false,
2618 ));
2619 let struct_field_h = Arc::new(Field::new_list(
2620 "h",
2621 Field::new_list_field(DataType::Int16, false),
2622 true,
2623 ));
2624 let struct_field_e = Arc::new(Field::new_struct(
2625 "e",
2626 vec![
2627 struct_field_f.clone(),
2628 struct_field_g.clone(),
2629 struct_field_h.clone(),
2630 ],
2631 false,
2632 ));
2633 let schema = Schema::new(vec![
2634 Field::new("a", DataType::Int32, false),
2635 Field::new("b", DataType::Int32, true),
2636 Field::new_struct(
2637 "c",
2638 vec![struct_field_d.clone(), struct_field_e.clone()],
2639 false,
2640 ),
2641 ]);
2642
2643 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2645 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2646 let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2647 let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2648
2649 let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2650
2651 let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2654
2655 let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2657 .len(5)
2658 .add_buffer(g_value_offsets.clone())
2659 .add_child_data(g_value.to_data())
2660 .build()
2661 .unwrap();
2662 let g = ListArray::from(g_list_data);
2663 let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2665 .len(5)
2666 .add_buffer(g_value_offsets)
2667 .add_child_data(g_value.to_data())
2668 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2669 .build()
2670 .unwrap();
2671 let h = ListArray::from(h_list_data);
2672
2673 let e = StructArray::from(vec![
2674 (struct_field_f, Arc::new(f) as ArrayRef),
2675 (struct_field_g, Arc::new(g) as ArrayRef),
2676 (struct_field_h, Arc::new(h) as ArrayRef),
2677 ]);
2678
2679 let c = StructArray::from(vec![
2680 (struct_field_d, Arc::new(d) as ArrayRef),
2681 (struct_field_e, Arc::new(e) as ArrayRef),
2682 ]);
2683
2684 let batch = RecordBatch::try_new(
2686 Arc::new(schema),
2687 vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2688 )
2689 .unwrap();
2690
2691 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2692 roundtrip(batch, Some(SMALL_SIZE / 3));
2693 }
2694
2695 #[test]
2696 fn arrow_writer_complex_mixed() {
2697 let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2702 let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2703 let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2704 let schema = Schema::new(vec![Field::new(
2705 "some_nested_object",
2706 DataType::Struct(Fields::from(vec![
2707 offset_field.clone(),
2708 partition_field.clone(),
2709 topic_field.clone(),
2710 ])),
2711 false,
2712 )]);
2713
2714 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2716 let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2717 let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2718
2719 let some_nested_object = StructArray::from(vec![
2720 (offset_field, Arc::new(offset) as ArrayRef),
2721 (partition_field, Arc::new(partition) as ArrayRef),
2722 (topic_field, Arc::new(topic) as ArrayRef),
2723 ]);
2724
2725 let batch =
2727 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2728
2729 roundtrip(batch, Some(SMALL_SIZE / 2));
2730 }
2731
2732 #[test]
2733 fn arrow_writer_map() {
2734 let json_content = r#"
2736 {"stocks":{"long": "$AAA", "short": "$BBB"}}
2737 {"stocks":{"long": null, "long": "$CCC", "short": null}}
2738 {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2739 "#;
2740 let entries_struct_type = DataType::Struct(Fields::from(vec![
2741 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2742 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2743 ]));
2744 let stocks_field = Field::new(
2745 "stocks",
2746 DataType::Map(
2747 Arc::new(Field::new(
2748 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2749 entries_struct_type,
2750 false,
2751 )),
2752 false,
2753 ),
2754 true,
2755 );
2756 let schema = Arc::new(Schema::new(vec![stocks_field]));
2757 let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2758 let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2759
2760 let batch = reader.next().unwrap().unwrap();
2761 roundtrip(batch, None);
2762 }
2763
2764 #[test]
2765 fn arrow_writer_2_level_struct() {
2766 let field_c = Field::new("c", DataType::Int32, true);
2768 let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2769 let type_a = DataType::Struct(vec![field_b.clone()].into());
2770 let field_a = Field::new("a", type_a, true);
2771 let schema = Schema::new(vec![field_a.clone()]);
2772
2773 let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2775 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2776 .len(6)
2777 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2778 .add_child_data(c.into_data())
2779 .build()
2780 .unwrap();
2781 let b = StructArray::from(b_data);
2782 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2783 .len(6)
2784 .null_bit_buffer(Some(Buffer::from([0b00101111])))
2785 .add_child_data(b.into_data())
2786 .build()
2787 .unwrap();
2788 let a = StructArray::from(a_data);
2789
2790 assert_eq!(a.null_count(), 1);
2791 assert_eq!(a.column(0).null_count(), 2);
2792
2793 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2795
2796 roundtrip(batch, Some(SMALL_SIZE / 2));
2797 }
2798
2799 #[test]
2800 fn arrow_writer_2_level_struct_non_null() {
2801 let field_c = Field::new("c", DataType::Int32, false);
2803 let type_b = DataType::Struct(vec![field_c].into());
2804 let field_b = Field::new("b", type_b.clone(), false);
2805 let type_a = DataType::Struct(vec![field_b].into());
2806 let field_a = Field::new("a", type_a.clone(), false);
2807 let schema = Schema::new(vec![field_a]);
2808
2809 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2811 let b_data = ArrayDataBuilder::new(type_b)
2812 .len(6)
2813 .add_child_data(c.into_data())
2814 .build()
2815 .unwrap();
2816 let b = StructArray::from(b_data);
2817 let a_data = ArrayDataBuilder::new(type_a)
2818 .len(6)
2819 .add_child_data(b.into_data())
2820 .build()
2821 .unwrap();
2822 let a = StructArray::from(a_data);
2823
2824 assert_eq!(a.null_count(), 0);
2825 assert_eq!(a.column(0).null_count(), 0);
2826
2827 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2829
2830 roundtrip(batch, Some(SMALL_SIZE / 2));
2831 }
2832
2833 #[test]
2834 fn arrow_writer_2_level_struct_mixed_null() {
2835 let field_c = Field::new("c", DataType::Int32, false);
2837 let type_b = DataType::Struct(vec![field_c].into());
2838 let field_b = Field::new("b", type_b.clone(), true);
2839 let type_a = DataType::Struct(vec![field_b].into());
2840 let field_a = Field::new("a", type_a.clone(), false);
2841 let schema = Schema::new(vec![field_a]);
2842
2843 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2845 let b_data = ArrayDataBuilder::new(type_b)
2846 .len(6)
2847 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2848 .add_child_data(c.into_data())
2849 .build()
2850 .unwrap();
2851 let b = StructArray::from(b_data);
2852 let a_data = ArrayDataBuilder::new(type_a)
2854 .len(6)
2855 .add_child_data(b.into_data())
2856 .build()
2857 .unwrap();
2858 let a = StructArray::from(a_data);
2859
2860 assert_eq!(a.null_count(), 0);
2861 assert_eq!(a.column(0).null_count(), 2);
2862
2863 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2865
2866 roundtrip(batch, Some(SMALL_SIZE / 2));
2867 }
2868
2869 #[test]
2870 fn arrow_writer_2_level_struct_mixed_null_2() {
2871 let field_c = Field::new("c", DataType::Int32, false);
2873 let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
2874 let field_e = Field::new(
2875 "e",
2876 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2877 false,
2878 );
2879
2880 let field_b = Field::new(
2881 "b",
2882 DataType::Struct(vec![field_c, field_d, field_e].into()),
2883 false,
2884 );
2885 let type_a = DataType::Struct(vec![field_b.clone()].into());
2886 let field_a = Field::new("a", type_a, true);
2887 let schema = Schema::new(vec![field_a.clone()]);
2888
2889 let c = Int32Array::from_iter_values(0..6);
2891 let d = FixedSizeBinaryArray::try_from_iter(
2892 ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
2893 )
2894 .expect("four byte values");
2895 let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
2896 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2897 .len(6)
2898 .add_child_data(c.into_data())
2899 .add_child_data(d.into_data())
2900 .add_child_data(e.into_data())
2901 .build()
2902 .unwrap();
2903 let b = StructArray::from(b_data);
2904 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2905 .len(6)
2906 .null_bit_buffer(Some(Buffer::from([0b00100101])))
2907 .add_child_data(b.into_data())
2908 .build()
2909 .unwrap();
2910 let a = StructArray::from(a_data);
2911
2912 assert_eq!(a.null_count(), 3);
2913 assert_eq!(a.column(0).null_count(), 0);
2914
2915 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2917
2918 roundtrip(batch, Some(SMALL_SIZE / 2));
2919 }
2920
2921 #[test]
2922 fn test_fixed_size_binary_in_dict() {
2923 fn test_fixed_size_binary_in_dict_inner<K>()
2924 where
2925 K: ArrowDictionaryKeyType,
2926 K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
2927 <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
2928 {
2929 let field = Field::new(
2930 "a",
2931 DataType::Dictionary(
2932 Box::new(K::DATA_TYPE),
2933 Box::new(DataType::FixedSizeBinary(4)),
2934 ),
2935 false,
2936 );
2937 let schema = Schema::new(vec![field]);
2938
2939 let keys: Vec<K::Native> = vec![
2940 K::Native::try_from(0u8).unwrap(),
2941 K::Native::try_from(0u8).unwrap(),
2942 K::Native::try_from(1u8).unwrap(),
2943 ];
2944 let keys = PrimitiveArray::<K>::from_iter_values(keys);
2945 let values = FixedSizeBinaryArray::try_from_iter(
2946 vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
2947 )
2948 .unwrap();
2949
2950 let data = DictionaryArray::<K>::new(keys, Arc::new(values));
2951 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
2952 roundtrip(batch, None);
2953 }
2954
2955 test_fixed_size_binary_in_dict_inner::<UInt8Type>();
2956 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2957 test_fixed_size_binary_in_dict_inner::<UInt32Type>();
2958 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2959 test_fixed_size_binary_in_dict_inner::<Int8Type>();
2960 test_fixed_size_binary_in_dict_inner::<Int16Type>();
2961 test_fixed_size_binary_in_dict_inner::<Int32Type>();
2962 test_fixed_size_binary_in_dict_inner::<Int64Type>();
2963 }
2964
2965 #[test]
2966 fn test_empty_dict() {
2967 let struct_fields = Fields::from(vec![Field::new(
2968 "dict",
2969 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2970 false,
2971 )]);
2972
2973 let schema = Schema::new(vec![Field::new_struct(
2974 "struct",
2975 struct_fields.clone(),
2976 true,
2977 )]);
2978 let dictionary = Arc::new(DictionaryArray::new(
2979 Int32Array::new_null(5),
2980 Arc::new(StringArray::new_null(0)),
2981 ));
2982
2983 let s = StructArray::new(
2984 struct_fields,
2985 vec![dictionary],
2986 Some(NullBuffer::new_null(5)),
2987 );
2988
2989 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
2990 roundtrip(batch, None);
2991 }
2992 #[test]
2993 fn arrow_writer_page_size() {
2994 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
2995
2996 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
2997
2998 for i in 0..10 {
3000 let value = i
3001 .to_string()
3002 .repeat(10)
3003 .chars()
3004 .take(10)
3005 .collect::<String>();
3006
3007 builder.append_value(value);
3008 }
3009
3010 let array = Arc::new(builder.finish());
3011
3012 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3013
3014 let file = tempfile::tempfile().unwrap();
3015
3016 let props = WriterProperties::builder()
3018 .set_data_page_size_limit(1)
3019 .set_dictionary_page_size_limit(1)
3020 .set_write_batch_size(1)
3021 .build();
3022
3023 let mut writer =
3024 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3025 .expect("Unable to write file");
3026 writer.write(&batch).unwrap();
3027 writer.close().unwrap();
3028
3029 let options = ReadOptionsBuilder::new().with_page_index().build();
3030 let reader =
3031 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3032
3033 let column = reader.metadata().row_group(0).columns();
3034
3035 assert_eq!(column.len(), 1);
3036
3037 assert!(
3040 column[0].dictionary_page_offset().is_some(),
3041 "Expected a dictionary page"
3042 );
3043
3044 assert!(reader.metadata().offset_index().is_some());
3045 let offset_indexes = &reader.metadata().offset_index().unwrap()[0];
3046
3047 let page_locations = offset_indexes[0].page_locations.clone();
3048
3049 assert_eq!(
3052 page_locations.len(),
3053 10,
3054 "Expected 10 pages but got {page_locations:#?}"
3055 );
3056 }
3057
3058 #[test]
3059 fn arrow_writer_float_nans() {
3060 let f16_field = Field::new("a", DataType::Float16, false);
3061 let f32_field = Field::new("b", DataType::Float32, false);
3062 let f64_field = Field::new("c", DataType::Float64, false);
3063 let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3064
3065 let f16_values = (0..MEDIUM_SIZE)
3066 .map(|i| {
3067 Some(if i % 2 == 0 {
3068 f16::NAN
3069 } else {
3070 f16::from_f32(i as f32)
3071 })
3072 })
3073 .collect::<Float16Array>();
3074
3075 let f32_values = (0..MEDIUM_SIZE)
3076 .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3077 .collect::<Float32Array>();
3078
3079 let f64_values = (0..MEDIUM_SIZE)
3080 .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3081 .collect::<Float64Array>();
3082
3083 let batch = RecordBatch::try_new(
3084 Arc::new(schema),
3085 vec![
3086 Arc::new(f16_values),
3087 Arc::new(f32_values),
3088 Arc::new(f64_values),
3089 ],
3090 )
3091 .unwrap();
3092
3093 roundtrip(batch, None);
3094 }
3095
3096 const SMALL_SIZE: usize = 7;
3097 const MEDIUM_SIZE: usize = 63;
3098
3099 fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3102 let mut files = vec![];
3103 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3104 let mut props = WriterProperties::builder().set_writer_version(version);
3105
3106 if let Some(size) = max_row_group_size {
3107 props = props.set_max_row_group_row_count(Some(size))
3108 }
3109
3110 let props = props.build();
3111 files.push(roundtrip_opts(&expected_batch, props))
3112 }
3113 files
3114 }
3115
3116 fn roundtrip_opts_with_array_validation<F>(
3120 expected_batch: &RecordBatch,
3121 props: WriterProperties,
3122 validate: F,
3123 ) -> Bytes
3124 where
3125 F: Fn(&ArrayData, &ArrayData),
3126 {
3127 let mut file = vec![];
3128
3129 let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3130 .expect("Unable to write file");
3131 writer.write(expected_batch).unwrap();
3132 writer.close().unwrap();
3133
3134 let file = Bytes::from(file);
3135 let mut record_batch_reader =
3136 ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3137
3138 let actual_batch = record_batch_reader
3139 .next()
3140 .expect("No batch found")
3141 .expect("Unable to get batch");
3142
3143 assert_eq!(expected_batch.schema(), actual_batch.schema());
3144 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3145 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3146 for i in 0..expected_batch.num_columns() {
3147 let expected_data = expected_batch.column(i).to_data();
3148 let actual_data = actual_batch.column(i).to_data();
3149 validate(&expected_data, &actual_data);
3150 }
3151
3152 file
3153 }
3154
3155 fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3156 roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3157 a.validate_full().expect("valid expected data");
3158 b.validate_full().expect("valid actual data");
3159 assert_eq!(a, b)
3160 })
3161 }
3162
3163 struct RoundTripOptions {
3164 values: ArrayRef,
3165 schema: SchemaRef,
3166 bloom_filter: bool,
3167 bloom_filter_ndv: Option<u64>,
3168 bloom_filter_position: BloomFilterPosition,
3169 }
3170
3171 impl RoundTripOptions {
3172 fn new(values: ArrayRef, nullable: bool) -> Self {
3173 let data_type = values.data_type().clone();
3174 let schema = Schema::new(vec![Field::new("col", data_type, nullable)]);
3175 Self {
3176 values,
3177 schema: Arc::new(schema),
3178 bloom_filter: false,
3179 bloom_filter_ndv: None,
3180 bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3181 }
3182 }
3183 }
3184
3185 fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec<Bytes> {
3186 one_column_roundtrip_with_options(RoundTripOptions::new(values, nullable))
3187 }
3188
3189 fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef) -> Vec<Bytes> {
3190 let mut options = RoundTripOptions::new(values, false);
3191 options.schema = schema;
3192 one_column_roundtrip_with_options(options)
3193 }
3194
3195 fn one_column_roundtrip_with_options(options: RoundTripOptions) -> Vec<Bytes> {
3196 let RoundTripOptions {
3197 values,
3198 schema,
3199 bloom_filter,
3200 bloom_filter_ndv,
3201 bloom_filter_position,
3202 } = options;
3203
3204 let encodings = match values.data_type() {
3205 DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3206 vec![
3207 Encoding::PLAIN,
3208 Encoding::DELTA_BYTE_ARRAY,
3209 Encoding::DELTA_LENGTH_BYTE_ARRAY,
3210 ]
3211 }
3212 DataType::Int64
3213 | DataType::Int32
3214 | DataType::Int16
3215 | DataType::Int8
3216 | DataType::UInt64
3217 | DataType::UInt32
3218 | DataType::UInt16
3219 | DataType::UInt8 => vec![
3220 Encoding::PLAIN,
3221 Encoding::DELTA_BINARY_PACKED,
3222 Encoding::BYTE_STREAM_SPLIT,
3223 ],
3224 DataType::Float32 | DataType::Float64 => {
3225 vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
3226 }
3227 _ => vec![Encoding::PLAIN],
3228 };
3229
3230 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3231
3232 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3233
3234 let mut files = vec![];
3235 for dictionary_size in [0, 1, 1024] {
3236 for encoding in &encodings {
3237 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3238 for row_group_size in row_group_sizes {
3239 let mut builder = WriterProperties::builder()
3240 .set_writer_version(version)
3241 .set_max_row_group_row_count(Some(row_group_size))
3242 .set_dictionary_enabled(dictionary_size != 0)
3243 .set_dictionary_page_size_limit(dictionary_size.max(1))
3244 .set_encoding(*encoding)
3245 .set_bloom_filter_enabled(bloom_filter)
3246 .set_bloom_filter_position(bloom_filter_position);
3247 if let Some(ndv) = bloom_filter_ndv {
3248 builder = builder.set_bloom_filter_max_ndv(ndv);
3249 }
3250 let props = builder.build();
3251
3252 files.push(roundtrip_opts(&expected_batch, props))
3253 }
3254 }
3255 }
3256 }
3257 files
3258 }
3259
3260 fn values_required<A, I>(iter: I) -> Vec<Bytes>
3261 where
3262 A: From<Vec<I::Item>> + Array + 'static,
3263 I: IntoIterator,
3264 {
3265 let raw_values: Vec<_> = iter.into_iter().collect();
3266 let values = Arc::new(A::from(raw_values));
3267 one_column_roundtrip(values, false)
3268 }
3269
3270 fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3271 where
3272 A: From<Vec<Option<I::Item>>> + Array + 'static,
3273 I: IntoIterator,
3274 {
3275 let optional_raw_values: Vec<_> = iter
3276 .into_iter()
3277 .enumerate()
3278 .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3279 .collect();
3280 let optional_values = Arc::new(A::from(optional_raw_values));
3281 one_column_roundtrip(optional_values, true)
3282 }
3283
3284 fn required_and_optional<A, I>(iter: I)
3285 where
3286 A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3287 I: IntoIterator + Clone,
3288 {
3289 values_required::<A, I>(iter.clone());
3290 values_optional::<A, I>(iter);
3291 }
3292
3293 fn check_bloom_filter<T: AsBytes>(
3294 files: Vec<Bytes>,
3295 file_column: String,
3296 positive_values: Vec<T>,
3297 negative_values: Vec<T>,
3298 ) {
3299 files.into_iter().take(1).for_each(|file| {
3300 let file_reader = SerializedFileReader::new_with_options(
3301 file,
3302 ReadOptionsBuilder::new()
3303 .with_reader_properties(
3304 ReaderProperties::builder()
3305 .set_read_bloom_filter(true)
3306 .build(),
3307 )
3308 .build(),
3309 )
3310 .expect("Unable to open file as Parquet");
3311 let metadata = file_reader.metadata();
3312
3313 let mut bloom_filters: Vec<_> = vec![];
3315 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3316 if let Some((column_index, _)) = row_group
3317 .columns()
3318 .iter()
3319 .enumerate()
3320 .find(|(_, column)| column.column_path().string() == file_column)
3321 {
3322 let row_group_reader = file_reader
3323 .get_row_group(ri)
3324 .expect("Unable to read row group");
3325 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3326 bloom_filters.push(sbbf.clone());
3327 } else {
3328 panic!("No bloom filter for column named {file_column} found");
3329 }
3330 } else {
3331 panic!("No column named {file_column} found");
3332 }
3333 }
3334
3335 positive_values.iter().for_each(|value| {
3336 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3337 assert!(
3338 found.is_some(),
3339 "{}",
3340 format!("Value {:?} should be in bloom filter", value.as_bytes())
3341 );
3342 });
3343
3344 negative_values.iter().for_each(|value| {
3345 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3346 assert!(
3347 found.is_none(),
3348 "{}",
3349 format!("Value {:?} should not be in bloom filter", value.as_bytes())
3350 );
3351 });
3352 });
3353 }
3354
3355 #[test]
3356 fn all_null_primitive_single_column() {
3357 let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3358 one_column_roundtrip(values, true);
3359 }
3360 #[test]
3361 fn null_single_column() {
3362 let values = Arc::new(NullArray::new(SMALL_SIZE));
3363 one_column_roundtrip(values, true);
3364 }
3366
3367 #[test]
3368 fn bool_single_column() {
3369 required_and_optional::<BooleanArray, _>(
3370 [true, false].iter().cycle().copied().take(SMALL_SIZE),
3371 );
3372 }
3373
3374 #[test]
3375 fn bool_large_single_column() {
3376 let values = Arc::new(
3377 [None, Some(true), Some(false)]
3378 .iter()
3379 .cycle()
3380 .copied()
3381 .take(200_000)
3382 .collect::<BooleanArray>(),
3383 );
3384 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3385 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3386 let file = tempfile::tempfile().unwrap();
3387
3388 let mut writer =
3389 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3390 .expect("Unable to write file");
3391 writer.write(&expected_batch).unwrap();
3392 writer.close().unwrap();
3393 }
3394
3395 #[test]
3396 fn check_page_offset_index_with_nan() {
3397 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3398 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3399 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3400
3401 let mut out = Vec::with_capacity(1024);
3402 let mut writer =
3403 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3404 writer.write(&batch).unwrap();
3405 let file_meta_data = writer.close().unwrap();
3406 for row_group in file_meta_data.row_groups() {
3407 for column in row_group.columns() {
3408 assert!(column.offset_index_offset().is_some());
3409 assert!(column.offset_index_length().is_some());
3410 assert!(column.column_index_offset().is_none());
3411 assert!(column.column_index_length().is_none());
3412 }
3413 }
3414 }
3415
3416 #[test]
3417 fn i8_single_column() {
3418 required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3419 }
3420
3421 #[test]
3422 fn i16_single_column() {
3423 required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3424 }
3425
3426 #[test]
3427 fn i32_single_column() {
3428 required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3429 }
3430
3431 #[test]
3432 fn i64_single_column() {
3433 required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3434 }
3435
3436 #[test]
3437 fn u8_single_column() {
3438 required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3439 }
3440
3441 #[test]
3442 fn u16_single_column() {
3443 required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3444 }
3445
3446 #[test]
3447 fn u32_single_column() {
3448 required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3449 }
3450
3451 #[test]
3452 fn u64_single_column() {
3453 required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3454 }
3455
3456 #[test]
3457 fn f32_single_column() {
3458 required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3459 }
3460
3461 #[test]
3462 fn f64_single_column() {
3463 required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3464 }
3465
3466 #[test]
3471 fn timestamp_second_single_column() {
3472 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3473 let values = Arc::new(TimestampSecondArray::from(raw_values));
3474
3475 one_column_roundtrip(values, false);
3476 }
3477
3478 #[test]
3479 fn timestamp_millisecond_single_column() {
3480 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3481 let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3482
3483 one_column_roundtrip(values, false);
3484 }
3485
3486 #[test]
3487 fn timestamp_microsecond_single_column() {
3488 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3489 let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3490
3491 one_column_roundtrip(values, false);
3492 }
3493
3494 #[test]
3495 fn timestamp_nanosecond_single_column() {
3496 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3497 let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3498
3499 one_column_roundtrip(values, false);
3500 }
3501
3502 #[test]
3503 fn date32_single_column() {
3504 required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3505 }
3506
3507 #[test]
3508 fn date64_single_column() {
3509 required_and_optional::<Date64Array, _>(
3511 (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3512 );
3513 }
3514
3515 #[test]
3516 fn time32_second_single_column() {
3517 required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3518 }
3519
3520 #[test]
3521 fn time32_millisecond_single_column() {
3522 required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3523 }
3524
3525 #[test]
3526 fn time64_microsecond_single_column() {
3527 required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3528 }
3529
3530 #[test]
3531 fn time64_nanosecond_single_column() {
3532 required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3533 }
3534
3535 #[test]
3536 fn duration_second_single_column() {
3537 required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3538 }
3539
3540 #[test]
3541 fn duration_millisecond_single_column() {
3542 required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3543 }
3544
3545 #[test]
3546 fn duration_microsecond_single_column() {
3547 required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3548 }
3549
3550 #[test]
3551 fn duration_nanosecond_single_column() {
3552 required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3553 }
3554
3555 #[test]
3556 fn interval_year_month_single_column() {
3557 required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3558 }
3559
3560 #[test]
3561 fn interval_day_time_single_column() {
3562 required_and_optional::<IntervalDayTimeArray, _>(vec![
3563 IntervalDayTime::new(0, 1),
3564 IntervalDayTime::new(0, 3),
3565 IntervalDayTime::new(3, -2),
3566 IntervalDayTime::new(-200, 4),
3567 ]);
3568 }
3569
3570 #[test]
3571 #[should_panic(
3572 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3573 )]
3574 fn interval_month_day_nano_single_column() {
3575 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3576 IntervalMonthDayNano::new(0, 1, 5),
3577 IntervalMonthDayNano::new(0, 3, 2),
3578 IntervalMonthDayNano::new(3, -2, -5),
3579 IntervalMonthDayNano::new(-200, 4, -1),
3580 ]);
3581 }
3582
3583 #[test]
3584 fn binary_single_column() {
3585 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3586 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3587 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3588
3589 values_required::<BinaryArray, _>(many_vecs_iter);
3591 }
3592
3593 #[test]
3594 fn binary_view_single_column() {
3595 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3596 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3597 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3598
3599 values_required::<BinaryViewArray, _>(many_vecs_iter);
3601 }
3602
3603 #[test]
3604 fn i32_column_bloom_filter_at_end() {
3605 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3606 let mut options = RoundTripOptions::new(array, false);
3607 options.bloom_filter = true;
3608 options.bloom_filter_position = BloomFilterPosition::End;
3609
3610 let files = one_column_roundtrip_with_options(options);
3611 check_bloom_filter(
3612 files,
3613 "col".to_string(),
3614 (0..SMALL_SIZE as i32).collect(),
3615 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3616 );
3617 }
3618
3619 #[test]
3620 fn i32_column_bloom_filter() {
3621 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3622 let mut options = RoundTripOptions::new(array, false);
3623 options.bloom_filter = true;
3624
3625 let files = one_column_roundtrip_with_options(options);
3626 check_bloom_filter(
3627 files,
3628 "col".to_string(),
3629 (0..SMALL_SIZE as i32).collect(),
3630 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3631 );
3632 }
3633
3634 #[test]
3639 fn i32_column_bloom_filter_fixed_ndv() {
3640 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3641
3642 let mut options = RoundTripOptions::new(array.clone(), false);
3644 options.bloom_filter = true;
3645 options.bloom_filter_ndv = Some(1_000_000);
3646
3647 let files = one_column_roundtrip_with_options(options);
3648 check_bloom_filter(
3649 files,
3650 "col".to_string(),
3651 (0..SMALL_SIZE as i32).collect(),
3652 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3653 );
3654
3655 let mut options = RoundTripOptions::new(array, false);
3657 options.bloom_filter = true;
3658 options.bloom_filter_ndv = Some(3);
3659
3660 let files = one_column_roundtrip_with_options(options);
3661 check_bloom_filter(
3662 files,
3663 "col".to_string(),
3664 (0..SMALL_SIZE as i32).collect(),
3665 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3666 );
3667 }
3668
3669 #[test]
3670 fn binary_column_bloom_filter() {
3671 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3672 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3673 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3674
3675 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
3676 let mut options = RoundTripOptions::new(array, false);
3677 options.bloom_filter = true;
3678
3679 let files = one_column_roundtrip_with_options(options);
3680 check_bloom_filter(
3681 files,
3682 "col".to_string(),
3683 many_vecs,
3684 vec![vec![(SMALL_SIZE + 1) as u8]],
3685 );
3686 }
3687
3688 #[test]
3689 fn empty_string_null_column_bloom_filter() {
3690 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3691 let raw_strs = raw_values.iter().map(|s| s.as_str());
3692
3693 let array = Arc::new(StringArray::from_iter_values(raw_strs));
3694 let mut options = RoundTripOptions::new(array, false);
3695 options.bloom_filter = true;
3696
3697 let files = one_column_roundtrip_with_options(options);
3698
3699 let optional_raw_values: Vec<_> = raw_values
3700 .iter()
3701 .enumerate()
3702 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
3703 .collect();
3704 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
3706 }
3707
3708 #[test]
3709 fn large_binary_single_column() {
3710 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3711 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3712 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3713
3714 values_required::<LargeBinaryArray, _>(many_vecs_iter);
3716 }
3717
3718 #[test]
3719 fn fixed_size_binary_single_column() {
3720 let mut builder = FixedSizeBinaryBuilder::new(4);
3721 builder.append_value(b"0123").unwrap();
3722 builder.append_null();
3723 builder.append_value(b"8910").unwrap();
3724 builder.append_value(b"1112").unwrap();
3725 let array = Arc::new(builder.finish());
3726
3727 one_column_roundtrip(array, true);
3728 }
3729
3730 #[test]
3731 fn string_single_column() {
3732 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3733 let raw_strs = raw_values.iter().map(|s| s.as_str());
3734
3735 required_and_optional::<StringArray, _>(raw_strs);
3736 }
3737
3738 #[test]
3739 fn large_string_single_column() {
3740 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3741 let raw_strs = raw_values.iter().map(|s| s.as_str());
3742
3743 required_and_optional::<LargeStringArray, _>(raw_strs);
3744 }
3745
3746 #[test]
3747 fn string_view_single_column() {
3748 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3749 let raw_strs = raw_values.iter().map(|s| s.as_str());
3750
3751 required_and_optional::<StringViewArray, _>(raw_strs);
3752 }
3753
3754 #[test]
3755 fn null_list_single_column() {
3756 let null_field = Field::new_list_field(DataType::Null, true);
3757 let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
3758
3759 let schema = Schema::new(vec![list_field]);
3760
3761 let a_values = NullArray::new(2);
3763 let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
3764 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3765 DataType::Null,
3766 true,
3767 ))))
3768 .len(3)
3769 .add_buffer(a_value_offsets)
3770 .null_bit_buffer(Some(Buffer::from([0b00000101])))
3771 .add_child_data(a_values.into_data())
3772 .build()
3773 .unwrap();
3774
3775 let a = ListArray::from(a_list_data);
3776
3777 assert!(a.is_valid(0));
3778 assert!(!a.is_valid(1));
3779 assert!(a.is_valid(2));
3780
3781 assert_eq!(a.value(0).len(), 0);
3782 assert_eq!(a.value(2).len(), 2);
3783 assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
3784
3785 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3786 roundtrip(batch, None);
3787 }
3788
3789 #[test]
3790 fn list_single_column() {
3791 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3792 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
3793 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3794 DataType::Int32,
3795 false,
3796 ))))
3797 .len(5)
3798 .add_buffer(a_value_offsets)
3799 .null_bit_buffer(Some(Buffer::from([0b00011011])))
3800 .add_child_data(a_values.into_data())
3801 .build()
3802 .unwrap();
3803
3804 assert_eq!(a_list_data.null_count(), 1);
3805
3806 let a = ListArray::from(a_list_data);
3807 let values = Arc::new(a);
3808
3809 one_column_roundtrip(values, true);
3810 }
3811
3812 #[test]
3813 fn large_list_single_column() {
3814 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3815 let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
3816 let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
3817 "large_item",
3818 DataType::Int32,
3819 true,
3820 ))))
3821 .len(5)
3822 .add_buffer(a_value_offsets)
3823 .add_child_data(a_values.into_data())
3824 .null_bit_buffer(Some(Buffer::from([0b00011011])))
3825 .build()
3826 .unwrap();
3827
3828 assert_eq!(a_list_data.null_count(), 1);
3830
3831 let a = LargeListArray::from(a_list_data);
3832 let values = Arc::new(a);
3833
3834 one_column_roundtrip(values, true);
3835 }
3836
3837 #[test]
3838 fn list_nested_nulls() {
3839 use arrow::datatypes::Int32Type;
3840 let data = vec![
3841 Some(vec![Some(1)]),
3842 Some(vec![Some(2), Some(3)]),
3843 None,
3844 Some(vec![Some(4), Some(5), None]),
3845 Some(vec![None]),
3846 Some(vec![Some(6), Some(7)]),
3847 ];
3848
3849 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
3850 one_column_roundtrip(Arc::new(list), true);
3851
3852 let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
3853 one_column_roundtrip(Arc::new(list), true);
3854 }
3855
3856 #[test]
3857 fn list_utf8_view_selective_padding_roundtrip() {
3858 let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
3859 let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
3860 builder.values().append_value("a");
3861 builder.values().append_null();
3862 builder.append(true);
3863 builder.append(false);
3866 builder.values().append_value("large payload over 12 bytes");
3868 builder.append(true);
3869
3870 one_column_roundtrip(Arc::new(builder.finish()), true);
3871 }
3872
3873 #[test]
3874 fn struct_single_column() {
3875 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3876 let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
3877 let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
3878
3879 let values = Arc::new(s);
3880 one_column_roundtrip(values, false);
3881 }
3882
3883 #[test]
3884 fn list_and_map_coerced_names() {
3885 let list_field =
3887 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
3888 let map_field = Field::new_map(
3889 "my_map",
3890 "my_entries",
3891 Field::new("my_keys", DataType::Int32, false),
3892 Field::new("my_values", DataType::Int32, true),
3893 false,
3894 true,
3895 );
3896
3897 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
3898 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
3899
3900 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
3901
3902 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
3904 let file = tempfile::tempfile().unwrap();
3905 let mut writer =
3906 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
3907
3908 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
3909 writer.write(&batch).unwrap();
3910 let file_metadata = writer.close().unwrap();
3911
3912 let schema = file_metadata.file_metadata().schema();
3913 let list_field = &schema.get_fields()[0].get_fields()[0];
3915 assert_eq!(list_field.get_fields()[0].name(), "element");
3916
3917 let map_field = &schema.get_fields()[1].get_fields()[0];
3918 assert_eq!(map_field.name(), "key_value");
3920 assert_eq!(map_field.get_fields()[0].name(), "key");
3922 assert_eq!(map_field.get_fields()[1].name(), "value");
3924
3925 let reader = SerializedFileReader::new(file).unwrap();
3927 let file_schema = reader.metadata().file_metadata().schema();
3928 let fields = file_schema.get_fields();
3929 let list_field = &fields[0].get_fields()[0];
3930 assert_eq!(list_field.get_fields()[0].name(), "element");
3931 let map_field = &fields[1].get_fields()[0];
3932 assert_eq!(map_field.name(), "key_value");
3933 assert_eq!(map_field.get_fields()[0].name(), "key");
3934 assert_eq!(map_field.get_fields()[1].name(), "value");
3935 }
3936
3937 #[test]
3938 fn fallback_flush_data_page() {
3939 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
3941 let values = Arc::new(StringArray::from(raw_values));
3942 let encodings = vec![
3943 Encoding::DELTA_BYTE_ARRAY,
3944 Encoding::DELTA_LENGTH_BYTE_ARRAY,
3945 ];
3946 let data_type = values.data_type().clone();
3947 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
3948 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3949
3950 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3951 let data_page_size_limit: usize = 32;
3952 let write_batch_size: usize = 16;
3953
3954 for encoding in &encodings {
3955 for row_group_size in row_group_sizes {
3956 let props = WriterProperties::builder()
3957 .set_writer_version(WriterVersion::PARQUET_2_0)
3958 .set_max_row_group_row_count(Some(row_group_size))
3959 .set_dictionary_enabled(false)
3960 .set_encoding(*encoding)
3961 .set_data_page_size_limit(data_page_size_limit)
3962 .set_write_batch_size(write_batch_size)
3963 .build();
3964
3965 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
3966 let string_array_a = StringArray::from(a.clone());
3967 let string_array_b = StringArray::from(b.clone());
3968 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
3969 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
3970 assert_eq!(
3971 vec_a, vec_b,
3972 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
3973 );
3974 });
3975 }
3976 }
3977 }
3978
3979 #[test]
3980 fn arrow_writer_string_dictionary() {
3981 #[allow(deprecated)]
3983 let schema = Arc::new(Schema::new(vec![Field::new_dict(
3984 "dictionary",
3985 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3986 true,
3987 42,
3988 true,
3989 )]));
3990
3991 let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
3993 .iter()
3994 .copied()
3995 .collect();
3996
3997 one_column_roundtrip_with_schema(Arc::new(d), schema);
3999 }
4000
4001 #[test]
4002 fn arrow_writer_test_type_compatibility() {
4003 fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4004 where
4005 T1: Array + 'static,
4006 T2: Array + 'static,
4007 {
4008 let schema1 = Arc::new(Schema::new(vec![Field::new(
4009 "a",
4010 array1.data_type().clone(),
4011 false,
4012 )]));
4013
4014 let file = tempfile().unwrap();
4015 let mut writer =
4016 ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4017
4018 let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4019 writer.write(&rb1).unwrap();
4020
4021 let schema2 = Arc::new(Schema::new(vec![Field::new(
4022 "a",
4023 array2.data_type().clone(),
4024 false,
4025 )]));
4026 let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4027 writer.write(&rb2).unwrap();
4028
4029 writer.close().unwrap();
4030
4031 let mut record_batch_reader =
4032 ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4033 let actual_batch = record_batch_reader.next().unwrap().unwrap();
4034
4035 let expected_batch =
4036 RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4037 assert_eq!(actual_batch, expected_batch);
4038 }
4039
4040 ensure_compatible_write(
4043 DictionaryArray::new(
4044 UInt8Array::from_iter_values(vec![0]),
4045 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4046 ),
4047 StringArray::from_iter_values(vec!["barquet"]),
4048 DictionaryArray::new(
4049 UInt8Array::from_iter_values(vec![0, 1]),
4050 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4051 ),
4052 );
4053
4054 ensure_compatible_write(
4055 StringArray::from_iter_values(vec!["parquet"]),
4056 DictionaryArray::new(
4057 UInt8Array::from_iter_values(vec![0]),
4058 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4059 ),
4060 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4061 );
4062
4063 ensure_compatible_write(
4066 DictionaryArray::new(
4067 UInt8Array::from_iter_values(vec![0]),
4068 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4069 ),
4070 DictionaryArray::new(
4071 UInt16Array::from_iter_values(vec![0]),
4072 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4073 ),
4074 DictionaryArray::new(
4075 UInt8Array::from_iter_values(vec![0, 1]),
4076 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4077 ),
4078 );
4079
4080 ensure_compatible_write(
4082 DictionaryArray::new(
4083 UInt8Array::from_iter_values(vec![0]),
4084 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4085 ),
4086 DictionaryArray::new(
4087 UInt8Array::from_iter_values(vec![0]),
4088 Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4089 ),
4090 DictionaryArray::new(
4091 UInt8Array::from_iter_values(vec![0, 1]),
4092 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4093 ),
4094 );
4095
4096 ensure_compatible_write(
4098 DictionaryArray::new(
4099 UInt8Array::from_iter_values(vec![0]),
4100 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4101 ),
4102 LargeStringArray::from_iter_values(vec!["barquet"]),
4103 DictionaryArray::new(
4104 UInt8Array::from_iter_values(vec![0, 1]),
4105 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4106 ),
4107 );
4108
4109 ensure_compatible_write(
4112 StringArray::from_iter_values(vec!["parquet"]),
4113 LargeStringArray::from_iter_values(vec!["barquet"]),
4114 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4115 );
4116
4117 ensure_compatible_write(
4118 LargeStringArray::from_iter_values(vec!["parquet"]),
4119 StringArray::from_iter_values(vec!["barquet"]),
4120 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4121 );
4122
4123 ensure_compatible_write(
4124 StringArray::from_iter_values(vec!["parquet"]),
4125 StringViewArray::from_iter_values(vec!["barquet"]),
4126 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4127 );
4128
4129 ensure_compatible_write(
4130 StringViewArray::from_iter_values(vec!["parquet"]),
4131 StringArray::from_iter_values(vec!["barquet"]),
4132 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4133 );
4134
4135 ensure_compatible_write(
4136 LargeStringArray::from_iter_values(vec!["parquet"]),
4137 StringViewArray::from_iter_values(vec!["barquet"]),
4138 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4139 );
4140
4141 ensure_compatible_write(
4142 StringViewArray::from_iter_values(vec!["parquet"]),
4143 LargeStringArray::from_iter_values(vec!["barquet"]),
4144 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4145 );
4146
4147 ensure_compatible_write(
4150 BinaryArray::from_iter_values(vec![b"parquet"]),
4151 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4152 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4153 );
4154
4155 ensure_compatible_write(
4156 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4157 BinaryArray::from_iter_values(vec![b"barquet"]),
4158 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4159 );
4160
4161 ensure_compatible_write(
4162 BinaryArray::from_iter_values(vec![b"parquet"]),
4163 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4164 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4165 );
4166
4167 ensure_compatible_write(
4168 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4169 BinaryArray::from_iter_values(vec![b"barquet"]),
4170 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4171 );
4172
4173 ensure_compatible_write(
4174 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4175 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4176 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4177 );
4178
4179 ensure_compatible_write(
4180 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4181 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4182 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4183 );
4184
4185 let list_field_metadata = HashMap::from_iter(vec![(
4188 PARQUET_FIELD_ID_META_KEY.to_string(),
4189 "1".to_string(),
4190 )]);
4191 let list_field = Field::new_list_field(DataType::Int32, false);
4192
4193 let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4194 let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4195
4196 let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4197 let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4198
4199 let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4200 let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4201
4202 ensure_compatible_write(
4203 ListArray::try_new(
4205 Arc::new(
4206 list_field
4207 .clone()
4208 .with_metadata(list_field_metadata.clone()),
4209 ),
4210 offsets1,
4211 values1,
4212 None,
4213 )
4214 .unwrap(),
4215 ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4217 ListArray::try_new(
4219 Arc::new(
4220 list_field
4221 .clone()
4222 .with_metadata(list_field_metadata.clone()),
4223 ),
4224 offsets_expected,
4225 values_expected,
4226 None,
4227 )
4228 .unwrap(),
4229 );
4230 }
4231
4232 #[test]
4233 fn arrow_writer_primitive_dictionary() {
4234 #[allow(deprecated)]
4236 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4237 "dictionary",
4238 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4239 true,
4240 42,
4241 true,
4242 )]));
4243
4244 let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4246 builder.append(12345678).unwrap();
4247 builder.append_null();
4248 builder.append(22345678).unwrap();
4249 builder.append(12345678).unwrap();
4250 let d = builder.finish();
4251
4252 one_column_roundtrip_with_schema(Arc::new(d), schema);
4253 }
4254
4255 #[test]
4256 fn arrow_writer_decimal32_dictionary() {
4257 let integers = vec![12345, 56789, 34567];
4258
4259 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4260
4261 let values = Decimal32Array::from(integers.clone())
4262 .with_precision_and_scale(5, 2)
4263 .unwrap();
4264
4265 let array = DictionaryArray::new(keys, Arc::new(values));
4266 one_column_roundtrip(Arc::new(array.clone()), true);
4267
4268 let values = Decimal32Array::from(integers)
4269 .with_precision_and_scale(9, 2)
4270 .unwrap();
4271
4272 let array = array.with_values(Arc::new(values));
4273 one_column_roundtrip(Arc::new(array), true);
4274 }
4275
4276 #[test]
4277 fn arrow_writer_decimal64_dictionary() {
4278 let integers = vec![12345, 56789, 34567];
4279
4280 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4281
4282 let values = Decimal64Array::from(integers.clone())
4283 .with_precision_and_scale(5, 2)
4284 .unwrap();
4285
4286 let array = DictionaryArray::new(keys, Arc::new(values));
4287 one_column_roundtrip(Arc::new(array.clone()), true);
4288
4289 let values = Decimal64Array::from(integers)
4290 .with_precision_and_scale(12, 2)
4291 .unwrap();
4292
4293 let array = array.with_values(Arc::new(values));
4294 one_column_roundtrip(Arc::new(array), true);
4295 }
4296
4297 #[test]
4298 fn arrow_writer_decimal128_dictionary() {
4299 let integers = vec![12345, 56789, 34567];
4300
4301 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4302
4303 let values = Decimal128Array::from(integers.clone())
4304 .with_precision_and_scale(5, 2)
4305 .unwrap();
4306
4307 let array = DictionaryArray::new(keys, Arc::new(values));
4308 one_column_roundtrip(Arc::new(array.clone()), true);
4309
4310 let values = Decimal128Array::from(integers)
4311 .with_precision_and_scale(12, 2)
4312 .unwrap();
4313
4314 let array = array.with_values(Arc::new(values));
4315 one_column_roundtrip(Arc::new(array), true);
4316 }
4317
4318 #[test]
4319 fn arrow_writer_decimal256_dictionary() {
4320 let integers = vec![
4321 i256::from_i128(12345),
4322 i256::from_i128(56789),
4323 i256::from_i128(34567),
4324 ];
4325
4326 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4327
4328 let values = Decimal256Array::from(integers.clone())
4329 .with_precision_and_scale(5, 2)
4330 .unwrap();
4331
4332 let array = DictionaryArray::new(keys, Arc::new(values));
4333 one_column_roundtrip(Arc::new(array.clone()), true);
4334
4335 let values = Decimal256Array::from(integers)
4336 .with_precision_and_scale(12, 2)
4337 .unwrap();
4338
4339 let array = array.with_values(Arc::new(values));
4340 one_column_roundtrip(Arc::new(array), true);
4341 }
4342
4343 #[test]
4344 fn arrow_writer_string_dictionary_unsigned_index() {
4345 #[allow(deprecated)]
4347 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4348 "dictionary",
4349 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4350 true,
4351 42,
4352 true,
4353 )]));
4354
4355 let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4357 .iter()
4358 .copied()
4359 .collect();
4360
4361 one_column_roundtrip_with_schema(Arc::new(d), schema);
4362 }
4363
4364 #[test]
4365 fn u32_min_max() {
4366 let src = [
4368 u32::MIN,
4369 1,
4370 (i32::MAX as u32) - 1,
4371 i32::MAX as u32,
4372 (i32::MAX as u32) + 1,
4373 u32::MAX - 1,
4374 u32::MAX,
4375 ];
4376 let values = Arc::new(UInt32Array::from_iter_values(src.iter().cloned()));
4377 let files = one_column_roundtrip(values, false);
4378
4379 for file in files {
4380 let reader = SerializedFileReader::new(file).unwrap();
4382 let metadata = reader.metadata();
4383
4384 let mut row_offset = 0;
4385 for row_group in metadata.row_groups() {
4386 assert_eq!(row_group.num_columns(), 1);
4387 let column = row_group.column(0);
4388
4389 let num_values = column.num_values() as usize;
4390 let src_slice = &src[row_offset..row_offset + num_values];
4391 row_offset += column.num_values() as usize;
4392
4393 let stats = column.statistics().unwrap();
4394 if let Statistics::Int32(stats) = stats {
4395 assert_eq!(
4396 *stats.min_opt().unwrap() as u32,
4397 *src_slice.iter().min().unwrap()
4398 );
4399 assert_eq!(
4400 *stats.max_opt().unwrap() as u32,
4401 *src_slice.iter().max().unwrap()
4402 );
4403 } else {
4404 panic!("Statistics::Int32 missing")
4405 }
4406 }
4407 }
4408 }
4409
4410 #[test]
4411 fn u64_min_max() {
4412 let src = [
4414 u64::MIN,
4415 1,
4416 (i64::MAX as u64) - 1,
4417 i64::MAX as u64,
4418 (i64::MAX as u64) + 1,
4419 u64::MAX - 1,
4420 u64::MAX,
4421 ];
4422 let values = Arc::new(UInt64Array::from_iter_values(src.iter().cloned()));
4423 let files = one_column_roundtrip(values, false);
4424
4425 for file in files {
4426 let reader = SerializedFileReader::new(file).unwrap();
4428 let metadata = reader.metadata();
4429
4430 let mut row_offset = 0;
4431 for row_group in metadata.row_groups() {
4432 assert_eq!(row_group.num_columns(), 1);
4433 let column = row_group.column(0);
4434
4435 let num_values = column.num_values() as usize;
4436 let src_slice = &src[row_offset..row_offset + num_values];
4437 row_offset += column.num_values() as usize;
4438
4439 let stats = column.statistics().unwrap();
4440 if let Statistics::Int64(stats) = stats {
4441 assert_eq!(
4442 *stats.min_opt().unwrap() as u64,
4443 *src_slice.iter().min().unwrap()
4444 );
4445 assert_eq!(
4446 *stats.max_opt().unwrap() as u64,
4447 *src_slice.iter().max().unwrap()
4448 );
4449 } else {
4450 panic!("Statistics::Int64 missing")
4451 }
4452 }
4453 }
4454 }
4455
4456 #[test]
4457 fn statistics_null_counts_only_nulls() {
4458 let values = Arc::new(UInt64Array::from(vec![None, None]));
4460 let files = one_column_roundtrip(values, true);
4461
4462 for file in files {
4463 let reader = SerializedFileReader::new(file).unwrap();
4465 let metadata = reader.metadata();
4466 assert_eq!(metadata.num_row_groups(), 1);
4467 let row_group = metadata.row_group(0);
4468 assert_eq!(row_group.num_columns(), 1);
4469 let column = row_group.column(0);
4470 let stats = column.statistics().unwrap();
4471 assert_eq!(stats.null_count_opt(), Some(2));
4472 }
4473 }
4474
4475 #[test]
4476 fn test_list_of_struct_roundtrip() {
4477 let int_field = Field::new("a", DataType::Int32, true);
4479 let int_field2 = Field::new("b", DataType::Int32, true);
4480
4481 let int_builder = Int32Builder::with_capacity(10);
4482 let int_builder2 = Int32Builder::with_capacity(10);
4483
4484 let struct_builder = StructBuilder::new(
4485 vec![int_field, int_field2],
4486 vec![Box::new(int_builder), Box::new(int_builder2)],
4487 );
4488 let mut list_builder = ListBuilder::new(struct_builder);
4489
4490 let values = list_builder.values();
4495 values
4496 .field_builder::<Int32Builder>(0)
4497 .unwrap()
4498 .append_value(1);
4499 values
4500 .field_builder::<Int32Builder>(1)
4501 .unwrap()
4502 .append_value(2);
4503 values.append(true);
4504 list_builder.append(true);
4505
4506 list_builder.append(true);
4508
4509 list_builder.append(false);
4511
4512 let values = list_builder.values();
4514 values
4515 .field_builder::<Int32Builder>(0)
4516 .unwrap()
4517 .append_null();
4518 values
4519 .field_builder::<Int32Builder>(1)
4520 .unwrap()
4521 .append_null();
4522 values.append(false);
4523 values
4524 .field_builder::<Int32Builder>(0)
4525 .unwrap()
4526 .append_null();
4527 values
4528 .field_builder::<Int32Builder>(1)
4529 .unwrap()
4530 .append_null();
4531 values.append(false);
4532 list_builder.append(true);
4533
4534 let values = list_builder.values();
4536 values
4537 .field_builder::<Int32Builder>(0)
4538 .unwrap()
4539 .append_null();
4540 values
4541 .field_builder::<Int32Builder>(1)
4542 .unwrap()
4543 .append_value(3);
4544 values.append(true);
4545 list_builder.append(true);
4546
4547 let values = list_builder.values();
4549 values
4550 .field_builder::<Int32Builder>(0)
4551 .unwrap()
4552 .append_value(2);
4553 values
4554 .field_builder::<Int32Builder>(1)
4555 .unwrap()
4556 .append_null();
4557 values.append(true);
4558 list_builder.append(true);
4559
4560 let array = Arc::new(list_builder.finish());
4561
4562 one_column_roundtrip(array, true);
4563 }
4564
4565 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4566 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4567 }
4568
4569 #[test]
4570 fn test_aggregates_records() {
4571 let arrays = [
4572 Int32Array::from((0..100).collect::<Vec<_>>()),
4573 Int32Array::from((0..50).collect::<Vec<_>>()),
4574 Int32Array::from((200..500).collect::<Vec<_>>()),
4575 ];
4576
4577 let schema = Arc::new(Schema::new(vec![Field::new(
4578 "int",
4579 ArrowDataType::Int32,
4580 false,
4581 )]));
4582
4583 let file = tempfile::tempfile().unwrap();
4584
4585 let props = WriterProperties::builder()
4586 .set_max_row_group_row_count(Some(200))
4587 .build();
4588
4589 let mut writer =
4590 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4591
4592 for array in arrays {
4593 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4594 writer.write(&batch).unwrap();
4595 }
4596
4597 writer.close().unwrap();
4598
4599 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4600 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4601
4602 let batches = builder
4603 .with_batch_size(100)
4604 .build()
4605 .unwrap()
4606 .collect::<ArrowResult<Vec<_>>>()
4607 .unwrap();
4608
4609 assert_eq!(batches.len(), 5);
4610 assert!(batches.iter().all(|x| x.num_columns() == 1));
4611
4612 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4613
4614 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4615
4616 let values: Vec<_> = batches
4617 .iter()
4618 .flat_map(|x| {
4619 x.column(0)
4620 .as_any()
4621 .downcast_ref::<Int32Array>()
4622 .unwrap()
4623 .values()
4624 .iter()
4625 .cloned()
4626 })
4627 .collect();
4628
4629 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4630 assert_eq!(&values, &expected_values)
4631 }
4632
4633 #[test]
4634 fn complex_aggregate() {
4635 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4637 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4638 let struct_a = Arc::new(Field::new(
4639 "struct_a",
4640 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4641 true,
4642 ));
4643
4644 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4645 let struct_b = Arc::new(Field::new(
4646 "struct_b",
4647 DataType::Struct(vec![list_a.clone()].into()),
4648 false,
4649 ));
4650
4651 let schema = Arc::new(Schema::new(vec![struct_b]));
4652
4653 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4655 let field_b_array =
4656 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4657
4658 let struct_a_array = StructArray::from(vec![
4659 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4660 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4661 ]);
4662
4663 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4664 .len(5)
4665 .add_buffer(Buffer::from_iter(vec![
4666 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4667 ]))
4668 .null_bit_buffer(Some(Buffer::from_iter(vec![
4669 true, false, true, false, true,
4670 ])))
4671 .child_data(vec![struct_a_array.into_data()])
4672 .build()
4673 .unwrap();
4674
4675 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4676 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4677
4678 let batch1 =
4679 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4680 .unwrap();
4681
4682 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4683 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4684
4685 let struct_a_array = StructArray::from(vec![
4686 (field_a, Arc::new(field_a_array) as ArrayRef),
4687 (field_b, Arc::new(field_b_array) as ArrayRef),
4688 ]);
4689
4690 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4691 .len(2)
4692 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4693 .child_data(vec![struct_a_array.into_data()])
4694 .build()
4695 .unwrap();
4696
4697 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4698 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4699
4700 let batch2 =
4701 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4702 .unwrap();
4703
4704 let batches = &[batch1, batch2];
4705
4706 let expected = r#"
4709 +-------------------------------------------------------------------------------------------------------+
4710 | struct_b |
4711 +-------------------------------------------------------------------------------------------------------+
4712 | {list: [{leaf_a: 1, leaf_b: 1}]} |
4713 | {list: } |
4714 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
4715 | {list: } |
4716 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
4717 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4718 | {list: [{leaf_a: 10, leaf_b: }]} |
4719 +-------------------------------------------------------------------------------------------------------+
4720 "#.trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4721
4722 let actual = pretty_format_batches(batches).unwrap().to_string();
4723 assert_eq!(actual, expected);
4724
4725 let file = tempfile::tempfile().unwrap();
4727 let props = WriterProperties::builder()
4728 .set_max_row_group_row_count(Some(6))
4729 .build();
4730
4731 let mut writer =
4732 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4733
4734 for batch in batches {
4735 writer.write(batch).unwrap();
4736 }
4737 writer.close().unwrap();
4738
4739 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4744 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4745
4746 let batches = builder
4747 .with_batch_size(2)
4748 .build()
4749 .unwrap()
4750 .collect::<ArrowResult<Vec<_>>>()
4751 .unwrap();
4752
4753 assert_eq!(batches.len(), 4);
4754 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4755 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4756
4757 let actual = pretty_format_batches(&batches).unwrap().to_string();
4758 assert_eq!(actual, expected);
4759 }
4760
4761 #[test]
4762 fn test_arrow_writer_metadata() {
4763 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4764 let file_schema = batch_schema.clone().with_metadata(
4765 vec![("foo".to_string(), "bar".to_string())]
4766 .into_iter()
4767 .collect(),
4768 );
4769
4770 let batch = RecordBatch::try_new(
4771 Arc::new(batch_schema),
4772 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4773 )
4774 .unwrap();
4775
4776 let mut buf = Vec::with_capacity(1024);
4777 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4778 writer.write(&batch).unwrap();
4779 writer.close().unwrap();
4780 }
4781
4782 #[test]
4783 fn test_arrow_writer_nullable() {
4784 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4785 let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
4786 let file_schema = Arc::new(file_schema);
4787
4788 let batch = RecordBatch::try_new(
4789 Arc::new(batch_schema),
4790 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4791 )
4792 .unwrap();
4793
4794 let mut buf = Vec::with_capacity(1024);
4795 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
4796 writer.write(&batch).unwrap();
4797 writer.close().unwrap();
4798
4799 let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
4800 let back = read.next().unwrap().unwrap();
4801 assert_eq!(back.schema(), file_schema);
4802 assert_ne!(back.schema(), batch.schema());
4803 assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
4804 }
4805
4806 #[test]
4807 fn in_progress_accounting() {
4808 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4810
4811 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4813
4814 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4816
4817 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4818
4819 assert_eq!(writer.in_progress_size(), 0);
4821 assert_eq!(writer.in_progress_rows(), 0);
4822 assert_eq!(writer.memory_size(), 0);
4823 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
4825
4826 let initial_size = writer.in_progress_size();
4828 assert!(initial_size > 0);
4829 assert_eq!(writer.in_progress_rows(), 5);
4830 let initial_memory = writer.memory_size();
4831 assert!(initial_memory > 0);
4832 assert!(
4834 initial_size <= initial_memory,
4835 "{initial_size} <= {initial_memory}"
4836 );
4837
4838 writer.write(&batch).unwrap();
4840 assert!(writer.in_progress_size() > initial_size);
4841 assert_eq!(writer.in_progress_rows(), 10);
4842 assert!(writer.memory_size() > initial_memory);
4843 assert!(
4844 writer.in_progress_size() <= writer.memory_size(),
4845 "in_progress_size {} <= memory_size {}",
4846 writer.in_progress_size(),
4847 writer.memory_size()
4848 );
4849
4850 let pre_flush_bytes_written = writer.bytes_written();
4852 writer.flush().unwrap();
4853 assert_eq!(writer.in_progress_size(), 0);
4854 assert_eq!(writer.memory_size(), 0);
4855 assert!(writer.bytes_written() > pre_flush_bytes_written);
4856
4857 writer.close().unwrap();
4858 }
4859
4860 #[test]
4861 fn test_writer_all_null() {
4862 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4863 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4864 let batch = RecordBatch::try_from_iter(vec![
4865 ("a", Arc::new(a) as ArrayRef),
4866 ("b", Arc::new(b) as ArrayRef),
4867 ])
4868 .unwrap();
4869
4870 let mut buf = Vec::with_capacity(1024);
4871 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4872 writer.write(&batch).unwrap();
4873 writer.close().unwrap();
4874
4875 let bytes = Bytes::from(buf);
4876 let options = ReadOptionsBuilder::new().with_page_index().build();
4877 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4878 let index = reader.metadata().offset_index().unwrap();
4879
4880 assert_eq!(index.len(), 1);
4881 assert_eq!(index[0].len(), 2); assert_eq!(index[0][0].page_locations().len(), 1); assert_eq!(index[0][1].page_locations().len(), 1); }
4885
4886 #[test]
4887 fn test_disabled_statistics_with_page() {
4888 let file_schema = Schema::new(vec![
4889 Field::new("a", DataType::Utf8, true),
4890 Field::new("b", DataType::Utf8, true),
4891 ]);
4892 let file_schema = Arc::new(file_schema);
4893
4894 let batch = RecordBatch::try_new(
4895 file_schema.clone(),
4896 vec![
4897 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
4898 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
4899 ],
4900 )
4901 .unwrap();
4902
4903 let props = WriterProperties::builder()
4904 .set_statistics_enabled(EnabledStatistics::None)
4905 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
4906 .build();
4907
4908 let mut buf = Vec::with_capacity(1024);
4909 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
4910 writer.write(&batch).unwrap();
4911
4912 let metadata = writer.close().unwrap();
4913 assert_eq!(metadata.num_row_groups(), 1);
4914 let row_group = metadata.row_group(0);
4915 assert_eq!(row_group.num_columns(), 2);
4916 assert!(row_group.column(0).offset_index_offset().is_some());
4918 assert!(row_group.column(0).column_index_offset().is_some());
4919 assert!(row_group.column(1).offset_index_offset().is_some());
4921 assert!(row_group.column(1).column_index_offset().is_none());
4922
4923 let options = ReadOptionsBuilder::new().with_page_index().build();
4924 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
4925
4926 let row_group = reader.get_row_group(0).unwrap();
4927 let a_col = row_group.metadata().column(0);
4928 let b_col = row_group.metadata().column(1);
4929
4930 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
4932 let min = byte_array_stats.min_opt().unwrap();
4933 let max = byte_array_stats.max_opt().unwrap();
4934
4935 assert_eq!(min.as_bytes(), b"a");
4936 assert_eq!(max.as_bytes(), b"d");
4937 } else {
4938 panic!("expecting Statistics::ByteArray");
4939 }
4940
4941 assert!(b_col.statistics().is_none());
4943
4944 let offset_index = reader.metadata().offset_index().unwrap();
4945 assert_eq!(offset_index.len(), 1); assert_eq!(offset_index[0].len(), 2); let column_index = reader.metadata().column_index().unwrap();
4949 assert_eq!(column_index.len(), 1); assert_eq!(column_index[0].len(), 2); let a_idx = &column_index[0][0];
4953 assert!(
4954 matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)),
4955 "{a_idx:?}"
4956 );
4957 let b_idx = &column_index[0][1];
4958 assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
4959 }
4960
4961 #[test]
4962 fn test_disabled_statistics_with_chunk() {
4963 let file_schema = Schema::new(vec![
4964 Field::new("a", DataType::Utf8, true),
4965 Field::new("b", DataType::Utf8, true),
4966 ]);
4967 let file_schema = Arc::new(file_schema);
4968
4969 let batch = RecordBatch::try_new(
4970 file_schema.clone(),
4971 vec![
4972 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
4973 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
4974 ],
4975 )
4976 .unwrap();
4977
4978 let props = WriterProperties::builder()
4979 .set_statistics_enabled(EnabledStatistics::None)
4980 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
4981 .build();
4982
4983 let mut buf = Vec::with_capacity(1024);
4984 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
4985 writer.write(&batch).unwrap();
4986
4987 let metadata = writer.close().unwrap();
4988 assert_eq!(metadata.num_row_groups(), 1);
4989 let row_group = metadata.row_group(0);
4990 assert_eq!(row_group.num_columns(), 2);
4991 assert!(row_group.column(0).offset_index_offset().is_some());
4993 assert!(row_group.column(0).column_index_offset().is_none());
4994 assert!(row_group.column(1).offset_index_offset().is_some());
4996 assert!(row_group.column(1).column_index_offset().is_none());
4997
4998 let options = ReadOptionsBuilder::new().with_page_index().build();
4999 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5000
5001 let row_group = reader.get_row_group(0).unwrap();
5002 let a_col = row_group.metadata().column(0);
5003 let b_col = row_group.metadata().column(1);
5004
5005 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5007 let min = byte_array_stats.min_opt().unwrap();
5008 let max = byte_array_stats.max_opt().unwrap();
5009
5010 assert_eq!(min.as_bytes(), b"a");
5011 assert_eq!(max.as_bytes(), b"d");
5012 } else {
5013 panic!("expecting Statistics::ByteArray");
5014 }
5015
5016 assert!(b_col.statistics().is_none());
5018
5019 let column_index = reader.metadata().column_index().unwrap();
5020 assert_eq!(column_index.len(), 1); assert_eq!(column_index[0].len(), 2); let a_idx = &column_index[0][0];
5024 assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}");
5025 let b_idx = &column_index[0][1];
5026 assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5027 }
5028
5029 #[test]
5030 fn test_arrow_writer_skip_metadata() {
5031 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5032 let file_schema = Arc::new(batch_schema.clone());
5033
5034 let batch = RecordBatch::try_new(
5035 Arc::new(batch_schema),
5036 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5037 )
5038 .unwrap();
5039 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5040
5041 let mut buf = Vec::with_capacity(1024);
5042 let mut writer =
5043 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5044 writer.write(&batch).unwrap();
5045 writer.close().unwrap();
5046
5047 let bytes = Bytes::from(buf);
5048 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5049 assert_eq!(file_schema, *reader_builder.schema());
5050 if let Some(key_value_metadata) = reader_builder
5051 .metadata()
5052 .file_metadata()
5053 .key_value_metadata()
5054 {
5055 assert!(
5056 !key_value_metadata
5057 .iter()
5058 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5059 );
5060 }
5061 }
5062
5063 #[test]
5064 fn test_arrow_writer_skip_path_in_schema() {
5065 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5066 let file_schema = Arc::new(batch_schema.clone());
5067
5068 let batch = RecordBatch::try_new(
5069 Arc::new(batch_schema),
5070 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5071 )
5072 .unwrap();
5073
5074 let skip_options = ArrowWriterOptions::new();
5076
5077 let mut buf = Vec::with_capacity(1024);
5078 let mut writer =
5079 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5080 writer.write(&batch).unwrap();
5081 writer.close().unwrap();
5082
5083 let skip_options = ArrowWriterOptions::new().with_properties(
5085 WriterProperties::builder()
5086 .set_write_path_in_schema(false)
5087 .build(),
5088 );
5089
5090 let mut buf2 = Vec::with_capacity(1024);
5091 let mut writer =
5092 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5093 .unwrap();
5094 writer.write(&batch).unwrap();
5095 writer.close().unwrap();
5096
5097 assert!(buf.len() > buf2.len());
5099 }
5100
5101 #[test]
5102 fn mismatched_schemas() {
5103 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5104 let file_schema = Arc::new(Schema::new(vec![Field::new(
5105 "temperature",
5106 DataType::Float64,
5107 false,
5108 )]));
5109
5110 let batch = RecordBatch::try_new(
5111 Arc::new(batch_schema),
5112 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5113 )
5114 .unwrap();
5115
5116 let mut buf = Vec::with_capacity(1024);
5117 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5118
5119 let err = writer.write(&batch).unwrap_err().to_string();
5120 assert_eq!(
5121 err,
5122 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5123 );
5124 }
5125
5126 #[test]
5127 fn test_roundtrip_empty_schema() {
5129 let empty_batch = RecordBatch::try_new_with_options(
5131 Arc::new(Schema::empty()),
5132 vec![],
5133 &RecordBatchOptions::default().with_row_count(Some(0)),
5134 )
5135 .unwrap();
5136
5137 let mut parquet_bytes: Vec<u8> = Vec::new();
5139 let mut writer =
5140 ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5141 writer.write(&empty_batch).unwrap();
5142 writer.close().unwrap();
5143
5144 let bytes = Bytes::from(parquet_bytes);
5146 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5147 assert_eq!(reader.schema(), &empty_batch.schema());
5148 let batches: Vec<_> = reader
5149 .build()
5150 .unwrap()
5151 .collect::<ArrowResult<Vec<_>>>()
5152 .unwrap();
5153 assert_eq!(batches.len(), 0);
5154 }
5155
5156 #[test]
5157 fn test_page_stats_not_written_by_default() {
5158 let string_field = Field::new("a", DataType::Utf8, false);
5159 let schema = Schema::new(vec![string_field]);
5160 let raw_string_values = vec!["Blart Versenwald III"];
5161 let string_values = StringArray::from(raw_string_values.clone());
5162 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5163
5164 let props = WriterProperties::builder()
5165 .set_statistics_enabled(EnabledStatistics::Page)
5166 .set_dictionary_enabled(false)
5167 .set_encoding(Encoding::PLAIN)
5168 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5169 .build();
5170
5171 let file = roundtrip_opts(&batch, props);
5172
5173 let first_page = &file[4..];
5178 let mut prot = ThriftSliceInputProtocol::new(first_page);
5179 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5180 let stats = hdr.data_page_header.unwrap().statistics;
5181
5182 assert!(stats.is_none());
5183 }
5184
5185 #[test]
5186 fn test_page_stats_when_enabled() {
5187 let string_field = Field::new("a", DataType::Utf8, false);
5188 let schema = Schema::new(vec![string_field]);
5189 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5190 let string_values = StringArray::from(raw_string_values.clone());
5191 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5192
5193 let props = WriterProperties::builder()
5194 .set_statistics_enabled(EnabledStatistics::Page)
5195 .set_dictionary_enabled(false)
5196 .set_encoding(Encoding::PLAIN)
5197 .set_write_page_header_statistics(true)
5198 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5199 .build();
5200
5201 let file = roundtrip_opts(&batch, props);
5202
5203 let first_page = &file[4..];
5208 let mut prot = ThriftSliceInputProtocol::new(first_page);
5209 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5210 let stats = hdr.data_page_header.unwrap().statistics;
5211
5212 let stats = stats.unwrap();
5213 assert!(stats.is_max_value_exact.unwrap());
5215 assert!(stats.is_min_value_exact.unwrap());
5216 assert_eq!(stats.max_value.unwrap(), "Blart Versenwald III".as_bytes());
5217 assert_eq!(stats.min_value.unwrap(), "Andrew Lamb".as_bytes());
5218 }
5219
5220 #[test]
5221 fn test_page_stats_truncation() {
5222 let string_field = Field::new("a", DataType::Utf8, false);
5223 let binary_field = Field::new("b", DataType::Binary, false);
5224 let schema = Schema::new(vec![string_field, binary_field]);
5225
5226 let raw_string_values = vec!["Blart Versenwald III"];
5227 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5228 let raw_binary_value_refs = raw_binary_values
5229 .iter()
5230 .map(|x| x.as_slice())
5231 .collect::<Vec<_>>();
5232
5233 let string_values = StringArray::from(raw_string_values.clone());
5234 let binary_values = BinaryArray::from(raw_binary_value_refs);
5235 let batch = RecordBatch::try_new(
5236 Arc::new(schema),
5237 vec![Arc::new(string_values), Arc::new(binary_values)],
5238 )
5239 .unwrap();
5240
5241 let props = WriterProperties::builder()
5242 .set_statistics_truncate_length(Some(2))
5243 .set_dictionary_enabled(false)
5244 .set_encoding(Encoding::PLAIN)
5245 .set_write_page_header_statistics(true)
5246 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5247 .build();
5248
5249 let file = roundtrip_opts(&batch, props);
5250
5251 let first_page = &file[4..];
5256 let mut prot = ThriftSliceInputProtocol::new(first_page);
5257 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5258 let stats = hdr.data_page_header.unwrap().statistics;
5259 assert!(stats.is_some());
5260 let stats = stats.unwrap();
5261 assert!(!stats.is_max_value_exact.unwrap());
5263 assert!(!stats.is_min_value_exact.unwrap());
5264 assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5265 assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5266
5267 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5269 let mut prot = ThriftSliceInputProtocol::new(second_page);
5270 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5271 let stats = hdr.data_page_header.unwrap().statistics;
5272 assert!(stats.is_some());
5273 let stats = stats.unwrap();
5274 assert!(!stats.is_max_value_exact.unwrap());
5276 assert!(!stats.is_min_value_exact.unwrap());
5277 assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5278 assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5279 }
5280
5281 #[test]
5282 fn test_page_encoding_statistics_roundtrip() {
5283 let batch_schema = Schema::new(vec![Field::new(
5284 "int32",
5285 arrow_schema::DataType::Int32,
5286 false,
5287 )]);
5288
5289 let batch = RecordBatch::try_new(
5290 Arc::new(batch_schema.clone()),
5291 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5292 )
5293 .unwrap();
5294
5295 let mut file: File = tempfile::tempfile().unwrap();
5296 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5297 writer.write(&batch).unwrap();
5298 let file_metadata = writer.close().unwrap();
5299
5300 assert_eq!(file_metadata.num_row_groups(), 1);
5301 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5302 assert!(
5303 file_metadata
5304 .row_group(0)
5305 .column(0)
5306 .page_encoding_stats()
5307 .is_some()
5308 );
5309 let chunk_page_stats = file_metadata
5310 .row_group(0)
5311 .column(0)
5312 .page_encoding_stats()
5313 .unwrap();
5314
5315 let options = ReadOptionsBuilder::new()
5317 .with_page_index()
5318 .with_encoding_stats_as_mask(false)
5319 .build();
5320 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5321
5322 let rowgroup = reader.get_row_group(0).expect("row group missing");
5323 assert_eq!(rowgroup.num_columns(), 1);
5324 let column = rowgroup.metadata().column(0);
5325 assert!(column.page_encoding_stats().is_some());
5326 let file_page_stats = column.page_encoding_stats().unwrap();
5327 assert_eq!(chunk_page_stats, file_page_stats);
5328 }
5329
5330 #[test]
5331 fn test_different_dict_page_size_limit() {
5332 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5333 let schema = Arc::new(Schema::new(vec![
5334 Field::new("col0", arrow_schema::DataType::Int64, false),
5335 Field::new("col1", arrow_schema::DataType::Int64, false),
5336 ]));
5337 let batch =
5338 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5339
5340 let props = WriterProperties::builder()
5341 .set_dictionary_page_size_limit(1024 * 1024)
5342 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5343 .build();
5344 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5345 writer.write(&batch).unwrap();
5346 let data = Bytes::from(writer.into_inner().unwrap());
5347
5348 let mut metadata = ParquetMetaDataReader::new();
5349 metadata.try_parse(&data).unwrap();
5350 let metadata = metadata.finish().unwrap();
5351 let col0_meta = metadata.row_group(0).column(0);
5352 let col1_meta = metadata.row_group(0).column(1);
5353
5354 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5355 let mut reader =
5356 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5357 let page = reader.get_next_page().unwrap().unwrap();
5358 match page {
5359 Page::DictionaryPage { buf, .. } => buf.len(),
5360 _ => panic!("expected DictionaryPage"),
5361 }
5362 };
5363
5364 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5365 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5366 }
5367
5368 #[test]
5369 fn test_arrow_writer_granular_mode_roundtrip() {
5370 let small = "tiny".to_string();
5379 let big = "x".repeat(64 * 1024);
5380 let strings: Vec<String> = (0..256)
5381 .map(|i| {
5382 if i % 16 == 0 {
5383 big.clone()
5384 } else {
5385 small.clone()
5386 }
5387 })
5388 .collect();
5389
5390 let schema = Arc::new(Schema::new(vec![Field::new(
5391 "col",
5392 ArrowDataType::Utf8,
5393 false,
5394 )]));
5395 let batch = RecordBatch::try_new(
5396 schema.clone(),
5397 vec![Arc::new(StringArray::from(strings.clone())) as _],
5398 )
5399 .unwrap();
5400
5401 let props = WriterProperties::builder()
5402 .set_dictionary_enabled(false)
5403 .set_data_page_size_limit(16 * 1024)
5404 .build();
5405 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5406 writer.write(&batch).unwrap();
5407 let data = Bytes::from(writer.into_inner().unwrap());
5408
5409 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5410 let read = reader.next().unwrap().unwrap();
5411 assert!(reader.next().is_none(), "expected one batch");
5412 let col = read
5413 .column(0)
5414 .as_any()
5415 .downcast_ref::<StringArray>()
5416 .unwrap();
5417 assert_eq!(col.len(), strings.len());
5418 for (i, expected) in strings.iter().enumerate() {
5419 assert_eq!(
5420 col.value(i),
5421 expected.as_str(),
5422 "value mismatch at index {i}"
5423 );
5424 }
5425 }
5426
5427 #[test]
5428 fn test_arrow_writer_all_null_string_column() {
5429 let num_rows = 1024;
5434 let schema = Arc::new(Schema::new(vec![Field::new(
5435 "col",
5436 ArrowDataType::Utf8,
5437 true,
5438 )]));
5439 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5440 let batch = RecordBatch::try_new(
5441 schema.clone(),
5442 vec![Arc::new(StringArray::from(nulls)) as _],
5443 )
5444 .unwrap();
5445
5446 let props = WriterProperties::builder()
5447 .set_dictionary_enabled(false)
5448 .set_data_page_size_limit(16 * 1024)
5449 .build();
5450 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5451 writer.write(&batch).unwrap();
5452 let data = Bytes::from(writer.into_inner().unwrap());
5453
5454 let mut metadata = ParquetMetaDataReader::new();
5457 metadata.try_parse(&data).unwrap();
5458 let metadata = metadata.finish().unwrap();
5459 let row_group = metadata.row_group(0);
5460 let col_meta = row_group.column(0);
5461 assert_eq!(row_group.num_rows() as usize, num_rows);
5462 if let Some(stats) = col_meta.statistics() {
5465 assert_eq!(
5466 stats.null_count_opt().unwrap_or(0) as usize,
5467 num_rows,
5468 "expected all-null column to report null_count = num_rows"
5469 );
5470 }
5471
5472 let mut reader =
5473 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5474 let mut total_values = 0u32;
5475 while let Some(page) = reader.get_next_page().unwrap() {
5476 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5477 total_values += page.num_values();
5478 }
5479 }
5480 assert_eq!(
5481 total_values as usize, num_rows,
5482 "expected every level position to be represented in some page"
5483 );
5484 }
5485
5486 struct WriteBatchesShape {
5487 num_batches: usize,
5488 rows_per_batch: usize,
5489 row_size: usize,
5490 }
5491
5492 fn write_batches(
5494 WriteBatchesShape {
5495 num_batches,
5496 rows_per_batch,
5497 row_size,
5498 }: WriteBatchesShape,
5499 props: WriterProperties,
5500 ) -> ParquetRecordBatchReaderBuilder<File> {
5501 let schema = Arc::new(Schema::new(vec![Field::new(
5502 "str",
5503 ArrowDataType::Utf8,
5504 false,
5505 )]));
5506 let file = tempfile::tempfile().unwrap();
5507 let mut writer =
5508 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5509
5510 for batch_idx in 0..num_batches {
5511 let strings: Vec<String> = (0..rows_per_batch)
5512 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5513 .collect();
5514 let array = StringArray::from(strings);
5515 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5516 writer.write(&batch).unwrap();
5517 }
5518 writer.close().unwrap();
5519 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5520 }
5521
5522 #[test]
5523 fn test_row_group_limit_none_writes_single_row_group() {
5525 let props = WriterProperties::builder()
5526 .set_max_row_group_row_count(None)
5527 .set_max_row_group_bytes(None)
5528 .build();
5529
5530 let builder = write_batches(
5531 WriteBatchesShape {
5532 num_batches: 1,
5533 rows_per_batch: 1000,
5534 row_size: 4,
5535 },
5536 props,
5537 );
5538
5539 assert_eq!(
5540 &row_group_sizes(builder.metadata()),
5541 &[1000],
5542 "With no limits, all rows should be in a single row group"
5543 );
5544 }
5545
5546 #[test]
5547 fn test_row_group_limit_rows_only() {
5549 let props = WriterProperties::builder()
5550 .set_max_row_group_row_count(Some(300))
5551 .set_max_row_group_bytes(None)
5552 .build();
5553
5554 let builder = write_batches(
5555 WriteBatchesShape {
5556 num_batches: 1,
5557 rows_per_batch: 1000,
5558 row_size: 4,
5559 },
5560 props,
5561 );
5562
5563 assert_eq!(
5564 &row_group_sizes(builder.metadata()),
5565 &[300, 300, 300, 100],
5566 "Row groups should be split by row count"
5567 );
5568 }
5569
5570 #[test]
5571 fn test_row_group_limit_bytes_only() {
5573 let props = WriterProperties::builder()
5574 .set_max_row_group_row_count(None)
5575 .set_max_row_group_bytes(Some(3500))
5577 .build();
5578
5579 let builder = write_batches(
5580 WriteBatchesShape {
5581 num_batches: 10,
5582 rows_per_batch: 10,
5583 row_size: 100,
5584 },
5585 props,
5586 );
5587
5588 let sizes = row_group_sizes(builder.metadata());
5589
5590 assert!(
5591 sizes.len() > 1,
5592 "Should have multiple row groups due to byte limit, got {sizes:?}",
5593 );
5594
5595 let total_rows: i64 = sizes.iter().sum();
5596 assert_eq!(total_rows, 100, "Total rows should be preserved");
5597 }
5598
5599 #[test]
5600 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5602 let schema = Arc::new(Schema::new(vec![Field::new(
5603 "str",
5604 ArrowDataType::Utf8,
5605 false,
5606 )]));
5607 let file = tempfile::tempfile().unwrap();
5608
5609 let props = WriterProperties::builder()
5611 .set_max_row_group_row_count(None)
5612 .set_max_row_group_bytes(None)
5613 .build();
5614 let mut writer =
5615 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5616
5617 let first_array = StringArray::from(
5618 (0..10)
5619 .map(|i| format!("{:0>100}", i))
5620 .collect::<Vec<String>>(),
5621 );
5622 let first_batch =
5623 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5624 writer.write(&first_batch).unwrap();
5625 assert_eq!(writer.in_progress_rows(), 10);
5626
5627 writer.max_row_group_bytes = Some(1);
5630
5631 let second_array = StringArray::from(vec!["x".to_string()]);
5632 let second_batch =
5633 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5634 writer.write(&second_batch).unwrap();
5635 writer.close().unwrap();
5636 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5637
5638 assert_eq!(
5639 &row_group_sizes(builder.metadata()),
5640 &[10, 1],
5641 "The second write should flush an oversized in-progress row group first",
5642 );
5643 }
5644
5645 #[test]
5646 fn test_row_group_limit_both_row_wins_single_batch() {
5648 let props = WriterProperties::builder()
5649 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
5652
5653 let builder = write_batches(
5654 WriteBatchesShape {
5655 num_batches: 1,
5656 row_size: 4,
5657 rows_per_batch: 1000,
5658 },
5659 props,
5660 );
5661
5662 assert_eq!(
5663 &row_group_sizes(builder.metadata()),
5664 &[200, 200, 200, 200, 200],
5665 "Row limit should trigger before byte limit"
5666 );
5667 }
5668
5669 #[test]
5670 fn test_row_group_limit_both_row_wins_multiple_batches() {
5672 let props = WriterProperties::builder()
5673 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
5676
5677 let builder = write_batches(
5678 WriteBatchesShape {
5679 num_batches: 10,
5680 rows_per_batch: 10,
5681 row_size: 100,
5682 },
5683 props,
5684 );
5685
5686 assert_eq!(
5687 &row_group_sizes(builder.metadata()),
5688 &[5; 20],
5689 "Row limit should trigger before byte limit"
5690 );
5691 }
5692
5693 #[test]
5694 fn test_row_group_limit_both_bytes_wins() {
5696 let props = WriterProperties::builder()
5697 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
5700
5701 let builder = write_batches(
5702 WriteBatchesShape {
5703 num_batches: 10,
5704 rows_per_batch: 10,
5705 row_size: 100,
5706 },
5707 props,
5708 );
5709
5710 let sizes = row_group_sizes(builder.metadata());
5711
5712 assert!(
5713 sizes.len() > 1,
5714 "Byte limit should trigger before row limit, got {sizes:?}",
5715 );
5716
5717 assert!(
5718 sizes.iter().all(|&s| s < 1000),
5719 "No row group should hit the row limit"
5720 );
5721
5722 let total_rows: i64 = sizes.iter().sum();
5723 assert_eq!(total_rows, 100, "Total rows should be preserved");
5724 }
5725
5726 #[test]
5727 fn arrow_column_chunk_close_mut_drops_column_index() {
5728 use crate::arrow::ArrowSchemaConverter;
5729 use crate::file::writer::SerializedFileWriter;
5730
5731 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5732 let props = Arc::new(
5733 WriterProperties::builder()
5734 .set_statistics_enabled(EnabledStatistics::Page)
5735 .build(),
5736 );
5737 let parquet_schema = ArrowSchemaConverter::new()
5738 .with_coerce_types(props.coerce_types())
5739 .convert(&schema)
5740 .unwrap();
5741
5742 let mut buf = Vec::with_capacity(1024);
5743 let mut writer =
5744 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5745 .unwrap();
5746
5747 let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
5748 let mut col_writers = factory.create_column_writers(0).unwrap();
5749 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5750 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5751 col_writers[0].write(&leaves).unwrap();
5752 }
5753 let mut chunk = col_writers.pop().unwrap().close().unwrap();
5754
5755 assert!(
5757 chunk.close().column_index.is_some(),
5758 "EnabledStatistics::Page should produce a column_index"
5759 );
5760
5761 chunk.close_mut().column_index = None;
5763 assert!(chunk.close().column_index.is_none());
5764
5765 let mut rg = writer.next_row_group().unwrap();
5766 chunk.append_to_row_group(&mut rg).unwrap();
5767 rg.close().unwrap();
5768 let file_meta = writer.close().unwrap();
5769
5770 let cc = file_meta.row_group(0).column(0);
5773 assert!(cc.column_index_range().is_none());
5774 }
5775
5776 fn write_column_to_bytes(array: ArrayRef) -> Bytes {
5778 let schema = Arc::new(Schema::new(vec![Field::new(
5779 "col",
5780 array.data_type().clone(),
5781 true,
5782 )]));
5783 let buf = get_bytes_after_close(
5784 schema.clone(),
5785 &RecordBatch::try_new(schema, vec![array]).unwrap(),
5786 );
5787 Bytes::from(buf)
5788 }
5789
5790 fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
5794 let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
5795 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
5796 .unwrap()
5797 .build()
5798 .unwrap()
5799 .next()
5800 .unwrap()
5801 .unwrap()
5802 .column(0)
5803 .clone()
5804 }
5805
5806 fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
5807 let flat_schema = Arc::new(Schema::new(vec![Field::new(
5808 "col",
5809 flat.data_type().clone(),
5810 true,
5811 )]));
5812 let ree_bytes = write_column_to_bytes(ree);
5813 let flat_bytes = write_column_to_bytes(flat.clone());
5814 assert_eq!(
5815 ree_bytes, flat_bytes,
5816 "REE and flat bytes should be identical"
5817 );
5818
5819 let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
5820 let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
5821
5822 assert_eq!(decoded_ree.as_ref(), flat.as_ref());
5823 assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
5824 }
5825
5826 #[test]
5827 fn ree_string() {
5828 let ree: ArrayRef = Arc::new(
5829 [Some("a"), Some("a"), None, Some("b"), Some("b")]
5830 .into_iter()
5831 .collect::<Int32RunArray>(),
5832 );
5833 let flat: ArrayRef = Arc::new(StringArray::from(vec![
5834 Some("a"),
5835 Some("a"),
5836 None,
5837 Some("b"),
5838 Some("b"),
5839 ]));
5840 ree_write_read_roundtrip(ree, flat);
5841 }
5842
5843 #[test]
5844 fn ree_int32() {
5845 let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
5846 for v in [Some(1), Some(1), None, Some(2), Some(2)] {
5847 b.append_option(v);
5848 }
5849 let ree: ArrayRef = Arc::new(b.finish());
5850 let flat: ArrayRef = Arc::new(Int32Array::from(vec![
5851 Some(1),
5852 Some(1),
5853 None,
5854 Some(2),
5855 Some(2),
5856 ]));
5857 ree_write_read_roundtrip(ree, flat);
5858 }
5859
5860 #[test]
5861 fn ree_bool() {
5862 let ree: ArrayRef = Arc::new(
5864 RunArray::try_new(
5865 &Int32Array::from(vec![3, 5, 7]),
5866 &BooleanArray::from(vec![Some(true), None, Some(false)]),
5867 )
5868 .unwrap(),
5869 );
5870 let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
5871 Some(true),
5872 Some(true),
5873 Some(true),
5874 None,
5875 None,
5876 Some(false),
5877 Some(false),
5878 ]));
5879 ree_write_read_roundtrip(ree, flat);
5880 }
5881
5882 #[test]
5883 fn ree_fixed_size_binary() {
5884 let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
5885 let mut b = FixedSizeBinaryBuilder::new(2);
5886 for v in vals {
5887 match v {
5888 Some(x) => b.append_value(x).unwrap(),
5889 None => b.append_null(),
5890 }
5891 }
5892 b.finish()
5893 };
5894 let ree: ArrayRef = Arc::new(
5896 RunArray::try_new(
5897 &Int32Array::from(vec![2, 4, 6]),
5898 &mk(&[Some(b"aa"), None, Some(b"bb")]),
5899 )
5900 .unwrap(),
5901 );
5902 let flat: ArrayRef = Arc::new(mk(&[
5903 Some(b"aa"),
5904 Some(b"aa"),
5905 None,
5906 None,
5907 Some(b"bb"),
5908 Some(b"bb"),
5909 ]));
5910 ree_write_read_roundtrip(ree, flat);
5911 }
5912
5913 #[test]
5914 fn ree_single_run() {
5915 let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
5916 let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
5917 ree_write_read_roundtrip(ree, flat);
5918 }
5919
5920 #[test]
5921 fn ree_float32() {
5922 let ree: ArrayRef = Arc::new(
5924 RunArray::try_new(
5925 &Int32Array::from(vec![2, 4, 5]),
5926 &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
5927 )
5928 .unwrap(),
5929 );
5930 let flat: ArrayRef = Arc::new(Float32Array::from(vec![
5931 Some(1.0_f32),
5932 Some(1.0_f32),
5933 None,
5934 None,
5935 Some(2.5_f32),
5936 ]));
5937 ree_write_read_roundtrip(ree, flat);
5938 }
5939
5940 #[test]
5941 fn ree_sliced() {
5942 let full: ArrayRef = Arc::new(
5947 RunArray::try_new(
5948 &Int32Array::from(vec![3, 5, 7]),
5949 &StringArray::from(vec!["a", "b", "c"]),
5950 )
5951 .unwrap(),
5952 );
5953 let sliced = full.slice(2, 5);
5954 let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
5955 ree_write_read_roundtrip(sliced, flat);
5956 }
5957
5958 #[test]
5959 fn ree_struct_with_ree_child() {
5960 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
5963
5964 let col_a: ArrayRef = Arc::new(
5965 RunArray::try_new(
5966 &run_ends,
5967 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
5968 )
5969 .unwrap(),
5970 );
5971 let col_b: ArrayRef = Arc::new(
5972 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
5973 );
5974
5975 let struct_array: ArrayRef = Arc::new(StructArray::new(
5976 Fields::from(vec![
5977 Field::new("a", col_a.data_type().clone(), true),
5978 Field::new("b", col_b.data_type().clone(), true),
5979 ]),
5980 vec![col_a, col_b],
5981 None,
5982 ));
5983
5984 let schema = Arc::new(Schema::new(vec![Field::new(
5985 "row",
5986 struct_array.data_type().clone(),
5987 true,
5988 )]));
5989 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
5990
5991 let mut buf = Vec::new();
5992 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
5993 writer.write(&batch).unwrap();
5994 let metadata = writer.close().unwrap();
5995
5996 let parquet_schema = metadata.file_metadata().schema_descr();
5997 assert_eq!(parquet_schema.num_columns(), 2);
5998 assert_eq!(
5999 parquet_schema.column(0).physical_type(),
6000 crate::basic::Type::BYTE_ARRAY
6001 );
6002 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6003 assert_eq!(
6004 parquet_schema.column(1).physical_type(),
6005 crate::basic::Type::INT32
6006 );
6007 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6008 }
6009}