1use crate::column::chunker::ContentDefinedChunker;
21
22use bytes::Bytes;
23use std::io::Write;
24use std::slice::Iter;
25use std::sync::{Arc, Mutex};
26use std::vec::IntoIter;
27
28use arrow_array::cast::AsArray;
29use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter};
30use arrow_array::{PrimitiveArray, types::*};
31use arrow_schema::{
32 ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit,
33};
34
35use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_precision};
36
37use crate::arrow::ArrowSchemaConverter;
38use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder;
39use crate::basic::PageType;
40use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter};
41use crate::column::page_encryption::PageEncryptor;
42use crate::column::writer::encoder::ColumnValueEncoder;
43use crate::column::writer::{
44 ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer,
45};
46use crate::data_type::{ByteArray, FixedLenByteArray};
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 && in_progress.buffered_rows + batch.num_rows() > max_rows
373 {
374 let to_write = max_rows - in_progress.buffered_rows;
375 let a = batch.slice(0, to_write);
376 let b = batch.slice(to_write, batch.num_rows() - to_write);
377 self.write(&a)?;
378 return self.write(&b);
379 }
380
381 if let Some(max_bytes) = self.max_row_group_bytes
384 && in_progress.buffered_rows > 0
385 {
386 let current_bytes = in_progress.get_estimated_total_bytes();
387
388 if current_bytes >= max_bytes {
389 self.flush()?;
390 return self.write(batch);
391 }
392
393 if let Some(avg_row_bytes) = current_bytes
394 .checked_div(in_progress.buffered_rows)
395 .filter(|avg_row_bytes| *avg_row_bytes > 0)
396 {
397 let remaining_bytes = max_bytes - current_bytes;
399 let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
400
401 if batch.num_rows() > rows_that_fit {
402 if rows_that_fit > 0 {
403 let a = batch.slice(0, rows_that_fit);
404 let b = batch.slice(rows_that_fit, batch.num_rows() - rows_that_fit);
405 self.write(&a)?;
406 return self.write(&b);
407 } else {
408 self.flush()?;
409 return self.write(batch);
410 }
411 }
412 }
413 }
414
415 match self.cdc_chunkers.as_mut() {
416 Some(chunkers) => in_progress.write_with_chunkers(batch, chunkers)?,
417 None => in_progress.write(batch)?,
418 }
419
420 let should_flush = self
421 .max_row_group_row_count
422 .is_some_and(|max| in_progress.buffered_rows >= max)
423 || self
424 .max_row_group_bytes
425 .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
426
427 if should_flush {
428 self.flush()?
429 }
430 Ok(())
431 }
432
433 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_array_slice(array, indices.iter().copied())
1755 }
1756 ArrowDataType::Decimal64(_, _) => {
1757 let array = column.as_primitive::<Decimal64Type>();
1758 get_decimal_array_slice(array, indices.iter().copied())
1759 }
1760 ArrowDataType::Decimal128(_, _) => {
1761 let array = column.as_primitive::<Decimal128Type>();
1762 get_decimal_array_slice(array, indices.iter().copied())
1763 }
1764 ArrowDataType::Decimal256(_, _) => {
1765 let array = column.as_primitive::<Decimal256Type>();
1766 get_decimal_array_slice(array, indices.iter().copied())
1767 }
1768 ArrowDataType::Float16 => {
1769 let array = column.as_primitive::<Float16Type>();
1770 get_float_16_array_slice(array, indices.iter().copied())
1771 }
1772 _ => {
1773 return Err(ParquetError::NYI(
1774 "Attempting to write an Arrow type that is not yet implemented".to_string(),
1775 ));
1776 }
1777 };
1778 typed.write_batch_internal(
1779 bytes.as_slice(),
1780 None,
1781 levels.def_level_data().as_ref(),
1782 levels.rep_level_data().as_ref(),
1783 None,
1784 None,
1785 None,
1786 )
1787 }
1788 }
1789}
1790
1791fn write_primitive<E: ColumnValueEncoder>(
1792 writer: &mut GenericColumnWriter<E>,
1793 values: &E::Values,
1794 levels: &ArrayLevels,
1795) -> Result<usize> {
1796 writer.write_batch_internal(
1797 values,
1798 Some(levels.non_null_indices()),
1799 levels.def_level_data().as_ref(),
1800 levels.rep_level_data().as_ref(),
1801 None,
1802 None,
1803 None,
1804 )
1805}
1806
1807fn get_bool_array_slice(
1808 array: &arrow_array::BooleanArray,
1809 indices: impl ExactSizeIterator<Item = usize>,
1810) -> Vec<bool> {
1811 let mut values = Vec::with_capacity(indices.len());
1812 for i in indices {
1813 values.push(array.value(i))
1814 }
1815 values
1816}
1817
1818fn get_interval_ym_array_slice(
1821 array: &arrow_array::IntervalYearMonthArray,
1822 indices: impl ExactSizeIterator<Item = usize>,
1823) -> Vec<FixedLenByteArray> {
1824 chunk_array_slice(12, indices, move |i, chunk| {
1825 let value = array.value(i);
1826 chunk[0..4].copy_from_slice(&value.to_le_bytes());
1827 })
1828}
1829
1830fn get_interval_dt_array_slice(
1833 array: &arrow_array::IntervalDayTimeArray,
1834 indices: impl ExactSizeIterator<Item = usize>,
1835) -> Vec<FixedLenByteArray> {
1836 chunk_array_slice(12, indices, move |i, chunk| {
1837 let value = array.value(i);
1838 chunk[4..8].copy_from_slice(&value.days.to_le_bytes());
1839 chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1840 })
1841}
1842
1843trait NativeDecimalType: DecimalType {
1844 type NativeBytes: AsRef<[u8]>;
1845
1846 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
1847}
1848impl NativeDecimalType for Decimal32Type {
1849 type NativeBytes = [u8; Self::BYTE_LENGTH];
1850
1851 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1852 value.to_be_bytes()
1853 }
1854}
1855impl NativeDecimalType for Decimal64Type {
1856 type NativeBytes = [u8; Self::BYTE_LENGTH];
1857
1858 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1859 value.to_be_bytes()
1860 }
1861}
1862impl NativeDecimalType for Decimal128Type {
1863 type NativeBytes = [u8; Self::BYTE_LENGTH];
1864
1865 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1866 value.to_be_bytes()
1867 }
1868}
1869impl NativeDecimalType for Decimal256Type {
1870 type NativeBytes = [u8; Self::BYTE_LENGTH];
1871
1872 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1873 value.to_be_bytes()
1874 }
1875}
1876
1877fn get_decimal_array_slice<T: NativeDecimalType>(
1878 array: &PrimitiveArray<T>,
1879 indices: impl ExactSizeIterator<Item = usize>,
1880) -> Vec<FixedLenByteArray> {
1881 let chunk_size = decimal_length_from_precision(array.precision());
1882 assert!(chunk_size <= T::BYTE_LENGTH);
1883
1884 if chunk_size == T::BYTE_LENGTH {
1885 chunk_array_slice(chunk_size, indices, move |i, chunk| {
1887 let as_be_bytes = T::to_be_bytes(array.value(i));
1888 chunk.copy_from_slice(as_be_bytes.as_ref());
1889 })
1890 } else {
1891 chunk_array_slice(chunk_size, indices, move |i, chunk| {
1892 let as_be_bytes = T::to_be_bytes(array.value(i));
1893 let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..];
1894 chunk.copy_from_slice(resized_value);
1895 })
1896 }
1897}
1898
1899fn get_float_16_array_slice(
1900 array: &arrow_array::Float16Array,
1901 indices: impl ExactSizeIterator<Item = usize>,
1902) -> Vec<FixedLenByteArray> {
1903 chunk_array_slice(2, indices, move |i, chunk| {
1904 let value = array.value(i).to_le_bytes();
1905 chunk.copy_from_slice(&value);
1906 })
1907}
1908
1909fn get_fsb_array_slice(
1910 array: &arrow_array::FixedSizeBinaryArray,
1911 indices: impl ExactSizeIterator<Item = usize>,
1912) -> Vec<FixedLenByteArray> {
1913 chunk_array_slice(array.value_size(), indices, move |i, chunk| {
1914 let value = array.value(i);
1915 chunk.copy_from_slice(value);
1916 })
1917}
1918
1919#[inline]
1920fn chunk_array_slice(
1921 chunk_size: usize,
1922 indices: impl ExactSizeIterator<Item = usize>,
1923 writer: impl Fn(usize, &mut [u8]),
1924) -> Vec<FixedLenByteArray> {
1925 let capacity = indices.len() * chunk_size;
1926 let mut arena = vec![0; capacity];
1929 for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
1930 writer(i, chunk);
1931 }
1932 chunk_contiguous_vec(arena, chunk_size)
1933}
1934
1935fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
1936 let mut values = Vec::with_capacity(arena.len() / chunk_size);
1937 let mut arena = Bytes::from(arena);
1938 while arena.len() >= chunk_size {
1939 let slice = arena.split_to(chunk_size);
1940 values.push(FixedLenByteArray::from(ByteArray::from(slice)));
1941 }
1942 values
1943}
1944
1945#[cfg(test)]
1946mod tests {
1947 use super::*;
1948 use std::cmp::Ordering;
1949 use std::collections::HashMap;
1950
1951 use std::fs::File;
1952
1953 use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1954 use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
1955 use crate::column::page::{Page, PageReader};
1956 use crate::file::metadata::thrift::PageHeader;
1957 use crate::file::page_index::column_index::ColumnIndexMetaData;
1958 use crate::file::reader::SerializedPageReader;
1959 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
1960 use crate::schema::types::ColumnPath;
1961 use arrow::datatypes::ToByteSlice;
1962 use arrow::datatypes::{DataType, Schema};
1963 use arrow::error::Result as ArrowResult;
1964 use arrow::util::data_gen::create_random_array;
1965 use arrow::util::pretty::pretty_format_batches;
1966 use arrow::{array::*, buffer::Buffer};
1967 use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
1968 use arrow_schema::Fields;
1969 use half::f16;
1970 use num_traits::{FromPrimitive, ToPrimitive};
1971 use tempfile::tempfile;
1972
1973 use crate::basic::Encoding;
1974 use crate::data_type::AsBytes;
1975 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
1976 use crate::file::properties::{
1977 BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
1978 };
1979 use crate::file::serialized_reader::ReadOptionsBuilder;
1980 use crate::file::{
1981 reader::{FileReader, SerializedFileReader},
1982 statistics::Statistics,
1983 };
1984
1985 #[derive(Debug, Default)]
1990 struct RecordingPageStore {
1991 next: u64,
1992 blobs: HashMap<u64, Bytes>,
1993 puts: Arc<std::sync::atomic::AtomicUsize>,
1994 }
1995
1996 impl PageStore for RecordingPageStore {
1997 fn put(&mut self, value: Bytes) -> Result<PageKey> {
1998 let id = 100 + self.next * 7;
2000 self.next += 1;
2001 self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2002 self.blobs.insert(id, value);
2003 Ok(PageKey::new(id))
2004 }
2005
2006 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2007 self.blobs
2008 .remove(&key.get())
2009 .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2010 }
2011 }
2012
2013 #[derive(Debug)]
2014 struct RecordingPageStoreFactory {
2015 puts: Arc<std::sync::atomic::AtomicUsize>,
2016 }
2017
2018 impl PageStoreFactory for RecordingPageStoreFactory {
2019 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2020 Ok(Box::new(RecordingPageStore {
2021 puts: self.puts.clone(),
2022 ..Default::default()
2023 }))
2024 }
2025 }
2026
2027 #[test]
2031 fn custom_page_store_is_byte_identical_to_default() {
2032 let schema = Arc::new(Schema::new(vec![
2033 Field::new("i", DataType::Int32, true),
2034 Field::new("s", DataType::Utf8, true),
2036 ]));
2037 let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2038 let s = StringArray::from(vec![
2039 Some("a"),
2040 Some("bb"),
2041 Some("a"),
2042 None,
2043 Some("bb"),
2044 Some("ccc"),
2045 ]);
2046 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2047
2048 let props = WriterProperties::builder()
2051 .set_max_row_group_row_count(Some(3))
2052 .build();
2053
2054 let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2055 let mut buffer = Vec::new();
2056 let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2057 if let Some(factory) = factory {
2058 opts = opts.with_page_store_factory(factory);
2059 }
2060 let mut writer =
2061 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2062 writer.write(&batch).unwrap();
2063 writer.close().unwrap();
2064 buffer
2065 };
2066
2067 let default_bytes = write(None);
2068
2069 let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2070 let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2071 puts: puts.clone(),
2072 })));
2073
2074 assert!(
2075 puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2076 "custom PageStore was never written to"
2077 );
2078 assert_eq!(
2079 default_bytes, custom_bytes,
2080 "a custom PageStore must produce byte-identical output to the default"
2081 );
2082 }
2083
2084 #[test]
2090 fn dictionary_column_round_trips_with_offset_index_disabled() {
2091 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2092
2093 let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2096 let array = Int32Array::from(values.clone());
2097 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2098
2099 let props = WriterProperties::builder()
2100 .set_offset_index_disabled(true)
2101 .set_data_page_row_count_limit(4096)
2102 .build();
2103 let opts = ArrowWriterOptions::new().with_properties(props);
2104
2105 let mut buffer = Vec::new();
2106 let mut writer =
2107 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2108 writer.write(&batch).unwrap();
2109 writer.close().unwrap();
2110
2111 let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2112 let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2113 let read_values: Vec<Option<i32>> = read
2114 .iter()
2115 .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2116 .collect();
2117 assert_eq!(read_values, values);
2118 }
2119
2120 #[test]
2125 fn dictionary_page_is_routed_through_the_store() {
2126 #[derive(Debug, Default)]
2128 struct SizeRecordingPageStore {
2129 blobs: Vec<Bytes>,
2130 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2131 }
2132 impl PageStore for SizeRecordingPageStore {
2133 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2134 self.bytes_put
2135 .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2136 let key = PageKey::new(self.blobs.len() as u64);
2137 self.blobs.push(value);
2138 Ok(key)
2139 }
2140 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2141 Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2142 }
2143 }
2144 #[derive(Debug)]
2145 struct Factory {
2146 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2147 }
2148 impl PageStoreFactory for Factory {
2149 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2150 Ok(Box::new(SizeRecordingPageStore {
2151 bytes_put: self.bytes_put.clone(),
2152 ..Default::default()
2153 }))
2154 }
2155 }
2156
2157 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2158 let values: Vec<&str> = (0..2048)
2161 .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2162 .collect();
2163 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2164 .unwrap();
2165
2166 let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2167 let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2168 bytes_put: bytes_put.clone(),
2169 }));
2170
2171 let mut buffer = Vec::new();
2174 let mut writer =
2175 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2176 writer.write(&batch).unwrap();
2177 writer.close().unwrap();
2178
2179 let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2180 let column = reader.metadata().row_group(0).column(0);
2181 assert!(
2182 column.dictionary_page_offset().is_some(),
2183 "expected the column to be dictionary-encoded"
2184 );
2185
2186 assert_eq!(
2190 bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2191 column.compressed_size(),
2192 "the dictionary page must pass through the store like any other page"
2193 );
2194 }
2195
2196 #[test]
2197 fn arrow_writer() {
2198 let schema = Schema::new(vec![
2200 Field::new("a", DataType::Int32, false),
2201 Field::new("b", DataType::Int32, true),
2202 ]);
2203
2204 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2206 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2207
2208 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2210
2211 roundtrip(batch, Some(SMALL_SIZE / 2));
2212 }
2213
2214 fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2215 let mut buffer = vec![];
2216
2217 let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2218 writer.write(expected_batch).unwrap();
2219 writer.close().unwrap();
2220
2221 buffer
2222 }
2223
2224 fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2225 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2226 writer.write(expected_batch).unwrap();
2227 writer.into_inner().unwrap()
2228 }
2229
2230 #[test]
2231 fn roundtrip_bytes() {
2232 let schema = Arc::new(Schema::new(vec![
2234 Field::new("a", DataType::Int32, false),
2235 Field::new("b", DataType::Int32, true),
2236 ]));
2237
2238 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2240 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2241
2242 let expected_batch =
2244 RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2245
2246 for buffer in [
2247 get_bytes_after_close(schema.clone(), &expected_batch),
2248 get_bytes_by_into_inner(schema, &expected_batch),
2249 ] {
2250 let cursor = Bytes::from(buffer);
2251 let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2252
2253 let actual_batch = record_batch_reader
2254 .next()
2255 .expect("No batch found")
2256 .expect("Unable to get batch");
2257
2258 assert_eq!(expected_batch.schema(), actual_batch.schema());
2259 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2260 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2261 for i in 0..expected_batch.num_columns() {
2262 let expected_data = expected_batch.column(i).to_data();
2263 let actual_data = actual_batch.column(i).to_data();
2264
2265 assert_eq!(expected_data, actual_data);
2266 }
2267 }
2268 }
2269
2270 #[test]
2271 fn arrow_writer_non_null() {
2272 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2274
2275 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2277
2278 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2280
2281 roundtrip(batch, Some(SMALL_SIZE / 2));
2282 }
2283
2284 #[test]
2285 fn arrow_writer_list() {
2286 let schema = Schema::new(vec![Field::new(
2288 "a",
2289 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2290 true,
2291 )]);
2292
2293 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2295
2296 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2299
2300 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2302 DataType::Int32,
2303 false,
2304 ))))
2305 .len(5)
2306 .add_buffer(a_value_offsets)
2307 .add_child_data(a_values.into_data())
2308 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2309 .build()
2310 .unwrap();
2311 let a = ListArray::from(a_list_data);
2312
2313 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2315
2316 assert_eq!(batch.column(0).null_count(), 1);
2317
2318 roundtrip(batch, None);
2321 }
2322
2323 #[test]
2324 fn arrow_writer_list_non_null() {
2325 let schema = Schema::new(vec![Field::new(
2327 "a",
2328 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2329 false,
2330 )]);
2331
2332 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2334
2335 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2338
2339 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2341 DataType::Int32,
2342 false,
2343 ))))
2344 .len(5)
2345 .add_buffer(a_value_offsets)
2346 .add_child_data(a_values.into_data())
2347 .build()
2348 .unwrap();
2349 let a = ListArray::from(a_list_data);
2350
2351 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2353
2354 assert_eq!(batch.column(0).null_count(), 0);
2357
2358 roundtrip(batch, None);
2359 }
2360
2361 #[test]
2362 fn arrow_writer_list_view() {
2363 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2364 let schema = Schema::new(vec![Field::new(
2365 "a",
2366 DataType::ListView(list_field.clone()),
2367 true,
2368 )]);
2369
2370 let a = ListViewArray::new(
2372 list_field,
2373 vec![0, 1, 0, 3, 6].into(),
2374 vec![1, 2, 0, 3, 4].into(),
2375 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2376 Some(vec![true, true, false, true, true].into()),
2377 );
2378
2379 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2380
2381 assert_eq!(batch.column(0).null_count(), 1);
2382
2383 roundtrip(batch, None);
2384 }
2385
2386 #[test]
2387 fn arrow_writer_list_view_non_null() {
2388 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2389 let schema = Schema::new(vec![Field::new(
2390 "a",
2391 DataType::ListView(list_field.clone()),
2392 false,
2393 )]);
2394
2395 let a = ListViewArray::new(
2397 list_field,
2398 vec![0, 1, 0, 3, 6].into(),
2399 vec![1, 2, 0, 3, 4].into(),
2400 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2401 None,
2402 );
2403
2404 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2405
2406 assert_eq!(batch.column(0).null_count(), 0);
2407
2408 roundtrip(batch, None);
2409 }
2410
2411 #[test]
2412 fn arrow_writer_list_view_out_of_order() {
2413 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2414 let schema = Schema::new(vec![Field::new(
2415 "a",
2416 DataType::ListView(list_field.clone()),
2417 false,
2418 )]);
2419
2420 let a = ListViewArray::new(
2422 list_field,
2423 vec![0, 1, 0, 6, 3].into(),
2424 vec![1, 2, 0, 4, 3].into(),
2425 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2426 None,
2427 );
2428
2429 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2430
2431 roundtrip(batch, None);
2432 }
2433
2434 #[test]
2435 fn arrow_writer_large_list_view() {
2436 let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2437 let schema = Schema::new(vec![Field::new(
2438 "a",
2439 DataType::LargeListView(list_field.clone()),
2440 true,
2441 )]);
2442
2443 let a = LargeListViewArray::new(
2445 list_field,
2446 vec![0i64, 1, 0, 3, 6].into(),
2447 vec![1i64, 2, 0, 3, 4].into(),
2448 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2449 Some(vec![true, true, false, true, true].into()),
2450 );
2451
2452 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2453
2454 assert_eq!(batch.column(0).null_count(), 1);
2455
2456 roundtrip(batch, None);
2457 }
2458
2459 #[test]
2460 fn arrow_writer_list_view_with_struct() {
2461 let struct_fields = Fields::from(vec![
2463 Field::new("id", DataType::Int32, false),
2464 Field::new("name", DataType::Utf8, false),
2465 ]);
2466 let struct_type = DataType::Struct(struct_fields.clone());
2467 let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2468
2469 let schema = Schema::new(vec![Field::new(
2470 "a",
2471 DataType::ListView(list_field.clone()),
2472 true,
2473 )]);
2474
2475 let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2477 let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2478 let struct_array = StructArray::new(
2479 struct_fields,
2480 vec![Arc::new(id_array), Arc::new(name_array)],
2481 None,
2482 );
2483
2484 let list_view = ListViewArray::new(
2486 list_field,
2487 vec![0, 2, 2].into(), vec![2, 0, 3].into(), Arc::new(struct_array),
2490 Some(vec![true, false, true].into()),
2491 );
2492
2493 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap();
2494
2495 roundtrip(batch, None);
2496 }
2497
2498 #[test]
2499 fn arrow_writer_binary() {
2500 let string_field = Field::new("a", DataType::Utf8, false);
2501 let binary_field = Field::new("b", DataType::Binary, false);
2502 let schema = Schema::new(vec![string_field, binary_field]);
2503
2504 let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2505 let raw_binary_values = [
2506 b"foo".to_vec(),
2507 b"bar".to_vec(),
2508 b"baz".to_vec(),
2509 b"quux".to_vec(),
2510 ];
2511 let raw_binary_value_refs = raw_binary_values
2512 .iter()
2513 .map(|x| x.as_slice())
2514 .collect::<Vec<_>>();
2515
2516 let string_values = StringArray::from(raw_string_values.clone());
2517 let binary_values = BinaryArray::from(raw_binary_value_refs);
2518 let batch = RecordBatch::try_new(
2519 Arc::new(schema),
2520 vec![Arc::new(string_values), Arc::new(binary_values)],
2521 )
2522 .unwrap();
2523
2524 roundtrip(batch, Some(SMALL_SIZE / 2));
2525 }
2526
2527 #[test]
2528 fn arrow_writer_binary_view() {
2529 let string_field = Field::new("a", DataType::Utf8View, false);
2530 let binary_field = Field::new("b", DataType::BinaryView, false);
2531 let nullable_string_field = Field::new("a", DataType::Utf8View, true);
2532 let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]);
2533
2534 let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2535 let raw_binary_values = vec![
2536 b"foo".to_vec(),
2537 b"bar".to_vec(),
2538 b"large payload over 12 bytes".to_vec(),
2539 b"lulu".to_vec(),
2540 ];
2541 let nullable_string_values =
2542 vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2543
2544 let string_view_values = StringViewArray::from(raw_string_values);
2545 let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2546 let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2547 let batch = RecordBatch::try_new(
2548 Arc::new(schema),
2549 vec![
2550 Arc::new(string_view_values),
2551 Arc::new(binary_view_values),
2552 Arc::new(nullable_string_view_values),
2553 ],
2554 )
2555 .unwrap();
2556
2557 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2558 roundtrip(batch, None);
2559 }
2560
2561 #[test]
2562 fn arrow_writer_binary_view_long_value() {
2563 let string_field = Field::new("a", DataType::Utf8View, false);
2564 let binary_field = Field::new("b", DataType::BinaryView, false);
2565 let schema = Schema::new(vec![string_field, binary_field]);
2566
2567 let long = "a".repeat(128);
2571 let raw_string_values = vec!["foo", long.as_str(), "bar"];
2572 let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2573
2574 let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2575 let binary_view_values: ArrayRef =
2576 Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2577
2578 one_column_roundtrip(Arc::clone(&string_view_values), false);
2579 one_column_roundtrip(Arc::clone(&binary_view_values), false);
2580
2581 let batch = RecordBatch::try_new(
2582 Arc::new(schema),
2583 vec![string_view_values, binary_view_values],
2584 )
2585 .unwrap();
2586
2587 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
2589 let props = WriterProperties::builder()
2590 .set_writer_version(version)
2591 .set_dictionary_enabled(false)
2592 .build();
2593 roundtrip_opts(&batch, props);
2594 }
2595 }
2596
2597 fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2598 let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2599 let schema = Schema::new(vec![decimal_field]);
2600
2601 let decimal_values = vec![10_000, 50_000, 0, -100]
2602 .into_iter()
2603 .map(Some)
2604 .collect::<Decimal128Array>()
2605 .with_precision_and_scale(precision, scale)
2606 .unwrap();
2607
2608 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2609 }
2610
2611 #[test]
2612 fn arrow_writer_decimal() {
2613 let batch_int32_decimal = get_decimal_batch(5, 2);
2615 roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2616 let batch_int64_decimal = get_decimal_batch(12, 2);
2618 roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2619 let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2621 roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2622 }
2623
2624 #[test]
2625 fn arrow_writer_complex() {
2626 let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2628 let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2629 let struct_field_g = Arc::new(Field::new_list(
2630 "g",
2631 Field::new_list_field(DataType::Int16, true),
2632 false,
2633 ));
2634 let struct_field_h = Arc::new(Field::new_list(
2635 "h",
2636 Field::new_list_field(DataType::Int16, false),
2637 true,
2638 ));
2639 let struct_field_e = Arc::new(Field::new_struct(
2640 "e",
2641 vec![
2642 struct_field_f.clone(),
2643 struct_field_g.clone(),
2644 struct_field_h.clone(),
2645 ],
2646 false,
2647 ));
2648 let schema = Schema::new(vec![
2649 Field::new("a", DataType::Int32, false),
2650 Field::new("b", DataType::Int32, true),
2651 Field::new_struct(
2652 "c",
2653 vec![struct_field_d.clone(), struct_field_e.clone()],
2654 false,
2655 ),
2656 ]);
2657
2658 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2660 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2661 let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2662 let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2663
2664 let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2665
2666 let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2669
2670 let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2672 .len(5)
2673 .add_buffer(g_value_offsets.clone())
2674 .add_child_data(g_value.to_data())
2675 .build()
2676 .unwrap();
2677 let g = ListArray::from(g_list_data);
2678 let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2680 .len(5)
2681 .add_buffer(g_value_offsets)
2682 .add_child_data(g_value.to_data())
2683 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2684 .build()
2685 .unwrap();
2686 let h = ListArray::from(h_list_data);
2687
2688 let e = StructArray::from(vec![
2689 (struct_field_f, Arc::new(f) as ArrayRef),
2690 (struct_field_g, Arc::new(g) as ArrayRef),
2691 (struct_field_h, Arc::new(h) as ArrayRef),
2692 ]);
2693
2694 let c = StructArray::from(vec![
2695 (struct_field_d, Arc::new(d) as ArrayRef),
2696 (struct_field_e, Arc::new(e) as ArrayRef),
2697 ]);
2698
2699 let batch = RecordBatch::try_new(
2701 Arc::new(schema),
2702 vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2703 )
2704 .unwrap();
2705
2706 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2707 roundtrip(batch, Some(SMALL_SIZE / 3));
2708 }
2709
2710 #[test]
2711 fn arrow_writer_complex_mixed() {
2712 let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2717 let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2718 let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2719 let schema = Schema::new(vec![Field::new(
2720 "some_nested_object",
2721 DataType::Struct(Fields::from(vec![
2722 offset_field.clone(),
2723 partition_field.clone(),
2724 topic_field.clone(),
2725 ])),
2726 false,
2727 )]);
2728
2729 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2731 let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2732 let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2733
2734 let some_nested_object = StructArray::from(vec![
2735 (offset_field, Arc::new(offset) as ArrayRef),
2736 (partition_field, Arc::new(partition) as ArrayRef),
2737 (topic_field, Arc::new(topic) as ArrayRef),
2738 ]);
2739
2740 let batch =
2742 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2743
2744 roundtrip(batch, Some(SMALL_SIZE / 2));
2745 }
2746
2747 #[test]
2748 fn arrow_writer_map() {
2749 let json_content = r#"
2751 {"stocks":{"long": "$AAA", "short": "$BBB"}}
2752 {"stocks":{"long": null, "long": "$CCC", "short": null}}
2753 {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2754 "#;
2755 let entries_struct_type = DataType::Struct(Fields::from(vec![
2756 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2757 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2758 ]));
2759 let stocks_field = Field::new(
2760 "stocks",
2761 DataType::Map(
2762 Arc::new(Field::new(
2763 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2764 entries_struct_type,
2765 false,
2766 )),
2767 false,
2768 ),
2769 true,
2770 );
2771 let schema = Arc::new(Schema::new(vec![stocks_field]));
2772 let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2773 let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2774
2775 let batch = reader.next().unwrap().unwrap();
2776 roundtrip(batch, None);
2777 }
2778
2779 #[test]
2780 fn arrow_writer_2_level_struct() {
2781 let field_c = Field::new("c", DataType::Int32, true);
2783 let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2784 let type_a = DataType::Struct(vec![field_b.clone()].into());
2785 let field_a = Field::new("a", type_a, true);
2786 let schema = Schema::new(vec![field_a.clone()]);
2787
2788 let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2790 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2791 .len(6)
2792 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2793 .add_child_data(c.into_data())
2794 .build()
2795 .unwrap();
2796 let b = StructArray::from(b_data);
2797 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2798 .len(6)
2799 .null_bit_buffer(Some(Buffer::from([0b00101111])))
2800 .add_child_data(b.into_data())
2801 .build()
2802 .unwrap();
2803 let a = StructArray::from(a_data);
2804
2805 assert_eq!(a.null_count(), 1);
2806 assert_eq!(a.column(0).null_count(), 2);
2807
2808 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2810
2811 roundtrip(batch, Some(SMALL_SIZE / 2));
2812 }
2813
2814 #[test]
2815 fn arrow_writer_2_level_struct_non_null() {
2816 let field_c = Field::new("c", DataType::Int32, false);
2818 let type_b = DataType::Struct(vec![field_c].into());
2819 let field_b = Field::new("b", type_b.clone(), false);
2820 let type_a = DataType::Struct(vec![field_b].into());
2821 let field_a = Field::new("a", type_a.clone(), false);
2822 let schema = Schema::new(vec![field_a]);
2823
2824 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2826 let b_data = ArrayDataBuilder::new(type_b)
2827 .len(6)
2828 .add_child_data(c.into_data())
2829 .build()
2830 .unwrap();
2831 let b = StructArray::from(b_data);
2832 let a_data = ArrayDataBuilder::new(type_a)
2833 .len(6)
2834 .add_child_data(b.into_data())
2835 .build()
2836 .unwrap();
2837 let a = StructArray::from(a_data);
2838
2839 assert_eq!(a.null_count(), 0);
2840 assert_eq!(a.column(0).null_count(), 0);
2841
2842 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2844
2845 roundtrip(batch, Some(SMALL_SIZE / 2));
2846 }
2847
2848 #[test]
2849 fn arrow_writer_2_level_struct_mixed_null() {
2850 let field_c = Field::new("c", DataType::Int32, false);
2852 let type_b = DataType::Struct(vec![field_c].into());
2853 let field_b = Field::new("b", type_b.clone(), true);
2854 let type_a = DataType::Struct(vec![field_b].into());
2855 let field_a = Field::new("a", type_a.clone(), false);
2856 let schema = Schema::new(vec![field_a]);
2857
2858 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2860 let b_data = ArrayDataBuilder::new(type_b)
2861 .len(6)
2862 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2863 .add_child_data(c.into_data())
2864 .build()
2865 .unwrap();
2866 let b = StructArray::from(b_data);
2867 let a_data = ArrayDataBuilder::new(type_a)
2869 .len(6)
2870 .add_child_data(b.into_data())
2871 .build()
2872 .unwrap();
2873 let a = StructArray::from(a_data);
2874
2875 assert_eq!(a.null_count(), 0);
2876 assert_eq!(a.column(0).null_count(), 2);
2877
2878 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2880
2881 roundtrip(batch, Some(SMALL_SIZE / 2));
2882 }
2883
2884 #[test]
2885 fn arrow_writer_2_level_struct_mixed_null_2() {
2886 let field_c = Field::new("c", DataType::Int32, false);
2888 let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
2889 let field_e = Field::new(
2890 "e",
2891 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2892 false,
2893 );
2894
2895 let field_b = Field::new(
2896 "b",
2897 DataType::Struct(vec![field_c, field_d, field_e].into()),
2898 false,
2899 );
2900 let type_a = DataType::Struct(vec![field_b.clone()].into());
2901 let field_a = Field::new("a", type_a, true);
2902 let schema = Schema::new(vec![field_a.clone()]);
2903
2904 let c = Int32Array::from_iter_values(0..6);
2906 let d = FixedSizeBinaryArray::try_from_iter(
2907 ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
2908 )
2909 .expect("four byte values");
2910 let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
2911 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2912 .len(6)
2913 .add_child_data(c.into_data())
2914 .add_child_data(d.into_data())
2915 .add_child_data(e.into_data())
2916 .build()
2917 .unwrap();
2918 let b = StructArray::from(b_data);
2919 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2920 .len(6)
2921 .null_bit_buffer(Some(Buffer::from([0b00100101])))
2922 .add_child_data(b.into_data())
2923 .build()
2924 .unwrap();
2925 let a = StructArray::from(a_data);
2926
2927 assert_eq!(a.null_count(), 3);
2928 assert_eq!(a.column(0).null_count(), 0);
2929
2930 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2932
2933 roundtrip(batch, Some(SMALL_SIZE / 2));
2934 }
2935
2936 #[test]
2937 fn test_fixed_size_binary_in_dict() {
2938 fn test_fixed_size_binary_in_dict_inner<K>()
2939 where
2940 K: ArrowDictionaryKeyType,
2941 K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
2942 <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
2943 {
2944 let field = Field::new(
2945 "a",
2946 DataType::Dictionary(
2947 Box::new(K::DATA_TYPE),
2948 Box::new(DataType::FixedSizeBinary(4)),
2949 ),
2950 false,
2951 );
2952 let schema = Schema::new(vec![field]);
2953
2954 let keys: Vec<K::Native> = vec![
2955 K::Native::try_from(0u8).unwrap(),
2956 K::Native::try_from(0u8).unwrap(),
2957 K::Native::try_from(1u8).unwrap(),
2958 ];
2959 let keys = PrimitiveArray::<K>::from_iter_values(keys);
2960 let values = FixedSizeBinaryArray::try_from_iter(
2961 vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
2962 )
2963 .unwrap();
2964
2965 let data = DictionaryArray::<K>::new(keys, Arc::new(values));
2966 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
2967 roundtrip(batch, None);
2968 }
2969
2970 test_fixed_size_binary_in_dict_inner::<UInt8Type>();
2971 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2972 test_fixed_size_binary_in_dict_inner::<UInt32Type>();
2973 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2974 test_fixed_size_binary_in_dict_inner::<Int8Type>();
2975 test_fixed_size_binary_in_dict_inner::<Int16Type>();
2976 test_fixed_size_binary_in_dict_inner::<Int32Type>();
2977 test_fixed_size_binary_in_dict_inner::<Int64Type>();
2978 }
2979
2980 #[test]
2981 fn test_empty_dict() {
2982 let struct_fields = Fields::from(vec![Field::new(
2983 "dict",
2984 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2985 false,
2986 )]);
2987
2988 let schema = Schema::new(vec![Field::new_struct(
2989 "struct",
2990 struct_fields.clone(),
2991 true,
2992 )]);
2993 let dictionary = Arc::new(DictionaryArray::new(
2994 Int32Array::new_null(5),
2995 Arc::new(StringArray::new_null(0)),
2996 ));
2997
2998 let s = StructArray::new(
2999 struct_fields,
3000 vec![dictionary],
3001 Some(NullBuffer::new_null(5)),
3002 );
3003
3004 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
3005 roundtrip(batch, None);
3006 }
3007 #[test]
3008 fn arrow_writer_page_size() {
3009 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3010
3011 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3012
3013 for i in 0..10 {
3015 let value = i
3016 .to_string()
3017 .repeat(10)
3018 .chars()
3019 .take(10)
3020 .collect::<String>();
3021
3022 builder.append_value(value);
3023 }
3024
3025 let array = Arc::new(builder.finish());
3026
3027 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3028
3029 let file = tempfile::tempfile().unwrap();
3030
3031 let props = WriterProperties::builder()
3033 .set_data_page_size_limit(1)
3034 .set_dictionary_page_size_limit(1)
3035 .set_write_batch_size(1)
3036 .build();
3037
3038 let mut writer =
3039 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3040 .expect("Unable to write file");
3041 writer.write(&batch).unwrap();
3042 writer.close().unwrap();
3043
3044 let options = ReadOptionsBuilder::new().with_page_index().build();
3045 let reader =
3046 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3047
3048 let column = reader.metadata().row_group(0).columns();
3049
3050 assert_eq!(column.len(), 1);
3051
3052 assert!(
3055 column[0].dictionary_page_offset().is_some(),
3056 "Expected a dictionary page"
3057 );
3058
3059 assert!(reader.metadata().offset_index().is_some());
3060 let offset_indexes = &reader.metadata().offset_index().unwrap()[0];
3061
3062 let page_locations = offset_indexes[0].page_locations.clone();
3063
3064 assert_eq!(
3067 page_locations.len(),
3068 10,
3069 "Expected 10 pages but got {page_locations:#?}"
3070 );
3071 }
3072
3073 #[test]
3074 fn arrow_writer_float_nans() {
3075 let f16_field = Field::new("a", DataType::Float16, false);
3076 let f32_field = Field::new("b", DataType::Float32, false);
3077 let f64_field = Field::new("c", DataType::Float64, false);
3078 let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3079
3080 let f16_values = (0..MEDIUM_SIZE)
3081 .map(|i| {
3082 Some(if i % 2 == 0 {
3083 f16::NAN
3084 } else {
3085 f16::from_f32(i as f32)
3086 })
3087 })
3088 .collect::<Float16Array>();
3089
3090 let f32_values = (0..MEDIUM_SIZE)
3091 .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3092 .collect::<Float32Array>();
3093
3094 let f64_values = (0..MEDIUM_SIZE)
3095 .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3096 .collect::<Float64Array>();
3097
3098 let batch = RecordBatch::try_new(
3099 Arc::new(schema),
3100 vec![
3101 Arc::new(f16_values),
3102 Arc::new(f32_values),
3103 Arc::new(f64_values),
3104 ],
3105 )
3106 .unwrap();
3107
3108 roundtrip(batch, None);
3109 }
3110
3111 const SMALL_SIZE: usize = 7;
3112 const MEDIUM_SIZE: usize = 63;
3113
3114 fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3117 let mut files = vec![];
3118 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3119 let mut props = WriterProperties::builder().set_writer_version(version);
3120
3121 if let Some(size) = max_row_group_size {
3122 props = props.set_max_row_group_row_count(Some(size))
3123 }
3124
3125 let props = props.build();
3126 files.push(roundtrip_opts(&expected_batch, props))
3127 }
3128 files
3129 }
3130
3131 fn roundtrip_opts_with_array_validation<F>(
3135 expected_batch: &RecordBatch,
3136 props: WriterProperties,
3137 validate: F,
3138 ) -> Bytes
3139 where
3140 F: Fn(&ArrayData, &ArrayData),
3141 {
3142 let mut file = vec![];
3143
3144 let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3145 .expect("Unable to write file");
3146 writer.write(expected_batch).unwrap();
3147 writer.close().unwrap();
3148
3149 let file = Bytes::from(file);
3150 let mut record_batch_reader =
3151 ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3152
3153 let actual_batch = record_batch_reader
3154 .next()
3155 .expect("No batch found")
3156 .expect("Unable to get batch");
3157
3158 assert_eq!(expected_batch.schema(), actual_batch.schema());
3159 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3160 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3161 for i in 0..expected_batch.num_columns() {
3162 let expected_data = expected_batch.column(i).to_data();
3163 let actual_data = actual_batch.column(i).to_data();
3164 validate(&expected_data, &actual_data);
3165 }
3166
3167 file
3168 }
3169
3170 fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3171 roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3172 a.validate_full().expect("valid expected data");
3173 b.validate_full().expect("valid actual data");
3174 assert_eq!(a, b)
3175 })
3176 }
3177
3178 struct RoundTripOptions {
3179 values: ArrayRef,
3180 schema: SchemaRef,
3181 bloom_filter: bool,
3182 bloom_filter_ndv: Option<u64>,
3183 bloom_filter_position: BloomFilterPosition,
3184 }
3185
3186 impl RoundTripOptions {
3187 fn new(values: ArrayRef, nullable: bool) -> Self {
3188 let data_type = values.data_type().clone();
3189 let schema = Schema::new(vec![Field::new("col", data_type, nullable)]);
3190 Self {
3191 values,
3192 schema: Arc::new(schema),
3193 bloom_filter: false,
3194 bloom_filter_ndv: None,
3195 bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3196 }
3197 }
3198 }
3199
3200 fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec<Bytes> {
3201 one_column_roundtrip_with_options(RoundTripOptions::new(values, nullable))
3202 }
3203
3204 fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef) -> Vec<Bytes> {
3205 let mut options = RoundTripOptions::new(values, false);
3206 options.schema = schema;
3207 one_column_roundtrip_with_options(options)
3208 }
3209
3210 fn one_column_roundtrip_with_options(options: RoundTripOptions) -> Vec<Bytes> {
3211 let RoundTripOptions {
3212 values,
3213 schema,
3214 bloom_filter,
3215 bloom_filter_ndv,
3216 bloom_filter_position,
3217 } = options;
3218
3219 let encodings = match values.data_type() {
3220 DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3221 vec![
3222 Encoding::PLAIN,
3223 Encoding::DELTA_BYTE_ARRAY,
3224 Encoding::DELTA_LENGTH_BYTE_ARRAY,
3225 ]
3226 }
3227 DataType::Int64
3228 | DataType::Int32
3229 | DataType::Int16
3230 | DataType::Int8
3231 | DataType::UInt64
3232 | DataType::UInt32
3233 | DataType::UInt16
3234 | DataType::UInt8 => vec![
3235 Encoding::PLAIN,
3236 Encoding::DELTA_BINARY_PACKED,
3237 Encoding::BYTE_STREAM_SPLIT,
3238 ],
3239 DataType::Float32 | DataType::Float64 => {
3240 vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
3241 }
3242 _ => vec![Encoding::PLAIN],
3243 };
3244
3245 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3246
3247 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3248
3249 let mut files = vec![];
3250 for dictionary_size in [0, 1, 1024] {
3251 for encoding in &encodings {
3252 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3253 for row_group_size in row_group_sizes {
3254 let mut builder = WriterProperties::builder()
3255 .set_writer_version(version)
3256 .set_max_row_group_row_count(Some(row_group_size))
3257 .set_dictionary_enabled(dictionary_size != 0)
3258 .set_dictionary_page_size_limit(dictionary_size.max(1))
3259 .set_encoding(*encoding)
3260 .set_bloom_filter_enabled(bloom_filter)
3261 .set_bloom_filter_position(bloom_filter_position);
3262 if let Some(ndv) = bloom_filter_ndv {
3263 builder = builder.set_bloom_filter_max_ndv(ndv);
3264 }
3265 let props = builder.build();
3266
3267 files.push(roundtrip_opts(&expected_batch, props))
3268 }
3269 }
3270 }
3271 }
3272 files
3273 }
3274
3275 fn values_required<A, I>(iter: I) -> Vec<Bytes>
3276 where
3277 A: From<Vec<I::Item>> + Array + 'static,
3278 I: IntoIterator,
3279 {
3280 let raw_values: Vec<_> = iter.into_iter().collect();
3281 let values = Arc::new(A::from(raw_values));
3282 one_column_roundtrip(values, false)
3283 }
3284
3285 fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3286 where
3287 A: From<Vec<Option<I::Item>>> + Array + 'static,
3288 I: IntoIterator,
3289 {
3290 let optional_raw_values: Vec<_> = iter
3291 .into_iter()
3292 .enumerate()
3293 .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3294 .collect();
3295 let optional_values = Arc::new(A::from(optional_raw_values));
3296 one_column_roundtrip(optional_values, true)
3297 }
3298
3299 fn required_and_optional<A, I>(iter: I)
3300 where
3301 A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3302 I: IntoIterator + Clone,
3303 {
3304 values_required::<A, I>(iter.clone());
3305 values_optional::<A, I>(iter);
3306 }
3307
3308 fn check_bloom_filter<T: AsBytes>(
3309 files: Vec<Bytes>,
3310 file_column: String,
3311 positive_values: Vec<T>,
3312 negative_values: Vec<T>,
3313 ) {
3314 files.into_iter().take(1).for_each(|file| {
3315 let file_reader = SerializedFileReader::new_with_options(
3316 file,
3317 ReadOptionsBuilder::new()
3318 .with_reader_properties(
3319 ReaderProperties::builder()
3320 .set_read_bloom_filter(true)
3321 .build(),
3322 )
3323 .build(),
3324 )
3325 .expect("Unable to open file as Parquet");
3326 let metadata = file_reader.metadata();
3327
3328 let mut bloom_filters: Vec<_> = vec![];
3330 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3331 if let Some((column_index, _)) = row_group
3332 .columns()
3333 .iter()
3334 .enumerate()
3335 .find(|(_, column)| column.column_path().string() == file_column)
3336 {
3337 let row_group_reader = file_reader
3338 .get_row_group(ri)
3339 .expect("Unable to read row group");
3340 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3341 bloom_filters.push(sbbf.clone());
3342 } else {
3343 panic!("No bloom filter for column named {file_column} found");
3344 }
3345 } else {
3346 panic!("No column named {file_column} found");
3347 }
3348 }
3349
3350 positive_values.iter().for_each(|value| {
3351 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3352 assert!(
3353 found.is_some(),
3354 "{}",
3355 format!("Value {:?} should be in bloom filter", value.as_bytes())
3356 );
3357 });
3358
3359 negative_values.iter().for_each(|value| {
3360 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3361 assert!(
3362 found.is_none(),
3363 "{}",
3364 format!("Value {:?} should not be in bloom filter", value.as_bytes())
3365 );
3366 });
3367 });
3368 }
3369
3370 #[test]
3371 fn all_null_primitive_single_column() {
3372 let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3373 one_column_roundtrip(values, true);
3374 }
3375 #[test]
3376 fn null_single_column() {
3377 let values = Arc::new(NullArray::new(SMALL_SIZE));
3378 one_column_roundtrip(values, true);
3379 }
3381
3382 #[test]
3383 fn bool_single_column() {
3384 required_and_optional::<BooleanArray, _>(
3385 [true, false].iter().cycle().copied().take(SMALL_SIZE),
3386 );
3387 }
3388
3389 #[test]
3390 fn bool_large_single_column() {
3391 let values = Arc::new(
3392 [None, Some(true), Some(false)]
3393 .iter()
3394 .cycle()
3395 .copied()
3396 .take(200_000)
3397 .collect::<BooleanArray>(),
3398 );
3399 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3400 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3401 let file = tempfile::tempfile().unwrap();
3402
3403 let mut writer =
3404 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3405 .expect("Unable to write file");
3406 writer.write(&expected_batch).unwrap();
3407 writer.close().unwrap();
3408 }
3409
3410 #[test]
3411 fn check_page_offset_index_with_nan() {
3412 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3413 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3414 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3415
3416 let mut out = Vec::with_capacity(1024);
3417 let mut writer =
3418 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3419 writer.write(&batch).unwrap();
3420 let file_meta_data = writer.close().unwrap();
3421 for row_group in file_meta_data.row_groups() {
3422 for column in row_group.columns() {
3423 assert!(column.offset_index_offset().is_some());
3424 assert!(column.offset_index_length().is_some());
3425 assert!(column.column_index_offset().is_some());
3426 assert!(column.column_index_length().is_some());
3427 }
3428 }
3429 assert!(file_meta_data.column_index().is_some());
3430 if let Some(col_indexes) = file_meta_data.column_index() {
3431 for rg_idx in col_indexes {
3432 for idx in rg_idx {
3433 assert!(idx.nan_counts().is_some());
3434 let float_idx = match idx {
3435 ColumnIndexMetaData::DOUBLE(idx) => idx,
3436 _ => panic!("expected double statistics"),
3437 };
3438 for i in 0..idx.num_pages() as usize {
3439 assert_eq!(float_idx.nan_count(i), Some(10));
3440 assert_eq!(
3441 f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3442 Ordering::Equal
3443 );
3444 assert_eq!(
3445 f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3446 Ordering::Equal
3447 );
3448 }
3449 }
3450 }
3451 }
3452 }
3453
3454 #[test]
3455 fn check_page_offset_index_with_mixed_nan() {
3456 let schema = Arc::new(Schema::new(vec![Field::new(
3457 "col",
3458 DataType::Float64,
3459 true,
3460 )]));
3461
3462 let mut out = Vec::with_capacity(1024);
3463 let props = WriterProperties::builder()
3464 .set_data_page_row_count_limit(10)
3465 .build();
3466 let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3467 .expect("Unable to write file");
3468
3469 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3471 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3472 writer.write(&batch).unwrap();
3473
3474 let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3476 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3477 writer.write(&batch).unwrap();
3478
3479 let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3481 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3482 writer.write(&batch).unwrap();
3483
3484 let values = Arc::new(Float64Array::from(vec![
3486 -1.0,
3487 0.0,
3488 f64::NAN,
3489 -f64::NAN,
3490 1.0,
3491 ]));
3492 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3493 writer.write(&batch).unwrap();
3494
3495 let file_meta_data = writer.close().unwrap();
3496
3497 let col_stats = file_meta_data
3499 .row_group(0)
3500 .column(0)
3501 .statistics()
3502 .expect("missing column chunk statistics");
3503
3504 assert_eq!(col_stats.nan_count_opt(), Some(22));
3505 assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3506 assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3507
3508 assert!(file_meta_data.column_index().is_some());
3509 let col_idx = &file_meta_data.column_index().as_ref().unwrap()[0][0];
3510 assert_eq!(col_idx.num_pages(), 4);
3511
3512 let float_idx = match col_idx {
3514 ColumnIndexMetaData::DOUBLE(idx) => idx,
3515 _ => panic!("expected double statistics"),
3516 };
3517
3518 assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3519 assert_eq!(
3520 f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3521 Ordering::Equal
3522 );
3523 assert_eq!(
3524 f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3525 Ordering::Equal
3526 );
3527 assert_eq!(
3528 (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3529 Ordering::Equal
3530 );
3531 assert_eq!(
3532 (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3533 Ordering::Equal
3534 );
3535 assert_eq!(float_idx.min_value(2), Some(&0.0));
3536 assert_eq!(float_idx.max_value(2), Some(&0.0));
3537 assert_eq!(float_idx.min_value(3), Some(&-1.0));
3538 assert_eq!(float_idx.max_value(3), Some(&1.0));
3539 }
3540
3541 #[test]
3542 fn i8_single_column() {
3543 required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3544 }
3545
3546 #[test]
3547 fn i16_single_column() {
3548 required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3549 }
3550
3551 #[test]
3552 fn i32_single_column() {
3553 required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3554 }
3555
3556 #[test]
3557 fn i64_single_column() {
3558 required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3559 }
3560
3561 #[test]
3562 fn u8_single_column() {
3563 required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3564 }
3565
3566 #[test]
3567 fn u16_single_column() {
3568 required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3569 }
3570
3571 #[test]
3572 fn u32_single_column() {
3573 required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3574 }
3575
3576 #[test]
3577 fn u64_single_column() {
3578 required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3579 }
3580
3581 #[test]
3582 fn f32_single_column() {
3583 required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3584 }
3585
3586 #[test]
3587 fn f64_single_column() {
3588 required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3589 }
3590
3591 #[test]
3596 fn timestamp_second_single_column() {
3597 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3598 let values = Arc::new(TimestampSecondArray::from(raw_values));
3599
3600 one_column_roundtrip(values, false);
3601 }
3602
3603 #[test]
3604 fn timestamp_millisecond_single_column() {
3605 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3606 let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3607
3608 one_column_roundtrip(values, false);
3609 }
3610
3611 #[test]
3612 fn timestamp_microsecond_single_column() {
3613 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3614 let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3615
3616 one_column_roundtrip(values, false);
3617 }
3618
3619 #[test]
3620 fn timestamp_nanosecond_single_column() {
3621 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3622 let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3623
3624 one_column_roundtrip(values, false);
3625 }
3626
3627 #[test]
3628 fn date32_single_column() {
3629 required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3630 }
3631
3632 #[test]
3633 fn date64_single_column() {
3634 required_and_optional::<Date64Array, _>(
3636 (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3637 );
3638 }
3639
3640 #[test]
3641 fn time32_second_single_column() {
3642 required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3643 }
3644
3645 #[test]
3646 fn time32_millisecond_single_column() {
3647 required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3648 }
3649
3650 #[test]
3651 fn time64_microsecond_single_column() {
3652 required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3653 }
3654
3655 #[test]
3656 fn time64_nanosecond_single_column() {
3657 required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3658 }
3659
3660 #[test]
3661 fn duration_second_single_column() {
3662 required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3663 }
3664
3665 #[test]
3666 fn duration_millisecond_single_column() {
3667 required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3668 }
3669
3670 #[test]
3671 fn duration_microsecond_single_column() {
3672 required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3673 }
3674
3675 #[test]
3676 fn duration_nanosecond_single_column() {
3677 required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3678 }
3679
3680 #[test]
3681 fn interval_year_month_single_column() {
3682 required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3683 }
3684
3685 #[test]
3686 fn interval_day_time_single_column() {
3687 required_and_optional::<IntervalDayTimeArray, _>(vec![
3688 IntervalDayTime::new(0, 1),
3689 IntervalDayTime::new(0, 3),
3690 IntervalDayTime::new(3, -2),
3691 IntervalDayTime::new(-200, 4),
3692 ]);
3693 }
3694
3695 #[test]
3696 #[should_panic(
3697 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3698 )]
3699 fn interval_month_day_nano_single_column() {
3700 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3701 IntervalMonthDayNano::new(0, 1, 5),
3702 IntervalMonthDayNano::new(0, 3, 2),
3703 IntervalMonthDayNano::new(3, -2, -5),
3704 IntervalMonthDayNano::new(-200, 4, -1),
3705 ]);
3706 }
3707
3708 #[test]
3709 fn binary_single_column() {
3710 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3711 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3712 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3713
3714 values_required::<BinaryArray, _>(many_vecs_iter);
3716 }
3717
3718 #[test]
3719 fn binary_view_single_column() {
3720 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3721 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3722 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3723
3724 values_required::<BinaryViewArray, _>(many_vecs_iter);
3726 }
3727
3728 #[test]
3729 fn i32_column_bloom_filter_at_end() {
3730 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3731 let mut options = RoundTripOptions::new(array, false);
3732 options.bloom_filter = true;
3733 options.bloom_filter_position = BloomFilterPosition::End;
3734
3735 let files = one_column_roundtrip_with_options(options);
3736 check_bloom_filter(
3737 files,
3738 "col".to_string(),
3739 (0..SMALL_SIZE as i32).collect(),
3740 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3741 );
3742 }
3743
3744 #[test]
3745 fn i32_column_bloom_filter() {
3746 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3747 let mut options = RoundTripOptions::new(array, false);
3748 options.bloom_filter = true;
3749
3750 let files = one_column_roundtrip_with_options(options);
3751 check_bloom_filter(
3752 files,
3753 "col".to_string(),
3754 (0..SMALL_SIZE as i32).collect(),
3755 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3756 );
3757 }
3758
3759 #[test]
3764 fn i32_column_bloom_filter_fixed_ndv() {
3765 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3766
3767 let mut options = RoundTripOptions::new(array.clone(), false);
3769 options.bloom_filter = true;
3770 options.bloom_filter_ndv = Some(1_000_000);
3771
3772 let files = one_column_roundtrip_with_options(options);
3773 check_bloom_filter(
3774 files,
3775 "col".to_string(),
3776 (0..SMALL_SIZE as i32).collect(),
3777 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3778 );
3779
3780 let mut options = RoundTripOptions::new(array, false);
3782 options.bloom_filter = true;
3783 options.bloom_filter_ndv = Some(3);
3784
3785 let files = one_column_roundtrip_with_options(options);
3786 check_bloom_filter(
3787 files,
3788 "col".to_string(),
3789 (0..SMALL_SIZE as i32).collect(),
3790 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3791 );
3792 }
3793
3794 #[test]
3795 fn binary_column_bloom_filter() {
3796 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3797 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3798 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3799
3800 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
3801 let mut options = RoundTripOptions::new(array, false);
3802 options.bloom_filter = true;
3803
3804 let files = one_column_roundtrip_with_options(options);
3805 check_bloom_filter(
3806 files,
3807 "col".to_string(),
3808 many_vecs,
3809 vec![vec![(SMALL_SIZE + 1) as u8]],
3810 );
3811 }
3812
3813 #[test]
3814 fn empty_string_null_column_bloom_filter() {
3815 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3816 let raw_strs = raw_values.iter().map(|s| s.as_str());
3817
3818 let array = Arc::new(StringArray::from_iter_values(raw_strs));
3819 let mut options = RoundTripOptions::new(array, false);
3820 options.bloom_filter = true;
3821
3822 let files = one_column_roundtrip_with_options(options);
3823
3824 let optional_raw_values: Vec<_> = raw_values
3825 .iter()
3826 .enumerate()
3827 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
3828 .collect();
3829 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
3831 }
3832
3833 #[test]
3834 fn large_binary_single_column() {
3835 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3836 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3837 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3838
3839 values_required::<LargeBinaryArray, _>(many_vecs_iter);
3841 }
3842
3843 #[test]
3844 fn fixed_size_binary_single_column() {
3845 let mut builder = FixedSizeBinaryBuilder::new(4);
3846 builder.append_value(b"0123").unwrap();
3847 builder.append_null();
3848 builder.append_value(b"8910").unwrap();
3849 builder.append_value(b"1112").unwrap();
3850 let array = Arc::new(builder.finish());
3851
3852 one_column_roundtrip(array, true);
3853 }
3854
3855 #[test]
3856 fn string_single_column() {
3857 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3858 let raw_strs = raw_values.iter().map(|s| s.as_str());
3859
3860 required_and_optional::<StringArray, _>(raw_strs);
3861 }
3862
3863 #[test]
3864 fn large_string_single_column() {
3865 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3866 let raw_strs = raw_values.iter().map(|s| s.as_str());
3867
3868 required_and_optional::<LargeStringArray, _>(raw_strs);
3869 }
3870
3871 #[test]
3872 fn string_view_single_column() {
3873 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3874 let raw_strs = raw_values.iter().map(|s| s.as_str());
3875
3876 required_and_optional::<StringViewArray, _>(raw_strs);
3877 }
3878
3879 #[test]
3880 fn null_list_single_column() {
3881 let null_field = Field::new_list_field(DataType::Null, true);
3882 let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
3883
3884 let schema = Schema::new(vec![list_field]);
3885
3886 let a_values = NullArray::new(2);
3888 let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
3889 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3890 DataType::Null,
3891 true,
3892 ))))
3893 .len(3)
3894 .add_buffer(a_value_offsets)
3895 .null_bit_buffer(Some(Buffer::from([0b00000101])))
3896 .add_child_data(a_values.into_data())
3897 .build()
3898 .unwrap();
3899
3900 let a = ListArray::from(a_list_data);
3901
3902 assert!(a.is_valid(0));
3903 assert!(!a.is_valid(1));
3904 assert!(a.is_valid(2));
3905
3906 assert_eq!(a.value(0).len(), 0);
3907 assert_eq!(a.value(2).len(), 2);
3908 assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
3909
3910 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3911 roundtrip(batch, None);
3912 }
3913
3914 #[test]
3915 fn list_single_column() {
3916 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3917 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
3918 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3919 DataType::Int32,
3920 false,
3921 ))))
3922 .len(5)
3923 .add_buffer(a_value_offsets)
3924 .null_bit_buffer(Some(Buffer::from([0b00011011])))
3925 .add_child_data(a_values.into_data())
3926 .build()
3927 .unwrap();
3928
3929 assert_eq!(a_list_data.null_count(), 1);
3930
3931 let a = ListArray::from(a_list_data);
3932 let values = Arc::new(a);
3933
3934 one_column_roundtrip(values, true);
3935 }
3936
3937 #[test]
3938 fn large_list_single_column() {
3939 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3940 let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
3941 let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
3942 "large_item",
3943 DataType::Int32,
3944 true,
3945 ))))
3946 .len(5)
3947 .add_buffer(a_value_offsets)
3948 .add_child_data(a_values.into_data())
3949 .null_bit_buffer(Some(Buffer::from([0b00011011])))
3950 .build()
3951 .unwrap();
3952
3953 assert_eq!(a_list_data.null_count(), 1);
3955
3956 let a = LargeListArray::from(a_list_data);
3957 let values = Arc::new(a);
3958
3959 one_column_roundtrip(values, true);
3960 }
3961
3962 #[test]
3963 fn list_nested_nulls() {
3964 use arrow::datatypes::Int32Type;
3965 let data = vec![
3966 Some(vec![Some(1)]),
3967 Some(vec![Some(2), Some(3)]),
3968 None,
3969 Some(vec![Some(4), Some(5), None]),
3970 Some(vec![None]),
3971 Some(vec![Some(6), Some(7)]),
3972 ];
3973
3974 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
3975 one_column_roundtrip(Arc::new(list), true);
3976
3977 let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
3978 one_column_roundtrip(Arc::new(list), true);
3979 }
3980
3981 #[test]
3982 fn list_utf8_view_selective_padding_roundtrip() {
3983 let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
3984 let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
3985 builder.values().append_value("a");
3986 builder.values().append_null();
3987 builder.append(true);
3988 builder.append(false);
3991 builder.values().append_value("large payload over 12 bytes");
3993 builder.append(true);
3994
3995 one_column_roundtrip(Arc::new(builder.finish()), true);
3996 }
3997
3998 #[test]
3999 fn struct_single_column() {
4000 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4001 let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4002 let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4003
4004 let values = Arc::new(s);
4005 one_column_roundtrip(values, false);
4006 }
4007
4008 #[test]
4009 fn list_and_map_coerced_names() {
4010 let list_field =
4012 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4013 let map_field = Field::new_map(
4014 "my_map",
4015 "my_entries",
4016 Field::new("my_keys", DataType::Int32, false),
4017 Field::new("my_values", DataType::Int32, true),
4018 false,
4019 true,
4020 );
4021
4022 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4023 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4024
4025 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4026
4027 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4029 let file = tempfile::tempfile().unwrap();
4030 let mut writer =
4031 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4032
4033 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4034 writer.write(&batch).unwrap();
4035 let file_metadata = writer.close().unwrap();
4036
4037 let schema = file_metadata.file_metadata().schema();
4038 let list_field = &schema.get_fields()[0].get_fields()[0];
4040 assert_eq!(list_field.get_fields()[0].name(), "element");
4041
4042 let map_field = &schema.get_fields()[1].get_fields()[0];
4043 assert_eq!(map_field.name(), "key_value");
4045 assert_eq!(map_field.get_fields()[0].name(), "key");
4047 assert_eq!(map_field.get_fields()[1].name(), "value");
4049
4050 let reader = SerializedFileReader::new(file).unwrap();
4052 let file_schema = reader.metadata().file_metadata().schema();
4053 let fields = file_schema.get_fields();
4054 let list_field = &fields[0].get_fields()[0];
4055 assert_eq!(list_field.get_fields()[0].name(), "element");
4056 let map_field = &fields[1].get_fields()[0];
4057 assert_eq!(map_field.name(), "key_value");
4058 assert_eq!(map_field.get_fields()[0].name(), "key");
4059 assert_eq!(map_field.get_fields()[1].name(), "value");
4060 }
4061
4062 #[test]
4063 fn fallback_flush_data_page() {
4064 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4066 let values = Arc::new(StringArray::from(raw_values));
4067 let encodings = vec![
4068 Encoding::DELTA_BYTE_ARRAY,
4069 Encoding::DELTA_LENGTH_BYTE_ARRAY,
4070 ];
4071 let data_type = values.data_type().clone();
4072 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4073 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4074
4075 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4076 let data_page_size_limit: usize = 32;
4077 let write_batch_size: usize = 16;
4078
4079 for encoding in &encodings {
4080 for row_group_size in row_group_sizes {
4081 let props = WriterProperties::builder()
4082 .set_writer_version(WriterVersion::PARQUET_2_0)
4083 .set_max_row_group_row_count(Some(row_group_size))
4084 .set_dictionary_enabled(false)
4085 .set_encoding(*encoding)
4086 .set_data_page_size_limit(data_page_size_limit)
4087 .set_write_batch_size(write_batch_size)
4088 .build();
4089
4090 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4091 let string_array_a = StringArray::from(a.clone());
4092 let string_array_b = StringArray::from(b.clone());
4093 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4094 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4095 assert_eq!(
4096 vec_a, vec_b,
4097 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4098 );
4099 });
4100 }
4101 }
4102 }
4103
4104 #[test]
4105 fn arrow_writer_string_dictionary() {
4106 #[allow(deprecated)]
4108 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4109 "dictionary",
4110 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4111 true,
4112 42,
4113 true,
4114 )]));
4115
4116 let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4118 .iter()
4119 .copied()
4120 .collect();
4121
4122 one_column_roundtrip_with_schema(Arc::new(d), schema);
4124 }
4125
4126 #[test]
4127 fn arrow_writer_test_type_compatibility() {
4128 fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4129 where
4130 T1: Array + 'static,
4131 T2: Array + 'static,
4132 {
4133 let schema1 = Arc::new(Schema::new(vec![Field::new(
4134 "a",
4135 array1.data_type().clone(),
4136 false,
4137 )]));
4138
4139 let file = tempfile().unwrap();
4140 let mut writer =
4141 ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4142
4143 let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4144 writer.write(&rb1).unwrap();
4145
4146 let schema2 = Arc::new(Schema::new(vec![Field::new(
4147 "a",
4148 array2.data_type().clone(),
4149 false,
4150 )]));
4151 let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4152 writer.write(&rb2).unwrap();
4153
4154 writer.close().unwrap();
4155
4156 let mut record_batch_reader =
4157 ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4158 let actual_batch = record_batch_reader.next().unwrap().unwrap();
4159
4160 let expected_batch =
4161 RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4162 assert_eq!(actual_batch, expected_batch);
4163 }
4164
4165 ensure_compatible_write(
4168 DictionaryArray::new(
4169 UInt8Array::from_iter_values(vec![0]),
4170 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4171 ),
4172 StringArray::from_iter_values(vec!["barquet"]),
4173 DictionaryArray::new(
4174 UInt8Array::from_iter_values(vec![0, 1]),
4175 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4176 ),
4177 );
4178
4179 ensure_compatible_write(
4180 StringArray::from_iter_values(vec!["parquet"]),
4181 DictionaryArray::new(
4182 UInt8Array::from_iter_values(vec![0]),
4183 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4184 ),
4185 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4186 );
4187
4188 ensure_compatible_write(
4191 DictionaryArray::new(
4192 UInt8Array::from_iter_values(vec![0]),
4193 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4194 ),
4195 DictionaryArray::new(
4196 UInt16Array::from_iter_values(vec![0]),
4197 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4198 ),
4199 DictionaryArray::new(
4200 UInt8Array::from_iter_values(vec![0, 1]),
4201 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4202 ),
4203 );
4204
4205 ensure_compatible_write(
4207 DictionaryArray::new(
4208 UInt8Array::from_iter_values(vec![0]),
4209 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4210 ),
4211 DictionaryArray::new(
4212 UInt8Array::from_iter_values(vec![0]),
4213 Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4214 ),
4215 DictionaryArray::new(
4216 UInt8Array::from_iter_values(vec![0, 1]),
4217 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4218 ),
4219 );
4220
4221 ensure_compatible_write(
4223 DictionaryArray::new(
4224 UInt8Array::from_iter_values(vec![0]),
4225 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4226 ),
4227 LargeStringArray::from_iter_values(vec!["barquet"]),
4228 DictionaryArray::new(
4229 UInt8Array::from_iter_values(vec![0, 1]),
4230 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4231 ),
4232 );
4233
4234 ensure_compatible_write(
4237 StringArray::from_iter_values(vec!["parquet"]),
4238 LargeStringArray::from_iter_values(vec!["barquet"]),
4239 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4240 );
4241
4242 ensure_compatible_write(
4243 LargeStringArray::from_iter_values(vec!["parquet"]),
4244 StringArray::from_iter_values(vec!["barquet"]),
4245 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4246 );
4247
4248 ensure_compatible_write(
4249 StringArray::from_iter_values(vec!["parquet"]),
4250 StringViewArray::from_iter_values(vec!["barquet"]),
4251 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4252 );
4253
4254 ensure_compatible_write(
4255 StringViewArray::from_iter_values(vec!["parquet"]),
4256 StringArray::from_iter_values(vec!["barquet"]),
4257 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4258 );
4259
4260 ensure_compatible_write(
4261 LargeStringArray::from_iter_values(vec!["parquet"]),
4262 StringViewArray::from_iter_values(vec!["barquet"]),
4263 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4264 );
4265
4266 ensure_compatible_write(
4267 StringViewArray::from_iter_values(vec!["parquet"]),
4268 LargeStringArray::from_iter_values(vec!["barquet"]),
4269 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4270 );
4271
4272 ensure_compatible_write(
4275 BinaryArray::from_iter_values(vec![b"parquet"]),
4276 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4277 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4278 );
4279
4280 ensure_compatible_write(
4281 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4282 BinaryArray::from_iter_values(vec![b"barquet"]),
4283 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4284 );
4285
4286 ensure_compatible_write(
4287 BinaryArray::from_iter_values(vec![b"parquet"]),
4288 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4289 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4290 );
4291
4292 ensure_compatible_write(
4293 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4294 BinaryArray::from_iter_values(vec![b"barquet"]),
4295 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4296 );
4297
4298 ensure_compatible_write(
4299 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4300 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4301 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4302 );
4303
4304 ensure_compatible_write(
4305 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4306 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4307 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4308 );
4309
4310 let list_field_metadata = HashMap::from_iter(vec![(
4313 PARQUET_FIELD_ID_META_KEY.to_string(),
4314 "1".to_string(),
4315 )]);
4316 let list_field = Field::new_list_field(DataType::Int32, false);
4317
4318 let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4319 let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4320
4321 let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4322 let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4323
4324 let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4325 let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4326
4327 ensure_compatible_write(
4328 ListArray::try_new(
4330 Arc::new(
4331 list_field
4332 .clone()
4333 .with_metadata(list_field_metadata.clone()),
4334 ),
4335 offsets1,
4336 values1,
4337 None,
4338 )
4339 .unwrap(),
4340 ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4342 ListArray::try_new(
4344 Arc::new(
4345 list_field
4346 .clone()
4347 .with_metadata(list_field_metadata.clone()),
4348 ),
4349 offsets_expected,
4350 values_expected,
4351 None,
4352 )
4353 .unwrap(),
4354 );
4355 }
4356
4357 #[test]
4358 fn arrow_writer_primitive_dictionary() {
4359 #[allow(deprecated)]
4361 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4362 "dictionary",
4363 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4364 true,
4365 42,
4366 true,
4367 )]));
4368
4369 let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4371 builder.append(12345678).unwrap();
4372 builder.append_null();
4373 builder.append(22345678).unwrap();
4374 builder.append(12345678).unwrap();
4375 let d = builder.finish();
4376
4377 one_column_roundtrip_with_schema(Arc::new(d), schema);
4378 }
4379
4380 #[test]
4381 fn arrow_writer_decimal32_dictionary() {
4382 let integers = vec![12345, 56789, 34567];
4383
4384 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4385
4386 let values = Decimal32Array::from(integers.clone())
4387 .with_precision_and_scale(5, 2)
4388 .unwrap();
4389
4390 let array = DictionaryArray::new(keys, Arc::new(values));
4391 one_column_roundtrip(Arc::new(array.clone()), true);
4392
4393 let values = Decimal32Array::from(integers)
4394 .with_precision_and_scale(9, 2)
4395 .unwrap();
4396
4397 let array = array.with_values(Arc::new(values));
4398 one_column_roundtrip(Arc::new(array), true);
4399 }
4400
4401 #[test]
4402 fn arrow_writer_decimal64_dictionary() {
4403 let integers = vec![12345, 56789, 34567];
4404
4405 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4406
4407 let values = Decimal64Array::from(integers.clone())
4408 .with_precision_and_scale(5, 2)
4409 .unwrap();
4410
4411 let array = DictionaryArray::new(keys, Arc::new(values));
4412 one_column_roundtrip(Arc::new(array.clone()), true);
4413
4414 let values = Decimal64Array::from(integers)
4415 .with_precision_and_scale(12, 2)
4416 .unwrap();
4417
4418 let array = array.with_values(Arc::new(values));
4419 one_column_roundtrip(Arc::new(array), true);
4420 }
4421
4422 #[test]
4423 fn arrow_writer_decimal128_dictionary() {
4424 let integers = vec![12345, 56789, 34567];
4425
4426 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4427
4428 let values = Decimal128Array::from(integers.clone())
4429 .with_precision_and_scale(5, 2)
4430 .unwrap();
4431
4432 let array = DictionaryArray::new(keys, Arc::new(values));
4433 one_column_roundtrip(Arc::new(array.clone()), true);
4434
4435 let values = Decimal128Array::from(integers)
4436 .with_precision_and_scale(12, 2)
4437 .unwrap();
4438
4439 let array = array.with_values(Arc::new(values));
4440 one_column_roundtrip(Arc::new(array), true);
4441 }
4442
4443 #[test]
4444 fn arrow_writer_decimal256_dictionary() {
4445 let integers = vec![
4446 i256::from_i128(12345),
4447 i256::from_i128(56789),
4448 i256::from_i128(34567),
4449 ];
4450
4451 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4452
4453 let values = Decimal256Array::from(integers.clone())
4454 .with_precision_and_scale(5, 2)
4455 .unwrap();
4456
4457 let array = DictionaryArray::new(keys, Arc::new(values));
4458 one_column_roundtrip(Arc::new(array.clone()), true);
4459
4460 let values = Decimal256Array::from(integers)
4461 .with_precision_and_scale(12, 2)
4462 .unwrap();
4463
4464 let array = array.with_values(Arc::new(values));
4465 one_column_roundtrip(Arc::new(array), true);
4466 }
4467
4468 #[test]
4469 fn arrow_writer_string_dictionary_unsigned_index() {
4470 #[allow(deprecated)]
4472 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4473 "dictionary",
4474 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4475 true,
4476 42,
4477 true,
4478 )]));
4479
4480 let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4482 .iter()
4483 .copied()
4484 .collect();
4485
4486 one_column_roundtrip_with_schema(Arc::new(d), schema);
4487 }
4488
4489 #[test]
4490 fn u32_min_max() {
4491 let src = [
4493 u32::MIN,
4494 1,
4495 (i32::MAX as u32) - 1,
4496 i32::MAX as u32,
4497 (i32::MAX as u32) + 1,
4498 u32::MAX - 1,
4499 u32::MAX,
4500 ];
4501 let values = Arc::new(UInt32Array::from_iter_values(src.iter().cloned()));
4502 let files = one_column_roundtrip(values, false);
4503
4504 for file in files {
4505 let reader = SerializedFileReader::new(file).unwrap();
4507 let metadata = reader.metadata();
4508
4509 let mut row_offset = 0;
4510 for row_group in metadata.row_groups() {
4511 assert_eq!(row_group.num_columns(), 1);
4512 let column = row_group.column(0);
4513
4514 let num_values = column.num_values() as usize;
4515 let src_slice = &src[row_offset..row_offset + num_values];
4516 row_offset += column.num_values() as usize;
4517
4518 let stats = column.statistics().unwrap();
4519 if let Statistics::Int32(stats) = stats {
4520 assert_eq!(
4521 *stats.min_opt().unwrap() as u32,
4522 *src_slice.iter().min().unwrap()
4523 );
4524 assert_eq!(
4525 *stats.max_opt().unwrap() as u32,
4526 *src_slice.iter().max().unwrap()
4527 );
4528 } else {
4529 panic!("Statistics::Int32 missing")
4530 }
4531 }
4532 }
4533 }
4534
4535 #[test]
4536 fn u64_min_max() {
4537 let src = [
4539 u64::MIN,
4540 1,
4541 (i64::MAX as u64) - 1,
4542 i64::MAX as u64,
4543 (i64::MAX as u64) + 1,
4544 u64::MAX - 1,
4545 u64::MAX,
4546 ];
4547 let values = Arc::new(UInt64Array::from_iter_values(src.iter().cloned()));
4548 let files = one_column_roundtrip(values, false);
4549
4550 for file in files {
4551 let reader = SerializedFileReader::new(file).unwrap();
4553 let metadata = reader.metadata();
4554
4555 let mut row_offset = 0;
4556 for row_group in metadata.row_groups() {
4557 assert_eq!(row_group.num_columns(), 1);
4558 let column = row_group.column(0);
4559
4560 let num_values = column.num_values() as usize;
4561 let src_slice = &src[row_offset..row_offset + num_values];
4562 row_offset += column.num_values() as usize;
4563
4564 let stats = column.statistics().unwrap();
4565 if let Statistics::Int64(stats) = stats {
4566 assert_eq!(
4567 *stats.min_opt().unwrap() as u64,
4568 *src_slice.iter().min().unwrap()
4569 );
4570 assert_eq!(
4571 *stats.max_opt().unwrap() as u64,
4572 *src_slice.iter().max().unwrap()
4573 );
4574 } else {
4575 panic!("Statistics::Int64 missing")
4576 }
4577 }
4578 }
4579 }
4580
4581 #[test]
4582 fn statistics_null_counts_only_nulls() {
4583 let values = Arc::new(UInt64Array::from(vec![None, None]));
4585 let files = one_column_roundtrip(values, true);
4586
4587 for file in files {
4588 let reader = SerializedFileReader::new(file).unwrap();
4590 let metadata = reader.metadata();
4591 assert_eq!(metadata.num_row_groups(), 1);
4592 let row_group = metadata.row_group(0);
4593 assert_eq!(row_group.num_columns(), 1);
4594 let column = row_group.column(0);
4595 let stats = column.statistics().unwrap();
4596 assert_eq!(stats.null_count_opt(), Some(2));
4597 }
4598 }
4599
4600 #[test]
4601 fn test_list_of_struct_roundtrip() {
4602 let int_field = Field::new("a", DataType::Int32, true);
4604 let int_field2 = Field::new("b", DataType::Int32, true);
4605
4606 let int_builder = Int32Builder::with_capacity(10);
4607 let int_builder2 = Int32Builder::with_capacity(10);
4608
4609 let struct_builder = StructBuilder::new(
4610 vec![int_field, int_field2],
4611 vec![Box::new(int_builder), Box::new(int_builder2)],
4612 );
4613 let mut list_builder = ListBuilder::new(struct_builder);
4614
4615 let values = list_builder.values();
4620 values
4621 .field_builder::<Int32Builder>(0)
4622 .unwrap()
4623 .append_value(1);
4624 values
4625 .field_builder::<Int32Builder>(1)
4626 .unwrap()
4627 .append_value(2);
4628 values.append(true);
4629 list_builder.append(true);
4630
4631 list_builder.append(true);
4633
4634 list_builder.append(false);
4636
4637 let values = list_builder.values();
4639 values
4640 .field_builder::<Int32Builder>(0)
4641 .unwrap()
4642 .append_null();
4643 values
4644 .field_builder::<Int32Builder>(1)
4645 .unwrap()
4646 .append_null();
4647 values.append(false);
4648 values
4649 .field_builder::<Int32Builder>(0)
4650 .unwrap()
4651 .append_null();
4652 values
4653 .field_builder::<Int32Builder>(1)
4654 .unwrap()
4655 .append_null();
4656 values.append(false);
4657 list_builder.append(true);
4658
4659 let values = list_builder.values();
4661 values
4662 .field_builder::<Int32Builder>(0)
4663 .unwrap()
4664 .append_null();
4665 values
4666 .field_builder::<Int32Builder>(1)
4667 .unwrap()
4668 .append_value(3);
4669 values.append(true);
4670 list_builder.append(true);
4671
4672 let values = list_builder.values();
4674 values
4675 .field_builder::<Int32Builder>(0)
4676 .unwrap()
4677 .append_value(2);
4678 values
4679 .field_builder::<Int32Builder>(1)
4680 .unwrap()
4681 .append_null();
4682 values.append(true);
4683 list_builder.append(true);
4684
4685 let array = Arc::new(list_builder.finish());
4686
4687 one_column_roundtrip(array, true);
4688 }
4689
4690 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4691 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4692 }
4693
4694 #[test]
4695 fn test_aggregates_records() {
4696 let arrays = [
4697 Int32Array::from((0..100).collect::<Vec<_>>()),
4698 Int32Array::from((0..50).collect::<Vec<_>>()),
4699 Int32Array::from((200..500).collect::<Vec<_>>()),
4700 ];
4701
4702 let schema = Arc::new(Schema::new(vec![Field::new(
4703 "int",
4704 ArrowDataType::Int32,
4705 false,
4706 )]));
4707
4708 let file = tempfile::tempfile().unwrap();
4709
4710 let props = WriterProperties::builder()
4711 .set_max_row_group_row_count(Some(200))
4712 .build();
4713
4714 let mut writer =
4715 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4716
4717 for array in arrays {
4718 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4719 writer.write(&batch).unwrap();
4720 }
4721
4722 writer.close().unwrap();
4723
4724 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4725 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4726
4727 let batches = builder
4728 .with_batch_size(100)
4729 .build()
4730 .unwrap()
4731 .collect::<ArrowResult<Vec<_>>>()
4732 .unwrap();
4733
4734 assert_eq!(batches.len(), 5);
4735 assert!(batches.iter().all(|x| x.num_columns() == 1));
4736
4737 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4738
4739 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4740
4741 let values: Vec<_> = batches
4742 .iter()
4743 .flat_map(|x| {
4744 x.column(0)
4745 .as_any()
4746 .downcast_ref::<Int32Array>()
4747 .unwrap()
4748 .values()
4749 .iter()
4750 .cloned()
4751 })
4752 .collect();
4753
4754 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4755 assert_eq!(&values, &expected_values)
4756 }
4757
4758 #[test]
4759 fn complex_aggregate() {
4760 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4762 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4763 let struct_a = Arc::new(Field::new(
4764 "struct_a",
4765 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4766 true,
4767 ));
4768
4769 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4770 let struct_b = Arc::new(Field::new(
4771 "struct_b",
4772 DataType::Struct(vec![list_a.clone()].into()),
4773 false,
4774 ));
4775
4776 let schema = Arc::new(Schema::new(vec![struct_b]));
4777
4778 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4780 let field_b_array =
4781 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4782
4783 let struct_a_array = StructArray::from(vec![
4784 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4785 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4786 ]);
4787
4788 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4789 .len(5)
4790 .add_buffer(Buffer::from_iter(vec![
4791 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4792 ]))
4793 .null_bit_buffer(Some(Buffer::from_iter(vec![
4794 true, false, true, false, true,
4795 ])))
4796 .child_data(vec![struct_a_array.into_data()])
4797 .build()
4798 .unwrap();
4799
4800 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4801 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4802
4803 let batch1 =
4804 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4805 .unwrap();
4806
4807 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4808 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4809
4810 let struct_a_array = StructArray::from(vec![
4811 (field_a, Arc::new(field_a_array) as ArrayRef),
4812 (field_b, Arc::new(field_b_array) as ArrayRef),
4813 ]);
4814
4815 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4816 .len(2)
4817 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4818 .child_data(vec![struct_a_array.into_data()])
4819 .build()
4820 .unwrap();
4821
4822 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4823 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4824
4825 let batch2 =
4826 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4827 .unwrap();
4828
4829 let batches = &[batch1, batch2];
4830
4831 let expected = r#"
4834 +-------------------------------------------------------------------------------------------------------+
4835 | struct_b |
4836 +-------------------------------------------------------------------------------------------------------+
4837 | {list: [{leaf_a: 1, leaf_b: 1}]} |
4838 | {list: } |
4839 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
4840 | {list: } |
4841 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
4842 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4843 | {list: [{leaf_a: 10, leaf_b: }]} |
4844 +-------------------------------------------------------------------------------------------------------+
4845 "#.trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4846
4847 let actual = pretty_format_batches(batches).unwrap().to_string();
4848 assert_eq!(actual, expected);
4849
4850 let file = tempfile::tempfile().unwrap();
4852 let props = WriterProperties::builder()
4853 .set_max_row_group_row_count(Some(6))
4854 .build();
4855
4856 let mut writer =
4857 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4858
4859 for batch in batches {
4860 writer.write(batch).unwrap();
4861 }
4862 writer.close().unwrap();
4863
4864 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4869 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4870
4871 let batches = builder
4872 .with_batch_size(2)
4873 .build()
4874 .unwrap()
4875 .collect::<ArrowResult<Vec<_>>>()
4876 .unwrap();
4877
4878 assert_eq!(batches.len(), 4);
4879 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4880 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4881
4882 let actual = pretty_format_batches(&batches).unwrap().to_string();
4883 assert_eq!(actual, expected);
4884 }
4885
4886 #[test]
4887 fn test_arrow_writer_metadata() {
4888 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4889 let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4890
4891 let batch = RecordBatch::try_new(
4892 Arc::new(batch_schema),
4893 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4894 )
4895 .unwrap();
4896
4897 let mut buf = Vec::with_capacity(1024);
4898 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4899 writer.write(&batch).unwrap();
4900 writer.close().unwrap();
4901 }
4902
4903 #[test]
4904 fn test_arrow_writer_nullable() {
4905 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4906 let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
4907 let file_schema = Arc::new(file_schema);
4908
4909 let batch = RecordBatch::try_new(
4910 Arc::new(batch_schema),
4911 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4912 )
4913 .unwrap();
4914
4915 let mut buf = Vec::with_capacity(1024);
4916 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
4917 writer.write(&batch).unwrap();
4918 writer.close().unwrap();
4919
4920 let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
4921 let back = read.next().unwrap().unwrap();
4922 assert_eq!(back.schema(), file_schema);
4923 assert_ne!(back.schema(), batch.schema());
4924 assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
4925 }
4926
4927 #[test]
4928 fn in_progress_accounting() {
4929 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4931
4932 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4934
4935 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4937
4938 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4939
4940 assert_eq!(writer.in_progress_size(), 0);
4942 assert_eq!(writer.in_progress_rows(), 0);
4943 assert_eq!(writer.memory_size(), 0);
4944 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
4946
4947 let initial_size = writer.in_progress_size();
4949 assert!(initial_size > 0);
4950 assert_eq!(writer.in_progress_rows(), 5);
4951 let initial_memory = writer.memory_size();
4952 assert!(initial_memory > 0);
4953 assert!(
4955 initial_size <= initial_memory,
4956 "{initial_size} <= {initial_memory}"
4957 );
4958
4959 writer.write(&batch).unwrap();
4961 assert!(writer.in_progress_size() > initial_size);
4962 assert_eq!(writer.in_progress_rows(), 10);
4963 assert!(writer.memory_size() > initial_memory);
4964 assert!(
4965 writer.in_progress_size() <= writer.memory_size(),
4966 "in_progress_size {} <= memory_size {}",
4967 writer.in_progress_size(),
4968 writer.memory_size()
4969 );
4970
4971 let pre_flush_bytes_written = writer.bytes_written();
4973 writer.flush().unwrap();
4974 assert_eq!(writer.in_progress_size(), 0);
4975 assert_eq!(writer.memory_size(), 0);
4976 assert!(writer.bytes_written() > pre_flush_bytes_written);
4977
4978 writer.close().unwrap();
4979 }
4980
4981 #[test]
4982 fn test_writer_all_null() {
4983 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4984 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4985 let batch = RecordBatch::try_from_iter(vec![
4986 ("a", Arc::new(a) as ArrayRef),
4987 ("b", Arc::new(b) as ArrayRef),
4988 ])
4989 .unwrap();
4990
4991 let mut buf = Vec::with_capacity(1024);
4992 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4993 writer.write(&batch).unwrap();
4994 writer.close().unwrap();
4995
4996 let bytes = Bytes::from(buf);
4997 let options = ReadOptionsBuilder::new().with_page_index().build();
4998 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4999 let index = reader.metadata().offset_index().unwrap();
5000
5001 assert_eq!(index.len(), 1);
5002 assert_eq!(index[0].len(), 2); assert_eq!(index[0][0].page_locations().len(), 1); assert_eq!(index[0][1].page_locations().len(), 1); }
5006
5007 #[test]
5008 fn test_disabled_statistics_with_page() {
5009 let file_schema = Schema::new(vec![
5010 Field::new("a", DataType::Utf8, true),
5011 Field::new("b", DataType::Utf8, true),
5012 ]);
5013 let file_schema = Arc::new(file_schema);
5014
5015 let batch = RecordBatch::try_new(
5016 file_schema.clone(),
5017 vec![
5018 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5019 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5020 ],
5021 )
5022 .unwrap();
5023
5024 let props = WriterProperties::builder()
5025 .set_statistics_enabled(EnabledStatistics::None)
5026 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5027 .build();
5028
5029 let mut buf = Vec::with_capacity(1024);
5030 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5031 writer.write(&batch).unwrap();
5032
5033 let metadata = writer.close().unwrap();
5034 assert_eq!(metadata.num_row_groups(), 1);
5035 let row_group = metadata.row_group(0);
5036 assert_eq!(row_group.num_columns(), 2);
5037 assert!(row_group.column(0).offset_index_offset().is_some());
5039 assert!(row_group.column(0).column_index_offset().is_some());
5040 assert!(row_group.column(1).offset_index_offset().is_some());
5042 assert!(row_group.column(1).column_index_offset().is_none());
5043
5044 let options = ReadOptionsBuilder::new().with_page_index().build();
5045 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5046
5047 let row_group = reader.get_row_group(0).unwrap();
5048 let a_col = row_group.metadata().column(0);
5049 let b_col = row_group.metadata().column(1);
5050
5051 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5053 let min = byte_array_stats.min_opt().unwrap();
5054 let max = byte_array_stats.max_opt().unwrap();
5055
5056 assert_eq!(min.as_bytes(), b"a");
5057 assert_eq!(max.as_bytes(), b"d");
5058 } else {
5059 panic!("expecting Statistics::ByteArray");
5060 }
5061
5062 assert!(b_col.statistics().is_none());
5064
5065 let offset_index = reader.metadata().offset_index().unwrap();
5066 assert_eq!(offset_index.len(), 1); assert_eq!(offset_index[0].len(), 2); let column_index = reader.metadata().column_index().unwrap();
5070 assert_eq!(column_index.len(), 1); assert_eq!(column_index[0].len(), 2); let a_idx = &column_index[0][0];
5074 assert!(
5075 matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)),
5076 "{a_idx:?}"
5077 );
5078 let b_idx = &column_index[0][1];
5079 assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5080 }
5081
5082 #[test]
5083 fn test_disabled_statistics_with_chunk() {
5084 let file_schema = Schema::new(vec![
5085 Field::new("a", DataType::Utf8, true),
5086 Field::new("b", DataType::Utf8, true),
5087 ]);
5088 let file_schema = Arc::new(file_schema);
5089
5090 let batch = RecordBatch::try_new(
5091 file_schema.clone(),
5092 vec![
5093 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5094 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5095 ],
5096 )
5097 .unwrap();
5098
5099 let props = WriterProperties::builder()
5100 .set_statistics_enabled(EnabledStatistics::None)
5101 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5102 .build();
5103
5104 let mut buf = Vec::with_capacity(1024);
5105 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5106 writer.write(&batch).unwrap();
5107
5108 let metadata = writer.close().unwrap();
5109 assert_eq!(metadata.num_row_groups(), 1);
5110 let row_group = metadata.row_group(0);
5111 assert_eq!(row_group.num_columns(), 2);
5112 assert!(row_group.column(0).offset_index_offset().is_some());
5114 assert!(row_group.column(0).column_index_offset().is_none());
5115 assert!(row_group.column(1).offset_index_offset().is_some());
5117 assert!(row_group.column(1).column_index_offset().is_none());
5118
5119 let options = ReadOptionsBuilder::new().with_page_index().build();
5120 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5121
5122 let row_group = reader.get_row_group(0).unwrap();
5123 let a_col = row_group.metadata().column(0);
5124 let b_col = row_group.metadata().column(1);
5125
5126 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5128 let min = byte_array_stats.min_opt().unwrap();
5129 let max = byte_array_stats.max_opt().unwrap();
5130
5131 assert_eq!(min.as_bytes(), b"a");
5132 assert_eq!(max.as_bytes(), b"d");
5133 } else {
5134 panic!("expecting Statistics::ByteArray");
5135 }
5136
5137 assert!(b_col.statistics().is_none());
5139
5140 let column_index = reader.metadata().column_index().unwrap();
5141 assert_eq!(column_index.len(), 1); assert_eq!(column_index[0].len(), 2); let a_idx = &column_index[0][0];
5145 assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}");
5146 let b_idx = &column_index[0][1];
5147 assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5148 }
5149
5150 #[test]
5151 fn test_arrow_writer_skip_metadata() {
5152 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5153 let file_schema = Arc::new(batch_schema.clone());
5154
5155 let batch = RecordBatch::try_new(
5156 Arc::new(batch_schema),
5157 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5158 )
5159 .unwrap();
5160 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5161
5162 let mut buf = Vec::with_capacity(1024);
5163 let mut writer =
5164 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5165 writer.write(&batch).unwrap();
5166 writer.close().unwrap();
5167
5168 let bytes = Bytes::from(buf);
5169 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5170 assert_eq!(file_schema, *reader_builder.schema());
5171 if let Some(key_value_metadata) = reader_builder
5172 .metadata()
5173 .file_metadata()
5174 .key_value_metadata()
5175 {
5176 assert!(
5177 !key_value_metadata
5178 .iter()
5179 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5180 );
5181 }
5182 }
5183
5184 #[test]
5185 fn test_arrow_writer_skip_path_in_schema() {
5186 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5187 let file_schema = Arc::new(batch_schema.clone());
5188
5189 let batch = RecordBatch::try_new(
5190 Arc::new(batch_schema),
5191 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5192 )
5193 .unwrap();
5194
5195 let skip_options = ArrowWriterOptions::new();
5197
5198 let mut buf = Vec::with_capacity(1024);
5199 let mut writer =
5200 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5201 writer.write(&batch).unwrap();
5202 writer.close().unwrap();
5203
5204 let skip_options = ArrowWriterOptions::new().with_properties(
5206 WriterProperties::builder()
5207 .set_write_path_in_schema(false)
5208 .build(),
5209 );
5210
5211 let mut buf2 = Vec::with_capacity(1024);
5212 let mut writer =
5213 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5214 .unwrap();
5215 writer.write(&batch).unwrap();
5216 writer.close().unwrap();
5217
5218 assert!(buf.len() > buf2.len());
5220 }
5221
5222 #[test]
5223 fn mismatched_schemas() {
5224 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5225 let file_schema = Arc::new(Schema::new(vec![Field::new(
5226 "temperature",
5227 DataType::Float64,
5228 false,
5229 )]));
5230
5231 let batch = RecordBatch::try_new(
5232 Arc::new(batch_schema),
5233 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5234 )
5235 .unwrap();
5236
5237 let mut buf = Vec::with_capacity(1024);
5238 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5239
5240 let err = writer.write(&batch).unwrap_err().to_string();
5241 assert_eq!(
5242 err,
5243 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5244 );
5245 }
5246
5247 #[test]
5248 fn test_roundtrip_empty_schema() {
5250 let empty_batch = RecordBatch::try_new_with_options(
5252 Arc::new(Schema::empty()),
5253 vec![],
5254 &RecordBatchOptions::default().with_row_count(Some(0)),
5255 )
5256 .unwrap();
5257
5258 let mut parquet_bytes: Vec<u8> = Vec::new();
5260 let mut writer =
5261 ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5262 writer.write(&empty_batch).unwrap();
5263 writer.close().unwrap();
5264
5265 let bytes = Bytes::from(parquet_bytes);
5267 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5268 assert_eq!(reader.schema(), &empty_batch.schema());
5269 let batches: Vec<_> = reader
5270 .build()
5271 .unwrap()
5272 .collect::<ArrowResult<Vec<_>>>()
5273 .unwrap();
5274 assert_eq!(batches.len(), 0);
5275 }
5276
5277 #[test]
5278 fn test_page_stats_not_written_by_default() {
5279 let string_field = Field::new("a", DataType::Utf8, false);
5280 let schema = Schema::new(vec![string_field]);
5281 let raw_string_values = vec!["Blart Versenwald III"];
5282 let string_values = StringArray::from(raw_string_values.clone());
5283 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5284
5285 let props = WriterProperties::builder()
5286 .set_statistics_enabled(EnabledStatistics::Page)
5287 .set_dictionary_enabled(false)
5288 .set_encoding(Encoding::PLAIN)
5289 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5290 .build();
5291
5292 let file = roundtrip_opts(&batch, props);
5293
5294 let first_page = &file[4..];
5299 let mut prot = ThriftSliceInputProtocol::new(first_page);
5300 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5301 let stats = hdr.data_page_header.unwrap().statistics;
5302
5303 assert!(stats.is_none());
5304 }
5305
5306 #[test]
5307 fn test_page_stats_when_enabled() {
5308 let string_field = Field::new("a", DataType::Utf8, false);
5309 let schema = Schema::new(vec![string_field]);
5310 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5311 let string_values = StringArray::from(raw_string_values.clone());
5312 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5313
5314 let props = WriterProperties::builder()
5315 .set_statistics_enabled(EnabledStatistics::Page)
5316 .set_dictionary_enabled(false)
5317 .set_encoding(Encoding::PLAIN)
5318 .set_write_page_header_statistics(true)
5319 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5320 .build();
5321
5322 let file = roundtrip_opts(&batch, props);
5323
5324 let first_page = &file[4..];
5329 let mut prot = ThriftSliceInputProtocol::new(first_page);
5330 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5331 let stats = hdr.data_page_header.unwrap().statistics;
5332
5333 let stats = stats.unwrap();
5334 assert!(stats.is_max_value_exact.unwrap());
5336 assert!(stats.is_min_value_exact.unwrap());
5337 assert_eq!(stats.max_value.unwrap(), "Blart Versenwald III".as_bytes());
5338 assert_eq!(stats.min_value.unwrap(), "Andrew Lamb".as_bytes());
5339 }
5340
5341 #[test]
5342 fn test_page_stats_truncation() {
5343 let string_field = Field::new("a", DataType::Utf8, false);
5344 let binary_field = Field::new("b", DataType::Binary, false);
5345 let schema = Schema::new(vec![string_field, binary_field]);
5346
5347 let raw_string_values = vec!["Blart Versenwald III"];
5348 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5349 let raw_binary_value_refs = raw_binary_values
5350 .iter()
5351 .map(|x| x.as_slice())
5352 .collect::<Vec<_>>();
5353
5354 let string_values = StringArray::from(raw_string_values.clone());
5355 let binary_values = BinaryArray::from(raw_binary_value_refs);
5356 let batch = RecordBatch::try_new(
5357 Arc::new(schema),
5358 vec![Arc::new(string_values), Arc::new(binary_values)],
5359 )
5360 .unwrap();
5361
5362 let props = WriterProperties::builder()
5363 .set_statistics_truncate_length(Some(2))
5364 .set_dictionary_enabled(false)
5365 .set_encoding(Encoding::PLAIN)
5366 .set_write_page_header_statistics(true)
5367 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5368 .build();
5369
5370 let file = roundtrip_opts(&batch, props);
5371
5372 let first_page = &file[4..];
5377 let mut prot = ThriftSliceInputProtocol::new(first_page);
5378 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5379 let stats = hdr.data_page_header.unwrap().statistics;
5380 assert!(stats.is_some());
5381 let stats = stats.unwrap();
5382 assert!(!stats.is_max_value_exact.unwrap());
5384 assert!(!stats.is_min_value_exact.unwrap());
5385 assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5386 assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5387
5388 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5390 let mut prot = ThriftSliceInputProtocol::new(second_page);
5391 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5392 let stats = hdr.data_page_header.unwrap().statistics;
5393 assert!(stats.is_some());
5394 let stats = stats.unwrap();
5395 assert!(!stats.is_max_value_exact.unwrap());
5397 assert!(!stats.is_min_value_exact.unwrap());
5398 assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5399 assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5400 }
5401
5402 #[test]
5403 fn test_page_encoding_statistics_roundtrip() {
5404 let batch_schema = Schema::new(vec![Field::new(
5405 "int32",
5406 arrow_schema::DataType::Int32,
5407 false,
5408 )]);
5409
5410 let batch = RecordBatch::try_new(
5411 Arc::new(batch_schema.clone()),
5412 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5413 )
5414 .unwrap();
5415
5416 let mut file: File = tempfile::tempfile().unwrap();
5417 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5418 writer.write(&batch).unwrap();
5419 let file_metadata = writer.close().unwrap();
5420
5421 assert_eq!(file_metadata.num_row_groups(), 1);
5422 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5423 assert!(
5424 file_metadata
5425 .row_group(0)
5426 .column(0)
5427 .page_encoding_stats()
5428 .is_some()
5429 );
5430 let chunk_page_stats = file_metadata
5431 .row_group(0)
5432 .column(0)
5433 .page_encoding_stats()
5434 .unwrap();
5435
5436 let options = ReadOptionsBuilder::new()
5438 .with_page_index()
5439 .with_encoding_stats_as_mask(false)
5440 .build();
5441 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5442
5443 let rowgroup = reader.get_row_group(0).expect("row group missing");
5444 assert_eq!(rowgroup.num_columns(), 1);
5445 let column = rowgroup.metadata().column(0);
5446 assert!(column.page_encoding_stats().is_some());
5447 let file_page_stats = column.page_encoding_stats().unwrap();
5448 assert_eq!(chunk_page_stats, file_page_stats);
5449 }
5450
5451 #[test]
5452 fn test_different_dict_page_size_limit() {
5453 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5454 let schema = Arc::new(Schema::new(vec![
5455 Field::new("col0", arrow_schema::DataType::Int64, false),
5456 Field::new("col1", arrow_schema::DataType::Int64, false),
5457 ]));
5458 let batch =
5459 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5460
5461 let props = WriterProperties::builder()
5462 .set_dictionary_page_size_limit(1024 * 1024)
5463 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5464 .build();
5465 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5466 writer.write(&batch).unwrap();
5467 let data = Bytes::from(writer.into_inner().unwrap());
5468
5469 let mut metadata = ParquetMetaDataReader::new();
5470 metadata.try_parse(&data).unwrap();
5471 let metadata = metadata.finish().unwrap();
5472 let col0_meta = metadata.row_group(0).column(0);
5473 let col1_meta = metadata.row_group(0).column(1);
5474
5475 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5476 let mut reader =
5477 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5478 let page = reader.get_next_page().unwrap().unwrap();
5479 match page {
5480 Page::DictionaryPage { buf, .. } => buf.len(),
5481 _ => panic!("expected DictionaryPage"),
5482 }
5483 };
5484
5485 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5486 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5487 }
5488
5489 #[test]
5490 fn test_arrow_writer_granular_mode_roundtrip() {
5491 let small = "tiny".to_string();
5500 let big = "x".repeat(64 * 1024);
5501 let strings: Vec<String> = (0..256)
5502 .map(|i| {
5503 if i % 16 == 0 {
5504 big.clone()
5505 } else {
5506 small.clone()
5507 }
5508 })
5509 .collect();
5510
5511 let schema = Arc::new(Schema::new(vec![Field::new(
5512 "col",
5513 ArrowDataType::Utf8,
5514 false,
5515 )]));
5516 let batch = RecordBatch::try_new(
5517 schema.clone(),
5518 vec![Arc::new(StringArray::from(strings.clone())) as _],
5519 )
5520 .unwrap();
5521
5522 let props = WriterProperties::builder()
5523 .set_dictionary_enabled(false)
5524 .set_data_page_size_limit(16 * 1024)
5525 .build();
5526 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5527 writer.write(&batch).unwrap();
5528 let data = Bytes::from(writer.into_inner().unwrap());
5529
5530 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5531 let read = reader.next().unwrap().unwrap();
5532 assert!(reader.next().is_none(), "expected one batch");
5533 let col = read
5534 .column(0)
5535 .as_any()
5536 .downcast_ref::<StringArray>()
5537 .unwrap();
5538 assert_eq!(col.len(), strings.len());
5539 for (i, expected) in strings.iter().enumerate() {
5540 assert_eq!(
5541 col.value(i),
5542 expected.as_str(),
5543 "value mismatch at index {i}"
5544 );
5545 }
5546 }
5547
5548 #[test]
5549 fn test_arrow_writer_all_null_string_column() {
5550 let num_rows = 1024;
5555 let schema = Arc::new(Schema::new(vec![Field::new(
5556 "col",
5557 ArrowDataType::Utf8,
5558 true,
5559 )]));
5560 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5561 let batch = RecordBatch::try_new(
5562 schema.clone(),
5563 vec![Arc::new(StringArray::from(nulls)) as _],
5564 )
5565 .unwrap();
5566
5567 let props = WriterProperties::builder()
5568 .set_dictionary_enabled(false)
5569 .set_data_page_size_limit(16 * 1024)
5570 .build();
5571 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5572 writer.write(&batch).unwrap();
5573 let data = Bytes::from(writer.into_inner().unwrap());
5574
5575 let mut metadata = ParquetMetaDataReader::new();
5578 metadata.try_parse(&data).unwrap();
5579 let metadata = metadata.finish().unwrap();
5580 let row_group = metadata.row_group(0);
5581 let col_meta = row_group.column(0);
5582 assert_eq!(row_group.num_rows() as usize, num_rows);
5583 if let Some(stats) = col_meta.statistics() {
5586 assert_eq!(
5587 stats.null_count_opt().unwrap_or(0) as usize,
5588 num_rows,
5589 "expected all-null column to report null_count = num_rows"
5590 );
5591 }
5592
5593 let mut reader =
5594 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5595 let mut total_values = 0u32;
5596 while let Some(page) = reader.get_next_page().unwrap() {
5597 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5598 total_values += page.num_values();
5599 }
5600 }
5601 assert_eq!(
5602 total_values as usize, num_rows,
5603 "expected every level position to be represented in some page"
5604 );
5605 }
5606
5607 struct WriteBatchesShape {
5608 num_batches: usize,
5609 rows_per_batch: usize,
5610 row_size: usize,
5611 }
5612
5613 fn write_batches(
5615 WriteBatchesShape {
5616 num_batches,
5617 rows_per_batch,
5618 row_size,
5619 }: WriteBatchesShape,
5620 props: WriterProperties,
5621 ) -> ParquetRecordBatchReaderBuilder<File> {
5622 let schema = Arc::new(Schema::new(vec![Field::new(
5623 "str",
5624 ArrowDataType::Utf8,
5625 false,
5626 )]));
5627 let file = tempfile::tempfile().unwrap();
5628 let mut writer =
5629 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5630
5631 for batch_idx in 0..num_batches {
5632 let strings: Vec<String> = (0..rows_per_batch)
5633 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5634 .collect();
5635 let array = StringArray::from(strings);
5636 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5637 writer.write(&batch).unwrap();
5638 }
5639 writer.close().unwrap();
5640 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5641 }
5642
5643 #[test]
5644 fn test_row_group_limit_none_writes_single_row_group() {
5646 let props = WriterProperties::builder()
5647 .set_max_row_group_row_count(None)
5648 .set_max_row_group_bytes(None)
5649 .build();
5650
5651 let builder = write_batches(
5652 WriteBatchesShape {
5653 num_batches: 1,
5654 rows_per_batch: 1000,
5655 row_size: 4,
5656 },
5657 props,
5658 );
5659
5660 assert_eq!(
5661 &row_group_sizes(builder.metadata()),
5662 &[1000],
5663 "With no limits, all rows should be in a single row group"
5664 );
5665 }
5666
5667 #[test]
5668 fn test_row_group_limit_rows_only() {
5670 let props = WriterProperties::builder()
5671 .set_max_row_group_row_count(Some(300))
5672 .set_max_row_group_bytes(None)
5673 .build();
5674
5675 let builder = write_batches(
5676 WriteBatchesShape {
5677 num_batches: 1,
5678 rows_per_batch: 1000,
5679 row_size: 4,
5680 },
5681 props,
5682 );
5683
5684 assert_eq!(
5685 &row_group_sizes(builder.metadata()),
5686 &[300, 300, 300, 100],
5687 "Row groups should be split by row count"
5688 );
5689 }
5690
5691 #[test]
5692 fn test_row_group_limit_bytes_only() {
5694 let props = WriterProperties::builder()
5695 .set_max_row_group_row_count(None)
5696 .set_max_row_group_bytes(Some(3500))
5698 .build();
5699
5700 let builder = write_batches(
5701 WriteBatchesShape {
5702 num_batches: 10,
5703 rows_per_batch: 10,
5704 row_size: 100,
5705 },
5706 props,
5707 );
5708
5709 let sizes = row_group_sizes(builder.metadata());
5710
5711 assert!(
5712 sizes.len() > 1,
5713 "Should have multiple row groups due to byte limit, got {sizes:?}",
5714 );
5715
5716 let total_rows: i64 = sizes.iter().sum();
5717 assert_eq!(total_rows, 100, "Total rows should be preserved");
5718 }
5719
5720 #[test]
5721 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5723 let schema = Arc::new(Schema::new(vec![Field::new(
5724 "str",
5725 ArrowDataType::Utf8,
5726 false,
5727 )]));
5728 let file = tempfile::tempfile().unwrap();
5729
5730 let props = WriterProperties::builder()
5732 .set_max_row_group_row_count(None)
5733 .set_max_row_group_bytes(None)
5734 .build();
5735 let mut writer =
5736 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5737
5738 let first_array = StringArray::from(
5739 (0..10)
5740 .map(|i| format!("{:0>100}", i))
5741 .collect::<Vec<String>>(),
5742 );
5743 let first_batch =
5744 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5745 writer.write(&first_batch).unwrap();
5746 assert_eq!(writer.in_progress_rows(), 10);
5747
5748 writer.max_row_group_bytes = Some(1);
5751
5752 let second_array = StringArray::from(vec!["x".to_string()]);
5753 let second_batch =
5754 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5755 writer.write(&second_batch).unwrap();
5756 writer.close().unwrap();
5757 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5758
5759 assert_eq!(
5760 &row_group_sizes(builder.metadata()),
5761 &[10, 1],
5762 "The second write should flush an oversized in-progress row group first",
5763 );
5764 }
5765
5766 #[test]
5767 fn test_row_group_limit_both_row_wins_single_batch() {
5769 let props = WriterProperties::builder()
5770 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
5773
5774 let builder = write_batches(
5775 WriteBatchesShape {
5776 num_batches: 1,
5777 row_size: 4,
5778 rows_per_batch: 1000,
5779 },
5780 props,
5781 );
5782
5783 assert_eq!(
5784 &row_group_sizes(builder.metadata()),
5785 &[200, 200, 200, 200, 200],
5786 "Row limit should trigger before byte limit"
5787 );
5788 }
5789
5790 #[test]
5791 fn test_row_group_limit_both_row_wins_multiple_batches() {
5793 let props = WriterProperties::builder()
5794 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
5797
5798 let builder = write_batches(
5799 WriteBatchesShape {
5800 num_batches: 10,
5801 rows_per_batch: 10,
5802 row_size: 100,
5803 },
5804 props,
5805 );
5806
5807 assert_eq!(
5808 &row_group_sizes(builder.metadata()),
5809 &[5; 20],
5810 "Row limit should trigger before byte limit"
5811 );
5812 }
5813
5814 #[test]
5815 fn test_row_group_limit_both_bytes_wins() {
5817 let props = WriterProperties::builder()
5818 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
5821
5822 let builder = write_batches(
5823 WriteBatchesShape {
5824 num_batches: 10,
5825 rows_per_batch: 10,
5826 row_size: 100,
5827 },
5828 props,
5829 );
5830
5831 let sizes = row_group_sizes(builder.metadata());
5832
5833 assert!(
5834 sizes.len() > 1,
5835 "Byte limit should trigger before row limit, got {sizes:?}",
5836 );
5837
5838 assert!(
5839 sizes.iter().all(|&s| s < 1000),
5840 "No row group should hit the row limit"
5841 );
5842
5843 let total_rows: i64 = sizes.iter().sum();
5844 assert_eq!(total_rows, 100, "Total rows should be preserved");
5845 }
5846
5847 #[test]
5848 fn arrow_column_chunk_close_mut_drops_column_index() {
5849 use crate::arrow::ArrowSchemaConverter;
5850 use crate::file::writer::SerializedFileWriter;
5851
5852 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5853 let props = Arc::new(
5854 WriterProperties::builder()
5855 .set_statistics_enabled(EnabledStatistics::Page)
5856 .build(),
5857 );
5858 let parquet_schema = ArrowSchemaConverter::new()
5859 .with_coerce_types(props.coerce_types())
5860 .convert(&schema)
5861 .unwrap();
5862
5863 let mut buf = Vec::with_capacity(1024);
5864 let mut writer =
5865 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5866 .unwrap();
5867
5868 let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
5869 let mut col_writers = factory.create_column_writers(0).unwrap();
5870 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5871 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5872 col_writers[0].write(&leaves).unwrap();
5873 }
5874 let mut chunk = col_writers.pop().unwrap().close().unwrap();
5875
5876 assert!(
5878 chunk.close().column_index.is_some(),
5879 "EnabledStatistics::Page should produce a column_index"
5880 );
5881
5882 chunk.close_mut().column_index = None;
5884 assert!(chunk.close().column_index.is_none());
5885
5886 let mut rg = writer.next_row_group().unwrap();
5887 chunk.append_to_row_group(&mut rg).unwrap();
5888 rg.close().unwrap();
5889 let file_meta = writer.close().unwrap();
5890
5891 let cc = file_meta.row_group(0).column(0);
5894 assert!(cc.column_index_range().is_none());
5895 }
5896
5897 fn write_column_to_bytes(array: ArrayRef) -> Bytes {
5899 let schema = Arc::new(Schema::new(vec![Field::new(
5900 "col",
5901 array.data_type().clone(),
5902 true,
5903 )]));
5904 let buf = get_bytes_after_close(
5905 schema.clone(),
5906 &RecordBatch::try_new(schema, vec![array]).unwrap(),
5907 );
5908 Bytes::from(buf)
5909 }
5910
5911 fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
5915 let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
5916 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
5917 .unwrap()
5918 .build()
5919 .unwrap()
5920 .next()
5921 .unwrap()
5922 .unwrap()
5923 .column(0)
5924 .clone()
5925 }
5926
5927 fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
5928 let flat_schema = Arc::new(Schema::new(vec![Field::new(
5929 "col",
5930 flat.data_type().clone(),
5931 true,
5932 )]));
5933 let ree_bytes = write_column_to_bytes(ree);
5934 let flat_bytes = write_column_to_bytes(flat.clone());
5935 assert_eq!(
5936 ree_bytes, flat_bytes,
5937 "REE and flat bytes should be identical"
5938 );
5939
5940 let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
5941 let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
5942
5943 assert_eq!(decoded_ree.as_ref(), flat.as_ref());
5944 assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
5945 }
5946
5947 #[test]
5948 fn ree_string() {
5949 let ree: ArrayRef = Arc::new(
5950 [Some("a"), Some("a"), None, Some("b"), Some("b")]
5951 .into_iter()
5952 .collect::<Int32RunArray>(),
5953 );
5954 let flat: ArrayRef = Arc::new(StringArray::from(vec![
5955 Some("a"),
5956 Some("a"),
5957 None,
5958 Some("b"),
5959 Some("b"),
5960 ]));
5961 ree_write_read_roundtrip(ree, flat);
5962 }
5963
5964 #[test]
5965 fn ree_int32() {
5966 let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
5967 for v in [Some(1), Some(1), None, Some(2), Some(2)] {
5968 b.append_option(v);
5969 }
5970 let ree: ArrayRef = Arc::new(b.finish());
5971 let flat: ArrayRef = Arc::new(Int32Array::from(vec![
5972 Some(1),
5973 Some(1),
5974 None,
5975 Some(2),
5976 Some(2),
5977 ]));
5978 ree_write_read_roundtrip(ree, flat);
5979 }
5980
5981 #[test]
5982 fn ree_bool() {
5983 let ree: ArrayRef = Arc::new(
5985 RunArray::try_new(
5986 &Int32Array::from(vec![3, 5, 7]),
5987 &BooleanArray::from(vec![Some(true), None, Some(false)]),
5988 )
5989 .unwrap(),
5990 );
5991 let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
5992 Some(true),
5993 Some(true),
5994 Some(true),
5995 None,
5996 None,
5997 Some(false),
5998 Some(false),
5999 ]));
6000 ree_write_read_roundtrip(ree, flat);
6001 }
6002
6003 #[test]
6004 fn ree_fixed_size_binary() {
6005 let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6006 let mut b = FixedSizeBinaryBuilder::new(2);
6007 for v in vals {
6008 match v {
6009 Some(x) => b.append_value(x).unwrap(),
6010 None => b.append_null(),
6011 }
6012 }
6013 b.finish()
6014 };
6015 let ree: ArrayRef = Arc::new(
6017 RunArray::try_new(
6018 &Int32Array::from(vec![2, 4, 6]),
6019 &mk(&[Some(b"aa"), None, Some(b"bb")]),
6020 )
6021 .unwrap(),
6022 );
6023 let flat: ArrayRef = Arc::new(mk(&[
6024 Some(b"aa"),
6025 Some(b"aa"),
6026 None,
6027 None,
6028 Some(b"bb"),
6029 Some(b"bb"),
6030 ]));
6031 ree_write_read_roundtrip(ree, flat);
6032 }
6033
6034 #[test]
6035 fn ree_single_run() {
6036 let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6037 let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6038 ree_write_read_roundtrip(ree, flat);
6039 }
6040
6041 #[test]
6042 fn ree_float32() {
6043 let ree: ArrayRef = Arc::new(
6045 RunArray::try_new(
6046 &Int32Array::from(vec![2, 4, 5]),
6047 &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6048 )
6049 .unwrap(),
6050 );
6051 let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6052 Some(1.0_f32),
6053 Some(1.0_f32),
6054 None,
6055 None,
6056 Some(2.5_f32),
6057 ]));
6058 ree_write_read_roundtrip(ree, flat);
6059 }
6060
6061 #[test]
6062 fn ree_sliced() {
6063 let full: ArrayRef = Arc::new(
6068 RunArray::try_new(
6069 &Int32Array::from(vec![3, 5, 7]),
6070 &StringArray::from(vec!["a", "b", "c"]),
6071 )
6072 .unwrap(),
6073 );
6074 let sliced = full.slice(2, 5);
6075 let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6076 ree_write_read_roundtrip(sliced, flat);
6077 }
6078
6079 #[test]
6080 fn ree_struct_with_ree_child() {
6081 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6084
6085 let col_a: ArrayRef = Arc::new(
6086 RunArray::try_new(
6087 &run_ends,
6088 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6089 )
6090 .unwrap(),
6091 );
6092 let col_b: ArrayRef = Arc::new(
6093 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6094 );
6095
6096 let struct_array: ArrayRef = Arc::new(StructArray::new(
6097 Fields::from(vec![
6098 Field::new("a", col_a.data_type().clone(), true),
6099 Field::new("b", col_b.data_type().clone(), true),
6100 ]),
6101 vec![col_a, col_b],
6102 None,
6103 ));
6104
6105 let schema = Arc::new(Schema::new(vec![Field::new(
6106 "row",
6107 struct_array.data_type().clone(),
6108 true,
6109 )]));
6110 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6111
6112 let mut buf = Vec::new();
6113 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6114 writer.write(&batch).unwrap();
6115 let metadata = writer.close().unwrap();
6116
6117 let parquet_schema = metadata.file_metadata().schema_descr();
6118 assert_eq!(parquet_schema.num_columns(), 2);
6119 assert_eq!(
6120 parquet_schema.column(0).physical_type(),
6121 crate::basic::Type::BYTE_ARRAY
6122 );
6123 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6124 assert_eq!(
6125 parquet_schema.column(1).physical_type(),
6126 crate::basic::Type::INT32
6127 );
6128 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6129 }
6130}