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 #[test]
2681 fn arrow_writer_string_view_dictionary() {
2682 let raw_string_values = vec!["a", "b", "large payload over 12 bytes"];
2683 let raw_binary_values = vec![
2684 b"a".to_vec(),
2685 b"b".to_vec(),
2686 b"large payload over 12 bytes".to_vec(),
2687 ];
2688
2689 let keys = UInt32Array::from(vec![Some(0), None, Some(2), Some(1), None]);
2690
2691 let string_view_values = Arc::new(StringViewArray::from(raw_string_values));
2692 let string_dict: ArrayRef = Arc::new(
2693 DictionaryArray::<UInt32Type>::try_new(keys.clone(), string_view_values).unwrap(),
2694 );
2695
2696 let binary_view_values = Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2697 let binary_dict: ArrayRef =
2698 Arc::new(DictionaryArray::<UInt32Type>::try_new(keys, binary_view_values).unwrap());
2699
2700 RoundTripTest::new(string_dict).run();
2701 RoundTripTest::new(binary_dict).run();
2702 }
2703
2704 fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2705 let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2706 let schema = Schema::new(vec![decimal_field]);
2707
2708 let decimal_values = vec![10_000, 50_000, 0, -100]
2709 .into_iter()
2710 .map(Some)
2711 .collect::<Decimal128Array>()
2712 .with_precision_and_scale(precision, scale)
2713 .unwrap();
2714
2715 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2716 }
2717
2718 #[test]
2719 fn arrow_writer_decimal() {
2720 let batch_int32_decimal = get_decimal_batch(5, 2);
2722 roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2723 let batch_int64_decimal = get_decimal_batch(12, 2);
2725 roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2726 let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2728 roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2729 }
2730
2731 #[test]
2732 fn arrow_writer_complex() {
2733 let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2735 let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2736 let struct_field_g = Arc::new(Field::new_list(
2737 "g",
2738 Field::new_list_field(DataType::Int16, true),
2739 false,
2740 ));
2741 let struct_field_h = Arc::new(Field::new_list(
2742 "h",
2743 Field::new_list_field(DataType::Int16, false),
2744 true,
2745 ));
2746 let struct_field_e = Arc::new(Field::new_struct(
2747 "e",
2748 vec![
2749 struct_field_f.clone(),
2750 struct_field_g.clone(),
2751 struct_field_h.clone(),
2752 ],
2753 false,
2754 ));
2755 let schema = Schema::new(vec![
2756 Field::new("a", DataType::Int32, false),
2757 Field::new("b", DataType::Int32, true),
2758 Field::new_struct(
2759 "c",
2760 vec![struct_field_d.clone(), struct_field_e.clone()],
2761 false,
2762 ),
2763 ]);
2764
2765 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2767 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2768 let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2769 let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2770
2771 let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2772
2773 let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2776
2777 let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2779 .len(5)
2780 .add_buffer(g_value_offsets.clone())
2781 .add_child_data(g_value.to_data())
2782 .build()
2783 .unwrap();
2784 let g = ListArray::from(g_list_data);
2785 let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2787 .len(5)
2788 .add_buffer(g_value_offsets)
2789 .add_child_data(g_value.to_data())
2790 .null_bit_buffer(Some(Buffer::from([0b00011011])))
2791 .build()
2792 .unwrap();
2793 let h = ListArray::from(h_list_data);
2794
2795 let e = StructArray::from(vec![
2796 (struct_field_f, Arc::new(f) as ArrayRef),
2797 (struct_field_g, Arc::new(g) as ArrayRef),
2798 (struct_field_h, Arc::new(h) as ArrayRef),
2799 ]);
2800
2801 let c = StructArray::from(vec![
2802 (struct_field_d, Arc::new(d) as ArrayRef),
2803 (struct_field_e, Arc::new(e) as ArrayRef),
2804 ]);
2805
2806 let batch = RecordBatch::try_new(
2808 Arc::new(schema),
2809 vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2810 )
2811 .unwrap();
2812
2813 roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2814 roundtrip(batch, Some(SMALL_SIZE / 3));
2815 }
2816
2817 #[test]
2818 fn arrow_writer_complex_mixed() {
2819 let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2824 let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2825 let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2826 let schema = Schema::new(vec![Field::new(
2827 "some_nested_object",
2828 DataType::Struct(Fields::from(vec![
2829 offset_field.clone(),
2830 partition_field.clone(),
2831 topic_field.clone(),
2832 ])),
2833 false,
2834 )]);
2835
2836 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2838 let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2839 let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2840
2841 let some_nested_object = StructArray::from(vec![
2842 (offset_field, Arc::new(offset) as ArrayRef),
2843 (partition_field, Arc::new(partition) as ArrayRef),
2844 (topic_field, Arc::new(topic) as ArrayRef),
2845 ]);
2846
2847 let batch =
2849 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2850
2851 roundtrip(batch, Some(SMALL_SIZE / 2));
2852 }
2853
2854 #[test]
2855 fn arrow_writer_map() {
2856 let json_content = r#"
2858 {"stocks":{"long": "$AAA", "short": "$BBB"}}
2859 {"stocks":{"long": null, "long": "$CCC", "short": null}}
2860 {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2861 "#;
2862 let entries_struct_type = DataType::Struct(Fields::from(vec![
2863 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2864 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2865 ]));
2866 let stocks_field = Field::new(
2867 "stocks",
2868 DataType::Map(
2869 Arc::new(Field::new(
2870 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2871 entries_struct_type,
2872 false,
2873 )),
2874 false,
2875 ),
2876 true,
2877 );
2878 let schema = Arc::new(Schema::new(vec![stocks_field]));
2879 let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2880 let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2881
2882 let batch = reader.next().unwrap().unwrap();
2883 roundtrip(batch, None);
2884 }
2885
2886 #[test]
2887 fn arrow_writer_2_level_struct() {
2888 let field_c = Field::new("c", DataType::Int32, true);
2890 let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2891 let type_a = DataType::Struct(vec![field_b.clone()].into());
2892 let field_a = Field::new("a", type_a, true);
2893 let schema = Schema::new(vec![field_a.clone()]);
2894
2895 let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2897 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2898 .len(6)
2899 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2900 .add_child_data(c.into_data())
2901 .build()
2902 .unwrap();
2903 let b = StructArray::from(b_data);
2904 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2905 .len(6)
2906 .null_bit_buffer(Some(Buffer::from([0b00101111])))
2907 .add_child_data(b.into_data())
2908 .build()
2909 .unwrap();
2910 let a = StructArray::from(a_data);
2911
2912 assert_eq!(a.null_count(), 1);
2913 assert_eq!(a.column(0).null_count(), 2);
2914
2915 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2917
2918 roundtrip(batch, Some(SMALL_SIZE / 2));
2919 }
2920
2921 #[test]
2922 fn arrow_writer_2_level_struct_non_null() {
2923 let field_c = Field::new("c", DataType::Int32, false);
2925 let type_b = DataType::Struct(vec![field_c].into());
2926 let field_b = Field::new("b", type_b.clone(), false);
2927 let type_a = DataType::Struct(vec![field_b].into());
2928 let field_a = Field::new("a", type_a.clone(), false);
2929 let schema = Schema::new(vec![field_a]);
2930
2931 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2933 let b_data = ArrayDataBuilder::new(type_b)
2934 .len(6)
2935 .add_child_data(c.into_data())
2936 .build()
2937 .unwrap();
2938 let b = StructArray::from(b_data);
2939 let a_data = ArrayDataBuilder::new(type_a)
2940 .len(6)
2941 .add_child_data(b.into_data())
2942 .build()
2943 .unwrap();
2944 let a = StructArray::from(a_data);
2945
2946 assert_eq!(a.null_count(), 0);
2947 assert_eq!(a.column(0).null_count(), 0);
2948
2949 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2951
2952 roundtrip(batch, Some(SMALL_SIZE / 2));
2953 }
2954
2955 #[test]
2956 fn arrow_writer_2_level_struct_mixed_null() {
2957 let field_c = Field::new("c", DataType::Int32, false);
2959 let type_b = DataType::Struct(vec![field_c].into());
2960 let field_b = Field::new("b", type_b.clone(), true);
2961 let type_a = DataType::Struct(vec![field_b].into());
2962 let field_a = Field::new("a", type_a.clone(), false);
2963 let schema = Schema::new(vec![field_a]);
2964
2965 let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2967 let b_data = ArrayDataBuilder::new(type_b)
2968 .len(6)
2969 .null_bit_buffer(Some(Buffer::from([0b00100111])))
2970 .add_child_data(c.into_data())
2971 .build()
2972 .unwrap();
2973 let b = StructArray::from(b_data);
2974 let a_data = ArrayDataBuilder::new(type_a)
2976 .len(6)
2977 .add_child_data(b.into_data())
2978 .build()
2979 .unwrap();
2980 let a = StructArray::from(a_data);
2981
2982 assert_eq!(a.null_count(), 0);
2983 assert_eq!(a.column(0).null_count(), 2);
2984
2985 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2987
2988 roundtrip(batch, Some(SMALL_SIZE / 2));
2989 }
2990
2991 #[test]
2992 fn arrow_writer_2_level_struct_mixed_null_2() {
2993 let field_c = Field::new("c", DataType::Int32, false);
2995 let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
2996 let field_e = Field::new(
2997 "e",
2998 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2999 false,
3000 );
3001
3002 let field_b = Field::new(
3003 "b",
3004 DataType::Struct(vec![field_c, field_d, field_e].into()),
3005 false,
3006 );
3007 let type_a = DataType::Struct(vec![field_b.clone()].into());
3008 let field_a = Field::new("a", type_a, true);
3009 let schema = Schema::new(vec![field_a.clone()]);
3010
3011 let c = Int32Array::from_iter_values(0..6);
3013 let d = FixedSizeBinaryArray::try_from_iter(
3014 ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
3015 )
3016 .expect("four byte values");
3017 let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
3018 let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
3019 .len(6)
3020 .add_child_data(c.into_data())
3021 .add_child_data(d.into_data())
3022 .add_child_data(e.into_data())
3023 .build()
3024 .unwrap();
3025 let b = StructArray::from(b_data);
3026 let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
3027 .len(6)
3028 .null_bit_buffer(Some(Buffer::from([0b00100101])))
3029 .add_child_data(b.into_data())
3030 .build()
3031 .unwrap();
3032 let a = StructArray::from(a_data);
3033
3034 assert_eq!(a.null_count(), 3);
3035 assert_eq!(a.column(0).null_count(), 0);
3036
3037 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3039
3040 roundtrip(batch, Some(SMALL_SIZE / 2));
3041 }
3042
3043 #[test]
3044 fn test_fixed_size_binary_in_dict() {
3045 fn test_fixed_size_binary_in_dict_inner<K>()
3046 where
3047 K: ArrowDictionaryKeyType,
3048 K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
3049 <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
3050 {
3051 let field = Field::new(
3052 "a",
3053 DataType::Dictionary(
3054 Box::new(K::DATA_TYPE),
3055 Box::new(DataType::FixedSizeBinary(4)),
3056 ),
3057 false,
3058 );
3059 let schema = Schema::new(vec![field]);
3060
3061 let keys: Vec<K::Native> = vec![
3062 K::Native::try_from(0u8).unwrap(),
3063 K::Native::try_from(0u8).unwrap(),
3064 K::Native::try_from(1u8).unwrap(),
3065 ];
3066 let keys = PrimitiveArray::<K>::from_iter_values(keys);
3067 let values = FixedSizeBinaryArray::try_from_iter(
3068 vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
3069 )
3070 .unwrap();
3071
3072 let data = DictionaryArray::<K>::new(keys, Arc::new(values));
3073 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
3074 roundtrip(batch, None);
3075 }
3076
3077 test_fixed_size_binary_in_dict_inner::<UInt8Type>();
3078 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3079 test_fixed_size_binary_in_dict_inner::<UInt32Type>();
3080 test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3081 test_fixed_size_binary_in_dict_inner::<Int8Type>();
3082 test_fixed_size_binary_in_dict_inner::<Int16Type>();
3083 test_fixed_size_binary_in_dict_inner::<Int32Type>();
3084 test_fixed_size_binary_in_dict_inner::<Int64Type>();
3085 }
3086
3087 #[test]
3088 fn test_empty_dict() {
3089 let struct_fields = Fields::from(vec![Field::new(
3090 "dict",
3091 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3092 false,
3093 )]);
3094
3095 let schema = Schema::new(vec![Field::new_struct(
3096 "struct",
3097 struct_fields.clone(),
3098 true,
3099 )]);
3100 let dictionary = Arc::new(DictionaryArray::new(
3101 Int32Array::new_null(5),
3102 Arc::new(StringArray::new_null(0)),
3103 ));
3104
3105 let s = StructArray::new(
3106 struct_fields,
3107 vec![dictionary],
3108 Some(NullBuffer::new_null(5)),
3109 );
3110
3111 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
3112 roundtrip(batch, None);
3113 }
3114 #[test]
3115 fn arrow_writer_page_size() {
3116 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3117
3118 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3119
3120 for i in 0..10 {
3122 let value = i
3123 .to_string()
3124 .repeat(10)
3125 .chars()
3126 .take(10)
3127 .collect::<String>();
3128
3129 builder.append_value(value);
3130 }
3131
3132 let array = Arc::new(builder.finish());
3133
3134 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3135
3136 let file = tempfile::tempfile().unwrap();
3137
3138 let props = WriterProperties::builder()
3140 .set_data_page_size_limit(1)
3141 .set_dictionary_page_size_limit(1)
3142 .set_write_batch_size(1)
3143 .build();
3144
3145 let mut writer =
3146 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3147 .expect("Unable to write file");
3148 writer.write(&batch).unwrap();
3149 writer.close().unwrap();
3150
3151 let options = ReadOptionsBuilder::new().with_page_index().build();
3152 let reader =
3153 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3154
3155 let column = reader.metadata().row_group(0).columns();
3156
3157 assert_eq!(column.len(), 1);
3158
3159 assert!(
3162 column[0].dictionary_page_offset().is_some(),
3163 "Expected a dictionary page"
3164 );
3165
3166 let page_index = reader
3167 .metadata()
3168 .page_index()
3169 .expect("page index should be present");
3170 let page_locations = page_index
3171 .page_locations(0, 0)
3172 .expect("page locations should exist");
3173
3174 assert_eq!(
3177 page_locations.len(),
3178 10,
3179 "Expected 10 pages but got {page_locations:#?}"
3180 );
3181 }
3182
3183 #[test]
3184 fn arrow_writer_float_nans() {
3185 let f16_field = Field::new("a", DataType::Float16, false);
3186 let f32_field = Field::new("b", DataType::Float32, false);
3187 let f64_field = Field::new("c", DataType::Float64, false);
3188 let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3189
3190 let f16_values = (0..MEDIUM_SIZE)
3191 .map(|i| {
3192 Some(if i % 2 == 0 {
3193 f16::NAN
3194 } else {
3195 f16::from_f32(i as f32)
3196 })
3197 })
3198 .collect::<Float16Array>();
3199
3200 let f32_values = (0..MEDIUM_SIZE)
3201 .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3202 .collect::<Float32Array>();
3203
3204 let f64_values = (0..MEDIUM_SIZE)
3205 .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3206 .collect::<Float64Array>();
3207
3208 let batch = RecordBatch::try_new(
3209 Arc::new(schema),
3210 vec![
3211 Arc::new(f16_values),
3212 Arc::new(f32_values),
3213 Arc::new(f64_values),
3214 ],
3215 )
3216 .unwrap();
3217
3218 roundtrip(batch, None);
3219 }
3220
3221 const SMALL_SIZE: usize = 7;
3222 const MEDIUM_SIZE: usize = 63;
3223
3224 fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3227 let mut files = vec![];
3228 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3229 let mut props = WriterProperties::builder().set_writer_version(version);
3230
3231 if let Some(size) = max_row_group_size {
3232 props = props.set_max_row_group_row_count(Some(size))
3233 }
3234
3235 let props = props.build();
3236 files.push(roundtrip_opts(&expected_batch, props))
3237 }
3238 files
3239 }
3240
3241 fn roundtrip_opts_with_array_validation<F>(
3245 expected_batch: &RecordBatch,
3246 props: WriterProperties,
3247 validate: F,
3248 ) -> Bytes
3249 where
3250 F: Fn(&ArrayData, &ArrayData),
3251 {
3252 let mut file = vec![];
3253
3254 let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3255 .expect("Unable to write file");
3256 writer.write(expected_batch).unwrap();
3257 writer.close().unwrap();
3258
3259 let file = Bytes::from(file);
3260 let mut record_batch_reader =
3261 ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3262
3263 let actual_batch = record_batch_reader
3264 .next()
3265 .expect("No batch found")
3266 .expect("Unable to get batch");
3267
3268 assert_eq!(expected_batch.schema(), actual_batch.schema());
3269 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3270 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3271 for i in 0..expected_batch.num_columns() {
3272 let expected_data = expected_batch.column(i).to_data();
3273 let actual_data = actual_batch.column(i).to_data();
3274 validate(&expected_data, &actual_data);
3275 }
3276
3277 file
3278 }
3279
3280 fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3281 roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3282 a.validate_full().expect("valid expected data");
3283 b.validate_full().expect("valid actual data");
3284 assert_eq!(a, b)
3285 })
3286 }
3287
3288 struct RoundTripTest {
3292 values: ArrayRef,
3293 schema: Option<SchemaRef>,
3295 nullable: bool,
3298 bloom_filter: bool,
3299 bloom_filter_ndv: Option<u64>,
3300 bloom_filter_position: BloomFilterPosition,
3301 }
3302
3303 impl RoundTripTest {
3304 fn new(values: ArrayRef) -> Self {
3306 Self {
3307 values,
3308 schema: None,
3309 nullable: true,
3310 bloom_filter: false,
3311 bloom_filter_ndv: None,
3312 bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3313 }
3314 }
3315
3316 fn with_schema(mut self, schema: SchemaRef) -> Self {
3318 self.schema = Some(schema);
3319 self
3320 }
3321
3322 fn with_nullable(mut self, nullable: bool) -> Self {
3324 self.nullable = nullable;
3325 self
3326 }
3327
3328 fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
3330 self.bloom_filter = bloom_filter;
3331 self
3332 }
3333
3334 fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) -> Self {
3336 self.bloom_filter_ndv = Some(bloom_filter_ndv);
3337 self
3338 }
3339
3340 fn with_bloom_filter_position(
3342 mut self,
3343 bloom_filter_position: BloomFilterPosition,
3344 ) -> Self {
3345 self.bloom_filter_position = bloom_filter_position;
3346 self
3347 }
3348
3349 fn run(self) -> Vec<Bytes> {
3351 let RoundTripTest {
3352 values,
3353 schema,
3354 nullable,
3355 bloom_filter,
3356 bloom_filter_ndv,
3357 bloom_filter_position,
3358 } = self;
3359
3360 let schema = schema.unwrap_or_else(|| {
3361 let data_type = values.data_type().clone();
3362 Arc::new(Schema::new(vec![Field::new("col", data_type, nullable)]))
3363 });
3364
3365 let encodings = match values.data_type() {
3366 DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3367 vec![
3368 Encoding::PLAIN,
3369 Encoding::DELTA_BYTE_ARRAY,
3370 Encoding::DELTA_LENGTH_BYTE_ARRAY,
3371 ]
3372 }
3373 DataType::Int64
3374 | DataType::Int32
3375 | DataType::Int16
3376 | DataType::Int8
3377 | DataType::UInt64
3378 | DataType::UInt32
3379 | DataType::UInt16
3380 | DataType::UInt8 => vec![
3381 Encoding::PLAIN,
3382 Encoding::DELTA_BINARY_PACKED,
3383 Encoding::BYTE_STREAM_SPLIT,
3384 ],
3385 DataType::Float32 | DataType::Float64 => {
3386 vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
3387 }
3388 _ => vec![Encoding::PLAIN],
3389 };
3390
3391 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3392
3393 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3394
3395 let mut files = vec![];
3396 for dictionary_size in [0, 1, 1024] {
3397 for encoding in &encodings {
3398 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3399 for row_group_size in row_group_sizes {
3400 let mut builder = WriterProperties::builder()
3401 .set_writer_version(version)
3402 .set_max_row_group_row_count(Some(row_group_size))
3403 .set_dictionary_enabled(dictionary_size != 0)
3404 .set_dictionary_page_size_limit(dictionary_size.max(1))
3405 .set_encoding(*encoding)
3406 .set_bloom_filter_enabled(bloom_filter)
3407 .set_bloom_filter_position(bloom_filter_position);
3408 if let Some(ndv) = bloom_filter_ndv {
3409 builder = builder.set_bloom_filter_max_ndv(ndv);
3410 }
3411 let props = builder.build();
3412
3413 files.push(roundtrip_opts(&expected_batch, props))
3414 }
3415 }
3416 }
3417 }
3418 files
3419 }
3420 }
3421
3422 fn values_required<A, I>(iter: I) -> Vec<Bytes>
3423 where
3424 A: From<Vec<I::Item>> + Array + 'static,
3425 I: IntoIterator,
3426 {
3427 let raw_values: Vec<_> = iter.into_iter().collect();
3428 let values = Arc::new(A::from(raw_values));
3429 RoundTripTest::new(values).with_nullable(false).run()
3430 }
3431
3432 fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3433 where
3434 A: From<Vec<Option<I::Item>>> + Array + 'static,
3435 I: IntoIterator,
3436 {
3437 let optional_raw_values: Vec<_> = iter
3438 .into_iter()
3439 .enumerate()
3440 .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3441 .collect();
3442 let optional_values = Arc::new(A::from(optional_raw_values));
3443 RoundTripTest::new(optional_values).run()
3444 }
3445
3446 fn required_and_optional<A, I>(iter: I)
3447 where
3448 A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3449 I: IntoIterator + Clone,
3450 {
3451 values_required::<A, I>(iter.clone());
3452 values_optional::<A, I>(iter);
3453 }
3454
3455 fn check_bloom_filter<T: AsBytes>(
3456 files: Vec<Bytes>,
3457 file_column: String,
3458 positive_values: Vec<T>,
3459 negative_values: Vec<T>,
3460 ) {
3461 files.into_iter().take(1).for_each(|file| {
3462 let file_reader = SerializedFileReader::new_with_options(
3463 file,
3464 ReadOptionsBuilder::new()
3465 .with_reader_properties(
3466 ReaderProperties::builder()
3467 .set_read_bloom_filter(true)
3468 .build(),
3469 )
3470 .build(),
3471 )
3472 .expect("Unable to open file as Parquet");
3473 let metadata = file_reader.metadata();
3474
3475 let mut bloom_filters: Vec<_> = vec![];
3477 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3478 if let Some((column_index, _)) = row_group
3479 .columns()
3480 .iter()
3481 .enumerate()
3482 .find(|(_, column)| column.column_path().string() == file_column)
3483 {
3484 let row_group_reader = file_reader
3485 .get_row_group(ri)
3486 .expect("Unable to read row group");
3487 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3488 bloom_filters.push(sbbf.clone());
3489 } else {
3490 panic!("No bloom filter for column named {file_column} found");
3491 }
3492 } else {
3493 panic!("No column named {file_column} found");
3494 }
3495 }
3496
3497 positive_values.iter().for_each(|value| {
3498 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3499 assert!(
3500 found.is_some(),
3501 "{}",
3502 format!("Value {:?} should be in bloom filter", value.as_bytes())
3503 );
3504 });
3505
3506 negative_values.iter().for_each(|value| {
3507 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3508 assert!(
3509 found.is_none(),
3510 "{}",
3511 format!("Value {:?} should not be in bloom filter", value.as_bytes())
3512 );
3513 });
3514 });
3515 }
3516
3517 #[test]
3518 fn all_null_primitive_single_column() {
3519 let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3520 RoundTripTest::new(values).run();
3521 }
3522 #[test]
3523 fn null_single_column() {
3524 let values = Arc::new(NullArray::new(SMALL_SIZE));
3525 RoundTripTest::new(values).run();
3526 }
3528
3529 #[test]
3530 fn bool_single_column() {
3531 required_and_optional::<BooleanArray, _>(
3532 [true, false].iter().cycle().copied().take(SMALL_SIZE),
3533 );
3534 }
3535
3536 #[test]
3537 fn bool_large_single_column() {
3538 let values = Arc::new(
3539 [None, Some(true), Some(false)]
3540 .iter()
3541 .cycle()
3542 .copied()
3543 .take(200_000)
3544 .collect::<BooleanArray>(),
3545 );
3546 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3547 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3548 let file = tempfile::tempfile().unwrap();
3549
3550 let mut writer =
3551 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3552 .expect("Unable to write file");
3553 writer.write(&expected_batch).unwrap();
3554 writer.close().unwrap();
3555 }
3556
3557 #[test]
3558 fn check_page_offset_index_with_nan() {
3559 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3560 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3561 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3562
3563 let mut out = Vec::with_capacity(1024);
3564 let mut writer =
3565 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3566 writer.write(&batch).unwrap();
3567 let file_meta_data = writer.close().unwrap();
3568 for row_group in file_meta_data.row_groups() {
3569 for column in row_group.columns() {
3570 assert!(column.offset_index_offset().is_some());
3571 assert!(column.offset_index_length().is_some());
3572 assert!(column.column_index_offset().is_some());
3573 assert!(column.column_index_length().is_some());
3574 }
3575 }
3576 if let Some(page_index) = file_meta_data.page_index() {
3577 for rg in 0..file_meta_data.num_row_groups() {
3578 for col in 0..file_meta_data.row_group(rg).num_columns() {
3579 let idx = page_index
3580 .column_index(rg, col)
3581 .expect("column index should exist");
3582 assert!(idx.nan_counts().is_some());
3583 let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
3584 panic!("expected double statistics")
3585 };
3586 for i in 0..idx.num_pages() as usize {
3587 assert_eq!(float_idx.nan_count(i), Some(10));
3588 assert_eq!(
3589 f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3590 Ordering::Equal
3591 );
3592 assert_eq!(
3593 f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3594 Ordering::Equal
3595 );
3596 }
3597 }
3598 }
3599 } else {
3600 panic!("page index should be present");
3601 }
3602 }
3603
3604 #[test]
3605 fn check_page_offset_index_with_mixed_nan() {
3606 let schema = Arc::new(Schema::new(vec![Field::new(
3607 "col",
3608 DataType::Float64,
3609 true,
3610 )]));
3611
3612 let mut out = Vec::with_capacity(1024);
3613 let props = WriterProperties::builder()
3614 .set_data_page_row_count_limit(10)
3615 .build();
3616 let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3617 .expect("Unable to write file");
3618
3619 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3621 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3622 writer.write(&batch).unwrap();
3623
3624 let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3626 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3627 writer.write(&batch).unwrap();
3628
3629 let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3631 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3632 writer.write(&batch).unwrap();
3633
3634 let values = Arc::new(Float64Array::from(vec![
3636 -1.0,
3637 0.0,
3638 f64::NAN,
3639 -f64::NAN,
3640 1.0,
3641 ]));
3642 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3643 writer.write(&batch).unwrap();
3644
3645 let file_meta_data = writer.close().unwrap();
3646
3647 let col_stats = file_meta_data
3649 .row_group(0)
3650 .column(0)
3651 .statistics()
3652 .expect("missing column chunk statistics");
3653
3654 assert_eq!(col_stats.nan_count_opt(), Some(22));
3655 assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3656 assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3657
3658 assert!(file_meta_data.page_index().is_some());
3659 let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
3660 assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
3661
3662 let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
3664 panic!("expected double statistics")
3665 };
3666
3667 assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3668 assert_eq!(
3669 f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3670 Ordering::Equal
3671 );
3672 assert_eq!(
3673 f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3674 Ordering::Equal
3675 );
3676 assert_eq!(
3677 (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3678 Ordering::Equal
3679 );
3680 assert_eq!(
3681 (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3682 Ordering::Equal
3683 );
3684 assert_eq!(float_idx.min_value(2), Some(&0.0));
3685 assert_eq!(float_idx.max_value(2), Some(&0.0));
3686 assert_eq!(float_idx.min_value(3), Some(&-1.0));
3687 assert_eq!(float_idx.max_value(3), Some(&1.0));
3688 }
3689
3690 #[test]
3691 fn i8_single_column() {
3692 required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3693 }
3694
3695 #[test]
3696 fn i16_single_column() {
3697 required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3698 }
3699
3700 #[test]
3701 fn i32_single_column() {
3702 required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3703 }
3704
3705 #[test]
3706 fn i64_single_column() {
3707 required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3708 }
3709
3710 #[test]
3711 fn u8_single_column() {
3712 required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3713 }
3714
3715 #[test]
3716 fn u16_single_column() {
3717 required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3718 }
3719
3720 #[test]
3721 fn u32_single_column() {
3722 required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3723 }
3724
3725 #[test]
3726 fn u64_single_column() {
3727 required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3728 }
3729
3730 #[test]
3731 fn f32_single_column() {
3732 required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3733 }
3734
3735 #[test]
3736 fn f64_single_column() {
3737 required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3738 }
3739
3740 #[test]
3745 fn timestamp_second_single_column() {
3746 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3747 let values = Arc::new(TimestampSecondArray::from(raw_values));
3748
3749 RoundTripTest::new(values).with_nullable(false).run();
3750 }
3751
3752 #[test]
3753 fn timestamp_millisecond_single_column() {
3754 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3755 let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3756
3757 RoundTripTest::new(values).with_nullable(false).run();
3758 }
3759
3760 #[test]
3761 fn timestamp_microsecond_single_column() {
3762 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3763 let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3764
3765 RoundTripTest::new(values).with_nullable(false).run();
3766 }
3767
3768 #[test]
3769 fn timestamp_nanosecond_single_column() {
3770 let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3771 let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3772
3773 RoundTripTest::new(values).with_nullable(false).run();
3774 }
3775
3776 #[test]
3777 fn date32_single_column() {
3778 required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3779 }
3780
3781 #[test]
3782 fn date64_single_column() {
3783 required_and_optional::<Date64Array, _>(
3785 (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3786 );
3787 }
3788
3789 #[test]
3790 fn time32_second_single_column() {
3791 required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3792 }
3793
3794 #[test]
3795 fn time32_millisecond_single_column() {
3796 required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3797 }
3798
3799 #[test]
3800 fn time64_microsecond_single_column() {
3801 required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3802 }
3803
3804 #[test]
3805 fn time64_nanosecond_single_column() {
3806 required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3807 }
3808
3809 #[test]
3810 fn duration_second_single_column() {
3811 required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3812 }
3813
3814 #[test]
3815 fn duration_millisecond_single_column() {
3816 required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3817 }
3818
3819 #[test]
3820 fn duration_microsecond_single_column() {
3821 required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3822 }
3823
3824 #[test]
3825 fn duration_nanosecond_single_column() {
3826 required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3827 }
3828
3829 #[test]
3830 fn interval_year_month_single_column() {
3831 required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3832 }
3833
3834 #[test]
3835 fn interval_day_time_single_column() {
3836 required_and_optional::<IntervalDayTimeArray, _>(vec![
3837 IntervalDayTime::new(0, 1),
3838 IntervalDayTime::new(0, 3),
3839 IntervalDayTime::new(3, -2),
3840 IntervalDayTime::new(-200, 4),
3841 ]);
3842 }
3843
3844 #[test]
3845 #[should_panic(
3846 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3847 )]
3848 fn interval_month_day_nano_single_column() {
3849 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3850 IntervalMonthDayNano::new(0, 1, 5),
3851 IntervalMonthDayNano::new(0, 3, 2),
3852 IntervalMonthDayNano::new(3, -2, -5),
3853 IntervalMonthDayNano::new(-200, 4, -1),
3854 ]);
3855 }
3856
3857 #[test]
3858 fn binary_single_column() {
3859 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3860 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3861 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3862
3863 values_required::<BinaryArray, _>(many_vecs_iter);
3865 }
3866
3867 #[test]
3868 fn binary_view_single_column() {
3869 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3870 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3871 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3872
3873 values_required::<BinaryViewArray, _>(many_vecs_iter);
3875 }
3876
3877 #[test]
3878 fn i32_column_bloom_filter_at_end() {
3879 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3880 let files = RoundTripTest::new(array)
3881 .with_nullable(false)
3882 .with_bloom_filter(true)
3883 .with_bloom_filter_position(BloomFilterPosition::End)
3884 .run();
3885
3886 check_bloom_filter(
3887 files,
3888 "col".to_string(),
3889 (0..SMALL_SIZE as i32).collect(),
3890 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3891 );
3892 }
3893
3894 #[test]
3895 fn i32_column_bloom_filter() {
3896 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3897 let files = RoundTripTest::new(array)
3898 .with_nullable(false)
3899 .with_bloom_filter(true)
3900 .run();
3901
3902 check_bloom_filter(
3903 files,
3904 "col".to_string(),
3905 (0..SMALL_SIZE as i32).collect(),
3906 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3907 );
3908 }
3909
3910 #[test]
3915 fn i32_column_bloom_filter_fixed_ndv() {
3916 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3917
3918 let files = RoundTripTest::new(array.clone())
3920 .with_nullable(false)
3921 .with_bloom_filter(true)
3922 .with_bloom_filter_ndv(1_000_000)
3923 .run();
3924
3925 check_bloom_filter(
3926 files,
3927 "col".to_string(),
3928 (0..SMALL_SIZE as i32).collect(),
3929 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3930 );
3931
3932 let files = RoundTripTest::new(array)
3934 .with_nullable(false)
3935 .with_bloom_filter(true)
3936 .with_bloom_filter_ndv(3)
3937 .run();
3938
3939 check_bloom_filter(
3940 files,
3941 "col".to_string(),
3942 (0..SMALL_SIZE as i32).collect(),
3943 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3944 );
3945 }
3946
3947 #[test]
3948 fn binary_column_bloom_filter() {
3949 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3950 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3951 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3952
3953 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
3954 let files = RoundTripTest::new(array)
3955 .with_nullable(false)
3956 .with_bloom_filter(true)
3957 .run();
3958
3959 check_bloom_filter(
3960 files,
3961 "col".to_string(),
3962 many_vecs,
3963 vec![vec![(SMALL_SIZE + 1) as u8]],
3964 );
3965 }
3966
3967 #[test]
3968 fn empty_string_null_column_bloom_filter() {
3969 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3970 let raw_strs = raw_values.iter().map(|s| s.as_str());
3971
3972 let array = Arc::new(StringArray::from_iter_values(raw_strs));
3973 let files = RoundTripTest::new(array)
3974 .with_nullable(false)
3975 .with_bloom_filter(true)
3976 .run();
3977
3978 let optional_raw_values: Vec<_> = raw_values
3979 .iter()
3980 .enumerate()
3981 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
3982 .collect();
3983 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
3985 }
3986
3987 #[test]
3988 fn large_binary_single_column() {
3989 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3990 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3991 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3992
3993 values_required::<LargeBinaryArray, _>(many_vecs_iter);
3995 }
3996
3997 #[test]
3998 fn fixed_size_binary_single_column() {
3999 let mut builder = FixedSizeBinaryBuilder::new(4);
4000 builder.append_value(b"0123").unwrap();
4001 builder.append_null();
4002 builder.append_value(b"8910").unwrap();
4003 builder.append_value(b"1112").unwrap();
4004 let array = Arc::new(builder.finish());
4005
4006 RoundTripTest::new(array).run();
4007 }
4008
4009 #[test]
4010 fn string_single_column() {
4011 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4012 let raw_strs = raw_values.iter().map(|s| s.as_str());
4013
4014 required_and_optional::<StringArray, _>(raw_strs);
4015 }
4016
4017 #[test]
4018 fn large_string_single_column() {
4019 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4020 let raw_strs = raw_values.iter().map(|s| s.as_str());
4021
4022 required_and_optional::<LargeStringArray, _>(raw_strs);
4023 }
4024
4025 #[test]
4026 fn string_view_single_column() {
4027 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4028 let raw_strs = raw_values.iter().map(|s| s.as_str());
4029
4030 required_and_optional::<StringViewArray, _>(raw_strs);
4031 }
4032
4033 #[test]
4034 fn null_list_single_column() {
4035 let null_field = Field::new_list_field(DataType::Null, true);
4036 let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
4037
4038 let schema = Schema::new(vec![list_field]);
4039
4040 let a_values = NullArray::new(2);
4042 let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
4043 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4044 DataType::Null,
4045 true,
4046 ))))
4047 .len(3)
4048 .add_buffer(a_value_offsets)
4049 .null_bit_buffer(Some(Buffer::from([0b00000101])))
4050 .add_child_data(a_values.into_data())
4051 .build()
4052 .unwrap();
4053
4054 let a = ListArray::from(a_list_data);
4055
4056 assert!(a.is_valid(0));
4057 assert!(!a.is_valid(1));
4058 assert!(a.is_valid(2));
4059
4060 assert_eq!(a.value(0).len(), 0);
4061 assert_eq!(a.value(2).len(), 2);
4062 assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
4063
4064 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4065 roundtrip(batch, None);
4066 }
4067
4068 #[test]
4069 fn list_single_column() {
4070 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4071 let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
4072 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4073 DataType::Int32,
4074 false,
4075 ))))
4076 .len(5)
4077 .add_buffer(a_value_offsets)
4078 .null_bit_buffer(Some(Buffer::from([0b00011011])))
4079 .add_child_data(a_values.into_data())
4080 .build()
4081 .unwrap();
4082
4083 assert_eq!(a_list_data.null_count(), 1);
4084
4085 let a = ListArray::from(a_list_data);
4086 let values = Arc::new(a);
4087
4088 RoundTripTest::new(values).run();
4089 }
4090
4091 #[test]
4092 fn large_list_single_column() {
4093 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4094 let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
4095 let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
4096 "large_item",
4097 DataType::Int32,
4098 true,
4099 ))))
4100 .len(5)
4101 .add_buffer(a_value_offsets)
4102 .add_child_data(a_values.into_data())
4103 .null_bit_buffer(Some(Buffer::from([0b00011011])))
4104 .build()
4105 .unwrap();
4106
4107 assert_eq!(a_list_data.null_count(), 1);
4109
4110 let a = LargeListArray::from(a_list_data);
4111 let values = Arc::new(a);
4112
4113 RoundTripTest::new(values).run();
4114 }
4115
4116 #[test]
4117 fn list_nested_nulls() {
4118 use arrow::datatypes::Int32Type;
4119 let data = vec![
4120 Some(vec![Some(1)]),
4121 Some(vec![Some(2), Some(3)]),
4122 None,
4123 Some(vec![Some(4), Some(5), None]),
4124 Some(vec![None]),
4125 Some(vec![Some(6), Some(7)]),
4126 ];
4127
4128 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
4129 RoundTripTest::new(Arc::new(list)).run();
4130
4131 let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
4132 RoundTripTest::new(Arc::new(list)).run();
4133 }
4134
4135 #[test]
4136 fn list_utf8_view_selective_padding_roundtrip() {
4137 let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
4138 let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
4139 builder.values().append_value("a");
4140 builder.values().append_null();
4141 builder.append(true);
4142 builder.append(false);
4145 builder.values().append_value("large payload over 12 bytes");
4147 builder.append(true);
4148
4149 RoundTripTest::new(Arc::new(builder.finish())).run();
4150 }
4151
4152 #[test]
4153 fn struct_single_column() {
4154 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4155 let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4156 let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4157
4158 let values = Arc::new(s);
4159 RoundTripTest::new(values).with_nullable(false).run();
4160 }
4161
4162 #[test]
4163 fn list_and_map_coerced_names() {
4164 let list_field =
4166 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4167 let map_field = Field::new_map(
4168 "my_map",
4169 "my_entries",
4170 Field::new("my_keys", DataType::Int32, false),
4171 Field::new("my_values", DataType::Int32, true),
4172 false,
4173 true,
4174 );
4175
4176 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4177 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4178
4179 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4180
4181 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4183 let file = tempfile::tempfile().unwrap();
4184 let mut writer =
4185 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4186
4187 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4188 writer.write(&batch).unwrap();
4189 let file_metadata = writer.close().unwrap();
4190
4191 let schema = file_metadata.file_metadata().schema();
4192 let list_field = &schema.get_fields()[0].get_fields()[0];
4194 assert_eq!(list_field.get_fields()[0].name(), "element");
4195
4196 let map_field = &schema.get_fields()[1].get_fields()[0];
4197 assert_eq!(map_field.name(), "key_value");
4199 assert_eq!(map_field.get_fields()[0].name(), "key");
4201 assert_eq!(map_field.get_fields()[1].name(), "value");
4203
4204 let reader = SerializedFileReader::new(file).unwrap();
4206 let file_schema = reader.metadata().file_metadata().schema();
4207 let fields = file_schema.get_fields();
4208 let list_field = &fields[0].get_fields()[0];
4209 assert_eq!(list_field.get_fields()[0].name(), "element");
4210 let map_field = &fields[1].get_fields()[0];
4211 assert_eq!(map_field.name(), "key_value");
4212 assert_eq!(map_field.get_fields()[0].name(), "key");
4213 assert_eq!(map_field.get_fields()[1].name(), "value");
4214 }
4215
4216 #[test]
4217 fn fallback_flush_data_page() {
4218 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4220 let values = Arc::new(StringArray::from(raw_values));
4221 let encodings = vec![
4222 Encoding::DELTA_BYTE_ARRAY,
4223 Encoding::DELTA_LENGTH_BYTE_ARRAY,
4224 ];
4225 let data_type = values.data_type().clone();
4226 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4227 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4228
4229 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4230 let data_page_size_limit: usize = 32;
4231 let write_batch_size: usize = 16;
4232
4233 for encoding in &encodings {
4234 for row_group_size in row_group_sizes {
4235 let props = WriterProperties::builder()
4236 .set_writer_version(WriterVersion::PARQUET_2_0)
4237 .set_max_row_group_row_count(Some(row_group_size))
4238 .set_dictionary_enabled(false)
4239 .set_encoding(*encoding)
4240 .set_data_page_size_limit(data_page_size_limit)
4241 .set_write_batch_size(write_batch_size)
4242 .build();
4243
4244 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4245 let string_array_a = StringArray::from(a.clone());
4246 let string_array_b = StringArray::from(b.clone());
4247 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4248 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4249 assert_eq!(
4250 vec_a, vec_b,
4251 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4252 );
4253 });
4254 }
4255 }
4256 }
4257
4258 #[test]
4259 fn arrow_writer_string_dictionary() {
4260 #[expect(deprecated)]
4262 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4263 "dictionary",
4264 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4265 true,
4266 42,
4267 true,
4268 )]));
4269
4270 let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4272 .iter()
4273 .copied()
4274 .collect();
4275
4276 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4278 }
4279
4280 #[test]
4281 fn arrow_writer_test_type_compatibility() {
4282 fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4283 where
4284 T1: Array + 'static,
4285 T2: Array + 'static,
4286 {
4287 let schema1 = Arc::new(Schema::new(vec![Field::new(
4288 "a",
4289 array1.data_type().clone(),
4290 false,
4291 )]));
4292
4293 let file = tempfile().unwrap();
4294 let mut writer =
4295 ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4296
4297 let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4298 writer.write(&rb1).unwrap();
4299
4300 let schema2 = Arc::new(Schema::new(vec![Field::new(
4301 "a",
4302 array2.data_type().clone(),
4303 false,
4304 )]));
4305 let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4306 writer.write(&rb2).unwrap();
4307
4308 writer.close().unwrap();
4309
4310 let mut record_batch_reader =
4311 ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4312 let actual_batch = record_batch_reader.next().unwrap().unwrap();
4313
4314 let expected_batch =
4315 RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4316 assert_eq!(actual_batch, expected_batch);
4317 }
4318
4319 ensure_compatible_write(
4322 DictionaryArray::new(
4323 UInt8Array::from_iter_values(vec![0]),
4324 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4325 ),
4326 StringArray::from_iter_values(vec!["barquet"]),
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(
4334 StringArray::from_iter_values(vec!["parquet"]),
4335 DictionaryArray::new(
4336 UInt8Array::from_iter_values(vec![0]),
4337 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4338 ),
4339 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4340 );
4341
4342 ensure_compatible_write(
4345 DictionaryArray::new(
4346 UInt8Array::from_iter_values(vec![0]),
4347 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4348 ),
4349 DictionaryArray::new(
4350 UInt16Array::from_iter_values(vec![0]),
4351 Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4352 ),
4353 DictionaryArray::new(
4354 UInt8Array::from_iter_values(vec![0, 1]),
4355 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4356 ),
4357 );
4358
4359 ensure_compatible_write(
4361 DictionaryArray::new(
4362 UInt8Array::from_iter_values(vec![0]),
4363 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4364 ),
4365 DictionaryArray::new(
4366 UInt8Array::from_iter_values(vec![0]),
4367 Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4368 ),
4369 DictionaryArray::new(
4370 UInt8Array::from_iter_values(vec![0, 1]),
4371 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4372 ),
4373 );
4374
4375 ensure_compatible_write(
4377 DictionaryArray::new(
4378 UInt8Array::from_iter_values(vec![0]),
4379 Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4380 ),
4381 LargeStringArray::from_iter_values(vec!["barquet"]),
4382 DictionaryArray::new(
4383 UInt8Array::from_iter_values(vec![0, 1]),
4384 Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4385 ),
4386 );
4387
4388 ensure_compatible_write(
4391 StringArray::from_iter_values(vec!["parquet"]),
4392 LargeStringArray::from_iter_values(vec!["barquet"]),
4393 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4394 );
4395
4396 ensure_compatible_write(
4397 LargeStringArray::from_iter_values(vec!["parquet"]),
4398 StringArray::from_iter_values(vec!["barquet"]),
4399 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4400 );
4401
4402 ensure_compatible_write(
4403 StringArray::from_iter_values(vec!["parquet"]),
4404 StringViewArray::from_iter_values(vec!["barquet"]),
4405 StringArray::from_iter_values(vec!["parquet", "barquet"]),
4406 );
4407
4408 ensure_compatible_write(
4409 StringViewArray::from_iter_values(vec!["parquet"]),
4410 StringArray::from_iter_values(vec!["barquet"]),
4411 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4412 );
4413
4414 ensure_compatible_write(
4415 LargeStringArray::from_iter_values(vec!["parquet"]),
4416 StringViewArray::from_iter_values(vec!["barquet"]),
4417 LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4418 );
4419
4420 ensure_compatible_write(
4421 StringViewArray::from_iter_values(vec!["parquet"]),
4422 LargeStringArray::from_iter_values(vec!["barquet"]),
4423 StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4424 );
4425
4426 ensure_compatible_write(
4429 BinaryArray::from_iter_values(vec![b"parquet"]),
4430 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4431 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4432 );
4433
4434 ensure_compatible_write(
4435 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4436 BinaryArray::from_iter_values(vec![b"barquet"]),
4437 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4438 );
4439
4440 ensure_compatible_write(
4441 BinaryArray::from_iter_values(vec![b"parquet"]),
4442 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4443 BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4444 );
4445
4446 ensure_compatible_write(
4447 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4448 BinaryArray::from_iter_values(vec![b"barquet"]),
4449 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4450 );
4451
4452 ensure_compatible_write(
4453 BinaryViewArray::from_iter_values(vec![b"parquet"]),
4454 LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4455 BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4456 );
4457
4458 ensure_compatible_write(
4459 LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4460 BinaryViewArray::from_iter_values(vec![b"barquet"]),
4461 LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4462 );
4463
4464 let list_field_metadata = HashMap::from_iter(vec![(
4467 PARQUET_FIELD_ID_META_KEY.to_string(),
4468 "1".to_string(),
4469 )]);
4470 let list_field = Field::new_list_field(DataType::Int32, false);
4471
4472 let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4473 let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4474
4475 let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4476 let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4477
4478 let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4479 let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4480
4481 ensure_compatible_write(
4482 ListArray::try_new(
4484 Arc::new(
4485 list_field
4486 .clone()
4487 .with_metadata(list_field_metadata.clone()),
4488 ),
4489 offsets1,
4490 values1,
4491 None,
4492 )
4493 .unwrap(),
4494 ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4496 ListArray::try_new(
4498 Arc::new(
4499 list_field
4500 .clone()
4501 .with_metadata(list_field_metadata.clone()),
4502 ),
4503 offsets_expected,
4504 values_expected,
4505 None,
4506 )
4507 .unwrap(),
4508 );
4509 }
4510
4511 #[test]
4512 fn arrow_writer_primitive_dictionary() {
4513 #[expect(deprecated)]
4515 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4516 "dictionary",
4517 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4518 true,
4519 42,
4520 true,
4521 )]));
4522
4523 let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4525 builder.append(12345678).unwrap();
4526 builder.append_null();
4527 builder.append(22345678).unwrap();
4528 builder.append(12345678).unwrap();
4529 let d = builder.finish();
4530
4531 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4532 }
4533
4534 #[test]
4535 fn arrow_writer_decimal32_dictionary() {
4536 let integers = vec![12345, 56789, 34567];
4537
4538 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4539
4540 let values = Decimal32Array::from(integers.clone())
4541 .with_precision_and_scale(5, 2)
4542 .unwrap();
4543
4544 let array = DictionaryArray::new(keys, Arc::new(values));
4545 RoundTripTest::new(Arc::new(array.clone())).run();
4546
4547 let values = Decimal32Array::from(integers)
4548 .with_precision_and_scale(9, 2)
4549 .unwrap();
4550
4551 let array = array.with_values(Arc::new(values));
4552 RoundTripTest::new(Arc::new(array)).run();
4553 }
4554
4555 #[test]
4556 fn arrow_writer_decimal64_dictionary() {
4557 let integers = vec![12345, 56789, 34567];
4558
4559 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4560
4561 let values = Decimal64Array::from(integers.clone())
4562 .with_precision_and_scale(5, 2)
4563 .unwrap();
4564
4565 let array = DictionaryArray::new(keys, Arc::new(values));
4566 RoundTripTest::new(Arc::new(array.clone())).run();
4567
4568 let values = Decimal64Array::from(integers)
4569 .with_precision_and_scale(12, 2)
4570 .unwrap();
4571
4572 let array = array.with_values(Arc::new(values));
4573 RoundTripTest::new(Arc::new(array)).run();
4574 }
4575
4576 #[test]
4577 fn arrow_writer_decimal128_dictionary() {
4578 let integers = vec![12345, 56789, 34567];
4579
4580 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4581
4582 let values = Decimal128Array::from(integers.clone())
4583 .with_precision_and_scale(5, 2)
4584 .unwrap();
4585
4586 let array = DictionaryArray::new(keys, Arc::new(values));
4587 RoundTripTest::new(Arc::new(array.clone())).run();
4588
4589 let values = Decimal128Array::from(integers)
4590 .with_precision_and_scale(12, 2)
4591 .unwrap();
4592
4593 let array = array.with_values(Arc::new(values));
4594 RoundTripTest::new(Arc::new(array)).run();
4595 }
4596
4597 #[test]
4598 fn arrow_writer_decimal256_dictionary() {
4599 let integers = vec![
4600 i256::from_i128(12345),
4601 i256::from_i128(56789),
4602 i256::from_i128(34567),
4603 ];
4604
4605 let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4606
4607 let values = Decimal256Array::from(integers.clone())
4608 .with_precision_and_scale(5, 2)
4609 .unwrap();
4610
4611 let array = DictionaryArray::new(keys, Arc::new(values));
4612 RoundTripTest::new(Arc::new(array.clone())).run();
4613
4614 let values = Decimal256Array::from(integers)
4615 .with_precision_and_scale(12, 2)
4616 .unwrap();
4617
4618 let array = array.with_values(Arc::new(values));
4619 RoundTripTest::new(Arc::new(array)).run();
4620 }
4621
4622 #[test]
4623 fn arrow_writer_string_dictionary_unsigned_index() {
4624 #[expect(deprecated)]
4626 let schema = Arc::new(Schema::new(vec![Field::new_dict(
4627 "dictionary",
4628 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4629 true,
4630 42,
4631 true,
4632 )]));
4633
4634 let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4636 .iter()
4637 .copied()
4638 .collect();
4639
4640 RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4641 }
4642
4643 #[test]
4644 fn u32_min_max() {
4645 let src = [
4647 u32::MIN,
4648 1,
4649 (i32::MAX as u32) - 1,
4650 i32::MAX as u32,
4651 (i32::MAX as u32) + 1,
4652 u32::MAX - 1,
4653 u32::MAX,
4654 ];
4655 let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4656 let files = RoundTripTest::new(values).with_nullable(false).run();
4657
4658 for file in files {
4659 let reader = SerializedFileReader::new(file).unwrap();
4661 let metadata = reader.metadata();
4662
4663 let mut row_offset = 0;
4664 for row_group in metadata.row_groups() {
4665 assert_eq!(row_group.num_columns(), 1);
4666 let column = row_group.column(0);
4667
4668 let num_values = column.num_values() as usize;
4669 let src_slice = &src[row_offset..row_offset + num_values];
4670 row_offset += column.num_values() as usize;
4671
4672 let stats = column.statistics().unwrap();
4673 if let Statistics::Int32(stats) = stats {
4674 assert_eq!(
4675 *stats.min_opt().unwrap() as u32,
4676 *src_slice.iter().min().unwrap()
4677 );
4678 assert_eq!(
4679 *stats.max_opt().unwrap() as u32,
4680 *src_slice.iter().max().unwrap()
4681 );
4682 } else {
4683 panic!("Statistics::Int32 missing")
4684 }
4685 }
4686 }
4687 }
4688
4689 #[test]
4690 fn u64_min_max() {
4691 let src = [
4693 u64::MIN,
4694 1,
4695 (i64::MAX as u64) - 1,
4696 i64::MAX as u64,
4697 (i64::MAX as u64) + 1,
4698 u64::MAX - 1,
4699 u64::MAX,
4700 ];
4701 let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4702 let files = RoundTripTest::new(values).with_nullable(false).run();
4703
4704 for file in files {
4705 let reader = SerializedFileReader::new(file).unwrap();
4707 let metadata = reader.metadata();
4708
4709 let mut row_offset = 0;
4710 for row_group in metadata.row_groups() {
4711 assert_eq!(row_group.num_columns(), 1);
4712 let column = row_group.column(0);
4713
4714 let num_values = column.num_values() as usize;
4715 let src_slice = &src[row_offset..row_offset + num_values];
4716 row_offset += column.num_values() as usize;
4717
4718 let stats = column.statistics().unwrap();
4719 if let Statistics::Int64(stats) = stats {
4720 assert_eq!(
4721 *stats.min_opt().unwrap() as u64,
4722 *src_slice.iter().min().unwrap()
4723 );
4724 assert_eq!(
4725 *stats.max_opt().unwrap() as u64,
4726 *src_slice.iter().max().unwrap()
4727 );
4728 } else {
4729 panic!("Statistics::Int64 missing")
4730 }
4731 }
4732 }
4733 }
4734
4735 #[test]
4736 fn statistics_null_counts_only_nulls() {
4737 let values = Arc::new(UInt64Array::from(vec![None, None]));
4739 let files = RoundTripTest::new(values).run();
4740
4741 for file in files {
4742 let reader = SerializedFileReader::new(file).unwrap();
4744 let metadata = reader.metadata();
4745 assert_eq!(metadata.num_row_groups(), 1);
4746 let row_group = metadata.row_group(0);
4747 assert_eq!(row_group.num_columns(), 1);
4748 let column = row_group.column(0);
4749 let stats = column.statistics().unwrap();
4750 assert_eq!(stats.null_count_opt(), Some(2));
4751 }
4752 }
4753
4754 #[test]
4755 fn test_list_of_struct_roundtrip() {
4756 let int_field = Field::new("a", DataType::Int32, true);
4758 let int_field2 = Field::new("b", DataType::Int32, true);
4759
4760 let int_builder = Int32Builder::with_capacity(10);
4761 let int_builder2 = Int32Builder::with_capacity(10);
4762
4763 let struct_builder = StructBuilder::new(
4764 vec![int_field, int_field2],
4765 vec![Box::new(int_builder), Box::new(int_builder2)],
4766 );
4767 let mut list_builder = ListBuilder::new(struct_builder);
4768
4769 let values = list_builder.values();
4774 values
4775 .field_builder::<Int32Builder>(0)
4776 .unwrap()
4777 .append_value(1);
4778 values
4779 .field_builder::<Int32Builder>(1)
4780 .unwrap()
4781 .append_value(2);
4782 values.append(true);
4783 list_builder.append(true);
4784
4785 list_builder.append(true);
4787
4788 list_builder.append(false);
4790
4791 let values = list_builder.values();
4793 values
4794 .field_builder::<Int32Builder>(0)
4795 .unwrap()
4796 .append_null();
4797 values
4798 .field_builder::<Int32Builder>(1)
4799 .unwrap()
4800 .append_null();
4801 values.append(false);
4802 values
4803 .field_builder::<Int32Builder>(0)
4804 .unwrap()
4805 .append_null();
4806 values
4807 .field_builder::<Int32Builder>(1)
4808 .unwrap()
4809 .append_null();
4810 values.append(false);
4811 list_builder.append(true);
4812
4813 let values = list_builder.values();
4815 values
4816 .field_builder::<Int32Builder>(0)
4817 .unwrap()
4818 .append_null();
4819 values
4820 .field_builder::<Int32Builder>(1)
4821 .unwrap()
4822 .append_value(3);
4823 values.append(true);
4824 list_builder.append(true);
4825
4826 let values = list_builder.values();
4828 values
4829 .field_builder::<Int32Builder>(0)
4830 .unwrap()
4831 .append_value(2);
4832 values
4833 .field_builder::<Int32Builder>(1)
4834 .unwrap()
4835 .append_null();
4836 values.append(true);
4837 list_builder.append(true);
4838
4839 let array = Arc::new(list_builder.finish());
4840
4841 RoundTripTest::new(array).run();
4842 }
4843
4844 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4845 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4846 }
4847
4848 #[test]
4849 fn test_aggregates_records() {
4850 let arrays = [
4851 Int32Array::from((0..100).collect::<Vec<_>>()),
4852 Int32Array::from((0..50).collect::<Vec<_>>()),
4853 Int32Array::from((200..500).collect::<Vec<_>>()),
4854 ];
4855
4856 let schema = Arc::new(Schema::new(vec![Field::new(
4857 "int",
4858 ArrowDataType::Int32,
4859 false,
4860 )]));
4861
4862 let file = tempfile::tempfile().unwrap();
4863
4864 let props = WriterProperties::builder()
4865 .set_max_row_group_row_count(Some(200))
4866 .build();
4867
4868 let mut writer =
4869 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4870
4871 for array in arrays {
4872 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4873 writer.write(&batch).unwrap();
4874 }
4875
4876 writer.close().unwrap();
4877
4878 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4879 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4880
4881 let batches = builder
4882 .with_batch_size(100)
4883 .build()
4884 .unwrap()
4885 .collect::<ArrowResult<Vec<_>>>()
4886 .unwrap();
4887
4888 assert_eq!(batches.len(), 5);
4889 assert!(batches.iter().all(|x| x.num_columns() == 1));
4890
4891 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4892
4893 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4894
4895 let values: Vec<_> = batches
4896 .iter()
4897 .flat_map(|x| {
4898 x.column(0)
4899 .as_any()
4900 .downcast_ref::<Int32Array>()
4901 .unwrap()
4902 .values()
4903 .iter()
4904 .copied()
4905 })
4906 .collect();
4907
4908 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4909 assert_eq!(&values, &expected_values)
4910 }
4911
4912 #[test]
4913 fn complex_aggregate() {
4914 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4916 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4917 let struct_a = Arc::new(Field::new(
4918 "struct_a",
4919 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4920 true,
4921 ));
4922
4923 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4924 let struct_b = Arc::new(Field::new(
4925 "struct_b",
4926 DataType::Struct(vec![list_a.clone()].into()),
4927 false,
4928 ));
4929
4930 let schema = Arc::new(Schema::new(vec![struct_b]));
4931
4932 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4934 let field_b_array =
4935 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4936
4937 let struct_a_array = StructArray::from(vec![
4938 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4939 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4940 ]);
4941
4942 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4943 .len(5)
4944 .add_buffer(Buffer::from_iter(vec![
4945 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4946 ]))
4947 .null_bit_buffer(Some(Buffer::from_iter(vec![
4948 true, false, true, false, true,
4949 ])))
4950 .child_data(vec![struct_a_array.into_data()])
4951 .build()
4952 .unwrap();
4953
4954 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4955 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4956
4957 let batch1 =
4958 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4959 .unwrap();
4960
4961 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4962 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4963
4964 let struct_a_array = StructArray::from(vec![
4965 (field_a, Arc::new(field_a_array) as ArrayRef),
4966 (field_b, Arc::new(field_b_array) as ArrayRef),
4967 ]);
4968
4969 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4970 .len(2)
4971 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4972 .child_data(vec![struct_a_array.into_data()])
4973 .build()
4974 .unwrap();
4975
4976 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4977 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4978
4979 let batch2 =
4980 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4981 .unwrap();
4982
4983 let batches = &[batch1, batch2];
4984
4985 let expected = r"
4988 +-------------------------------------------------------------------------------------------------------+
4989 | struct_b |
4990 +-------------------------------------------------------------------------------------------------------+
4991 | {list: [{leaf_a: 1, leaf_b: 1}]} |
4992 | {list: } |
4993 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
4994 | {list: } |
4995 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
4996 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4997 | {list: [{leaf_a: 10, leaf_b: }]} |
4998 +-------------------------------------------------------------------------------------------------------+
4999 ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
5000
5001 let actual = pretty_format_batches(batches).unwrap().to_string();
5002 assert_eq!(actual, expected);
5003
5004 let file = tempfile::tempfile().unwrap();
5006 let props = WriterProperties::builder()
5007 .set_max_row_group_row_count(Some(6))
5008 .build();
5009
5010 let mut writer =
5011 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
5012
5013 for batch in batches {
5014 writer.write(batch).unwrap();
5015 }
5016 writer.close().unwrap();
5017
5018 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5023 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
5024
5025 let batches = builder
5026 .with_batch_size(2)
5027 .build()
5028 .unwrap()
5029 .collect::<ArrowResult<Vec<_>>>()
5030 .unwrap();
5031
5032 assert_eq!(batches.len(), 4);
5033 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
5034 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
5035
5036 let actual = pretty_format_batches(&batches).unwrap().to_string();
5037 assert_eq!(actual, expected);
5038 }
5039
5040 #[test]
5041 fn test_arrow_writer_metadata() {
5042 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5043 let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
5044
5045 let batch = RecordBatch::try_new(
5046 Arc::new(batch_schema),
5047 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5048 )
5049 .unwrap();
5050
5051 let mut buf = Vec::with_capacity(1024);
5052 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
5053 writer.write(&batch).unwrap();
5054 writer.close().unwrap();
5055 }
5056
5057 #[test]
5058 fn test_arrow_writer_nullable() {
5059 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5060 let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
5061 let file_schema = Arc::new(file_schema);
5062
5063 let batch = RecordBatch::try_new(
5064 Arc::new(batch_schema),
5065 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5066 )
5067 .unwrap();
5068
5069 let mut buf = Vec::with_capacity(1024);
5070 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5071 writer.write(&batch).unwrap();
5072 writer.close().unwrap();
5073
5074 let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
5075 let back = read.next().unwrap().unwrap();
5076 assert_eq!(back.schema(), file_schema);
5077 assert_ne!(back.schema(), batch.schema());
5078 assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
5079 }
5080
5081 #[test]
5082 fn in_progress_accounting() {
5083 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
5085
5086 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5088
5089 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
5091
5092 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
5093
5094 assert_eq!(writer.in_progress_size(), 0);
5096 assert_eq!(writer.in_progress_rows(), 0);
5097 assert_eq!(writer.memory_size(), 0);
5098 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
5100
5101 let initial_size = writer.in_progress_size();
5103 assert!(initial_size > 0);
5104 assert_eq!(writer.in_progress_rows(), 5);
5105 let initial_memory = writer.memory_size();
5106 assert!(initial_memory > 0);
5107 assert!(
5109 initial_size <= initial_memory,
5110 "{initial_size} <= {initial_memory}"
5111 );
5112
5113 writer.write(&batch).unwrap();
5115 assert!(writer.in_progress_size() > initial_size);
5116 assert_eq!(writer.in_progress_rows(), 10);
5117 assert!(writer.memory_size() > initial_memory);
5118 assert!(
5119 writer.in_progress_size() <= writer.memory_size(),
5120 "in_progress_size {} <= memory_size {}",
5121 writer.in_progress_size(),
5122 writer.memory_size()
5123 );
5124
5125 let pre_flush_bytes_written = writer.bytes_written();
5127 writer.flush().unwrap();
5128 assert_eq!(writer.in_progress_size(), 0);
5129 assert_eq!(writer.memory_size(), 0);
5130 assert!(writer.bytes_written() > pre_flush_bytes_written);
5131
5132 writer.close().unwrap();
5133 }
5134
5135 #[test]
5136 fn test_writer_all_null() {
5137 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5138 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
5139 let batch = RecordBatch::try_from_iter(vec![
5140 ("a", Arc::new(a) as ArrayRef),
5141 ("b", Arc::new(b) as ArrayRef),
5142 ])
5143 .unwrap();
5144
5145 let mut buf = Vec::with_capacity(1024);
5146 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5147 writer.write(&batch).unwrap();
5148 writer.close().unwrap();
5149
5150 let bytes = Bytes::from(buf);
5151 let options = ReadOptionsBuilder::new().with_page_index().build();
5152 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
5153 let index = reader.metadata().page_index().unwrap();
5154
5155 assert_eq!(index.num_data_pages(0, 0), Some(1)); assert_eq!(index.num_data_pages(0, 1), Some(1)); }
5158
5159 #[test]
5160 fn test_disabled_statistics_with_page() {
5161 let file_schema = Schema::new(vec![
5162 Field::new("a", DataType::Utf8, true),
5163 Field::new("b", DataType::Utf8, true),
5164 ]);
5165 let file_schema = Arc::new(file_schema);
5166
5167 let batch = RecordBatch::try_new(
5168 file_schema.clone(),
5169 vec![
5170 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5171 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5172 ],
5173 )
5174 .unwrap();
5175
5176 let props = WriterProperties::builder()
5177 .set_statistics_enabled(EnabledStatistics::None)
5178 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5179 .build();
5180
5181 let mut buf = Vec::with_capacity(1024);
5182 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5183 writer.write(&batch).unwrap();
5184
5185 let metadata = writer.close().unwrap();
5186 assert_eq!(metadata.num_row_groups(), 1);
5187 let row_group = metadata.row_group(0);
5188 assert_eq!(row_group.num_columns(), 2);
5189 assert!(row_group.column(0).offset_index_offset().is_some());
5191 assert!(row_group.column(0).column_index_offset().is_some());
5192 assert!(row_group.column(1).offset_index_offset().is_some());
5194 assert!(row_group.column(1).column_index_offset().is_none());
5195
5196 let options = ReadOptionsBuilder::new().with_page_index().build();
5197 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5198
5199 let row_group = reader.get_row_group(0).unwrap();
5200 let a_col = row_group.metadata().column(0);
5201 let b_col = row_group.metadata().column(1);
5202
5203 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5205 let min = byte_array_stats.min_opt().unwrap();
5206 let max = byte_array_stats.max_opt().unwrap();
5207
5208 assert_eq!(min.as_bytes(), b"a");
5209 assert_eq!(max.as_bytes(), b"d");
5210 } else {
5211 panic!("expecting Statistics::ByteArray");
5212 }
5213
5214 assert!(b_col.statistics().is_none());
5216
5217 let page_index = reader.metadata().page_index().unwrap();
5218
5219 let a_idx = page_index.column_index(0, 0);
5220 assert!(
5221 matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
5222 "{a_idx:?}"
5223 );
5224 let b_idx = page_index.column_index(0, 1);
5225 assert!(b_idx.is_none(), "{b_idx:?}");
5226 }
5227
5228 #[test]
5229 fn test_disabled_statistics_with_chunk() {
5230 let file_schema = Schema::new(vec![
5231 Field::new("a", DataType::Utf8, true),
5232 Field::new("b", DataType::Utf8, true),
5233 ]);
5234 let file_schema = Arc::new(file_schema);
5235
5236 let batch = RecordBatch::try_new(
5237 file_schema.clone(),
5238 vec![
5239 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5240 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5241 ],
5242 )
5243 .unwrap();
5244
5245 let props = WriterProperties::builder()
5246 .set_statistics_enabled(EnabledStatistics::None)
5247 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5248 .build();
5249
5250 let mut buf = Vec::with_capacity(1024);
5251 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5252 writer.write(&batch).unwrap();
5253
5254 let metadata = writer.close().unwrap();
5255 assert_eq!(metadata.num_row_groups(), 1);
5256 let row_group = metadata.row_group(0);
5257 assert_eq!(row_group.num_columns(), 2);
5258 assert!(row_group.column(0).offset_index_offset().is_some());
5260 assert!(row_group.column(0).column_index_offset().is_none());
5261 assert!(row_group.column(1).offset_index_offset().is_some());
5263 assert!(row_group.column(1).column_index_offset().is_none());
5264
5265 let options = ReadOptionsBuilder::new().with_page_index().build();
5266 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5267
5268 let row_group = reader.get_row_group(0).unwrap();
5269 let a_col = row_group.metadata().column(0);
5270 let b_col = row_group.metadata().column(1);
5271
5272 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5274 let min = byte_array_stats.min_opt().unwrap();
5275 let max = byte_array_stats.max_opt().unwrap();
5276
5277 assert_eq!(min.as_bytes(), b"a");
5278 assert_eq!(max.as_bytes(), b"d");
5279 } else {
5280 panic!("expecting Statistics::ByteArray");
5281 }
5282
5283 assert!(b_col.statistics().is_none());
5285
5286 let page_index = reader.metadata().page_index().unwrap();
5287
5288 let a_idx = page_index.column_index(0, 0);
5289 assert!(a_idx.is_none(), "{a_idx:?}");
5290 let b_idx = page_index.column_index(0, 1);
5291 assert!(b_idx.is_none(), "{b_idx:?}");
5292 }
5293
5294 #[test]
5295 fn test_arrow_writer_skip_metadata() {
5296 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5297 let file_schema = Arc::new(batch_schema.clone());
5298
5299 let batch = RecordBatch::try_new(
5300 Arc::new(batch_schema),
5301 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5302 )
5303 .unwrap();
5304 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5305
5306 let mut buf = Vec::with_capacity(1024);
5307 let mut writer =
5308 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5309 writer.write(&batch).unwrap();
5310 writer.close().unwrap();
5311
5312 let bytes = Bytes::from(buf);
5313 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5314 assert_eq!(file_schema, *reader_builder.schema());
5315 if let Some(key_value_metadata) = reader_builder
5316 .metadata()
5317 .file_metadata()
5318 .key_value_metadata()
5319 {
5320 assert!(
5321 !key_value_metadata
5322 .iter()
5323 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5324 );
5325 }
5326 }
5327
5328 #[test]
5329 fn test_arrow_writer_skip_path_in_schema() {
5330 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5331 let file_schema = Arc::new(batch_schema.clone());
5332
5333 let batch = RecordBatch::try_new(
5334 Arc::new(batch_schema),
5335 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5336 )
5337 .unwrap();
5338
5339 let skip_options = ArrowWriterOptions::new();
5341
5342 let mut buf = Vec::with_capacity(1024);
5343 let mut writer =
5344 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5345 writer.write(&batch).unwrap();
5346 writer.close().unwrap();
5347
5348 let skip_options = ArrowWriterOptions::new().with_properties(
5350 WriterProperties::builder()
5351 .set_write_path_in_schema(false)
5352 .build(),
5353 );
5354
5355 let mut buf2 = Vec::with_capacity(1024);
5356 let mut writer =
5357 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5358 .unwrap();
5359 writer.write(&batch).unwrap();
5360 writer.close().unwrap();
5361
5362 assert!(buf.len() > buf2.len());
5364 }
5365
5366 #[test]
5367 fn mismatched_schemas() {
5368 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5369 let file_schema = Arc::new(Schema::new(vec![Field::new(
5370 "temperature",
5371 DataType::Float64,
5372 false,
5373 )]));
5374
5375 let batch = RecordBatch::try_new(
5376 Arc::new(batch_schema),
5377 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5378 )
5379 .unwrap();
5380
5381 let mut buf = Vec::with_capacity(1024);
5382 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5383
5384 let err = writer.write(&batch).unwrap_err().to_string();
5385 assert_eq!(
5386 err,
5387 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5388 );
5389 }
5390
5391 #[test]
5392 fn test_roundtrip_empty_schema() {
5394 let empty_batch = RecordBatch::try_new_with_options(
5396 Arc::new(Schema::empty()),
5397 vec![],
5398 &RecordBatchOptions::default().with_row_count(Some(0)),
5399 )
5400 .unwrap();
5401
5402 let mut parquet_bytes: Vec<u8> = Vec::new();
5404 let mut writer =
5405 ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5406 writer.write(&empty_batch).unwrap();
5407 writer.close().unwrap();
5408
5409 let bytes = Bytes::from(parquet_bytes);
5411 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5412 assert_eq!(reader.schema(), &empty_batch.schema());
5413 let batches: Vec<_> = reader
5414 .build()
5415 .unwrap()
5416 .collect::<ArrowResult<Vec<_>>>()
5417 .unwrap();
5418 assert_eq!(batches.len(), 0);
5419 }
5420
5421 #[test]
5422 fn test_page_stats_not_written_by_default() {
5423 let string_field = Field::new("a", DataType::Utf8, false);
5424 let schema = Schema::new(vec![string_field]);
5425 let raw_string_values = vec!["Blart Versenwald III"];
5426 let string_values = StringArray::from(raw_string_values.clone());
5427 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5428
5429 let props = WriterProperties::builder()
5430 .set_statistics_enabled(EnabledStatistics::Page)
5431 .set_dictionary_enabled(false)
5432 .set_encoding(Encoding::PLAIN)
5433 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5434 .build();
5435
5436 let file = roundtrip_opts(&batch, props);
5437
5438 let first_page = &file[4..];
5443 let mut prot = ThriftSliceInputProtocol::new(first_page);
5444 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5445 let stats = hdr.data_page_header.unwrap().statistics;
5446
5447 assert!(stats.is_none());
5448 }
5449
5450 #[test]
5451 fn test_page_stats_when_enabled() {
5452 let string_field = Field::new("a", DataType::Utf8, false);
5453 let schema = Schema::new(vec![string_field]);
5454 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5455 let string_values = StringArray::from(raw_string_values.clone());
5456 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5457
5458 let props = WriterProperties::builder()
5459 .set_statistics_enabled(EnabledStatistics::Page)
5460 .set_dictionary_enabled(false)
5461 .set_encoding(Encoding::PLAIN)
5462 .set_write_page_header_statistics(true)
5463 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5464 .build();
5465
5466 let file = roundtrip_opts(&batch, props);
5467
5468 let first_page = &file[4..];
5473 let mut prot = ThriftSliceInputProtocol::new(first_page);
5474 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5475 let stats = hdr.data_page_header.unwrap().statistics;
5476
5477 let stats = stats.unwrap();
5478 assert!(stats.is_max_value_exact.unwrap());
5480 assert!(stats.is_min_value_exact.unwrap());
5481 assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5482 assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5483 }
5484
5485 #[test]
5486 fn test_page_stats_truncation() {
5487 let string_field = Field::new("a", DataType::Utf8, false);
5488 let binary_field = Field::new("b", DataType::Binary, false);
5489 let schema = Schema::new(vec![string_field, binary_field]);
5490
5491 let raw_string_values = vec!["Blart Versenwald III"];
5492 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5493 let raw_binary_value_refs = raw_binary_values
5494 .iter()
5495 .map(|x| x.as_slice())
5496 .collect::<Vec<_>>();
5497
5498 let string_values = StringArray::from(raw_string_values.clone());
5499 let binary_values = BinaryArray::from(raw_binary_value_refs);
5500 let batch = RecordBatch::try_new(
5501 Arc::new(schema),
5502 vec![Arc::new(string_values), Arc::new(binary_values)],
5503 )
5504 .unwrap();
5505
5506 let props = WriterProperties::builder()
5507 .set_statistics_truncate_length(Some(2))
5508 .set_dictionary_enabled(false)
5509 .set_encoding(Encoding::PLAIN)
5510 .set_write_page_header_statistics(true)
5511 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5512 .build();
5513
5514 let file = roundtrip_opts(&batch, props);
5515
5516 let first_page = &file[4..];
5521 let mut prot = ThriftSliceInputProtocol::new(first_page);
5522 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5523 let stats = hdr.data_page_header.unwrap().statistics;
5524 assert!(stats.is_some());
5525 let stats = stats.unwrap();
5526 assert!(!stats.is_max_value_exact.unwrap());
5528 assert!(!stats.is_min_value_exact.unwrap());
5529 assert_eq!(stats.max_value.unwrap(), b"Bm");
5530 assert_eq!(stats.min_value.unwrap(), b"Bl");
5531
5532 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5534 let mut prot = ThriftSliceInputProtocol::new(second_page);
5535 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5536 let stats = hdr.data_page_header.unwrap().statistics;
5537 assert!(stats.is_some());
5538 let stats = stats.unwrap();
5539 assert!(!stats.is_max_value_exact.unwrap());
5541 assert!(!stats.is_min_value_exact.unwrap());
5542 assert_eq!(stats.max_value.unwrap(), b"Bm");
5543 assert_eq!(stats.min_value.unwrap(), b"Bl");
5544 }
5545
5546 #[test]
5547 fn test_page_encoding_statistics_roundtrip() {
5548 let batch_schema = Schema::new(vec![Field::new(
5549 "int32",
5550 arrow_schema::DataType::Int32,
5551 false,
5552 )]);
5553
5554 let batch = RecordBatch::try_new(
5555 Arc::new(batch_schema.clone()),
5556 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5557 )
5558 .unwrap();
5559
5560 let mut file: File = tempfile::tempfile().unwrap();
5561 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5562 writer.write(&batch).unwrap();
5563 let file_metadata = writer.close().unwrap();
5564
5565 assert_eq!(file_metadata.num_row_groups(), 1);
5566 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5567 assert!(
5568 file_metadata
5569 .row_group(0)
5570 .column(0)
5571 .page_encoding_stats()
5572 .is_some()
5573 );
5574 let chunk_page_stats = file_metadata
5575 .row_group(0)
5576 .column(0)
5577 .page_encoding_stats()
5578 .unwrap();
5579
5580 let options = ReadOptionsBuilder::new()
5582 .with_page_index()
5583 .with_encoding_stats_as_mask(false)
5584 .build();
5585 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5586
5587 let rowgroup = reader.get_row_group(0).expect("row group missing");
5588 assert_eq!(rowgroup.num_columns(), 1);
5589 let column = rowgroup.metadata().column(0);
5590 assert!(column.page_encoding_stats().is_some());
5591 let file_page_stats = column.page_encoding_stats().unwrap();
5592 assert_eq!(chunk_page_stats, file_page_stats);
5593 }
5594
5595 #[test]
5596 fn test_different_dict_page_size_limit() {
5597 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5598 let schema = Arc::new(Schema::new(vec![
5599 Field::new("col0", arrow_schema::DataType::Int64, false),
5600 Field::new("col1", arrow_schema::DataType::Int64, false),
5601 ]));
5602 let batch =
5603 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5604
5605 let props = WriterProperties::builder()
5606 .set_dictionary_page_size_limit(1024 * 1024)
5607 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5608 .build();
5609 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5610 writer.write(&batch).unwrap();
5611 let data = Bytes::from(writer.into_inner().unwrap());
5612
5613 let mut metadata = ParquetMetaDataReader::new();
5614 metadata.try_parse(&data).unwrap();
5615 let metadata = metadata.finish().unwrap();
5616 let col0_meta = metadata.row_group(0).column(0);
5617 let col1_meta = metadata.row_group(0).column(1);
5618
5619 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5620 let mut reader =
5621 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5622 let page = reader.get_next_page().unwrap().unwrap();
5623 match page {
5624 Page::DictionaryPage { buf, .. } => buf.len(),
5625 _ => panic!("expected DictionaryPage"),
5626 }
5627 };
5628
5629 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5630 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5631 }
5632
5633 #[test]
5634 fn test_arrow_writer_granular_mode_roundtrip() {
5635 let small = "tiny".to_string();
5644 let big = "x".repeat(64 * 1024);
5645 let strings: Vec<String> = (0..256)
5646 .map(|i| {
5647 if i % 16 == 0 {
5648 big.clone()
5649 } else {
5650 small.clone()
5651 }
5652 })
5653 .collect();
5654
5655 let schema = Arc::new(Schema::new(vec![Field::new(
5656 "col",
5657 ArrowDataType::Utf8,
5658 false,
5659 )]));
5660 let batch = RecordBatch::try_new(
5661 schema.clone(),
5662 vec![Arc::new(StringArray::from(strings.clone())) as _],
5663 )
5664 .unwrap();
5665
5666 let props = WriterProperties::builder()
5667 .set_dictionary_enabled(false)
5668 .set_data_page_size_limit(16 * 1024)
5669 .build();
5670 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5671 writer.write(&batch).unwrap();
5672 let data = Bytes::from(writer.into_inner().unwrap());
5673
5674 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5675 let read = reader.next().unwrap().unwrap();
5676 assert!(reader.next().is_none(), "expected one batch");
5677 let col = read
5678 .column(0)
5679 .as_any()
5680 .downcast_ref::<StringArray>()
5681 .unwrap();
5682 assert_eq!(col.len(), strings.len());
5683 for (i, expected) in strings.iter().enumerate() {
5684 assert_eq!(
5685 col.value(i),
5686 expected.as_str(),
5687 "value mismatch at index {i}"
5688 );
5689 }
5690 }
5691
5692 #[test]
5693 fn test_arrow_writer_all_null_string_column() {
5694 let num_rows = 1024;
5699 let schema = Arc::new(Schema::new(vec![Field::new(
5700 "col",
5701 ArrowDataType::Utf8,
5702 true,
5703 )]));
5704 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5705 let batch = RecordBatch::try_new(
5706 schema.clone(),
5707 vec![Arc::new(StringArray::from(nulls)) as _],
5708 )
5709 .unwrap();
5710
5711 let props = WriterProperties::builder()
5712 .set_dictionary_enabled(false)
5713 .set_data_page_size_limit(16 * 1024)
5714 .build();
5715 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5716 writer.write(&batch).unwrap();
5717 let data = Bytes::from(writer.into_inner().unwrap());
5718
5719 let mut metadata = ParquetMetaDataReader::new();
5722 metadata.try_parse(&data).unwrap();
5723 let metadata = metadata.finish().unwrap();
5724 let row_group = metadata.row_group(0);
5725 let col_meta = row_group.column(0);
5726 assert_eq!(row_group.num_rows() as usize, num_rows);
5727 if let Some(stats) = col_meta.statistics() {
5730 assert_eq!(
5731 stats.null_count_opt().unwrap_or(0) as usize,
5732 num_rows,
5733 "expected all-null column to report null_count = num_rows"
5734 );
5735 }
5736
5737 let mut reader =
5738 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5739 let mut total_values = 0u32;
5740 while let Some(page) = reader.get_next_page().unwrap() {
5741 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5742 total_values += page.num_values();
5743 }
5744 }
5745 assert_eq!(
5746 total_values as usize, num_rows,
5747 "expected every level position to be represented in some page"
5748 );
5749 }
5750
5751 struct WriteBatchesShape {
5752 num_batches: usize,
5753 rows_per_batch: usize,
5754 row_size: usize,
5755 }
5756
5757 fn write_batches(
5759 WriteBatchesShape {
5760 num_batches,
5761 rows_per_batch,
5762 row_size,
5763 }: WriteBatchesShape,
5764 props: WriterProperties,
5765 ) -> ParquetRecordBatchReaderBuilder<File> {
5766 let schema = Arc::new(Schema::new(vec![Field::new(
5767 "str",
5768 ArrowDataType::Utf8,
5769 false,
5770 )]));
5771 let file = tempfile::tempfile().unwrap();
5772 let mut writer =
5773 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5774
5775 for batch_idx in 0..num_batches {
5776 let strings: Vec<String> = (0..rows_per_batch)
5777 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5778 .collect();
5779 let array = StringArray::from(strings);
5780 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5781 writer.write(&batch).unwrap();
5782 }
5783 writer.close().unwrap();
5784 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5785 }
5786
5787 #[test]
5788 fn test_row_group_limit_none_writes_single_row_group() {
5790 let props = WriterProperties::builder()
5791 .set_max_row_group_row_count(None)
5792 .set_max_row_group_bytes(None)
5793 .build();
5794
5795 let builder = write_batches(
5796 WriteBatchesShape {
5797 num_batches: 1,
5798 rows_per_batch: 1000,
5799 row_size: 4,
5800 },
5801 props,
5802 );
5803
5804 assert_eq!(
5805 &row_group_sizes(builder.metadata()),
5806 &[1000],
5807 "With no limits, all rows should be in a single row group"
5808 );
5809 }
5810
5811 #[test]
5812 fn test_row_group_limit_rows_only() {
5814 let props = WriterProperties::builder()
5815 .set_max_row_group_row_count(Some(300))
5816 .set_max_row_group_bytes(None)
5817 .build();
5818
5819 let builder = write_batches(
5820 WriteBatchesShape {
5821 num_batches: 1,
5822 rows_per_batch: 1000,
5823 row_size: 4,
5824 },
5825 props,
5826 );
5827
5828 assert_eq!(
5829 &row_group_sizes(builder.metadata()),
5830 &[300, 300, 300, 100],
5831 "Row groups should be split by row count"
5832 );
5833 }
5834
5835 #[test]
5836 fn test_row_group_limit_rows_only_many_splits() {
5839 let props = WriterProperties::builder()
5840 .set_max_row_group_row_count(Some(1))
5841 .set_max_row_group_bytes(None)
5842 .build();
5843
5844 let rows = 50_000;
5845 let builder = write_batches(
5846 WriteBatchesShape {
5847 num_batches: 1,
5848 rows_per_batch: rows,
5849 row_size: 4,
5850 },
5851 props,
5852 );
5853
5854 let sizes = row_group_sizes(builder.metadata());
5855 assert_eq!(sizes.len(), rows, "Every row should get its own row group");
5856 assert_eq!(
5857 sizes.iter().sum::<i64>(),
5858 rows as i64,
5859 "Total rows should be preserved"
5860 );
5861 }
5862
5863 #[test]
5864 fn test_row_group_limit_bytes_only() {
5866 let props = WriterProperties::builder()
5867 .set_max_row_group_row_count(None)
5868 .set_max_row_group_bytes(Some(3500))
5870 .build();
5871
5872 let builder = write_batches(
5873 WriteBatchesShape {
5874 num_batches: 10,
5875 rows_per_batch: 10,
5876 row_size: 100,
5877 },
5878 props,
5879 );
5880
5881 let sizes = row_group_sizes(builder.metadata());
5882
5883 assert!(
5884 sizes.len() > 1,
5885 "Should have multiple row groups due to byte limit, got {sizes:?}",
5886 );
5887
5888 let total_rows: i64 = sizes.iter().sum();
5889 assert_eq!(total_rows, 100, "Total rows should be preserved");
5890 }
5891
5892 #[test]
5893 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5895 let schema = Arc::new(Schema::new(vec![Field::new(
5896 "str",
5897 ArrowDataType::Utf8,
5898 false,
5899 )]));
5900 let file = tempfile::tempfile().unwrap();
5901
5902 let props = WriterProperties::builder()
5904 .set_max_row_group_row_count(None)
5905 .set_max_row_group_bytes(None)
5906 .build();
5907 let mut writer =
5908 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5909
5910 let first_array = StringArray::from(
5911 (0..10)
5912 .map(|i| format!("{i:0>100}"))
5913 .collect::<Vec<String>>(),
5914 );
5915 let first_batch =
5916 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5917 writer.write(&first_batch).unwrap();
5918 assert_eq!(writer.in_progress_rows(), 10);
5919
5920 writer.max_row_group_bytes = Some(1);
5923
5924 let second_array = StringArray::from(vec!["x".to_string()]);
5925 let second_batch =
5926 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5927 writer.write(&second_batch).unwrap();
5928 writer.close().unwrap();
5929 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5930
5931 assert_eq!(
5932 &row_group_sizes(builder.metadata()),
5933 &[10, 1],
5934 "The second write should flush an oversized in-progress row group first",
5935 );
5936 }
5937
5938 #[test]
5939 fn test_row_group_limit_both_row_wins_single_batch() {
5941 let props = WriterProperties::builder()
5942 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
5945
5946 let builder = write_batches(
5947 WriteBatchesShape {
5948 num_batches: 1,
5949 row_size: 4,
5950 rows_per_batch: 1000,
5951 },
5952 props,
5953 );
5954
5955 assert_eq!(
5956 &row_group_sizes(builder.metadata()),
5957 &[200, 200, 200, 200, 200],
5958 "Row limit should trigger before byte limit"
5959 );
5960 }
5961
5962 #[test]
5963 fn test_row_group_limit_both_row_wins_multiple_batches() {
5965 let props = WriterProperties::builder()
5966 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
5969
5970 let builder = write_batches(
5971 WriteBatchesShape {
5972 num_batches: 10,
5973 rows_per_batch: 10,
5974 row_size: 100,
5975 },
5976 props,
5977 );
5978
5979 assert_eq!(
5980 &row_group_sizes(builder.metadata()),
5981 &[5; 20],
5982 "Row limit should trigger before byte limit"
5983 );
5984 }
5985
5986 #[test]
5987 fn test_row_group_limit_both_bytes_wins() {
5989 let props = WriterProperties::builder()
5990 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
5993
5994 let builder = write_batches(
5995 WriteBatchesShape {
5996 num_batches: 10,
5997 rows_per_batch: 10,
5998 row_size: 100,
5999 },
6000 props,
6001 );
6002
6003 let sizes = row_group_sizes(builder.metadata());
6004
6005 assert!(
6006 sizes.len() > 1,
6007 "Byte limit should trigger before row limit, got {sizes:?}",
6008 );
6009
6010 assert!(
6011 sizes.iter().all(|&s| s < 1000),
6012 "No row group should hit the row limit"
6013 );
6014
6015 let total_rows: i64 = sizes.iter().sum();
6016 assert_eq!(total_rows, 100, "Total rows should be preserved");
6017 }
6018
6019 #[test]
6020 fn test_row_group_limit_both_apply_to_same_batch() {
6023 let props = WriterProperties::builder()
6024 .set_max_row_group_row_count(Some(15))
6025 .set_max_row_group_bytes(Some(1500))
6026 .build();
6027
6028 let builder = write_batches(
6029 WriteBatchesShape {
6030 num_batches: 2,
6031 rows_per_batch: 10,
6032 row_size: 100,
6033 },
6034 props,
6035 );
6036
6037 assert_eq!(
6038 &row_group_sizes(builder.metadata()),
6039 &[14, 6],
6040 "Byte limit should still apply to a batch the row limit already split"
6041 );
6042 }
6043
6044 #[test]
6045 fn arrow_column_chunk_close_mut_drops_column_index() {
6046 use crate::arrow::ArrowSchemaConverter;
6047 use crate::file::writer::SerializedFileWriter;
6048
6049 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
6050 let props = Arc::new(
6051 WriterProperties::builder()
6052 .set_statistics_enabled(EnabledStatistics::Page)
6053 .build(),
6054 );
6055 let parquet_schema = ArrowSchemaConverter::new()
6056 .with_coerce_types(props.coerce_types())
6057 .convert(&schema)
6058 .unwrap();
6059
6060 let mut buf = Vec::with_capacity(1024);
6061 let mut writer =
6062 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
6063 .unwrap();
6064
6065 let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
6066 let mut col_writers = factory.create_column_writers(0).unwrap();
6067 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
6068 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
6069 col_writers[0].write(&leaves).unwrap();
6070 }
6071 let mut chunk = col_writers.pop().unwrap().close().unwrap();
6072
6073 assert!(
6075 chunk.close().column_index.is_some(),
6076 "EnabledStatistics::Page should produce a column_index"
6077 );
6078
6079 chunk.close_mut().column_index = None;
6081 assert!(chunk.close().column_index.is_none());
6082
6083 let mut rg = writer.next_row_group().unwrap();
6084 chunk.append_to_row_group(&mut rg).unwrap();
6085 rg.close().unwrap();
6086 let file_meta = writer.close().unwrap();
6087
6088 let cc = file_meta.row_group(0).column(0);
6091 assert!(cc.column_index_range().is_none());
6092 }
6093
6094 fn write_column_to_bytes(array: ArrayRef) -> Bytes {
6096 let schema = Arc::new(Schema::new(vec![Field::new(
6097 "col",
6098 array.data_type().clone(),
6099 true,
6100 )]));
6101 let buf = get_bytes_after_close(
6102 schema.clone(),
6103 &RecordBatch::try_new(schema, vec![array]).unwrap(),
6104 );
6105 Bytes::from(buf)
6106 }
6107
6108 fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
6112 let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
6113 ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
6114 .unwrap()
6115 .build()
6116 .unwrap()
6117 .next()
6118 .unwrap()
6119 .unwrap()
6120 .column(0)
6121 .clone()
6122 }
6123
6124 fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
6125 let flat_schema = Arc::new(Schema::new(vec![Field::new(
6126 "col",
6127 flat.data_type().clone(),
6128 true,
6129 )]));
6130 let ree_bytes = write_column_to_bytes(ree);
6131 let flat_bytes = write_column_to_bytes(flat.clone());
6132 assert_eq!(
6133 ree_bytes, flat_bytes,
6134 "REE and flat bytes should be identical"
6135 );
6136
6137 let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
6138 let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
6139
6140 assert_eq!(decoded_ree.as_ref(), flat.as_ref());
6141 assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
6142 }
6143
6144 #[test]
6145 fn ree_string() {
6146 let ree: ArrayRef = Arc::new(
6147 [Some("a"), Some("a"), None, Some("b"), Some("b")]
6148 .into_iter()
6149 .collect::<Int32RunArray>(),
6150 );
6151 let flat: ArrayRef = Arc::new(StringArray::from(vec![
6152 Some("a"),
6153 Some("a"),
6154 None,
6155 Some("b"),
6156 Some("b"),
6157 ]));
6158 ree_write_read_roundtrip(ree, flat);
6159 }
6160
6161 #[test]
6162 fn ree_int32() {
6163 let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
6164 for v in [Some(1), Some(1), None, Some(2), Some(2)] {
6165 b.append_option(v);
6166 }
6167 let ree: ArrayRef = Arc::new(b.finish());
6168 let flat: ArrayRef = Arc::new(Int32Array::from(vec![
6169 Some(1),
6170 Some(1),
6171 None,
6172 Some(2),
6173 Some(2),
6174 ]));
6175 ree_write_read_roundtrip(ree, flat);
6176 }
6177
6178 #[test]
6179 fn ree_bool() {
6180 let ree: ArrayRef = Arc::new(
6182 RunArray::try_new(
6183 &Int32Array::from(vec![3, 5, 7]),
6184 &BooleanArray::from(vec![Some(true), None, Some(false)]),
6185 )
6186 .unwrap(),
6187 );
6188 let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
6189 Some(true),
6190 Some(true),
6191 Some(true),
6192 None,
6193 None,
6194 Some(false),
6195 Some(false),
6196 ]));
6197 ree_write_read_roundtrip(ree, flat);
6198 }
6199
6200 #[test]
6201 fn ree_fixed_size_binary() {
6202 let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6203 let mut b = FixedSizeBinaryBuilder::new(2);
6204 for v in vals {
6205 match v {
6206 Some(x) => b.append_value(x).unwrap(),
6207 None => b.append_null(),
6208 }
6209 }
6210 b.finish()
6211 };
6212 let ree: ArrayRef = Arc::new(
6214 RunArray::try_new(
6215 &Int32Array::from(vec![2, 4, 6]),
6216 &mk(&[Some(b"aa"), None, Some(b"bb")]),
6217 )
6218 .unwrap(),
6219 );
6220 let flat: ArrayRef = Arc::new(mk(&[
6221 Some(b"aa"),
6222 Some(b"aa"),
6223 None,
6224 None,
6225 Some(b"bb"),
6226 Some(b"bb"),
6227 ]));
6228 ree_write_read_roundtrip(ree, flat);
6229 }
6230
6231 #[test]
6232 fn ree_single_run() {
6233 let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6234 let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6235 ree_write_read_roundtrip(ree, flat);
6236 }
6237
6238 #[test]
6239 fn ree_float32() {
6240 let ree: ArrayRef = Arc::new(
6242 RunArray::try_new(
6243 &Int32Array::from(vec![2, 4, 5]),
6244 &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6245 )
6246 .unwrap(),
6247 );
6248 let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6249 Some(1.0_f32),
6250 Some(1.0_f32),
6251 None,
6252 None,
6253 Some(2.5_f32),
6254 ]));
6255 ree_write_read_roundtrip(ree, flat);
6256 }
6257
6258 #[test]
6259 fn ree_sliced() {
6260 let full: ArrayRef = Arc::new(
6265 RunArray::try_new(
6266 &Int32Array::from(vec![3, 5, 7]),
6267 &StringArray::from(vec!["a", "b", "c"]),
6268 )
6269 .unwrap(),
6270 );
6271 let sliced = full.slice(2, 5);
6272 let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6273 ree_write_read_roundtrip(sliced, flat);
6274 }
6275
6276 #[test]
6277 fn test_number_distinct_values_exact_count() {
6278 let cardinality = 50u32;
6281 let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
6282 if i % 7 == 0 {
6283 None
6284 } else {
6285 Some((i % cardinality) as i32)
6286 }
6287 })));
6288 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
6289 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6290
6291 let props = WriterProperties::builder()
6292 .set_write_row_group_number_distinct_values(true)
6293 .build();
6294 let mut buf = Vec::new();
6295 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6296 writer.write(&batch).unwrap();
6297 let metadata = writer.close().unwrap();
6298
6299 let count = metadata
6300 .row_group(0)
6301 .column(0)
6302 .statistics()
6303 .and_then(|s| s.distinct_count_opt())
6304 .expect("distinct_count should be set");
6305 assert_eq!(count, cardinality as u64);
6307 }
6308
6309 #[test]
6310 fn test_number_distinct_values_not_written_by_default() {
6311 let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
6312 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
6313 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6314
6315 let mut buf = Vec::new();
6316 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
6317 writer.write(&batch).unwrap();
6318 let metadata = writer.close().unwrap();
6319
6320 let count = metadata
6321 .row_group(0)
6322 .column(0)
6323 .statistics()
6324 .and_then(|s| s.distinct_count_opt());
6325 assert!(count.is_none());
6326 }
6327
6328 #[test]
6329 fn ree_struct_with_ree_child() {
6330 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6333
6334 let col_a: ArrayRef = Arc::new(
6335 RunArray::try_new(
6336 &run_ends,
6337 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6338 )
6339 .unwrap(),
6340 );
6341 let col_b: ArrayRef = Arc::new(
6342 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6343 );
6344
6345 let struct_array: ArrayRef = Arc::new(StructArray::new(
6346 Fields::from(vec![
6347 Field::new("a", col_a.data_type().clone(), true),
6348 Field::new("b", col_b.data_type().clone(), true),
6349 ]),
6350 vec![col_a, col_b],
6351 None,
6352 ));
6353
6354 let schema = Arc::new(Schema::new(vec![Field::new(
6355 "row",
6356 struct_array.data_type().clone(),
6357 true,
6358 )]));
6359 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6360
6361 let mut buf = Vec::new();
6362 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6363 writer.write(&batch).unwrap();
6364 let metadata = writer.close().unwrap();
6365
6366 let parquet_schema = metadata.file_metadata().schema_descr();
6367 assert_eq!(parquet_schema.num_columns(), 2);
6368 assert_eq!(
6369 parquet_schema.column(0).physical_type(),
6370 crate::basic::Type::BYTE_ARRAY
6371 );
6372 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6373 assert_eq!(
6374 parquet_schema.column(1).physical_type(),
6375 crate::basic::Type::INT32
6376 );
6377 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6378 }
6379}