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