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