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