1use crate::column::chunker::ContentDefinedChunker;
8
9use bytes::Bytes;
10use std::io::Write;
11use std::ops::Range;
12use std::slice::Iter;
13use std::sync::{Arc, Mutex};
14use std::vec::IntoIter;
15
16use arrow_array::cast::AsArray;
17use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter, new_empty_array};
18use arrow_array::{PrimitiveArray, types::*};
19use arrow_schema::{
20 ArrowError, DataType as ArrowDataType, Field, FieldRef, IntervalUnit, SchemaRef, TimeUnit,
21};
22
23use super::schema::{
24 add_encoded_arrow_schema_to_metadata, decimal_length_from_precision, validate_map_key_type,
25};
26
27use crate::arrow::ArrowSchemaConverter;
28use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder;
29use crate::basic::PageType;
30use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter};
31use crate::column::page_encryption::PageEncryptor;
32use crate::column::writer::encoder::ColumnValueEncoder;
33use crate::column::writer::{
34 ColumnCloseResult, ColumnWriter, ColumnWriterImpl, GenericColumnWriter, get_column_writer,
35};
36use crate::data_type::{ByteArray, FixedLenByteArray, FixedLenByteArrayType};
37use std::collections::HashSet;
38type DistinctValuesSet = HashSet<u64>;
39#[cfg(feature = "encryption")]
40use crate::encryption::encrypt::FileEncryptor;
41use crate::errors::{ParquetError, Result};
42use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData};
43use crate::file::properties::{WriterProperties, WriterPropertiesPtr};
44use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
45use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
46use crate::schema::types::{ColumnDescPtr, SchemaDescriptor};
47use levels::{ArrayLevels, calculate_array_levels};
48
49mod byte_array;
50mod levels;
51
52#[doc(inline)]
53pub use crate::column::page_store::{
54 InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs,
55 PageStoreFactory,
56};
57
58pub struct ArrowWriter<W: Write> {
175 writer: SerializedFileWriter<W>,
177
178 in_progress: Option<ArrowRowGroupWriter>,
180
181 arrow_schema: SchemaRef,
185
186 row_group_writer_factory: ArrowRowGroupWriterFactory,
188
189 max_row_group_row_count: Option<usize>,
191
192 max_row_group_bytes: Option<usize>,
194
195 cdc_chunkers: Option<Vec<ContentDefinedChunker>>,
197}
198
199impl<W: Write + Send> std::fmt::Debug for ArrowWriter<W> {
200 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201 let buffered_memory = self.in_progress_size();
202 f.debug_struct("ArrowWriter")
203 .field("writer", &self.writer)
204 .field("in_progress_size", &format_args!("{buffered_memory} bytes"))
205 .field("in_progress_rows", &self.in_progress_rows())
206 .field("arrow_schema", &self.arrow_schema)
207 .field("max_row_group_row_count", &self.max_row_group_row_count)
208 .field("max_row_group_bytes", &self.max_row_group_bytes)
209 .finish()
210 }
211}
212
213impl<W: Write + Send> ArrowWriter<W> {
214 pub fn try_new(
220 writer: W,
221 arrow_schema: SchemaRef,
222 props: Option<WriterProperties>,
223 ) -> Result<Self> {
224 let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
225 Self::try_new_with_options(writer, arrow_schema, options)
226 }
227
228 pub fn try_new_with_options(
234 writer: W,
235 arrow_schema: SchemaRef,
236 options: ArrowWriterOptions,
237 ) -> Result<Self> {
238 for field in arrow_schema.fields() {
239 validate_map_key_type(field.data_type())?;
240 }
241 let mut props = options.properties;
242
243 let schema = if let Some(parquet_schema) = options.schema_descr {
244 parquet_schema.clone()
245 } else {
246 let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types());
247 if let Some(schema_root) = &options.schema_root {
248 converter = converter.schema_root(schema_root);
249 }
250
251 converter.convert(&arrow_schema)?
252 };
253
254 if !options.skip_arrow_metadata {
255 add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
257 }
258
259 let max_row_group_row_count = props.max_row_group_row_count();
260 let max_row_group_bytes = props.max_row_group_bytes();
261
262 let props_ptr = Arc::new(props);
263 let file_writer =
264 SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
265
266 let mut row_group_writer_factory =
267 ArrowRowGroupWriterFactory::try_new(&file_writer, arrow_schema.clone())?;
268 if let Some(page_store_factory) = options.page_store_factory {
269 row_group_writer_factory =
270 row_group_writer_factory.with_page_store_factory(page_store_factory);
271 }
272
273 let cdc_chunkers = props_ptr
274 .content_defined_chunking()
275 .map(|opts| {
276 file_writer
277 .schema_descr()
278 .columns()
279 .iter()
280 .map(|desc| ContentDefinedChunker::new(desc, opts))
281 .collect::<Result<Vec<_>>>()
282 })
283 .transpose()?;
284
285 Ok(Self {
286 writer: file_writer,
287 in_progress: None,
288 arrow_schema,
289 row_group_writer_factory,
290 max_row_group_row_count,
291 max_row_group_bytes,
292 cdc_chunkers,
293 })
294 }
295
296 pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
298 self.writer.flushed_row_groups()
299 }
300
301 pub fn memory_size(&self) -> usize {
306 match &self.in_progress {
307 Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(),
308 None => 0,
309 }
310 }
311
312 pub fn in_progress_size(&self) -> usize {
319 match &self.in_progress {
320 Some(in_progress) => in_progress
321 .writers
322 .iter()
323 .map(|x| x.get_estimated_total_bytes())
324 .sum(),
325 None => 0,
326 }
327 }
328
329 pub fn in_progress_rows(&self) -> usize {
331 self.in_progress
332 .as_ref()
333 .map(|x| x.buffered_rows)
334 .unwrap_or_default()
335 }
336
337 pub fn bytes_written(&self) -> usize {
339 self.writer.bytes_written()
340 }
341
342 pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
356 if batch.num_rows() == 0 {
357 return Ok(());
358 }
359
360 let mut remaining = batch.clone();
363
364 loop {
365 let in_progress = match &mut self.in_progress {
366 Some(in_progress) => in_progress,
367 x => x.insert(
368 self.row_group_writer_factory
369 .create_row_group_writer(self.writer.flushed_row_groups().len())?,
370 ),
371 };
372 let buffered_rows = in_progress.buffered_rows;
373
374 let mut split_at = match self.max_row_group_row_count {
377 Some(max_rows) if buffered_rows + remaining.num_rows() > max_rows => {
378 Some(max_rows - buffered_rows)
379 }
380 _ => None,
381 };
382
383 let candidate_rows = split_at.unwrap_or_else(|| remaining.num_rows());
388
389 if self.max_row_group_bytes.is_some() && buffered_rows == 0 {
390 if candidate_rows > 1 {
394 split_at = Some(1);
395 }
396 } else if let Some(max_bytes) = self.max_row_group_bytes {
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 let avg_row_bytes = current_bytes / buffered_rows;
405 if let Some(rows_that_fit) = (max_bytes - current_bytes).checked_div(avg_row_bytes)
406 {
407 if candidate_rows > rows_that_fit {
409 if rows_that_fit > 0 {
410 split_at = Some(rows_that_fit);
411 } else {
412 self.flush()?;
413 continue;
414 }
415 }
416 }
417 }
418
419 let rest = split_at.map(|to_write| {
420 let rest = remaining.slice(to_write, remaining.num_rows() - to_write);
421 remaining = remaining.slice(0, to_write);
422 rest
423 });
424
425 let in_progress = self.in_progress.as_mut().unwrap();
426 match self.cdc_chunkers.as_mut() {
427 Some(chunkers) => in_progress.write_with_chunkers(&remaining, chunkers)?,
428 None => in_progress.write(&remaining)?,
429 }
430
431 let should_flush = self
432 .max_row_group_row_count
433 .is_some_and(|max| in_progress.buffered_rows >= max)
434 || self
435 .max_row_group_bytes
436 .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
437
438 if should_flush {
439 self.flush()?
440 }
441
442 match rest {
443 Some(rest) => remaining = rest,
444 None => return Ok(()),
445 }
446 }
447 }
448
449 pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
454 self.writer.write_all(buf)
455 }
456
457 pub fn sync(&mut self) -> std::io::Result<()> {
459 self.writer.flush()
460 }
461
462 pub fn flush(&mut self) -> Result<()> {
467 let Some(in_progress) = self.in_progress.take() else {
468 return Ok(());
469 };
470
471 let mut row_group_writer = self.writer.next_row_group()?;
472 for chunk in in_progress.close()? {
473 chunk.append_to_row_group(&mut row_group_writer)?;
474 }
475 row_group_writer.close()?;
476 Ok(())
477 }
478
479 pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
483 self.writer.append_key_value_metadata(kv_metadata)
484 }
485
486 pub fn inner(&self) -> &W {
488 self.writer.inner()
489 }
490
491 pub fn inner_mut(&mut self) -> &mut W {
500 self.writer.inner_mut()
501 }
502
503 pub fn into_inner(mut self) -> Result<W> {
505 self.flush()?;
506 self.writer.into_inner()
507 }
508
509 pub fn finish(&mut self) -> Result<ParquetMetaData> {
515 self.flush()?;
516 self.writer.finish()
517 }
518
519 pub fn close(mut self) -> Result<ParquetMetaData> {
521 self.finish()
522 }
523
524 pub fn into_serialized_writer(
531 mut self,
532 ) -> Result<(SerializedFileWriter<W>, ArrowRowGroupWriterFactory)> {
533 self.flush()?;
534 Ok((self.writer, self.row_group_writer_factory))
535 }
536}
537
538impl<W: Write + Send> RecordBatchWriter for ArrowWriter<W> {
539 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
540 self.write(batch).map_err(|e| e.into())
541 }
542
543 fn close(self) -> std::result::Result<(), ArrowError> {
544 self.close()?;
545 Ok(())
546 }
547}
548
549#[derive(Debug, Clone, Default)]
553pub struct ArrowWriterOptions {
554 properties: WriterProperties,
555 skip_arrow_metadata: bool,
556 schema_root: Option<String>,
557 schema_descr: Option<SchemaDescriptor>,
558 page_store_factory: Option<Arc<dyn PageStoreFactory>>,
559}
560
561impl ArrowWriterOptions {
562 pub fn new() -> Self {
564 Self::default()
565 }
566
567 pub fn with_properties(self, properties: WriterProperties) -> Self {
569 Self { properties, ..self }
570 }
571
572 pub fn with_page_store_factory(self, page_store_factory: Arc<dyn PageStoreFactory>) -> Self {
658 Self {
659 page_store_factory: Some(page_store_factory),
660 ..self
661 }
662 }
663
664 pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
671 Self {
672 skip_arrow_metadata,
673 ..self
674 }
675 }
676
677 pub fn with_schema_root(self, schema_root: String) -> Self {
679 Self {
680 schema_root: Some(schema_root),
681 ..self
682 }
683 }
684
685 pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self {
691 Self {
692 schema_descr: Some(schema_descr),
693 ..self
694 }
695 }
696}
697
698struct ArrowColumnChunkData {
704 length: usize,
705 store: Box<dyn PageStore>,
706 keys: Vec<PageKey>,
707 dictionary_keys: Vec<PageKey>,
718 dictionary_len: usize,
722}
723
724impl ArrowColumnChunkData {
725 fn new(store: Box<dyn PageStore>) -> Self {
726 Self {
727 length: 0,
728 store,
729 keys: Vec::new(),
730 dictionary_keys: Vec::new(),
731 dictionary_len: 0,
732 }
733 }
734
735 fn push(&mut self, value: Bytes) -> Result<()> {
738 let key = self.store.put(value)?;
739 self.keys.push(key);
740 Ok(())
741 }
742
743 fn push_dictionary(&mut self, value: Bytes) -> Result<()> {
747 self.dictionary_len += value.len();
748 let key = self.store.put(value)?;
749 self.dictionary_keys.push(key);
750 Ok(())
751 }
752
753 fn memory_size(&self) -> usize {
756 self.store.memory_size()
757 }
758}
759
760struct StreamingColumnChunkPages {
769 store: Box<dyn PageStore>,
770 keys: IntoIter<PageKey>,
773}
774
775impl StreamingColumnChunkPages {
776 fn new(data: ArrowColumnChunkData) -> Self {
777 let keys = if data.dictionary_keys.is_empty() {
780 data.keys
781 } else {
782 let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len());
783 keys.extend(data.dictionary_keys);
784 keys.extend(data.keys);
785 keys
786 };
787 Self {
788 store: data.store,
789 keys: keys.into_iter(),
790 }
791 }
792}
793
794impl Iterator for StreamingColumnChunkPages {
795 type Item = Result<Bytes>;
796
797 fn next(&mut self) -> Option<Self::Item> {
798 let key = self.keys.next()?;
799 Some(self.store.take(key))
800 }
801}
802
803type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
808
809struct ArrowPageWriter {
810 buffer: SharedColumnChunk,
811 #[cfg(feature = "encryption")]
812 page_encryptor: Option<PageEncryptor>,
813}
814
815impl ArrowPageWriter {
816 fn new(store: Box<dyn PageStore>) -> Self {
818 Self {
819 buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))),
820 #[cfg(feature = "encryption")]
821 page_encryptor: None,
822 }
823 }
824
825 #[cfg(feature = "encryption")]
826 pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
827 self.page_encryptor = page_encryptor;
828 self
829 }
830
831 #[cfg(feature = "encryption")]
832 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
833 self.page_encryptor.as_mut()
834 }
835
836 #[cfg(not(feature = "encryption"))]
839 #[expect(
840 clippy::needless_pass_by_ref_mut,
841 reason = "mirrors the encryption-enabled signature"
842 )]
843 fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
844 None
845 }
846}
847
848impl PageWriter for ArrowPageWriter {
849 fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
850 let page = match self.page_encryptor_mut() {
851 Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
852 None => page,
853 };
854
855 let page_header = page.to_thrift_header()?;
856 let header = {
857 let mut header = Vec::with_capacity(1024);
858
859 match self.page_encryptor_mut() {
860 Some(page_encryptor) => {
861 page_encryptor.encrypt_page_header(&page_header, &mut header)?;
862 if page.compressed_page().is_data_page() {
863 page_encryptor.increment_page();
864 }
865 }
866 None => {
867 let mut protocol = ThriftCompactOutputProtocol::new(&mut header);
868 page_header.write_thrift(&mut protocol)?;
869 }
870 }
871
872 Bytes::from(header)
873 };
874
875 let mut buf = self.buffer.try_lock().unwrap();
876
877 let data = page.compressed_page().buffer().clone();
878 let compressed_size = data.len() + header.len();
879
880 let mut spec = PageWriteSpec::new();
881 spec.page_type = page.page_type();
882 spec.num_values = page.num_values();
883 spec.uncompressed_size = page.uncompressed_size() + header.len();
884 spec.offset = buf.length as u64;
885 spec.compressed_size = compressed_size;
886 spec.bytes_written = compressed_size as u64;
887
888 buf.length += compressed_size;
889 if spec.page_type == PageType::DICTIONARY_PAGE {
890 buf.push_dictionary(header)?;
893 buf.push_dictionary(data)?;
894 } else {
895 buf.push(header)?;
896 buf.push(data)?;
897 }
898
899 Ok(spec)
900 }
901
902 fn defers_dictionary_ordering(&self) -> bool {
903 true
908 }
909
910 fn buffered_memory_size(&self) -> usize {
911 self.buffer.try_lock().unwrap().memory_size()
914 }
915
916 fn close(&mut self) -> Result<()> {
917 Ok(())
918 }
919}
920
921#[derive(Debug)]
923pub struct ArrowLeafColumn(ArrayLevels);
924
925pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result<Vec<ArrowLeafColumn>> {
931 let levels = calculate_array_levels(array, field)?;
932 Ok(levels.into_iter().map(ArrowLeafColumn).collect())
933}
934
935pub struct ArrowColumnChunk {
937 data: ArrowColumnChunkData,
938 close: ColumnCloseResult,
939}
940
941impl std::fmt::Debug for ArrowColumnChunk {
942 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
943 f.debug_struct("ArrowColumnChunk")
944 .field("length", &self.data.length)
945 .finish_non_exhaustive()
946 }
947}
948
949impl ArrowColumnChunk {
950 pub fn close(&self) -> &ColumnCloseResult {
957 &self.close
958 }
959
960 pub fn close_mut(&mut self) -> &mut ColumnCloseResult {
967 &mut self.close
968 }
969
970 pub fn append_to_row_group<W: Write + Send>(
973 self,
974 writer: &mut SerializedRowGroupWriter<'_, W>,
975 ) -> Result<()> {
976 let ArrowColumnChunk { data, close } = self;
977
978 let close = close.update_dictionary_location(data.dictionary_len)?;
982
983 let pages = StreamingColumnChunkPages::new(data);
984 writer.append_column_from_pages(pages, close)
985 }
986}
987
988pub struct ArrowColumnWriter {
1086 writer: ArrowColumnWriterImpl,
1087 chunk: SharedColumnChunk,
1088 distinct_values_seen: Option<DistinctValuesSet>,
1091}
1092
1093impl std::fmt::Debug for ArrowColumnWriter {
1094 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1095 f.debug_struct("ArrowColumnWriter").finish_non_exhaustive()
1096 }
1097}
1098
1099enum ArrowColumnWriterImpl {
1100 ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
1101 Column(ColumnWriter<'static>),
1102}
1103
1104impl ArrowColumnWriter {
1105 pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> {
1107 col.0.validate()?;
1108 self.write_internal(&col.0)
1109 }
1110
1111 fn write_with_chunker(
1113 &mut self,
1114 col: &ArrowLeafColumn,
1115 chunker: &mut ContentDefinedChunker,
1116 ) -> Result<()> {
1117 let levels = &col.0;
1118 levels.validate()?;
1119 let chunks = chunker.get_arrow_chunks(
1120 levels.def_level_data().as_ref(),
1121 levels.rep_level_data().as_ref(),
1122 levels.array(),
1123 )?;
1124
1125 let num_chunks = chunks.len();
1126 for (i, chunk) in chunks.iter().enumerate() {
1127 let chunk_levels = levels.slice_for_chunk(chunk);
1128 self.write_internal(&chunk_levels)?;
1129
1130 if i + 1 < num_chunks {
1132 match &mut self.writer {
1133 ArrowColumnWriterImpl::Column(c) => c.add_data_page()?,
1134 ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?,
1135 }
1136 }
1137 }
1138 Ok(())
1139 }
1140
1141 fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
1142 if let Some(seen) = &mut self.distinct_values_seen {
1143 let array = levels.array();
1144 let non_null = levels.non_null_indices();
1145 match array.as_any_dictionary_opt() {
1146 Some(dict) => {
1147 let values = dict.values();
1150 let keys = dict.normalized_keys();
1151 let referenced_value_indices: Vec<usize> = non_null
1152 .iter()
1153 .map(|&pos| keys[pos])
1154 .filter(|&val_idx| values.is_valid(val_idx))
1155 .collect();
1156 update_distinct_values_seen(values.as_ref(), &referenced_value_indices, seen);
1157 }
1158 None => update_distinct_values_seen(array.as_ref(), non_null, seen),
1160 }
1161 }
1162
1163 match &mut self.writer {
1164 ArrowColumnWriterImpl::Column(c) => {
1165 let leaf = levels.array();
1166 match leaf.as_any_dictionary_opt() {
1167 Some(dictionary) => {
1168 let materialized =
1169 arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?;
1170 write_leaf(c, &materialized, levels)?
1171 }
1172 None => write_leaf(c, leaf, levels)?,
1173 };
1174 }
1175 ArrowColumnWriterImpl::ByteArray(c) => {
1176 write_primitive(c, levels.array().as_ref(), levels)?;
1177 }
1178 }
1179 Ok(())
1180 }
1181
1182 pub fn close(self) -> Result<ArrowColumnChunk> {
1189 let distinct_count = self
1190 .distinct_values_seen
1191 .as_ref()
1192 .filter(|s| !s.is_empty())
1193 .map(|s| s.len() as u64);
1194 let close = match self.writer {
1195 ArrowColumnWriterImpl::ByteArray(mut c) => {
1196 if let Some(count) = distinct_count {
1197 c.set_distinct_count_override(count);
1198 }
1199 c.close()?
1200 }
1201 ArrowColumnWriterImpl::Column(mut c) => {
1202 if let Some(count) = distinct_count {
1203 c.set_distinct_count_override(count);
1204 }
1205 c.close()?
1206 }
1207 };
1208 let chunk = Arc::try_unwrap(self.chunk)
1210 .map_err(|_| general_err!("Internal Error: the column chunk is still shared"))?;
1211 let data = chunk
1212 .into_inner()
1213 .map_err(|_| general_err!("The column chunk lock is poisoned"))?;
1214 Ok(ArrowColumnChunk { data, close })
1215 }
1216
1217 pub fn memory_size(&self) -> usize {
1228 match &self.writer {
1229 ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1230 ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1231 }
1232 }
1233
1234 pub fn get_estimated_total_bytes(&self) -> usize {
1242 match &self.writer {
1243 ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1244 ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1245 }
1246 }
1247}
1248
1249#[derive(Debug)]
1251struct ArrowTopLevelWriterSpec {
1252 field: FieldRef,
1256 leaf_range: Range<usize>,
1257}
1258
1259#[derive(Debug)]
1264enum ArrowWriteSchemaPlanError {
1265 General(String),
1266 Nyi(String),
1267 ArrowError(String),
1268}
1269
1270impl ArrowWriteSchemaPlanError {
1271 fn to_parquet_error(&self) -> ParquetError {
1272 match self {
1273 Self::General(message) => ParquetError::General(message.clone()),
1274 Self::Nyi(message) => ParquetError::NYI(message.clone()),
1275 Self::ArrowError(message) => ParquetError::ArrowError(message.clone()),
1276 }
1277 }
1278
1279 fn from_parquet_error(error: ParquetError) -> Self {
1280 match error {
1281 ParquetError::NYI(message) => Self::Nyi(message),
1282 ParquetError::General(message) => Self::General(message),
1283 ParquetError::ArrowError(message) => Self::ArrowError(message),
1284 error => Self::General(error.to_string()),
1285 }
1286 }
1287}
1288
1289#[derive(Debug)]
1294struct ArrowWriteSchemaPlan {
1295 fields: Box<[ArrowTopLevelWriterSpec]>,
1296 leaves: Box<[ColumnDescPtr]>,
1297}
1298
1299impl ArrowWriteSchemaPlan {
1300 fn try_new(parquet: &SchemaDescriptor, arrow: &SchemaRef) -> Result<Self> {
1301 let parquet_leaves = parquet.columns();
1302 let mut leaf_count = 0;
1303 let mut fields = Vec::with_capacity(arrow.fields.len());
1304
1305 for field in &arrow.fields {
1306 validate_map_key_type(field.data_type())?;
1307 let start = leaf_count;
1308 leaf_count += compute_leaves(field, &new_empty_array(field.data_type()))?.len();
1309 fields.push(ArrowTopLevelWriterSpec {
1310 field: Arc::clone(field),
1311 leaf_range: start..leaf_count,
1312 });
1313 }
1314
1315 if leaf_count != parquet_leaves.len() {
1316 return Err(ParquetError::General(format!(
1317 "Arrow schema maps to {} leaf columns but Parquet schema has {}",
1318 leaf_count,
1319 parquet_leaves.len()
1320 )));
1321 }
1322
1323 Ok(Self {
1324 fields: fields.into_boxed_slice(),
1325 leaves: parquet_leaves.to_vec().into_boxed_slice(),
1326 })
1327 }
1328}
1329
1330#[derive(Debug)]
1337struct ArrowRowGroupWriter {
1338 writers: Vec<ArrowColumnWriter>,
1339 schema_plan: Arc<ArrowWriteSchemaPlan>,
1340 buffered_rows: usize,
1341}
1342
1343impl ArrowRowGroupWriter {
1344 fn new(writers: Vec<ArrowColumnWriter>, schema_plan: Arc<ArrowWriteSchemaPlan>) -> Self {
1345 debug_assert_eq!(writers.len(), schema_plan.leaves.len());
1346 Self {
1347 writers,
1348 schema_plan,
1349 buffered_rows: 0,
1350 }
1351 }
1352
1353 fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1354 self.validate_batch_shape(batch, None)?;
1355 self.buffered_rows += batch.num_rows();
1356
1357 for (column_idx, field) in self.schema_plan.fields.iter().enumerate() {
1358 let leaves = compute_leaves(field.field.as_ref(), batch.column(column_idx))?;
1359 self.validate_leaf_count(field, leaves.len())?;
1360 for (offset, leaf) in leaves.into_iter().enumerate() {
1361 self.writers[field.leaf_range.start + offset].write(&leaf)?;
1362 }
1363 }
1364 Ok(())
1365 }
1366
1367 fn write_with_chunkers(
1368 &mut self,
1369 batch: &RecordBatch,
1370 chunkers: &mut [ContentDefinedChunker],
1371 ) -> Result<()> {
1372 self.validate_batch_shape(batch, Some(chunkers.len()))?;
1373 self.buffered_rows += batch.num_rows();
1374
1375 for (column_idx, field) in self.schema_plan.fields.iter().enumerate() {
1376 let leaves = compute_leaves(field.field.as_ref(), batch.column(column_idx))?;
1377 self.validate_leaf_count(field, leaves.len())?;
1378 for (offset, leaf) in leaves.into_iter().enumerate() {
1379 let leaf_idx = field.leaf_range.start + offset;
1380 self.writers[leaf_idx].write_with_chunker(&leaf, &mut chunkers[leaf_idx])?;
1381 }
1382 }
1383 Ok(())
1384 }
1385
1386 fn validate_batch_shape(
1387 &self,
1388 batch: &RecordBatch,
1389 chunker_count: Option<usize>,
1390 ) -> Result<()> {
1391 if batch.num_columns() != self.schema_plan.fields.len() {
1392 return Err(ParquetError::ArrowError(format!(
1393 "Incompatible schema: writer has {} top-level fields but batch has {} columns",
1394 self.schema_plan.fields.len(),
1395 batch.num_columns()
1396 )));
1397 }
1398 if self.writers.len() != self.schema_plan.leaves.len() {
1399 return Err(ParquetError::General(format!(
1400 "Arrow row-group writer has {} column writers for {} planned leaves",
1401 self.writers.len(),
1402 self.schema_plan.leaves.len()
1403 )));
1404 }
1405 if let Some(actual) = chunker_count
1406 && actual != self.schema_plan.leaves.len()
1407 {
1408 return Err(ParquetError::General(format!(
1409 "content-defined chunking has {actual} chunkers for {} planned leaves",
1410 self.schema_plan.leaves.len()
1411 )));
1412 }
1413 Ok(())
1414 }
1415
1416 fn validate_leaf_count(&self, field: &ArrowTopLevelWriterSpec, actual: usize) -> Result<()> {
1417 let expected = field.leaf_range.end - field.leaf_range.start;
1418 if actual != expected {
1419 return Err(ParquetError::ArrowError(format!(
1420 "Incompatible schema: field '{}' produced {actual} leaf columns but writer expects {expected}",
1421 field.field.name()
1422 )));
1423 }
1424 Ok(())
1425 }
1426
1427 fn get_estimated_total_bytes(&self) -> usize {
1429 self.writers
1430 .iter()
1431 .map(|x| x.get_estimated_total_bytes())
1432 .sum()
1433 }
1434
1435 fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1436 self.writers
1437 .into_iter()
1438 .map(|writer| writer.close())
1439 .collect()
1440 }
1441}
1442
1443#[derive(Debug)]
1448pub struct ArrowRowGroupWriterFactory {
1449 schema_plan: std::result::Result<Arc<ArrowWriteSchemaPlan>, ArrowWriteSchemaPlanError>,
1454 props: WriterPropertiesPtr,
1455 page_store_factory: Arc<dyn PageStoreFactory>,
1456 #[cfg(feature = "encryption")]
1457 file_encryptor: Option<Arc<FileEncryptor>>,
1458}
1459
1460impl ArrowRowGroupWriterFactory {
1461 #[deprecated(
1466 since = "60.2.0",
1467 note = "Use `ArrowRowGroupWriterFactory::try_new` instead"
1468 )]
1469 pub fn new<W: Write + Send>(
1470 file_writer: &SerializedFileWriter<W>,
1471 arrow_schema: SchemaRef,
1472 ) -> Self {
1473 let schema_plan =
1474 ArrowWriteSchemaPlan::try_new(file_writer.schema_descr_ptr().as_ref(), &arrow_schema)
1475 .map(Arc::new)
1476 .map_err(ArrowWriteSchemaPlanError::from_parquet_error);
1477 Self::from_schema_plan(file_writer, schema_plan)
1478 }
1479
1480 pub fn try_new<W: Write + Send>(
1489 file_writer: &SerializedFileWriter<W>,
1490 arrow_schema: SchemaRef,
1491 ) -> Result<Self> {
1492 let schema_plan =
1493 ArrowWriteSchemaPlan::try_new(file_writer.schema_descr_ptr().as_ref(), &arrow_schema)?;
1494 Ok(Self::from_schema_plan(
1495 file_writer,
1496 Ok(Arc::new(schema_plan)),
1497 ))
1498 }
1499
1500 fn from_schema_plan<W: Write + Send>(
1501 file_writer: &SerializedFileWriter<W>,
1502 schema_plan: std::result::Result<Arc<ArrowWriteSchemaPlan>, ArrowWriteSchemaPlanError>,
1503 ) -> Self {
1504 let props = Arc::clone(file_writer.properties());
1505 Self {
1506 schema_plan,
1507 props,
1508 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1509 #[cfg(feature = "encryption")]
1510 file_encryptor: file_writer.file_encryptor(),
1511 }
1512 }
1513
1514 pub fn with_page_store_factory(
1518 mut self,
1519 page_store_factory: Arc<dyn PageStoreFactory>,
1520 ) -> Self {
1521 self.page_store_factory = page_store_factory;
1522 self
1523 }
1524
1525 fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1526 let schema_plan = Arc::clone(self.schema_plan()?);
1527 let writers = self.create_column_writers_from_plan(row_group_index, &schema_plan)?;
1528 Ok(ArrowRowGroupWriter::new(writers, schema_plan))
1529 }
1530
1531 pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1533 self.create_column_writers_from_plan(row_group_index, self.schema_plan()?)
1534 }
1535
1536 fn schema_plan(&self) -> Result<&Arc<ArrowWriteSchemaPlan>> {
1537 match &self.schema_plan {
1538 Ok(plan) => Ok(plan),
1539 Err(error) => Err(error.to_parquet_error()),
1540 }
1541 }
1542
1543 fn create_column_writers_from_plan(
1544 &self,
1545 row_group_index: usize,
1546 schema_plan: &ArrowWriteSchemaPlan,
1547 ) -> Result<Vec<ArrowColumnWriter>> {
1548 let mut writers = Vec::with_capacity(schema_plan.leaves.len());
1549 let mut leaves = schema_plan.leaves.iter();
1550 let column_factory = self.column_writer_factory(row_group_index);
1551 for field in &schema_plan.fields {
1552 column_factory.get_arrow_column_writer(
1553 field.field.data_type(),
1554 &self.props,
1555 &mut leaves,
1556 &mut writers,
1557 )?;
1558 }
1559 Ok(writers)
1560 }
1561
1562 #[cfg(feature = "encryption")]
1563 fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1564 ArrowColumnWriterFactory::new()
1565 .with_page_store_factory(self.page_store_factory.clone())
1566 .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1567 }
1568
1569 #[cfg(not(feature = "encryption"))]
1570 fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1571 ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1572 }
1573}
1574
1575struct ArrowColumnWriterFactory {
1577 page_store_factory: Arc<dyn PageStoreFactory>,
1579 #[cfg(feature = "encryption")]
1580 row_group_index: usize,
1581 #[cfg(feature = "encryption")]
1582 file_encryptor: Option<Arc<FileEncryptor>>,
1583}
1584
1585impl ArrowColumnWriterFactory {
1586 pub fn new() -> Self {
1587 Self {
1588 page_store_factory: Arc::new(InMemoryPageStoreFactory),
1589 #[cfg(feature = "encryption")]
1590 row_group_index: 0,
1591 #[cfg(feature = "encryption")]
1592 file_encryptor: None,
1593 }
1594 }
1595
1596 pub fn with_page_store_factory(
1598 mut self,
1599 page_store_factory: Arc<dyn PageStoreFactory>,
1600 ) -> Self {
1601 self.page_store_factory = page_store_factory;
1602 self
1603 }
1604
1605 #[cfg(feature = "encryption")]
1606 pub fn with_file_encryptor(
1607 mut self,
1608 row_group_index: usize,
1609 file_encryptor: Option<Arc<FileEncryptor>>,
1610 ) -> Self {
1611 self.row_group_index = row_group_index;
1612 self.file_encryptor = file_encryptor;
1613 self
1614 }
1615
1616 #[cfg(feature = "encryption")]
1617 fn create_page_writer(
1618 &self,
1619 column_descriptor: &ColumnDescPtr,
1620 column_index: usize,
1621 ) -> Result<Box<ArrowPageWriter>> {
1622 let column_path = column_descriptor.path().string();
1623 let page_encryptor = PageEncryptor::create_if_column_encrypted(
1624 self.file_encryptor.as_ref(),
1625 self.row_group_index,
1626 column_index,
1627 &column_path,
1628 )?;
1629 let args = PageStoreArgs::new(column_index, column_descriptor);
1630 let store = self.page_store_factory.create(&args)?;
1631 Ok(Box::new(
1632 ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1633 ))
1634 }
1635
1636 #[cfg(not(feature = "encryption"))]
1637 fn create_page_writer(
1638 &self,
1639 column_descriptor: &ColumnDescPtr,
1640 column_index: usize,
1641 ) -> Result<Box<ArrowPageWriter>> {
1642 let args = PageStoreArgs::new(column_index, column_descriptor);
1643 let store = self.page_store_factory.create(&args)?;
1644 Ok(Box::new(ArrowPageWriter::new(store)))
1645 }
1646
1647 fn get_arrow_column_writer(
1650 &self,
1651 data_type: &ArrowDataType,
1652 props: &WriterPropertiesPtr,
1653 leaves: &mut Iter<'_, ColumnDescPtr>,
1654 out: &mut Vec<ArrowColumnWriter>,
1655 ) -> Result<()> {
1656 let write_distinct_values = props.write_row_group_number_distinct_values();
1657
1658 let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1660 let page_writer = self.create_page_writer(desc, out.len())?;
1661 let chunk = page_writer.buffer.clone();
1662 let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1663 Ok(ArrowColumnWriter {
1664 chunk,
1665 writer: ArrowColumnWriterImpl::Column(writer),
1666 distinct_values_seen: write_distinct_values.then(HashSet::new),
1667 })
1668 };
1669
1670 let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1672 let page_writer = self.create_page_writer(desc, out.len())?;
1673 let chunk = page_writer.buffer.clone();
1674 let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1675 Ok(ArrowColumnWriter {
1676 chunk,
1677 writer: ArrowColumnWriterImpl::ByteArray(writer),
1678 distinct_values_seen: write_distinct_values.then(HashSet::new),
1679 })
1680 };
1681
1682 match data_type {
1683 _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1684 ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1685 out.push(col(leaves.next().unwrap())?)
1686 }
1687 ArrowDataType::LargeBinary
1688 | ArrowDataType::Binary
1689 | ArrowDataType::Utf8
1690 | ArrowDataType::LargeUtf8
1691 | ArrowDataType::BinaryView
1692 | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1693 ArrowDataType::List(f)
1694 | ArrowDataType::LargeList(f)
1695 | ArrowDataType::FixedSizeList(f, _)
1696 | ArrowDataType::ListView(f)
1697 | ArrowDataType::LargeListView(f) => {
1698 self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1699 }
1700 ArrowDataType::Struct(fields) => {
1701 for field in fields {
1702 self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1703 }
1704 }
1705 ArrowDataType::Map(f, _) => match f.data_type() {
1706 ArrowDataType::Struct(f) => {
1707 self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1708 self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1709 }
1710 _ => unreachable!("invalid map type"),
1711 },
1712 ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1713 ArrowDataType::Utf8
1714 | ArrowDataType::LargeUtf8
1715 | ArrowDataType::Binary
1716 | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1717 ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1718 out.push(bytes(leaves.next().unwrap())?)
1719 }
1720 _ => self.get_arrow_column_writer(value_type, props, leaves, out)?,
1721 },
1722 ArrowDataType::RunEndEncoded(_, value_field) => {
1723 self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1724 }
1725 _ => {
1726 return Err(ParquetError::NYI(format!(
1727 "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1728 )));
1729 }
1730 }
1731 Ok(())
1732 }
1733}
1734
1735fn write_leaf(
1736 writer: &mut ColumnWriter<'_>,
1737 column: &dyn arrow_array::Array,
1738 levels: &ArrayLevels,
1739) -> Result<usize> {
1740 let indices = levels.non_null_indices();
1741
1742 match writer {
1743 ColumnWriter::Int32ColumnWriter(typed) => visit_i32(
1745 column,
1746 WritePrimitive {
1747 writer: typed,
1748 levels,
1749 },
1750 ),
1751 ColumnWriter::BoolColumnWriter(typed) => {
1752 let array = column.as_boolean();
1753 let values = get_bool_array_slice(array, indices.iter().copied());
1754 typed.write_batch_internal(
1755 values.as_slice(),
1756 None,
1757 levels.def_level_data().as_ref(),
1758 levels.rep_level_data().as_ref(),
1759 None,
1760 None,
1761 None,
1762 )
1763 }
1764 ColumnWriter::Int64ColumnWriter(typed) => visit_i64(
1765 column,
1766 WritePrimitive {
1767 writer: typed,
1768 levels,
1769 },
1770 ),
1771 ColumnWriter::Int96ColumnWriter(_typed) => {
1772 unreachable!("Currently unreachable because data type not supported")
1773 }
1774 ColumnWriter::FloatColumnWriter(typed) => {
1775 let array = column.as_primitive::<Float32Type>();
1776 write_primitive(typed, array.values(), levels)
1777 }
1778 ColumnWriter::DoubleColumnWriter(typed) => {
1779 let array = column.as_primitive::<Float64Type>();
1780 write_primitive(typed, array.values(), levels)
1781 }
1782 ColumnWriter::ByteArrayColumnWriter(_) => {
1783 unreachable!("should use ByteArrayWriter")
1784 }
1785 ColumnWriter::FixedLenByteArrayColumnWriter(typed) => visit_fixed_len_byte_array(
1786 column,
1787 WriteFixedLenByteArray {
1788 writer: typed,
1789 levels,
1790 },
1791 ),
1792 }
1793}
1794
1795pub(crate) trait PrimitiveVisitor<T> {
1804 type Output;
1805
1806 fn slice(self, values: &[T]) -> Result<Self::Output>;
1808
1809 fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<Self::Output>;
1811}
1812
1813pub(crate) fn visit_i32<V: PrimitiveVisitor<i32>>(
1815 column: &dyn arrow_array::Array,
1816 visitor: V,
1817) -> Result<V::Output> {
1818 match column.data_type() {
1820 ArrowDataType::Null => visitor.slice(Int32Array::new_null(column.len()).values()),
1821 ArrowDataType::Int8 => {
1822 visitor.map(column.as_primitive::<Int8Type>().values(), |x| x as i32)
1823 }
1824 ArrowDataType::Int16 => {
1825 visitor.map(column.as_primitive::<Int16Type>().values(), |x| x as i32)
1826 }
1827 ArrowDataType::Int32 => visitor.slice(column.as_primitive::<Int32Type>().values()),
1828 ArrowDataType::UInt8 => {
1829 visitor.map(column.as_primitive::<UInt8Type>().values(), |x| x as i32)
1830 }
1831 ArrowDataType::UInt16 => {
1832 visitor.map(column.as_primitive::<UInt16Type>().values(), |x| x as i32)
1833 }
1834 ArrowDataType::UInt32 => {
1835 let values = column.as_primitive::<UInt32Type>().values();
1838 visitor.slice(values.inner().typed_data())
1839 }
1840 ArrowDataType::Date32 => visitor.slice(column.as_primitive::<Date32Type>().values()),
1841 ArrowDataType::Time32(TimeUnit::Second) => {
1842 visitor.slice(column.as_primitive::<Time32SecondType>().values())
1843 }
1844 ArrowDataType::Time32(TimeUnit::Millisecond) => {
1845 visitor.slice(column.as_primitive::<Time32MillisecondType>().values())
1846 }
1847 ArrowDataType::Date64 => {
1848 visitor.map(column.as_primitive::<Date64Type>().values(), |x| {
1850 (x / 86_400_000) as i32
1851 })
1852 }
1853 ArrowDataType::Decimal32(_, _) => {
1854 visitor.slice(column.as_primitive::<Decimal32Type>().values())
1855 }
1856 ArrowDataType::Decimal64(_, _) => {
1857 visitor.map(column.as_primitive::<Decimal64Type>().values(), |v| {
1859 v as i32
1860 })
1861 }
1862 ArrowDataType::Decimal128(_, _) => {
1863 visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1865 v as i32
1866 })
1867 }
1868 ArrowDataType::Decimal256(_, _) => {
1869 visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1871 v.as_i128() as i32
1872 })
1873 }
1874 d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1875 }
1876}
1877
1878pub(crate) fn visit_i64<V: PrimitiveVisitor<i64>>(
1880 column: &dyn arrow_array::Array,
1881 visitor: V,
1882) -> Result<V::Output> {
1883 match column.data_type() {
1885 ArrowDataType::Date64 => visitor.slice(column.as_primitive::<Date64Type>().values()),
1886 ArrowDataType::Int64 => visitor.slice(column.as_primitive::<Int64Type>().values()),
1887 ArrowDataType::UInt64 => {
1888 let values = column.as_primitive::<UInt64Type>().values();
1891 visitor.slice(values.inner().typed_data())
1892 }
1893 ArrowDataType::Time64(TimeUnit::Microsecond) => {
1894 visitor.slice(column.as_primitive::<Time64MicrosecondType>().values())
1895 }
1896 ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1897 visitor.slice(column.as_primitive::<Time64NanosecondType>().values())
1898 }
1899 ArrowDataType::Timestamp(unit, _) => match unit {
1900 TimeUnit::Second => {
1901 visitor.slice(column.as_primitive::<TimestampSecondType>().values())
1902 }
1903 TimeUnit::Millisecond => {
1904 visitor.slice(column.as_primitive::<TimestampMillisecondType>().values())
1905 }
1906 TimeUnit::Microsecond => {
1907 visitor.slice(column.as_primitive::<TimestampMicrosecondType>().values())
1908 }
1909 TimeUnit::Nanosecond => {
1910 visitor.slice(column.as_primitive::<TimestampNanosecondType>().values())
1911 }
1912 },
1913 ArrowDataType::Duration(unit) => match unit {
1914 TimeUnit::Second => visitor.slice(column.as_primitive::<DurationSecondType>().values()),
1915 TimeUnit::Millisecond => {
1916 visitor.slice(column.as_primitive::<DurationMillisecondType>().values())
1917 }
1918 TimeUnit::Microsecond => {
1919 visitor.slice(column.as_primitive::<DurationMicrosecondType>().values())
1920 }
1921 TimeUnit::Nanosecond => {
1922 visitor.slice(column.as_primitive::<DurationNanosecondType>().values())
1923 }
1924 },
1925 ArrowDataType::Decimal64(_, _) => {
1926 visitor.slice(column.as_primitive::<Decimal64Type>().values())
1927 }
1928 ArrowDataType::Decimal128(_, _) => {
1929 visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1931 v as i64
1932 })
1933 }
1934 ArrowDataType::Decimal256(_, _) => {
1935 visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1937 v.as_i128() as i64
1938 })
1939 }
1940 d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1941 }
1942}
1943
1944struct WritePrimitive<'a, 'b, E: ColumnValueEncoder> {
1946 writer: &'a mut GenericColumnWriter<'b, E>,
1947 levels: &'a ArrayLevels,
1948}
1949
1950impl<T, E> PrimitiveVisitor<T> for WritePrimitive<'_, '_, E>
1951where
1952 E: ColumnValueEncoder<Values = [T]>,
1953{
1954 type Output = usize;
1955
1956 fn slice(self, values: &[T]) -> Result<usize> {
1957 write_primitive(self.writer, values, self.levels)
1958 }
1959
1960 fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<usize> {
1961 let values: Vec<T> = values.iter().map(|v| convert(*v)).collect();
1962 write_primitive(self.writer, &values, self.levels)
1963 }
1964}
1965
1966pub(crate) trait FixedLenByteArrayVisitor {
1972 type Output;
1973
1974 fn decimal_length(&self, precision: u8, native_length: usize) -> usize;
1979
1980 fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<Self::Output>;
1985}
1986
1987pub(crate) fn visit_fixed_len_byte_array<V: FixedLenByteArrayVisitor>(
1990 column: &dyn arrow_array::Array,
1991 visitor: V,
1992) -> Result<V::Output> {
1993 match column.data_type() {
1994 ArrowDataType::Interval(interval_unit) => match interval_unit {
1995 IntervalUnit::YearMonth => {
1996 let array = column.as_primitive::<IntervalYearMonthType>();
1999 visitor.values(12, |i, value| {
2000 value[0..4].copy_from_slice(&array.value(i).to_le_bytes());
2001 })
2002 }
2003 IntervalUnit::DayTime => {
2004 let array = column.as_primitive::<IntervalDayTimeType>();
2008 visitor.values(12, |i, value| {
2009 let interval = array.value(i);
2010 value[4..8].copy_from_slice(&interval.days.to_le_bytes());
2011 value[8..12].copy_from_slice(&interval.milliseconds.to_le_bytes());
2012 })
2013 }
2014 IntervalUnit::MonthDayNano => Err(ParquetError::NYI(format!(
2015 "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
2016 ))),
2017 },
2018 ArrowDataType::FixedSizeBinary(_) => {
2019 let array = column.as_fixed_size_binary();
2020 visitor.values(array.value_size(), |i, value| {
2021 value.copy_from_slice(array.value(i));
2022 })
2023 }
2024 ArrowDataType::Decimal32(_, _) => {
2025 visit_decimal(column.as_primitive::<Decimal32Type>(), visitor)
2026 }
2027 ArrowDataType::Decimal64(_, _) => {
2028 visit_decimal(column.as_primitive::<Decimal64Type>(), visitor)
2029 }
2030 ArrowDataType::Decimal128(_, _) => {
2031 visit_decimal(column.as_primitive::<Decimal128Type>(), visitor)
2032 }
2033 ArrowDataType::Decimal256(_, _) => {
2034 visit_decimal(column.as_primitive::<Decimal256Type>(), visitor)
2035 }
2036 ArrowDataType::Float16 => {
2037 let array = column.as_primitive::<Float16Type>();
2038 visitor.values(2, |i, value| {
2039 value.copy_from_slice(&array.value(i).to_le_bytes());
2040 })
2041 }
2042 _ => Err(ParquetError::NYI(
2043 "Attempting to write an Arrow type that is not yet implemented".to_string(),
2044 )),
2045 }
2046}
2047
2048fn visit_decimal<T: NativeDecimalType, V: FixedLenByteArrayVisitor>(
2055 array: &PrimitiveArray<T>,
2056 visitor: V,
2057) -> Result<V::Output> {
2058 let length = visitor.decimal_length(array.precision(), T::BYTE_LENGTH);
2059 let be_bytes = |i| T::to_be_bytes(array.value(i));
2060 if length == T::BYTE_LENGTH {
2061 visitor.values(length, |i, value| {
2063 value.copy_from_slice(be_bytes(i).as_ref());
2064 })
2065 } else {
2066 visitor.values(length, |i, value| {
2067 value.copy_from_slice(&be_bytes(i).as_ref()[T::BYTE_LENGTH - length..]);
2068 })
2069 }
2070}
2071
2072struct WriteFixedLenByteArray<'a, 'b> {
2074 writer: &'a mut ColumnWriterImpl<'b, FixedLenByteArrayType>,
2075 levels: &'a ArrayLevels,
2076}
2077
2078impl FixedLenByteArrayVisitor for WriteFixedLenByteArray<'_, '_> {
2079 type Output = usize;
2080
2081 fn decimal_length(&self, precision: u8, native_length: usize) -> usize {
2082 let length = decimal_length_from_precision(precision);
2083 assert!(length <= native_length);
2085 length
2086 }
2087
2088 fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<usize> {
2089 let indices = self.levels.non_null_indices().iter().copied();
2090 let values = chunk_array_slice(length, indices, fill);
2091 self.writer.write_batch_internal(
2092 values.as_slice(),
2093 None,
2094 self.levels.def_level_data().as_ref(),
2095 self.levels.rep_level_data().as_ref(),
2096 None,
2097 None,
2098 None,
2099 )
2100 }
2101}
2102
2103fn write_primitive<E: ColumnValueEncoder>(
2104 writer: &mut GenericColumnWriter<E>,
2105 values: &E::Values,
2106 levels: &ArrayLevels,
2107) -> Result<usize> {
2108 writer.write_batch_internal(
2109 values,
2110 Some(levels.non_null_indices()),
2111 levels.def_level_data().as_ref(),
2112 levels.rep_level_data().as_ref(),
2113 None,
2114 None,
2115 None,
2116 )
2117}
2118
2119fn get_bool_array_slice(
2120 array: &arrow_array::BooleanArray,
2121 indices: impl ExactSizeIterator<Item = usize>,
2122) -> Vec<bool> {
2123 let mut values = Vec::with_capacity(indices.len());
2124 for i in indices {
2125 values.push(array.value(i))
2126 }
2127 values
2128}
2129
2130trait NativeDecimalType: DecimalType {
2131 type NativeBytes: AsRef<[u8]>;
2132
2133 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
2134}
2135impl NativeDecimalType for Decimal32Type {
2136 type NativeBytes = [u8; Self::BYTE_LENGTH];
2137
2138 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2139 value.to_be_bytes()
2140 }
2141}
2142impl NativeDecimalType for Decimal64Type {
2143 type NativeBytes = [u8; Self::BYTE_LENGTH];
2144
2145 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2146 value.to_be_bytes()
2147 }
2148}
2149impl NativeDecimalType for Decimal128Type {
2150 type NativeBytes = [u8; Self::BYTE_LENGTH];
2151
2152 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2153 value.to_be_bytes()
2154 }
2155}
2156impl NativeDecimalType for Decimal256Type {
2157 type NativeBytes = [u8; Self::BYTE_LENGTH];
2158
2159 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2160 value.to_be_bytes()
2161 }
2162}
2163
2164#[inline]
2165fn chunk_array_slice(
2166 chunk_size: usize,
2167 indices: impl ExactSizeIterator<Item = usize>,
2168 writer: impl Fn(usize, &mut [u8]),
2169) -> Vec<FixedLenByteArray> {
2170 let capacity = indices.len() * chunk_size;
2171 let mut arena = vec![0; capacity];
2174 for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
2175 writer(i, chunk);
2176 }
2177 chunk_contiguous_vec(arena, chunk_size)
2178}
2179
2180fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
2181 let mut values = Vec::with_capacity(arena.len() / chunk_size);
2182 let mut arena = Bytes::from(arena);
2183 while arena.len() >= chunk_size {
2184 let slice = arena.split_to(chunk_size);
2185 values.push(FixedLenByteArray::from(ByteArray::from(slice)));
2186 }
2187 values
2188}
2189
2190#[inline]
2192fn hash_bytes(bytes: &[u8]) -> u64 {
2193 twox_hash::XxHash64::oneshot(0, bytes)
2194}
2195
2196fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
2198 use ArrowDataType::*;
2199 match dt {
2200 Int8 | UInt8 => Some(1),
2201 Int16 | UInt16 | Float16 => Some(2),
2202 Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) => Some(4),
2203 Int64
2204 | UInt64
2205 | Float64
2206 | Date64
2207 | Time64(_)
2208 | Timestamp(_, _)
2209 | Duration(_)
2210 | Decimal64(_, _) => Some(8),
2211 Interval(IntervalUnit::YearMonth) => Some(4),
2212 Interval(IntervalUnit::DayTime) => Some(8),
2213 Interval(IntervalUnit::MonthDayNano) => Some(16),
2214 Decimal128(_, _) => Some(16),
2215 Decimal256(_, _) => Some(32),
2216 _ => None,
2217 }
2218}
2219
2220fn update_distinct_values_seen(
2226 array: &dyn arrow_array::Array,
2227 non_null_indices: &[usize],
2228 seen: &mut DistinctValuesSet,
2229) {
2230 let data = array.to_data();
2231 let offset = data.offset();
2232
2233 match array.data_type() {
2234 ArrowDataType::Boolean => {
2235 let arr = array
2236 .as_any()
2237 .downcast_ref::<arrow_array::BooleanArray>()
2238 .unwrap();
2239 for &row in non_null_indices {
2240 seen.insert(arr.value(row) as u64);
2241 }
2242 }
2243 ArrowDataType::Utf8 | ArrowDataType::Binary => {
2244 let offsets = data.buffers()[0].typed_data::<i32>();
2245 let values = data.buffers()[1].as_slice();
2246 for &row in non_null_indices {
2247 let start = offsets[offset + row] as usize;
2248 let end = offsets[offset + row + 1] as usize;
2249 seen.insert(hash_bytes(&values[start..end]));
2250 }
2251 }
2252 ArrowDataType::LargeUtf8 | ArrowDataType::LargeBinary => {
2253 let offsets = data.buffers()[0].typed_data::<i64>();
2254 let values = data.buffers()[1].as_slice();
2255 for &row in non_null_indices {
2256 let start = offsets[offset + row] as usize;
2257 let end = offsets[offset + row + 1] as usize;
2258 seen.insert(hash_bytes(&values[start..end]));
2259 }
2260 }
2261 ArrowDataType::FixedSizeBinary(byte_width) => {
2262 let byte_width = *byte_width as usize;
2263 let buffer = data.buffers()[0].as_slice();
2264 for &row in non_null_indices {
2265 let start = (offset + row) * byte_width;
2266 seen.insert(hash_bytes(&buffer[start..start + byte_width]));
2267 }
2268 }
2269 ArrowDataType::Utf8View => {
2270 let string_view_array = array.as_string_view();
2271 for &row in non_null_indices {
2272 seen.insert(hash_bytes(string_view_array.value(row).as_bytes()));
2273 }
2274 }
2275 ArrowDataType::BinaryView => {
2276 let binary_view_array = array.as_binary_view();
2277 for &row in non_null_indices {
2278 seen.insert(hash_bytes(binary_view_array.value(row)));
2279 }
2280 }
2281 data_type => {
2282 if let Some(width) = fixed_byte_width(data_type) {
2283 let buffer = data.buffers()[0].as_slice();
2284 for &row in non_null_indices {
2285 let pos = (offset + row) * width;
2286 seen.insert(hash_bytes(&buffer[pos..pos + width]));
2287 }
2288 }
2289 }
2291 }
2292}
2293
2294#[cfg(test)]
2296use crate as parquet_crate;
2297
2298#[cfg(test)]
2299#[path = "../../../tests/arrow_writer/roundtrip_helpers.rs"]
2300mod roundtrip_helpers;
2301
2302#[cfg(test)]
2303mod tests {
2304 use super::roundtrip_helpers::{
2305 RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip_opts,
2306 roundtrip_opts_with_array_validation,
2307 };
2308 use super::*;
2309 use std::cmp::Ordering;
2310 use std::collections::HashMap;
2311
2312 use std::fs::File;
2313
2314 use crate::arrow::ARROW_SCHEMA_META_KEY;
2315 use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
2316 use crate::column::page::{Page, PageReader};
2317 use crate::file::metadata::thrift::PageHeader;
2318 use crate::file::page_index::column_index::ColumnIndexMetaData;
2319 use crate::file::reader::SerializedPageReader;
2320 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
2321 use crate::schema::types::ColumnPath;
2322 use arrow::datatypes::{DataType, Schema};
2323 use arrow::error::Result as ArrowResult;
2324 use arrow::util::data_gen::create_random_array;
2325 use arrow::util::pretty::pretty_format_batches;
2326 use arrow::{array::*, buffer::Buffer, compute::cast};
2327 use arrow_buffer::{IntervalMonthDayNano, NullBuffer, OffsetBuffer};
2328 use arrow_schema::Fields;
2329 use arrow_select::concat::concat;
2330
2331 use crate::basic::{Encoding, EncodingMask};
2332 use crate::data_type::AsBytes;
2333 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
2334 use crate::file::properties::{
2335 BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
2336 };
2337 use crate::file::serialized_reader::ReadOptionsBuilder;
2338 use crate::file::{
2339 reader::{FileReader, SerializedFileReader},
2340 statistics::Statistics,
2341 };
2342
2343 #[derive(Debug, Default)]
2348 struct RecordingPageStore {
2349 next: u64,
2350 blobs: HashMap<u64, Bytes>,
2351 puts: Arc<std::sync::atomic::AtomicUsize>,
2352 }
2353
2354 impl PageStore for RecordingPageStore {
2355 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2356 let id = 100 + self.next * 7;
2358 self.next += 1;
2359 self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2360 self.blobs.insert(id, value);
2361 Ok(PageKey::new(id))
2362 }
2363
2364 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2365 self.blobs
2366 .remove(&key.get())
2367 .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2368 }
2369 }
2370
2371 #[derive(Debug)]
2372 struct RecordingPageStoreFactory {
2373 puts: Arc<std::sync::atomic::AtomicUsize>,
2374 }
2375
2376 impl PageStoreFactory for RecordingPageStoreFactory {
2377 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2378 Ok(Box::new(RecordingPageStore {
2379 puts: self.puts.clone(),
2380 ..Default::default()
2381 }))
2382 }
2383 }
2384
2385 #[test]
2389 fn custom_page_store_is_byte_identical_to_default() {
2390 let schema = Arc::new(Schema::new(vec![
2391 Field::new("i", DataType::Int32, true),
2392 Field::new("s", DataType::Utf8, true),
2394 ]));
2395 let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2396 let s = StringArray::from(vec![
2397 Some("a"),
2398 Some("bb"),
2399 Some("a"),
2400 None,
2401 Some("bb"),
2402 Some("ccc"),
2403 ]);
2404 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2405
2406 let props = WriterProperties::builder()
2409 .set_max_row_group_row_count(Some(3))
2410 .build();
2411
2412 let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2413 let mut buffer = Vec::new();
2414 let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2415 if let Some(factory) = factory {
2416 opts = opts.with_page_store_factory(factory);
2417 }
2418 let mut writer =
2419 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2420 writer.write(&batch).unwrap();
2421 writer.close().unwrap();
2422 buffer
2423 };
2424
2425 let default_bytes = write(None);
2426
2427 let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2428 let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2429 puts: puts.clone(),
2430 })));
2431
2432 assert!(
2433 puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2434 "custom PageStore was never written to"
2435 );
2436 assert_eq!(
2437 default_bytes, custom_bytes,
2438 "a custom PageStore must produce byte-identical output to the default"
2439 );
2440 }
2441
2442 #[test]
2448 #[cfg_attr(miri, ignore)] fn dictionary_column_round_trips_with_offset_index_disabled() {
2450 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2451
2452 let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2455 let array = Int32Array::from(values.clone());
2456 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2457
2458 let props = WriterProperties::builder()
2459 .set_offset_index_disabled(true)
2460 .set_data_page_row_count_limit(4096)
2461 .build();
2462 let opts = ArrowWriterOptions::new().with_properties(props);
2463
2464 let mut buffer = Vec::new();
2465 let mut writer =
2466 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2467 writer.write(&batch).unwrap();
2468 writer.close().unwrap();
2469
2470 let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2471 let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2472 let read_values: Vec<Option<i32>> = read
2473 .iter()
2474 .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2475 .collect();
2476 assert_eq!(read_values, values);
2477 }
2478
2479 #[test]
2484 fn dictionary_page_is_routed_through_the_store() {
2485 #[derive(Debug, Default)]
2487 struct SizeRecordingPageStore {
2488 blobs: Vec<Bytes>,
2489 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2490 }
2491 impl PageStore for SizeRecordingPageStore {
2492 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2493 self.bytes_put
2494 .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2495 let key = PageKey::new(self.blobs.len() as u64);
2496 self.blobs.push(value);
2497 Ok(key)
2498 }
2499 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2500 Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2501 }
2502 }
2503 #[derive(Debug)]
2504 struct Factory {
2505 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2506 }
2507 impl PageStoreFactory for Factory {
2508 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2509 Ok(Box::new(SizeRecordingPageStore {
2510 bytes_put: self.bytes_put.clone(),
2511 ..Default::default()
2512 }))
2513 }
2514 }
2515
2516 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2517 let values: Vec<&str> = (0..2048)
2520 .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2521 .collect();
2522 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2523 .unwrap();
2524
2525 let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2526 let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2527 bytes_put: bytes_put.clone(),
2528 }));
2529
2530 let mut buffer = Vec::new();
2533 let mut writer =
2534 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2535 writer.write(&batch).unwrap();
2536 writer.close().unwrap();
2537
2538 let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2539 let column = reader.metadata().row_group(0).column(0);
2540 assert!(
2541 column.dictionary_page_offset().is_some(),
2542 "expected the column to be dictionary-encoded"
2543 );
2544
2545 assert_eq!(
2549 bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2550 column.compressed_size(),
2551 "the dictionary page must pass through the store like any other page"
2552 );
2553 }
2554
2555 #[test]
2560 fn dictionary_page_encoding_stats_lists_dictionary_first() {
2561 use crate::basic::PageType;
2562 use crate::file::metadata::ParquetMetaDataOptions;
2563
2564 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2567 let values: Vec<Option<i32>> = (0..4096).map(|i| Some(i % 4)).collect();
2568 let array = Int32Array::from(values);
2569 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2570
2571 let mut buffer = Vec::new();
2572 let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None).unwrap();
2573 writer.write(&batch).unwrap();
2574 writer.close().unwrap();
2575 let buffer = Bytes::from(buffer);
2576
2577 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2581 let metadata = ParquetMetaDataReader::new()
2582 .with_metadata_options(Some(options))
2583 .parse_and_finish(&buffer)
2584 .unwrap();
2585
2586 let col = metadata.row_group(0).column(0);
2587 assert!(
2588 col.dictionary_page_offset().is_some(),
2589 "column should be dictionary encoded"
2590 );
2591 let stats = col
2592 .page_encoding_stats()
2593 .expect("full (non-mask) page encoding stats should be present");
2594
2595 assert_eq!(
2598 stats[0].page_type,
2599 PageType::DICTIONARY_PAGE,
2600 "the DICTIONARY_PAGE encoding stat must come first, got {stats:?}"
2601 );
2602 assert!(
2603 stats[1..]
2604 .iter()
2605 .all(|s| s.page_type == PageType::DATA_PAGE)
2606 );
2607 }
2608
2609 fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2610 let mut buffer = vec![];
2611
2612 let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2613 writer.write(expected_batch).unwrap();
2614 writer.close().unwrap();
2615
2616 buffer
2617 }
2618
2619 fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2620 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2621 writer.write(expected_batch).unwrap();
2622 writer.into_inner().unwrap()
2623 }
2624
2625 fn write_string_dictionaries(arrays: &[ArrayRef]) -> usize {
2628 let schema = Arc::new(Schema::new(vec![Field::new(
2629 "col",
2630 arrays[0].data_type().clone(),
2631 true,
2632 )]));
2633 let props = WriterProperties::builder().set_write_batch_size(2).build();
2634 let mut file = vec![];
2635 let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), Some(props)).unwrap();
2636 for array in arrays {
2637 let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();
2638 writer.write(&batch).unwrap();
2639 }
2640 writer.close().unwrap();
2641
2642 let reader = SerializedFileReader::new(Bytes::from(file.clone())).unwrap();
2643 let mut pages = reader
2644 .get_row_group(0)
2645 .unwrap()
2646 .get_column_page_reader(0)
2647 .unwrap();
2648 let dictionary_page = pages.get_next_page().unwrap().unwrap();
2649 assert!(matches!(dictionary_page, Page::DictionaryPage { .. }));
2650
2651 let expected: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
2652 let expected = concat(&expected).unwrap();
2653 let actual = read_column(file);
2654 assert_eq!(
2655 cast(&actual, &DataType::Utf8).unwrap().as_ref(),
2656 cast(&expected, &DataType::Utf8).unwrap().as_ref()
2657 );
2658
2659 dictionary_page.num_values() as usize
2660 }
2661
2662 #[test]
2663 fn arrow_writer_string_dictionary_cache_reused_for_same_values() {
2664 let shared: ArrayRef = Arc::new(StringArray::from(vec![
2666 "a",
2667 "bb",
2668 "longer than eight bytes",
2669 "dup",
2670 "dup",
2671 ]));
2672 let keys = Int32Array::from(vec![
2673 Some(0),
2674 Some(1),
2675 None,
2676 Some(2),
2677 Some(3),
2678 Some(4),
2679 Some(0),
2680 ]);
2681 let first = DictionaryArray::new(keys, shared.clone());
2682 let keys = Int32Array::from(vec![Some(4), Some(4), Some(1), None, Some(2), Some(0)]);
2685 let second = DictionaryArray::new(keys, shared).slice(1, 4);
2686
2687 let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second)];
2688 assert_eq!(write_string_dictionaries(&arrays), 4);
2690 }
2691
2692 #[test]
2693 fn arrow_writer_string_dictionary_cache_replaced_for_new_values() {
2694 let first_values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
2695 let first = DictionaryArray::new(Int32Array::from(vec![0, 1, 2, 0]), first_values.clone());
2696 let second_values: ArrayRef = Arc::new(StringArray::from(vec!["zz", "a", "new value!!"]));
2699 let second = DictionaryArray::new(
2700 Int32Array::from(vec![Some(2), Some(1), None, Some(0), Some(1)]),
2701 second_values,
2702 );
2703 let third = DictionaryArray::new(Int32Array::from(vec![2, 1, 0]), first_values);
2705
2706 let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second), Arc::new(third)];
2707 assert_eq!(write_string_dictionaries(&arrays), 5);
2709 }
2710
2711 #[test]
2712 fn roundtrip_bytes() {
2713 let schema = Arc::new(Schema::new(vec![
2715 Field::new("a", DataType::Int32, false),
2716 Field::new("b", DataType::Int32, true),
2717 ]));
2718
2719 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2721 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2722
2723 let expected_batch =
2725 RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2726
2727 for buffer in [
2728 get_bytes_after_close(schema.clone(), &expected_batch),
2729 get_bytes_by_into_inner(schema, &expected_batch),
2730 ] {
2731 let cursor = Bytes::from(buffer);
2732 let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2733
2734 let actual_batch = record_batch_reader
2735 .next()
2736 .expect("No batch found")
2737 .expect("Unable to get batch");
2738
2739 assert_eq!(expected_batch.schema(), actual_batch.schema());
2740 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2741 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2742 for i in 0..expected_batch.num_columns() {
2743 let expected_data = expected_batch.column(i).to_data();
2744 let actual_data = actual_batch.column(i).to_data();
2745
2746 assert_eq!(expected_data, actual_data);
2747 }
2748 }
2749 }
2750
2751 fn read_column(file: Vec<u8>) -> ArrayRef {
2752 let reader = ParquetRecordBatchReader::try_new(Bytes::from(file), 4096).unwrap();
2753 let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2754 let arrays: Vec<&dyn Array> = batches.iter().map(|b| b.column(0).as_ref()).collect();
2755 concat(&arrays).unwrap()
2756 }
2757
2758 fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
2759 let writer_schema = Arc::new(Schema::new(vec![field]));
2760 let batch_schema = Arc::new(Schema::new(vec![Field::new(
2761 "c",
2762 col.data_type().clone(),
2763 col.logical_null_count() != 0,
2764 )]));
2765 let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
2766 let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
2769 let mut file = vec![];
2770 let mut writer =
2771 ArrowWriter::try_new_with_options(&mut file, writer_schema, options).unwrap();
2772 writer.write(&batch).unwrap();
2773 writer.close().unwrap();
2774 read_column(file)
2775 }
2776
2777 fn assert_compatible_roundtrip(field: Field, actual: ArrayRef, expected: &ArrayRef) {
2778 assert!(!compute_leaves(&field, &actual).unwrap().is_empty());
2779 assert_eq!(
2780 roundtrip_compatible_column(field, actual).as_ref(),
2781 expected.as_ref()
2782 );
2783 }
2784
2785 fn run_end_type(value_type: DataType, nullable_values: bool) -> DataType {
2786 DataType::RunEndEncoded(
2787 Arc::new(Field::new(
2788 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2789 DataType::Int32,
2790 false,
2791 )),
2792 Arc::new(Field::new(
2793 Field::REE_VALUES_FIELD_DEFAULT_NAME,
2794 value_type,
2795 nullable_values,
2796 )),
2797 )
2798 }
2799
2800 #[test]
2801 fn arrow_writer_required_ree_schema_rejects_nullable_dense_input() {
2802 let writer_schema = Arc::new(Schema::new(vec![Field::new(
2803 "c",
2804 run_end_type(DataType::Int32, false),
2805 false,
2806 )]));
2807 let batch_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2808 let batch = RecordBatch::try_new(
2809 batch_schema.clone(),
2810 vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
2811 )
2812 .unwrap();
2813
2814 for skip_metadata in [false, true] {
2815 for cdc in [false, true] {
2816 let props = WriterProperties::builder()
2817 .set_content_defined_chunking(cdc.then(Default::default))
2818 .build();
2819 let options = ArrowWriterOptions::new()
2820 .with_properties(props)
2821 .with_skip_arrow_metadata(skip_metadata);
2822 let mut writer = ArrowWriter::try_new_with_options(
2823 Vec::new(),
2824 writer_schema.clone(),
2825 options.clone(),
2826 )
2827 .unwrap();
2828 let error = writer.write(&batch).unwrap_err().to_string();
2829 assert!(
2830 error.contains("Found null at index 1 for required field 'c'"),
2831 "skip_metadata={skip_metadata}, cdc={cdc}: {error}"
2832 );
2833
2834 let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 8, 9]));
2836 let batch =
2837 RecordBatch::try_new(batch_schema.clone(), vec![values.clone()]).unwrap();
2838 let mut file = vec![];
2839 let mut writer =
2840 ArrowWriter::try_new_with_options(&mut file, writer_schema.clone(), options)
2841 .unwrap();
2842 writer.write(&batch).unwrap();
2843 writer.close().unwrap();
2844 assert_eq!(read_column(file).as_ref(), values.as_ref());
2845 }
2846 }
2847 }
2848
2849 #[test]
2850 fn ree_nulls_respect_dense_schema_nullability() {
2851 let run_ends = Int32Array::from(vec![2, 4, 5]);
2852 let values = Int32Array::from(vec![Some(1), None, Some(2)]);
2853 let ree: ArrayRef = Arc::new(Int32RunArray::try_new(&run_ends, &values).unwrap());
2854 let ree_batch = RecordBatch::try_new(
2855 Arc::new(Schema::new(vec![Field::new(
2856 "c",
2857 ree.data_type().clone(),
2858 true,
2859 )])),
2860 vec![ree],
2861 )
2862 .unwrap();
2863
2864 let nullable_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2865 let mut file = vec![];
2866 let mut writer = ArrowWriter::try_new(&mut file, nullable_schema, None).unwrap();
2867 writer.write(&ree_batch).unwrap();
2868 writer.close().unwrap();
2869 assert_eq!(
2870 read_column(file).as_primitive::<Int32Type>(),
2871 &Int32Array::from(vec![Some(1), Some(1), None, None, Some(2)])
2872 );
2873
2874 let required_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
2875 let mut writer = ArrowWriter::try_new(Vec::new(), required_schema, None).unwrap();
2876 let error = writer.write(&ree_batch).unwrap_err().to_string();
2877 assert!(error.contains("Found null") && error.contains("required field 'c'"));
2878 }
2879
2880 #[test]
2881 fn required_null_validation_respects_reachability_and_write_timing() {
2882 let required = Field::new("c", DataType::Int32, false);
2883 let nullable: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None]));
2884 let leaves = compute_leaves(&required, &nullable).unwrap();
2885 let schema = Arc::new(Schema::new(vec![required.clone()]));
2886 let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2887 let (_file, factory) = writer.into_serialized_writer().unwrap();
2888 let mut columns = factory.create_column_writers(0).unwrap();
2889 assert!(
2890 columns[0]
2891 .write(&leaves[0])
2892 .unwrap_err()
2893 .to_string()
2894 .contains("required field 'c'")
2895 );
2896
2897 let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None]));
2899 let dict: ArrayRef = Arc::new(DictionaryArray::new(
2900 Int8Array::from(vec![0, 0]),
2901 values.clone(),
2902 ));
2903 assert_eq!(
2904 roundtrip_compatible_column(required.clone(), dict).as_primitive::<Int32Type>(),
2905 &Int32Array::from(vec![7, 7])
2906 );
2907 let dict: ArrayRef = Arc::new(DictionaryArray::new(Int8Array::from(vec![0, 1]), values));
2908 let leaves = compute_leaves(&required, &dict).unwrap();
2909 assert!(leaves[0].0.validate().is_err());
2910 let dict: ArrayRef = Arc::new(DictionaryArray::new(
2911 Int8Array::from(vec![Some(0), None]),
2912 Arc::new(Int32Array::from(vec![7])),
2913 ));
2914 let leaves = compute_leaves(&required, &dict).unwrap();
2915 assert!(leaves[0].0.validate().is_err());
2916
2917 let actual_fields = Fields::from(vec![Field::new("source_child", DataType::Int32, true)]);
2920 let array: ArrayRef = Arc::new(StructArray::new(
2921 actual_fields.clone(),
2922 vec![nullable.clone()],
2923 Some(NullBuffer::from(vec![true, false])),
2924 ));
2925 let unmasked: ArrayRef = Arc::new(StructArray::new(
2926 actual_fields,
2927 vec![nullable.clone()],
2928 None,
2929 ));
2930 for child_type in [DataType::Int32, run_end_type(DataType::Int32, false)] {
2931 let target = Field::new(
2932 "c",
2933 DataType::Struct(Fields::from(vec![Field::new(
2934 "target_child",
2935 child_type,
2936 false,
2937 )])),
2938 true,
2939 );
2940 let leaves = compute_leaves(&target, &unmasked).unwrap();
2941 assert!(leaves[0].0.validate().is_err());
2942 let leaves = compute_leaves(&target, &array).unwrap();
2943 assert!(leaves[0].0.validate().is_ok());
2944 let result = roundtrip_compatible_column(target, array.clone());
2945 assert_eq!(result.as_struct().fields()[0].name(), "target_child");
2946 assert!(result.is_null(1));
2947 }
2948
2949 let values: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(9), None]));
2951 let list: ArrayRef = Arc::new(ListArray::new(
2952 Arc::new(Field::new_list_field(DataType::Int32, true)),
2953 OffsetBuffer::new(vec![1_i32, 2].into()),
2954 values,
2955 None,
2956 ));
2957 let target = Field::new(
2958 "c",
2959 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2960 false,
2961 );
2962 assert!(
2963 compute_leaves(&target, &list).unwrap()[0]
2964 .0
2965 .validate()
2966 .is_ok()
2967 );
2968 assert_eq!(
2969 roundtrip_compatible_column(required, nullable.slice(0, 1)).as_primitive::<Int32Type>(),
2970 &Int32Array::from(vec![1])
2971 );
2972 }
2973
2974 fn check_required_list_validation(list_type: fn(Arc<Field>) -> DataType) {
2975 let make_array = |values: Vec<Option<i32>>,
2976 offsets: Vec<i32>,
2977 sizes: Vec<i32>,
2978 nulls: Option<NullBuffer>| {
2979 let child = Arc::new(Field::new("source_child", DataType::Int32, true));
2980 let values: ArrayRef = Arc::new(Int32Array::from(values));
2981 match list_type(child.clone()) {
2982 DataType::List(_) => Arc::new(ListArray::new(
2983 child,
2984 OffsetBuffer::new(offsets.into()),
2985 values,
2986 nulls,
2987 )) as ArrayRef,
2988 DataType::LargeList(_) => Arc::new(LargeListArray::new(
2989 child,
2990 OffsetBuffer::new(
2991 offsets
2992 .into_iter()
2993 .map(i64::from)
2994 .collect::<Vec<_>>()
2995 .into(),
2996 ),
2997 values,
2998 nulls,
2999 )) as ArrayRef,
3000 DataType::ListView(_) => Arc::new(ListViewArray::new(
3001 child,
3002 offsets.into(),
3003 sizes.into(),
3004 values,
3005 nulls,
3006 )) as ArrayRef,
3007 DataType::LargeListView(_) => Arc::new(LargeListViewArray::new(
3008 child,
3009 offsets
3010 .into_iter()
3011 .map(i64::from)
3012 .collect::<Vec<_>>()
3013 .into(),
3014 sizes.into_iter().map(i64::from).collect::<Vec<_>>().into(),
3015 values,
3016 nulls,
3017 )) as ArrayRef,
3018 _ => unreachable!(),
3019 }
3020 };
3021 let is_view = matches!(
3022 list_type(Arc::new(Field::new("item", DataType::Int32, true))),
3023 DataType::ListView(_) | DataType::LargeListView(_)
3024 );
3025 let offsets = |starts: Vec<i32>, end: i32| {
3026 let mut result = starts;
3027 if !is_view {
3028 result.push(end);
3029 }
3030 result
3031 };
3032 let target = Field::new(
3033 "c",
3034 list_type(Arc::new(Field::new("target_child", DataType::Int32, false))),
3035 true,
3036 );
3037 let invalid = make_array(vec![None, Some(7)], offsets(vec![0], 2), vec![2], None);
3038 let valid = make_array(vec![Some(7), Some(8)], offsets(vec![0], 2), vec![2], None);
3039 let masked = make_array(
3040 vec![None, Some(7)],
3041 offsets(vec![0, 1], 2),
3042 vec![1, 1],
3043 Some(NullBuffer::from(vec![false, true])),
3044 );
3045 let unused = make_array(
3046 vec![None, Some(7), None],
3047 offsets(vec![1], 2),
3048 vec![1],
3049 None,
3050 );
3051 let empty = make_array(vec![None], offsets(vec![0], 0), vec![0], None);
3052 let sliced = make_array(
3053 vec![None, Some(7)],
3054 offsets(vec![0, 1], 2),
3055 vec![1, 1],
3056 None,
3057 )
3058 .slice(1, 1);
3059 let schema = Arc::new(Schema::new(vec![target.clone()]));
3060
3061 for (name, array, invalid) in [
3063 ("reachable null", invalid, true),
3064 ("valid required child", valid.clone(), false),
3065 ("null ancestor", masked, false),
3066 ("unused child slots", unused, false),
3067 ("empty list", empty, false),
3068 ("out-of-slice null", sliced, false),
3069 ] {
3070 let leaves = compute_leaves(&target, &array).unwrap();
3071 assert_eq!(leaves.len(), 1);
3072 for cdc in [false, true] {
3073 let props = WriterProperties::builder()
3074 .set_content_defined_chunking(cdc.then(Default::default))
3075 .build();
3076 let mut writer =
3077 ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3078 let mut chunkers = writer.cdc_chunkers.take();
3079 let (_file, factory) = writer.into_serialized_writer().unwrap();
3080 let mut columns = factory.create_column_writers(0).unwrap();
3081 let before = (
3082 columns[0].memory_size(),
3083 columns[0].get_estimated_total_bytes(),
3084 );
3085 let result = match &mut chunkers {
3086 Some(chunkers) => columns[0].write_with_chunker(&leaves[0], &mut chunkers[0]),
3087 None => columns[0].write(&leaves[0]),
3088 };
3089 if invalid {
3090 let error = result.unwrap_err().to_string();
3091 assert!(
3092 error.contains("Found null at index 0 for required field 'target_child'"),
3093 "{name}, cdc={cdc}: {error}"
3094 );
3095 assert_eq!(
3096 before,
3097 (
3098 columns[0].memory_size(),
3099 columns[0].get_estimated_total_bytes()
3100 ),
3101 "invalid leaf must be rejected before encoding"
3102 );
3103 let valid_leaves = compute_leaves(&target, &valid).unwrap();
3105 match &mut chunkers {
3106 Some(chunkers) => columns[0]
3107 .write_with_chunker(&valid_leaves[0], &mut chunkers[0])
3108 .unwrap(),
3109 None => columns[0].write(&valid_leaves[0]).unwrap(),
3110 }
3111 assert_eq!(
3112 columns
3113 .pop()
3114 .unwrap()
3115 .close()
3116 .unwrap()
3117 .close
3118 .metadata
3119 .num_values(),
3120 2
3121 );
3122 } else {
3123 result.unwrap();
3124 columns.pop().unwrap().close().unwrap();
3125 }
3126
3127 let batch = RecordBatch::try_new(
3129 Arc::new(Schema::new(vec![Field::new(
3130 "c",
3131 array.data_type().clone(),
3132 true,
3133 )])),
3134 vec![array.clone()],
3135 )
3136 .unwrap();
3137 let props = WriterProperties::builder()
3138 .set_content_defined_chunking(cdc.then(Default::default))
3139 .build();
3140 let mut writer =
3141 ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3142 let result = writer.write(&batch);
3143 if invalid {
3144 assert!(
3145 result
3146 .unwrap_err()
3147 .to_string()
3148 .contains("required field 'target_child'")
3149 );
3150 } else {
3151 result.unwrap();
3152 writer.close().unwrap();
3153 }
3154 }
3155 }
3156 }
3157
3158 #[test]
3159 fn required_list_validation_preserves_level_generation() {
3160 check_required_list_validation(DataType::List);
3161 }
3162
3163 #[test]
3164 fn required_large_list_validation_preserves_level_generation() {
3165 check_required_list_validation(DataType::LargeList);
3166 }
3167
3168 #[test]
3169 fn required_list_view_validation_preserves_level_generation() {
3170 check_required_list_validation(DataType::ListView);
3171 }
3172
3173 #[test]
3174 fn required_large_list_view_validation_preserves_level_generation() {
3175 check_required_list_validation(DataType::LargeListView);
3176 }
3177
3178 fn nullable_int32_list() -> ArrayRef {
3179 Arc::new(ListArray::new(
3180 Arc::new(Field::new_list_field(DataType::Int32, true)),
3181 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3182 Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
3183 None,
3184 ))
3185 }
3186
3187 #[test]
3188 fn arrow_writer_dense_list_under_ree_schema() {
3189 let list = nullable_int32_list();
3190 let field = Field::new("c", run_end_type(list.data_type().clone(), false), false);
3192 assert_compatible_roundtrip(field, list.clone(), &list);
3193 }
3194
3195 #[test]
3196 fn arrow_writer_dense_list_under_dictionary_schema() {
3197 let list = nullable_int32_list();
3198 let field = Field::new(
3199 "c",
3200 DataType::Dictionary(Box::new(DataType::Int8), Box::new(list.data_type().clone())),
3201 false,
3202 );
3203 assert_compatible_roundtrip(field, list.clone(), &list);
3204 }
3205
3206 #[test]
3207 fn arrow_writer_dense_values_under_nested_ree_schema() {
3208 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3209 let field = Field::new(
3210 "c",
3211 run_end_type(run_end_type(DataType::Int32, true), false),
3212 false,
3213 );
3214 assert_compatible_roundtrip(field, dense.clone(), &dense);
3215 }
3216
3217 #[test]
3218 fn arrow_writer_dense_values_under_dictionary_ree_schema() {
3219 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3220 let field = Field::new(
3221 "c",
3222 DataType::Dictionary(
3223 Box::new(DataType::Int8),
3224 Box::new(run_end_type(DataType::Int32, true)),
3225 ),
3226 false,
3227 );
3228 assert_compatible_roundtrip(field, dense.clone(), &dense);
3229 }
3230
3231 #[test]
3232 fn arrow_writer_dense_values_under_ree_dictionary_schema() {
3233 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3234 let field = Field::new(
3237 "c",
3238 run_end_type(
3239 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3240 true,
3241 ),
3242 false,
3243 );
3244 assert_compatible_roundtrip(field, dense.clone(), &dense);
3245 }
3246
3247 #[test]
3248 fn arrow_writer_dense_struct_under_ree_child_schema() {
3249 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3250 let struct_fields = Fields::from(vec![Field::new("value", DataType::Int32, true)]);
3251 let dense_struct: ArrayRef = Arc::new(StructArray::new(struct_fields, vec![dense], None));
3252 let field = Field::new(
3254 "c",
3255 DataType::Struct(Fields::from(vec![Field::new(
3256 "value",
3257 run_end_type(DataType::Int32, true),
3258 false,
3259 )])),
3260 false,
3261 );
3262 assert_compatible_roundtrip(field, dense_struct.clone(), &dense_struct);
3263 }
3264
3265 #[test]
3266 fn arrow_writer_dense_list_under_ree_child_schema() {
3267 let list = nullable_int32_list();
3268 let field = Field::new(
3269 "c",
3270 DataType::List(Arc::new(Field::new(
3271 "item",
3272 run_end_type(DataType::Int32, true),
3273 false,
3274 ))),
3275 false,
3276 );
3277 assert_compatible_roundtrip(field, list.clone(), &list);
3278 }
3279
3280 #[test]
3281 fn arrow_writer_dense_list_under_dictionary_child_schema() {
3282 let list = nullable_int32_list();
3283 let field = Field::new(
3284 "c",
3285 DataType::List(Arc::new(Field::new_list_field(
3286 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3287 true,
3288 ))),
3289 false,
3290 );
3291 assert_compatible_roundtrip(field, list.clone(), &list);
3292 }
3293
3294 #[test]
3295 fn arrow_writer_dictionary_list_under_dense_child_schema() {
3296 let dict: ArrayRef = Arc::new(DictionaryArray::new(
3297 Int8Array::from(vec![Some(0), None, Some(1)]),
3298 Arc::new(Int32Array::from(vec![1, 3])),
3299 ));
3300 let dict_list: ArrayRef = Arc::new(ListArray::new(
3301 Arc::new(Field::new_list_field(dict.data_type().clone(), true)),
3302 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3303 dict,
3304 None,
3305 ));
3306 let expected = nullable_int32_list();
3307 assert_compatible_roundtrip(
3308 Field::new("c", expected.data_type().clone(), false),
3309 dict_list,
3310 &expected,
3311 );
3312 }
3313
3314 #[test]
3315 fn arrow_writer_dense_map_under_ree_value_schema() {
3316 let key_field = Arc::new(Field::new("keys", DataType::Utf8, false));
3317 let entries = StructArray::new(
3318 Fields::from(vec![
3319 key_field.clone(),
3320 Arc::new(Field::new("values", DataType::Int32, true)),
3321 ]),
3322 vec![
3323 Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
3324 Arc::new(Int32Array::from(vec![Some(7), None, Some(9)])),
3325 ],
3326 None,
3327 );
3328 let map: ArrayRef = Arc::new(MapArray::new(
3329 Arc::new(Field::new("entries", entries.data_type().clone(), false)),
3330 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3331 entries,
3332 None,
3333 false,
3334 ));
3335 let field = Field::new(
3337 "c",
3338 DataType::Map(
3339 Arc::new(Field::new(
3340 "entries",
3341 DataType::Struct(Fields::from(vec![
3342 key_field,
3343 Arc::new(Field::new(
3344 "values",
3345 run_end_type(DataType::Int32, true),
3346 false,
3347 )),
3348 ])),
3349 false,
3350 )),
3351 false,
3352 ),
3353 false,
3354 );
3355 assert_compatible_roundtrip(field, map.clone(), &map);
3356 }
3357
3358 #[test]
3359 fn arrow_writer_nested_dictionary_metadata_roundtrips_with_default_reader() {
3360 let dense: ArrayRef = Arc::new(StructArray::new(
3361 Fields::from(vec![Field::new("item", DataType::Int32, true)]),
3362 vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
3363 None,
3364 ));
3365 let field = Field::new(
3366 "value",
3367 DataType::Dictionary(
3368 Box::new(DataType::Int8),
3369 Box::new(dense.data_type().clone()),
3370 ),
3371 false,
3372 )
3373 .with_metadata(HashMap::from([(
3374 "description".to_string(),
3375 "nested dictionary".to_string(),
3376 )]));
3377 let writer_schema = Arc::new(Schema::new(vec![field.clone()]));
3378 let expected_schema = Arc::new(Schema::new(vec![
3379 field.with_data_type(dense.data_type().clone()),
3380 ]));
3381 let batch = RecordBatch::try_new(expected_schema.clone(), vec![dense]).unwrap();
3382
3383 let mut file = vec![];
3386 let mut writer = ArrowWriter::try_new(&mut file, writer_schema, None).unwrap();
3387 writer.write(&batch).unwrap();
3388 writer.close().unwrap();
3389
3390 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file)).unwrap();
3391 assert!(
3392 builder
3393 .metadata()
3394 .file_metadata()
3395 .key_value_metadata()
3396 .unwrap()
3397 .iter()
3398 .any(|kv| kv.key == ARROW_SCHEMA_META_KEY)
3399 );
3400 assert_eq!(builder.schema(), &expected_schema);
3401 let decoded = builder
3402 .build()
3403 .unwrap()
3404 .collect::<ArrowResult<Vec<_>>>()
3405 .unwrap();
3406 assert_eq!(decoded, vec![batch]);
3407 }
3408
3409 #[test]
3410 fn arrow_writer_rejects_missing_or_extra_top_level_columns_cleanly() {
3411 let one_field = Arc::new(Schema::new(vec![Field::new(
3412 "a",
3413 ArrowDataType::Int32,
3414 false,
3415 )]));
3416 let two_fields = Arc::new(Schema::new(vec![
3417 Field::new("a", ArrowDataType::Int32, false),
3418 Field::new("b", ArrowDataType::Int32, false),
3419 ]));
3420 let one_column = RecordBatch::try_new(
3421 Arc::clone(&one_field),
3422 vec![Arc::new(Int32Array::from(vec![1]))],
3423 )
3424 .unwrap();
3425 let two_columns = RecordBatch::try_new(
3426 Arc::clone(&two_fields),
3427 vec![
3428 Arc::new(Int32Array::from(vec![1])),
3429 Arc::new(Int32Array::from(vec![2])),
3430 ],
3431 )
3432 .unwrap();
3433
3434 let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&two_fields), None).unwrap();
3435 let err = writer.write(&one_column).unwrap_err();
3436 assert!(
3437 err.to_string()
3438 .contains("writer has 2 top-level fields but batch has 1 columns"),
3439 "{err}"
3440 );
3441 assert_eq!(writer.in_progress_rows(), 0);
3442
3443 let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&one_field), None).unwrap();
3444 let err = writer.write(&two_columns).unwrap_err();
3445 assert!(
3446 err.to_string()
3447 .contains("writer has 1 top-level fields but batch has 2 columns"),
3448 "{err}"
3449 );
3450 assert_eq!(writer.in_progress_rows(), 0);
3451 }
3452
3453 #[test]
3454 fn arrow_row_group_factory_rejects_parquet_leaf_count_mismatch_cleanly() {
3455 let one_field = Arc::new(Schema::new(vec![Field::new(
3456 "a",
3457 ArrowDataType::Int32,
3458 false,
3459 )]));
3460 let two_fields = Arc::new(Schema::new(vec![
3461 Field::new("a", ArrowDataType::Int32, false),
3462 Field::new("b", ArrowDataType::Int32, false),
3463 ]));
3464 let parquet = ArrowSchemaConverter::new().convert(&one_field).unwrap();
3465 let props = Arc::new(WriterProperties::default());
3466 let file_writer =
3467 SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), Arc::clone(&props))
3468 .unwrap();
3469
3470 let eager_error =
3471 ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&two_fields)).unwrap_err();
3472 assert!(matches!(&eager_error, ParquetError::General(_)));
3473 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&two_fields));
3475 let Err(err) = factory.create_column_writers(0) else {
3476 panic!("mismatched schemas should not create column writers")
3477 };
3478 assert_eq!(err.to_string(), eager_error.to_string());
3479 assert!(
3480 err.to_string()
3481 .contains("Arrow schema maps to 2 leaf columns but Parquet schema has 1"),
3482 "{err}"
3483 );
3484
3485 let parquet = ArrowSchemaConverter::new().convert(&two_fields).unwrap();
3486 let file_writer =
3487 SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), props).unwrap();
3488 let eager_error =
3489 ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&one_field)).unwrap_err();
3490 assert!(matches!(&eager_error, ParquetError::General(_)));
3491 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file_writer, one_field);
3493 let Err(err) = factory.create_column_writers(0) else {
3494 panic!("mismatched schemas should not create column writers")
3495 };
3496 assert_eq!(err.to_string(), eager_error.to_string());
3497 assert!(
3498 err.to_string()
3499 .contains("Arrow schema maps to 1 leaf columns but Parquet schema has 2"),
3500 "{err}"
3501 );
3502 }
3503
3504 #[test]
3505 fn arrow_writer_required_map_keys_reject_only_reachable_nulls() {
3506 let keys: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
3507 Int32Array::from(vec![1, 0, 1]),
3508 Arc::new(Int32Array::from(vec![None, Some(7)])),
3509 ));
3510 let values = Int32Array::from(vec![Some(10), Some(20), None]);
3511 let fields = Fields::from(vec![
3512 Field::new("key", keys.data_type().clone(), false),
3513 Field::new("value", DataType::Int32, true),
3514 ]);
3515 assert!(
3516 StructArray::try_new(
3517 fields.clone(),
3518 vec![keys.clone(), Arc::new(values.clone())],
3519 None
3520 )
3521 .is_err()
3522 );
3523
3524 let data = ArrayData::builder(DataType::Struct(fields))
3527 .len(3)
3528 .add_child_data(keys.to_data())
3529 .add_child_data(values.to_data())
3530 .build()
3531 .unwrap();
3532 data.validate_full().unwrap();
3533 let entries = StructArray::from(data);
3534 let entries_field = Arc::new(Field::new("entries", entries.data_type().clone(), false));
3535 let offsets = OffsetBuffer::new(vec![0i32, 1, 2, 3].into());
3536 let map = MapArray::try_new(
3537 entries_field.clone(),
3538 offsets.clone(),
3539 entries.clone(),
3540 None,
3541 false,
3542 )
3543 .unwrap();
3544 assert_eq!(map.keys().logical_null_count(), 1);
3545 let masked = MapArray::try_new(
3546 entries_field,
3547 offsets,
3548 entries,
3549 Some(NullBuffer::from(vec![true, false, true])),
3550 false,
3551 )
3552 .unwrap();
3553 let writer_schema = Arc::new(Schema::new(vec![Field::new(
3554 "map",
3555 DataType::Map(
3556 Arc::new(Field::new(
3557 "entries",
3558 DataType::Struct(
3559 vec![
3560 Field::new("key", DataType::Int32, false),
3561 Field::new("value", DataType::Int32, true),
3562 ]
3563 .into(),
3564 ),
3565 false,
3566 )),
3567 false,
3568 ),
3569 true,
3570 )]));
3571 for cdc in [false, true] {
3572 let props = WriterProperties::builder()
3573 .set_content_defined_chunking(cdc.then(Default::default))
3574 .build();
3575 for (array, expected_keys, expected_null_maps) in [
3576 (map.clone(), None, 0),
3577 (map.slice(0, 1), Some(1), 0),
3578 (map.slice(2, 1), Some(1), 0),
3579 (masked.clone(), Some(2), 1),
3580 ] {
3581 let schema = Arc::new(Schema::new(vec![Field::new(
3582 "map",
3583 array.data_type().clone(),
3584 true,
3585 )]));
3586 let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
3587 let mut bytes = Vec::new();
3588 let mut writer =
3589 ArrowWriter::try_new(&mut bytes, writer_schema.clone(), Some(props.clone()))
3590 .unwrap();
3591 let result = writer.write(&batch);
3592 let Some(expected_keys) = expected_keys else {
3593 let err = result.unwrap_err();
3594 assert!(
3595 err.to_string().contains("Found null")
3596 && err.to_string().contains("required field 'key'"),
3597 "{err}"
3598 );
3599 continue;
3600 };
3601 result.unwrap();
3602 writer.close().unwrap();
3603 let actual = ParquetRecordBatchReader::try_new(Bytes::from(bytes), 1024)
3604 .unwrap()
3605 .next()
3606 .unwrap()
3607 .unwrap();
3608 let actual = actual.column(0).as_map();
3609 assert_eq!(actual.null_count(), expected_null_maps);
3610 assert_eq!(
3611 actual.keys().as_primitive::<Int32Type>().values().as_ref(),
3612 vec![7; expected_keys]
3613 );
3614 }
3615 }
3616 }
3617
3618 #[test]
3619 fn arrow_writer_rejects_nullable_map_key_schemas_early() {
3620 let schema_with_key = |key| {
3621 Arc::new(Schema::new(vec![Field::new(
3622 "map",
3623 DataType::Map(
3624 Arc::new(Field::new(
3625 "entries",
3626 DataType::Struct(
3627 vec![key, Field::new("value", DataType::Int32, true)].into(),
3628 ),
3629 false,
3630 )),
3631 false,
3632 ),
3633 true,
3634 )]))
3635 };
3636 let required = schema_with_key(Field::new("key", DataType::Int32, false));
3637 let parquet = ArrowSchemaConverter::new().convert(&required).unwrap();
3638 let ree = DataType::RunEndEncoded(
3639 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3640 Arc::new(Field::new("values", DataType::Int32, true)),
3641 );
3642 for key in [
3643 Field::new("key", DataType::Int32, true),
3644 Field::new("key", ree.clone(), false),
3645 Field::new(
3646 "key",
3647 DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
3648 false,
3649 ),
3650 ] {
3651 let schema = schema_with_key(key);
3652 for explicit_schema in [false, true] {
3653 for skip_metadata in [false, true] {
3654 let mut options =
3655 ArrowWriterOptions::new().with_skip_arrow_metadata(skip_metadata);
3656 if explicit_schema {
3657 options = options.with_parquet_schema(parquet.clone());
3658 }
3659 let mut bytes = Vec::new();
3660 let err =
3661 ArrowWriter::try_new_with_options(&mut bytes, schema.clone(), options)
3662 .unwrap_err();
3663 assert!(
3664 err.to_string()
3665 .contains("Map key field 'key' must be non-nullable"),
3666 "{err}"
3667 );
3668 assert!(
3669 bytes.is_empty(),
3670 "schema rejection must precede file initialization"
3671 );
3672 }
3673 }
3674 let file = SerializedFileWriter::new(
3675 Vec::new(),
3676 parquet.root_schema_ptr(),
3677 Arc::new(WriterProperties::default()),
3678 )
3679 .unwrap();
3680 let eager_error =
3681 ArrowRowGroupWriterFactory::try_new(&file, Arc::clone(&schema)).unwrap_err();
3682 assert!(matches!(&eager_error, ParquetError::ArrowError(_)));
3683 assert_eq!(
3684 eager_error.to_string(),
3685 "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3686 );
3687 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file, schema);
3689 for row_group_index in 0..2 {
3690 let Err(err) = factory.create_column_writers(row_group_index) else {
3691 panic!("nullable map-key schema should not create column writers")
3692 };
3693 assert!(matches!(&err, ParquetError::ArrowError(_)), "{err:?}");
3694 assert_eq!(
3695 err.to_string(),
3696 "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3697 );
3698 }
3699 }
3700 }
3701
3702 #[test]
3703 fn eager_row_group_validation_rejects_writer_and_chunker_count_mismatches() {
3704 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
3705 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))])
3706 .unwrap();
3707 let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3708 let (_file, factory) = writer.into_serialized_writer().unwrap();
3709 let mut group = factory.create_row_group_writer(0).unwrap();
3710 assert!(
3711 group
3712 .write_with_chunkers(&batch, &mut [])
3713 .unwrap_err()
3714 .to_string()
3715 .contains("chunkers")
3716 );
3717 assert_eq!(group.buffered_rows, 0);
3718 group.writers.clear();
3719 assert!(
3720 group
3721 .write(&batch)
3722 .unwrap_err()
3723 .to_string()
3724 .contains("column writers")
3725 );
3726 assert_eq!(group.buffered_rows, 0);
3727 }
3728
3729 #[test]
3730 fn arrow_write_schema_plan_preserves_error_variants() {
3731 for error in [
3732 ParquetError::General("schema mismatch".to_owned()),
3733 ParquetError::NYI("unsupported type".to_owned()),
3734 ParquetError::ArrowError("invalid schema".to_owned()),
3735 ] {
3736 let expected_variant = std::mem::discriminant(&error);
3737 let expected_message = error.to_string();
3738 let cached = ArrowWriteSchemaPlanError::from_parquet_error(error);
3739 for _ in 0..2 {
3740 let error = cached.to_parquet_error();
3741 assert_eq!(std::mem::discriminant(&error), expected_variant);
3742 assert_eq!(error.to_string(), expected_message);
3743 }
3744 }
3745 }
3746
3747 #[test]
3748 fn arrow_write_schema_plan_caches_nested_leaf_ranges() {
3749 let nested = ArrowDataType::Struct(Fields::from(vec![
3750 Field::new("i", ArrowDataType::Int32, false),
3751 Field::new("s", ArrowDataType::Utf8, true),
3752 ]));
3753 let list = ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Boolean, true)));
3754 let ree = ArrowDataType::RunEndEncoded(
3755 Arc::new(Field::new("run_ends", ArrowDataType::Int32, false)),
3756 Arc::new(Field::new("values", ArrowDataType::Int64, false)),
3757 );
3758 let dictionary_fsb = ArrowDataType::Dictionary(
3759 Box::new(ArrowDataType::Int8),
3760 Box::new(ArrowDataType::FixedSizeBinary(2)),
3761 );
3762 let schema = Arc::new(Schema::new(vec![
3763 Field::new("nested", nested, true),
3764 Field::new("list", list, true),
3765 Field::new("ree", ree, false),
3766 Field::new("dictionary_fsb", dictionary_fsb, true),
3767 ]));
3768 let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
3769
3770 let plan = ArrowWriteSchemaPlan::try_new(&parquet, &schema).unwrap();
3771 let ranges = plan
3772 .fields
3773 .iter()
3774 .map(|field| field.leaf_range.clone())
3775 .collect::<Vec<_>>();
3776 assert_eq!(ranges, vec![0..2, 2..3, 3..4, 4..5]);
3777 assert_eq!(plan.leaves.len(), parquet.num_columns());
3778 assert_eq!(
3779 plan.leaves[4].physical_type(),
3780 crate::basic::Type::FIXED_LEN_BYTE_ARRAY
3781 );
3782 }
3783
3784 #[test]
3785 fn cached_schema_plan_preserves_compatible_physical_layout_alternation() {
3786 let writer_schema = Arc::new(Schema::new(vec![
3787 Field::new("number", ArrowDataType::Int32, false),
3788 Field::new("text", ArrowDataType::Utf8, false),
3789 Field::new("bytes", ArrowDataType::Binary, false),
3790 ]));
3791
3792 let dense = RecordBatch::try_new(
3793 Arc::clone(&writer_schema),
3794 vec![
3795 Arc::new(Int32Array::from(vec![1])),
3796 Arc::new(StringArray::from(vec!["a"])),
3797 Arc::new(BinaryArray::from_iter_values([b"a".as_slice()])),
3798 ],
3799 )
3800 .unwrap();
3801
3802 let number_dictionary = DictionaryArray::<Int8Type>::try_new(
3803 Int8Array::from(vec![0]),
3804 Arc::new(Int32Array::from(vec![2])),
3805 )
3806 .unwrap();
3807 let text_dictionary = DictionaryArray::<Int8Type>::try_new(
3808 Int8Array::from(vec![0]),
3809 Arc::new(StringArray::from(vec!["b"])),
3810 )
3811 .unwrap();
3812 let bytes_dictionary = DictionaryArray::<Int8Type>::try_new(
3813 Int8Array::from(vec![0]),
3814 Arc::new(BinaryArray::from_iter_values([b"b".as_slice()])),
3815 )
3816 .unwrap();
3817 let dictionary = RecordBatch::try_from_iter(vec![
3818 ("number", Arc::new(number_dictionary) as ArrayRef),
3819 ("text", Arc::new(text_dictionary) as ArrayRef),
3820 ("bytes", Arc::new(bytes_dictionary) as ArrayRef),
3821 ])
3822 .unwrap();
3823
3824 let run_ends = Int32Array::from(vec![1]);
3825 let number_ree: ArrayRef =
3826 Arc::new(Int32RunArray::try_new(&run_ends, &Int32Array::from(vec![3])).unwrap());
3827 let text_ree: ArrayRef =
3828 Arc::new(Int32RunArray::try_new(&run_ends, &StringArray::from(vec!["c"])).unwrap());
3829 let bytes_ree: ArrayRef = Arc::new(
3830 Int32RunArray::try_new(&run_ends, &BinaryArray::from_iter_values([b"c".as_slice()]))
3831 .unwrap(),
3832 );
3833 let ree = RecordBatch::try_from_iter(vec![
3834 ("number", number_ree),
3835 ("text", text_ree),
3836 ("bytes", bytes_ree),
3837 ])
3838 .unwrap();
3839
3840 let alternate = RecordBatch::try_from_iter(vec![
3841 ("number", Arc::new(Int32Array::from(vec![4])) as ArrayRef),
3842 (
3843 "text",
3844 Arc::new(LargeStringArray::from(vec!["d"])) as ArrayRef,
3845 ),
3846 (
3847 "bytes",
3848 Arc::new(BinaryViewArray::from_iter_values([b"d".as_slice()])) as ArrayRef,
3849 ),
3850 ])
3851 .unwrap();
3852
3853 let props = WriterProperties::builder()
3854 .set_max_row_group_row_count(Some(1))
3855 .build();
3856 let mut out = Vec::new();
3857 let mut writer =
3858 ArrowWriter::try_new(&mut out, Arc::clone(&writer_schema), Some(props)).unwrap();
3859 for batch in [&dense, &dictionary, &ree, &alternate] {
3860 writer.write(batch).unwrap();
3861 }
3862 writer.close().unwrap();
3863
3864 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(out)).unwrap();
3865 assert_eq!(builder.metadata().num_row_groups(), 4);
3866 let mut reader = builder.build().unwrap();
3867 let actual = reader.next().unwrap().unwrap();
3868 assert_eq!(
3869 actual.column(0).as_primitive::<Int32Type>(),
3870 &Int32Array::from(vec![1, 2, 3, 4])
3871 );
3872 assert_eq!(
3873 actual.column(1).as_string::<i32>(),
3874 &StringArray::from(vec!["a", "b", "c", "d"])
3875 );
3876 assert_eq!(
3877 actual.column(2).as_binary::<i32>(),
3878 &BinaryArray::from_iter_values([
3879 b"a".as_slice(),
3880 b"b".as_slice(),
3881 b"c".as_slice(),
3882 b"d".as_slice(),
3883 ])
3884 );
3885 }
3886
3887 #[test]
3888 fn cached_schema_plan_does_not_retain_batch_arrays() {
3889 let schema = Arc::new(Schema::new(vec![Field::new(
3890 "a",
3891 ArrowDataType::Int32,
3892 false,
3893 )]));
3894 let array = Arc::new(Int32Array::from(vec![1, 2, 3]));
3895 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap();
3896 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3897
3898 writer.write(&batch).unwrap();
3899 drop(batch);
3900 assert_eq!(
3901 Arc::strong_count(&array),
3902 1,
3903 "schema and row-group caches must not retain input ArrayRefs"
3904 );
3905 writer.close().unwrap();
3906 }
3907
3908 #[test]
3909 fn arrow_writer_page_size() {
3910 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3911
3912 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3913
3914 for i in 0..10 {
3916 let value = i
3917 .to_string()
3918 .repeat(10)
3919 .chars()
3920 .take(10)
3921 .collect::<String>();
3922
3923 builder.append_value(value);
3924 }
3925
3926 let array = Arc::new(builder.finish());
3927
3928 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3929
3930 let file = tempfile::tempfile().unwrap();
3931
3932 let props = WriterProperties::builder()
3934 .set_data_page_size_limit(1)
3935 .set_dictionary_page_size_limit(1)
3936 .set_write_batch_size(1)
3937 .build();
3938
3939 let mut writer =
3940 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3941 .expect("Unable to write file");
3942 writer.write(&batch).unwrap();
3943 writer.close().unwrap();
3944
3945 let options = ReadOptionsBuilder::new().with_page_index().build();
3946 let reader =
3947 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3948
3949 let column = reader.metadata().row_group(0).columns();
3950
3951 assert_eq!(column.len(), 1);
3952
3953 assert!(
3956 column[0].dictionary_page_offset().is_some(),
3957 "Expected a dictionary page"
3958 );
3959
3960 let page_index = reader
3961 .metadata()
3962 .page_index()
3963 .expect("page index should be present");
3964 let page_locations = page_index
3965 .page_locations(0, 0)
3966 .expect("page locations should exist");
3967
3968 assert_eq!(
3971 page_locations.len(),
3972 10,
3973 "Expected 10 pages but got {page_locations:#?}"
3974 );
3975 }
3976
3977 const MEDIUM_SIZE: usize = 63;
3978
3979 fn check_bloom_filter<T: AsBytes>(
3980 files: Vec<Bytes>,
3981 file_column: String,
3982 positive_values: Vec<T>,
3983 negative_values: Vec<T>,
3984 ) {
3985 files.into_iter().take(1).for_each(|file| {
3986 let file_reader = SerializedFileReader::new_with_options(
3987 file,
3988 ReadOptionsBuilder::new()
3989 .with_reader_properties(
3990 ReaderProperties::builder()
3991 .set_read_bloom_filter(true)
3992 .build(),
3993 )
3994 .build(),
3995 )
3996 .expect("Unable to open file as Parquet");
3997 let metadata = file_reader.metadata();
3998
3999 let mut bloom_filters: Vec<_> = vec![];
4001 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
4002 if let Some((column_index, _)) = row_group
4003 .columns()
4004 .iter()
4005 .enumerate()
4006 .find(|(_, column)| column.column_path().string() == file_column)
4007 {
4008 let row_group_reader = file_reader
4009 .get_row_group(ri)
4010 .expect("Unable to read row group");
4011 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
4012 bloom_filters.push(sbbf.clone());
4013 } else {
4014 panic!("No bloom filter for column named {file_column} found");
4015 }
4016 } else {
4017 panic!("No column named {file_column} found");
4018 }
4019 }
4020
4021 positive_values.iter().for_each(|value| {
4022 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4023 assert!(
4024 found.is_some(),
4025 "{}",
4026 format!("Value {:?} should be in bloom filter", value.as_bytes())
4027 );
4028 });
4029
4030 negative_values.iter().for_each(|value| {
4031 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4032 assert!(
4033 found.is_none(),
4034 "{}",
4035 format!("Value {:?} should not be in bloom filter", value.as_bytes())
4036 );
4037 });
4038 });
4039 }
4040
4041 #[test]
4042 #[cfg_attr(miri, ignore)] fn bool_large_single_column() {
4044 let values = Arc::new(
4045 [None, Some(true), Some(false)]
4046 .iter()
4047 .cycle()
4048 .copied()
4049 .take(200_000)
4050 .collect::<BooleanArray>(),
4051 );
4052 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
4053 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4054 let file = tempfile::tempfile().unwrap();
4055
4056 let mut writer =
4057 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
4058 .expect("Unable to write file");
4059 writer.write(&expected_batch).unwrap();
4060 writer.close().unwrap();
4061 }
4062
4063 #[test]
4064 fn check_page_offset_index_with_nan() {
4065 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4066 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
4067 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4068
4069 let mut out = Vec::with_capacity(1024);
4070 let mut writer =
4071 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
4072 writer.write(&batch).unwrap();
4073 let file_meta_data = writer.close().unwrap();
4074 for row_group in file_meta_data.row_groups() {
4075 for column in row_group.columns() {
4076 assert!(column.offset_index_offset().is_some());
4077 assert!(column.offset_index_length().is_some());
4078 assert!(column.column_index_offset().is_some());
4079 assert!(column.column_index_length().is_some());
4080 }
4081 }
4082 if let Some(page_index) = file_meta_data.page_index() {
4083 for rg in 0..file_meta_data.num_row_groups() {
4084 for col in 0..file_meta_data.row_group(rg).num_columns() {
4085 let idx = page_index
4086 .column_index(rg, col)
4087 .expect("column index should exist");
4088 assert!(idx.nan_counts().is_some());
4089 let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
4090 panic!("expected double statistics")
4091 };
4092 for i in 0..idx.num_pages() as usize {
4093 assert_eq!(float_idx.nan_count(i), Some(10));
4094 assert_eq!(
4095 f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
4096 Ordering::Equal
4097 );
4098 assert_eq!(
4099 f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
4100 Ordering::Equal
4101 );
4102 }
4103 }
4104 }
4105 } else {
4106 panic!("page index should be present");
4107 }
4108 }
4109
4110 #[test]
4111 fn check_page_offset_index_with_mixed_nan() {
4112 let schema = Arc::new(Schema::new(vec![Field::new(
4113 "col",
4114 DataType::Float64,
4115 true,
4116 )]));
4117
4118 let mut out = Vec::with_capacity(1024);
4119 let props = WriterProperties::builder()
4120 .set_data_page_row_count_limit(10)
4121 .build();
4122 let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
4123 .expect("Unable to write file");
4124
4125 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4127 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4128 writer.write(&batch).unwrap();
4129
4130 let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
4132 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4133 writer.write(&batch).unwrap();
4134
4135 let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
4137 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4138 writer.write(&batch).unwrap();
4139
4140 let values = Arc::new(Float64Array::from(vec![
4142 -1.0,
4143 0.0,
4144 f64::NAN,
4145 -f64::NAN,
4146 1.0,
4147 ]));
4148 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4149 writer.write(&batch).unwrap();
4150
4151 let file_meta_data = writer.close().unwrap();
4152
4153 let col_stats = file_meta_data
4155 .row_group(0)
4156 .column(0)
4157 .statistics()
4158 .expect("missing column chunk statistics");
4159
4160 assert_eq!(col_stats.nan_count_opt(), Some(22));
4161 assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
4162 assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
4163
4164 assert!(file_meta_data.page_index().is_some());
4165 let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
4166 assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
4167
4168 let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
4170 panic!("expected double statistics")
4171 };
4172
4173 assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
4174 assert_eq!(
4175 f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
4176 Ordering::Equal
4177 );
4178 assert_eq!(
4179 f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
4180 Ordering::Equal
4181 );
4182 assert_eq!(
4183 (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
4184 Ordering::Equal
4185 );
4186 assert_eq!(
4187 (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
4188 Ordering::Equal
4189 );
4190 assert_eq!(float_idx.min_value(2), Some(&0.0));
4191 assert_eq!(float_idx.max_value(2), Some(&0.0));
4192 assert_eq!(float_idx.min_value(3), Some(&-1.0));
4193 assert_eq!(float_idx.max_value(3), Some(&1.0));
4194 }
4195
4196 #[test]
4197 #[should_panic(
4198 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
4199 )]
4200 fn interval_month_day_nano_single_column() {
4201 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
4202 IntervalMonthDayNano::new(0, 1, 5),
4203 IntervalMonthDayNano::new(0, 3, 2),
4204 IntervalMonthDayNano::new(3, -2, -5),
4205 IntervalMonthDayNano::new(-200, 4, -1),
4206 ]);
4207 }
4208
4209 #[test]
4210 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_at_end() {
4212 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4213 let files = RoundTripTest::new(array)
4214 .with_nullable(false)
4215 .with_bloom_filter(true)
4216 .with_bloom_filter_position(BloomFilterPosition::End)
4217 .run();
4218
4219 check_bloom_filter(
4220 files,
4221 "col".to_string(),
4222 (0..SMALL_SIZE as i32).collect(),
4223 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4224 );
4225 }
4226
4227 #[test]
4228 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter() {
4230 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4231 let files = RoundTripTest::new(array)
4232 .with_nullable(false)
4233 .with_bloom_filter(true)
4234 .run();
4235
4236 check_bloom_filter(
4237 files,
4238 "col".to_string(),
4239 (0..SMALL_SIZE as i32).collect(),
4240 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4241 );
4242 }
4243
4244 fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
4245 let schema = Arc::new(Schema::new(vec![Field::new(
4246 "col",
4247 array.data_type().clone(),
4248 false,
4249 )]));
4250 let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
4251 let props = WriterProperties::builder()
4252 .set_dictionary_enabled(true)
4253 .set_dictionary_page_size_limit(dictionary_page_size_limit)
4254 .set_write_batch_size(256)
4255 .set_bloom_filter_enabled(true)
4256 .build();
4257 let mut buf = Vec::new();
4258 let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
4259 writer.write(&batch).unwrap();
4260 writer.close().unwrap();
4261 Bytes::from(buf)
4262 }
4263
4264 fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4265 let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4266 *metadata
4267 .row_group(0)
4268 .column(0)
4269 .page_encoding_stats_mask()
4270 .unwrap()
4271 }
4272
4273 #[test]
4276 fn string_column_bloom_filter_populated_from_dictionary() {
4277 let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4278 let array = Arc::new(StringArray::from_iter_values(&values));
4279 let file = write_with_bloom_filter(array, 1024 * 1024);
4280 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4281
4282 check_bloom_filter(
4283 vec![file],
4284 "col".to_string(),
4285 (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4286 (10..20)
4287 .map(|i| format!("value-{i}").into_bytes())
4288 .collect(),
4289 );
4290 }
4291
4292 #[test]
4295 fn string_column_bloom_filter_across_dictionary_fallback() {
4296 let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4297 let array = Arc::new(StringArray::from_iter_values(&values));
4298 let file = write_with_bloom_filter(array, 1024);
4299 let encodings = data_page_encoding_mask(&file);
4300 assert!(
4301 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4302 "expected dictionary and plain data pages, got {encodings:?}"
4303 );
4304
4305 check_bloom_filter(
4306 vec![file],
4307 "col".to_string(),
4308 values.into_iter().map(String::into_bytes).collect(),
4309 (2000..2010)
4310 .map(|i| format!("value-{i}").into_bytes())
4311 .collect(),
4312 );
4313 }
4314
4315 #[test]
4316 fn i64_column_bloom_filter_populated_from_dictionary() {
4317 let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4318 let file = write_with_bloom_filter(array, 1024 * 1024);
4319 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4320
4321 check_bloom_filter(
4322 vec![file],
4323 "col".to_string(),
4324 (0..10i64).collect(),
4325 (10..20i64).collect(),
4326 );
4327 }
4328
4329 #[test]
4330 fn i64_column_bloom_filter_across_dictionary_fallback() {
4331 let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4332 let file = write_with_bloom_filter(array, 1024);
4333 let encodings = data_page_encoding_mask(&file);
4334 assert!(
4335 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4336 "expected dictionary and plain data pages, got {encodings:?}"
4337 );
4338
4339 check_bloom_filter(
4340 vec![file],
4341 "col".to_string(),
4342 (0..2000i64).collect(),
4343 (2000..2010i64).collect(),
4344 );
4345 }
4346
4347 #[test]
4352 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_fixed_ndv() {
4354 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4355
4356 let files = RoundTripTest::new(array.clone())
4358 .with_nullable(false)
4359 .with_bloom_filter(true)
4360 .with_bloom_filter_ndv(1_000_000)
4361 .run();
4362
4363 check_bloom_filter(
4364 files,
4365 "col".to_string(),
4366 (0..SMALL_SIZE as i32).collect(),
4367 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4368 );
4369
4370 let files = RoundTripTest::new(array)
4372 .with_nullable(false)
4373 .with_bloom_filter(true)
4374 .with_bloom_filter_ndv(3)
4375 .run();
4376
4377 check_bloom_filter(
4378 files,
4379 "col".to_string(),
4380 (0..SMALL_SIZE as i32).collect(),
4381 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4382 );
4383 }
4384
4385 #[test]
4386 #[cfg_attr(miri, ignore)] fn binary_column_bloom_filter() {
4388 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4389 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4390 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4391
4392 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4393 let files = RoundTripTest::new(array)
4394 .with_nullable(false)
4395 .with_bloom_filter(true)
4396 .run();
4397
4398 check_bloom_filter(
4399 files,
4400 "col".to_string(),
4401 many_vecs,
4402 vec![vec![(SMALL_SIZE + 1) as u8]],
4403 );
4404 }
4405
4406 #[test]
4407 #[cfg_attr(miri, ignore)] fn empty_string_null_column_bloom_filter() {
4409 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4410 let raw_strs = raw_values.iter().map(|s| s.as_str());
4411
4412 let array = Arc::new(StringArray::from_iter_values(raw_strs));
4413 let files = RoundTripTest::new(array)
4414 .with_nullable(false)
4415 .with_bloom_filter(true)
4416 .run();
4417
4418 let optional_raw_values: Vec<_> = raw_values
4419 .iter()
4420 .enumerate()
4421 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4422 .collect();
4423 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4425 }
4426
4427 #[test]
4428 fn list_and_map_coerced_names() {
4429 let list_field =
4431 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4432 let map_field = Field::new_map(
4433 "my_map",
4434 "my_entries",
4435 Field::new("my_keys", DataType::Int32, false),
4436 Field::new("my_values", DataType::Int32, true),
4437 false,
4438 true,
4439 );
4440
4441 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4442 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4443
4444 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4445
4446 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4448 let file = tempfile::tempfile().unwrap();
4449 let mut writer =
4450 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4451
4452 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4453 writer.write(&batch).unwrap();
4454 let file_metadata = writer.close().unwrap();
4455
4456 let schema = file_metadata.file_metadata().schema();
4457 let list_field = &schema.get_fields()[0].get_fields()[0];
4459 assert_eq!(list_field.get_fields()[0].name(), "element");
4460
4461 let map_field = &schema.get_fields()[1].get_fields()[0];
4462 assert_eq!(map_field.name(), "key_value");
4464 assert_eq!(map_field.get_fields()[0].name(), "key");
4466 assert_eq!(map_field.get_fields()[1].name(), "value");
4468
4469 let reader = SerializedFileReader::new(file).unwrap();
4471 let file_schema = reader.metadata().file_metadata().schema();
4472 let fields = file_schema.get_fields();
4473 let list_field = &fields[0].get_fields()[0];
4474 assert_eq!(list_field.get_fields()[0].name(), "element");
4475 let map_field = &fields[1].get_fields()[0];
4476 assert_eq!(map_field.name(), "key_value");
4477 assert_eq!(map_field.get_fields()[0].name(), "key");
4478 assert_eq!(map_field.get_fields()[1].name(), "value");
4479 }
4480
4481 #[test]
4482 #[cfg_attr(miri, ignore)] fn fallback_flush_data_page() {
4484 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4486 let values = Arc::new(StringArray::from(raw_values));
4487 let encodings = vec![
4488 Encoding::DELTA_BYTE_ARRAY,
4489 Encoding::DELTA_LENGTH_BYTE_ARRAY,
4490 ];
4491 let data_type = values.data_type().clone();
4492 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4493 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4494
4495 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4496 let data_page_size_limit: usize = 32;
4497 let write_batch_size: usize = 16;
4498
4499 for encoding in &encodings {
4500 for row_group_size in row_group_sizes {
4501 let props = WriterProperties::builder()
4502 .set_writer_version(WriterVersion::PARQUET_2_0)
4503 .set_max_row_group_row_count(Some(row_group_size))
4504 .set_dictionary_enabled(false)
4505 .set_encoding(*encoding)
4506 .set_data_page_size_limit(data_page_size_limit)
4507 .set_write_batch_size(write_batch_size)
4508 .build();
4509
4510 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4511 let string_array_a = StringArray::from(a.clone());
4512 let string_array_b = StringArray::from(b.clone());
4513 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4514 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4515 assert_eq!(
4516 vec_a, vec_b,
4517 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4518 );
4519 });
4520 }
4521 }
4522 }
4523
4524 #[test]
4525 #[cfg_attr(miri, ignore)] fn u32_min_max() {
4527 let src = [
4529 u32::MIN,
4530 1,
4531 (i32::MAX as u32) - 1,
4532 i32::MAX as u32,
4533 (i32::MAX as u32) + 1,
4534 u32::MAX - 1,
4535 u32::MAX,
4536 ];
4537 let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4538 let files = RoundTripTest::new(values).with_nullable(false).run();
4539
4540 for file in files {
4541 let reader = SerializedFileReader::new(file).unwrap();
4543 let metadata = reader.metadata();
4544
4545 let mut row_offset = 0;
4546 for row_group in metadata.row_groups() {
4547 assert_eq!(row_group.num_columns(), 1);
4548 let column = row_group.column(0);
4549
4550 let num_values = column.num_values() as usize;
4551 let src_slice = &src[row_offset..row_offset + num_values];
4552 row_offset += column.num_values() as usize;
4553
4554 let stats = column.statistics().unwrap();
4555 if let Statistics::Int32(stats) = stats {
4556 assert_eq!(
4557 *stats.min_opt().unwrap() as u32,
4558 *src_slice.iter().min().unwrap()
4559 );
4560 assert_eq!(
4561 *stats.max_opt().unwrap() as u32,
4562 *src_slice.iter().max().unwrap()
4563 );
4564 } else {
4565 panic!("Statistics::Int32 missing")
4566 }
4567 }
4568 }
4569 }
4570
4571 #[test]
4572 #[cfg_attr(miri, ignore)] fn u64_min_max() {
4574 let src = [
4576 u64::MIN,
4577 1,
4578 (i64::MAX as u64) - 1,
4579 i64::MAX as u64,
4580 (i64::MAX as u64) + 1,
4581 u64::MAX - 1,
4582 u64::MAX,
4583 ];
4584 let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4585 let files = RoundTripTest::new(values).with_nullable(false).run();
4586
4587 for file in files {
4588 let reader = SerializedFileReader::new(file).unwrap();
4590 let metadata = reader.metadata();
4591
4592 let mut row_offset = 0;
4593 for row_group in metadata.row_groups() {
4594 assert_eq!(row_group.num_columns(), 1);
4595 let column = row_group.column(0);
4596
4597 let num_values = column.num_values() as usize;
4598 let src_slice = &src[row_offset..row_offset + num_values];
4599 row_offset += column.num_values() as usize;
4600
4601 let stats = column.statistics().unwrap();
4602 if let Statistics::Int64(stats) = stats {
4603 assert_eq!(
4604 *stats.min_opt().unwrap() as u64,
4605 *src_slice.iter().min().unwrap()
4606 );
4607 assert_eq!(
4608 *stats.max_opt().unwrap() as u64,
4609 *src_slice.iter().max().unwrap()
4610 );
4611 } else {
4612 panic!("Statistics::Int64 missing")
4613 }
4614 }
4615 }
4616 }
4617
4618 #[test]
4619 #[cfg_attr(miri, ignore)] fn statistics_null_counts_only_nulls() {
4621 let values = Arc::new(UInt64Array::from(vec![None, None]));
4623 let files = RoundTripTest::new(values).run();
4624
4625 for file in files {
4626 let reader = SerializedFileReader::new(file).unwrap();
4628 let metadata = reader.metadata();
4629 assert_eq!(metadata.num_row_groups(), 1);
4630 let row_group = metadata.row_group(0);
4631 assert_eq!(row_group.num_columns(), 1);
4632 let column = row_group.column(0);
4633 let stats = column.statistics().unwrap();
4634 assert_eq!(stats.null_count_opt(), Some(2));
4635 }
4636 }
4637
4638 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4639 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4640 }
4641
4642 #[test]
4643 fn test_aggregates_records() {
4644 let arrays = [
4645 Int32Array::from((0..100).collect::<Vec<_>>()),
4646 Int32Array::from((0..50).collect::<Vec<_>>()),
4647 Int32Array::from((200..500).collect::<Vec<_>>()),
4648 ];
4649
4650 let schema = Arc::new(Schema::new(vec![Field::new(
4651 "int",
4652 ArrowDataType::Int32,
4653 false,
4654 )]));
4655
4656 let file = tempfile::tempfile().unwrap();
4657
4658 let props = WriterProperties::builder()
4659 .set_max_row_group_row_count(Some(200))
4660 .build();
4661
4662 let mut writer =
4663 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4664
4665 for array in arrays {
4666 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4667 writer.write(&batch).unwrap();
4668 }
4669
4670 writer.close().unwrap();
4671
4672 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4673 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4674
4675 let batches = builder
4676 .with_batch_size(100)
4677 .build()
4678 .unwrap()
4679 .collect::<ArrowResult<Vec<_>>>()
4680 .unwrap();
4681
4682 assert_eq!(batches.len(), 5);
4683 assert!(batches.iter().all(|x| x.num_columns() == 1));
4684
4685 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4686
4687 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4688
4689 let values: Vec<_> = batches
4690 .iter()
4691 .flat_map(|x| {
4692 x.column(0)
4693 .as_any()
4694 .downcast_ref::<Int32Array>()
4695 .unwrap()
4696 .values()
4697 .iter()
4698 .copied()
4699 })
4700 .collect();
4701
4702 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4703 assert_eq!(&values, &expected_values)
4704 }
4705
4706 #[test]
4707 fn complex_aggregate() {
4708 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4710 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4711 let struct_a = Arc::new(Field::new(
4712 "struct_a",
4713 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4714 true,
4715 ));
4716
4717 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4718 let struct_b = Arc::new(Field::new(
4719 "struct_b",
4720 DataType::Struct(vec![list_a.clone()].into()),
4721 false,
4722 ));
4723
4724 let schema = Arc::new(Schema::new(vec![struct_b]));
4725
4726 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4728 let field_b_array =
4729 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4730
4731 let struct_a_array = StructArray::from(vec![
4732 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4733 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4734 ]);
4735
4736 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4737 .len(5)
4738 .add_buffer(Buffer::from_iter(vec![
4739 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4740 ]))
4741 .null_bit_buffer(Some(Buffer::from_iter(vec![
4742 true, false, true, false, true,
4743 ])))
4744 .child_data(vec![struct_a_array.into_data()])
4745 .build()
4746 .unwrap();
4747
4748 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4749 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4750
4751 let batch1 =
4752 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4753 .unwrap();
4754
4755 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4756 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4757
4758 let struct_a_array = StructArray::from(vec![
4759 (field_a, Arc::new(field_a_array) as ArrayRef),
4760 (field_b, Arc::new(field_b_array) as ArrayRef),
4761 ]);
4762
4763 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4764 .len(2)
4765 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4766 .child_data(vec![struct_a_array.into_data()])
4767 .build()
4768 .unwrap();
4769
4770 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4771 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4772
4773 let batch2 =
4774 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4775 .unwrap();
4776
4777 let batches = &[batch1, batch2];
4778
4779 let expected = r"
4782 +-------------------------------------------------------------------------------------------------------+
4783 | struct_b |
4784 +-------------------------------------------------------------------------------------------------------+
4785 | {list: [{leaf_a: 1, leaf_b: 1}]} |
4786 | {list: } |
4787 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
4788 | {list: } |
4789 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
4790 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4791 | {list: [{leaf_a: 10, leaf_b: }]} |
4792 +-------------------------------------------------------------------------------------------------------+
4793 ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4794
4795 let actual = pretty_format_batches(batches).unwrap().to_string();
4796 assert_eq!(actual, expected);
4797
4798 let file = tempfile::tempfile().unwrap();
4800 let props = WriterProperties::builder()
4801 .set_max_row_group_row_count(Some(6))
4802 .build();
4803
4804 let mut writer =
4805 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4806
4807 for batch in batches {
4808 writer.write(batch).unwrap();
4809 }
4810 writer.close().unwrap();
4811
4812 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4817 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4818
4819 let batches = builder
4820 .with_batch_size(2)
4821 .build()
4822 .unwrap()
4823 .collect::<ArrowResult<Vec<_>>>()
4824 .unwrap();
4825
4826 assert_eq!(batches.len(), 4);
4827 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4828 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4829
4830 let actual = pretty_format_batches(&batches).unwrap().to_string();
4831 assert_eq!(actual, expected);
4832 }
4833
4834 #[test]
4835 fn test_arrow_writer_metadata() {
4836 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4837 let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4838
4839 let batch = RecordBatch::try_new(
4840 Arc::new(batch_schema),
4841 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4842 )
4843 .unwrap();
4844
4845 let mut buf = Vec::with_capacity(1024);
4846 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4847 writer.write(&batch).unwrap();
4848 writer.close().unwrap();
4849 }
4850
4851 #[test]
4852 fn in_progress_accounting() {
4853 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4855
4856 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4858
4859 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4861
4862 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4863
4864 assert_eq!(writer.in_progress_size(), 0);
4866 assert_eq!(writer.in_progress_rows(), 0);
4867 assert_eq!(writer.memory_size(), 0);
4868 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
4870
4871 let initial_size = writer.in_progress_size();
4873 assert!(initial_size > 0);
4874 assert_eq!(writer.in_progress_rows(), 5);
4875 let initial_memory = writer.memory_size();
4876 assert!(initial_memory > 0);
4877 assert!(
4879 initial_size <= initial_memory,
4880 "{initial_size} <= {initial_memory}"
4881 );
4882
4883 writer.write(&batch).unwrap();
4885 assert!(writer.in_progress_size() > initial_size);
4886 assert_eq!(writer.in_progress_rows(), 10);
4887 assert!(writer.memory_size() > initial_memory);
4888 assert!(
4889 writer.in_progress_size() <= writer.memory_size(),
4890 "in_progress_size {} <= memory_size {}",
4891 writer.in_progress_size(),
4892 writer.memory_size()
4893 );
4894
4895 let pre_flush_bytes_written = writer.bytes_written();
4897 writer.flush().unwrap();
4898 assert_eq!(writer.in_progress_size(), 0);
4899 assert_eq!(writer.memory_size(), 0);
4900 assert!(writer.bytes_written() > pre_flush_bytes_written);
4901
4902 writer.close().unwrap();
4903 }
4904
4905 #[test]
4906 fn test_writer_all_null() {
4907 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4908 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4909 let batch = RecordBatch::try_from_iter(vec![
4910 ("a", Arc::new(a) as ArrayRef),
4911 ("b", Arc::new(b) as ArrayRef),
4912 ])
4913 .unwrap();
4914
4915 let mut buf = Vec::with_capacity(1024);
4916 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4917 writer.write(&batch).unwrap();
4918 writer.close().unwrap();
4919
4920 let bytes = Bytes::from(buf);
4921 let options = ReadOptionsBuilder::new().with_page_index().build();
4922 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4923 let index = reader.metadata().page_index().unwrap();
4924
4925 assert_eq!(index.num_data_pages(0, 0), Some(1)); assert_eq!(index.num_data_pages(0, 1), Some(1)); }
4928
4929 #[test]
4930 fn test_disabled_statistics_with_page() {
4931 let file_schema = Schema::new(vec![
4932 Field::new("a", DataType::Utf8, true),
4933 Field::new("b", DataType::Utf8, true),
4934 ]);
4935 let file_schema = Arc::new(file_schema);
4936
4937 let batch = RecordBatch::try_new(
4938 file_schema.clone(),
4939 vec![
4940 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
4941 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
4942 ],
4943 )
4944 .unwrap();
4945
4946 let props = WriterProperties::builder()
4947 .set_statistics_enabled(EnabledStatistics::None)
4948 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
4949 .build();
4950
4951 let mut buf = Vec::with_capacity(1024);
4952 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
4953 writer.write(&batch).unwrap();
4954
4955 let metadata = writer.close().unwrap();
4956 assert_eq!(metadata.num_row_groups(), 1);
4957 let row_group = metadata.row_group(0);
4958 assert_eq!(row_group.num_columns(), 2);
4959 assert!(row_group.column(0).offset_index_offset().is_some());
4961 assert!(row_group.column(0).column_index_offset().is_some());
4962 assert!(row_group.column(1).offset_index_offset().is_some());
4964 assert!(row_group.column(1).column_index_offset().is_none());
4965
4966 let options = ReadOptionsBuilder::new().with_page_index().build();
4967 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
4968
4969 let row_group = reader.get_row_group(0).unwrap();
4970 let a_col = row_group.metadata().column(0);
4971 let b_col = row_group.metadata().column(1);
4972
4973 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
4975 let min = byte_array_stats.min_opt().unwrap();
4976 let max = byte_array_stats.max_opt().unwrap();
4977
4978 assert_eq!(min.as_bytes(), b"a");
4979 assert_eq!(max.as_bytes(), b"d");
4980 } else {
4981 panic!("expecting Statistics::ByteArray");
4982 }
4983
4984 assert!(b_col.statistics().is_none());
4986
4987 let page_index = reader.metadata().page_index().unwrap();
4988
4989 let a_idx = page_index.column_index(0, 0);
4990 assert!(
4991 matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
4992 "{a_idx:?}"
4993 );
4994 let b_idx = page_index.column_index(0, 1);
4995 assert!(b_idx.is_none(), "{b_idx:?}");
4996 }
4997
4998 #[test]
4999 fn test_disabled_statistics_with_chunk() {
5000 let file_schema = Schema::new(vec![
5001 Field::new("a", DataType::Utf8, true),
5002 Field::new("b", DataType::Utf8, true),
5003 ]);
5004 let file_schema = Arc::new(file_schema);
5005
5006 let batch = RecordBatch::try_new(
5007 file_schema.clone(),
5008 vec![
5009 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5010 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5011 ],
5012 )
5013 .unwrap();
5014
5015 let props = WriterProperties::builder()
5016 .set_statistics_enabled(EnabledStatistics::None)
5017 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5018 .build();
5019
5020 let mut buf = Vec::with_capacity(1024);
5021 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5022 writer.write(&batch).unwrap();
5023
5024 let metadata = writer.close().unwrap();
5025 assert_eq!(metadata.num_row_groups(), 1);
5026 let row_group = metadata.row_group(0);
5027 assert_eq!(row_group.num_columns(), 2);
5028 assert!(row_group.column(0).offset_index_offset().is_some());
5030 assert!(row_group.column(0).column_index_offset().is_none());
5031 assert!(row_group.column(1).offset_index_offset().is_some());
5033 assert!(row_group.column(1).column_index_offset().is_none());
5034
5035 let options = ReadOptionsBuilder::new().with_page_index().build();
5036 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5037
5038 let row_group = reader.get_row_group(0).unwrap();
5039 let a_col = row_group.metadata().column(0);
5040 let b_col = row_group.metadata().column(1);
5041
5042 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5044 let min = byte_array_stats.min_opt().unwrap();
5045 let max = byte_array_stats.max_opt().unwrap();
5046
5047 assert_eq!(min.as_bytes(), b"a");
5048 assert_eq!(max.as_bytes(), b"d");
5049 } else {
5050 panic!("expecting Statistics::ByteArray");
5051 }
5052
5053 assert!(b_col.statistics().is_none());
5055
5056 let page_index = reader.metadata().page_index().unwrap();
5057
5058 let a_idx = page_index.column_index(0, 0);
5059 assert!(a_idx.is_none(), "{a_idx:?}");
5060 let b_idx = page_index.column_index(0, 1);
5061 assert!(b_idx.is_none(), "{b_idx:?}");
5062 }
5063
5064 #[test]
5065 fn test_arrow_writer_skip_metadata() {
5066 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5067 let file_schema = Arc::new(batch_schema.clone());
5068
5069 let batch = RecordBatch::try_new(
5070 Arc::new(batch_schema),
5071 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5072 )
5073 .unwrap();
5074 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5075
5076 let mut buf = Vec::with_capacity(1024);
5077 let mut writer =
5078 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5079 writer.write(&batch).unwrap();
5080 writer.close().unwrap();
5081
5082 let bytes = Bytes::from(buf);
5083 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5084 assert_eq!(file_schema, *reader_builder.schema());
5085 if let Some(key_value_metadata) = reader_builder
5086 .metadata()
5087 .file_metadata()
5088 .key_value_metadata()
5089 {
5090 assert!(
5091 !key_value_metadata
5092 .iter()
5093 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5094 );
5095 }
5096 }
5097
5098 #[test]
5099 fn test_arrow_writer_skip_path_in_schema() {
5100 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5101 let file_schema = Arc::new(batch_schema.clone());
5102
5103 let batch = RecordBatch::try_new(
5104 Arc::new(batch_schema),
5105 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5106 )
5107 .unwrap();
5108
5109 let skip_options = ArrowWriterOptions::new();
5111
5112 let mut buf = Vec::with_capacity(1024);
5113 let mut writer =
5114 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5115 writer.write(&batch).unwrap();
5116 writer.close().unwrap();
5117
5118 let skip_options = ArrowWriterOptions::new().with_properties(
5120 WriterProperties::builder()
5121 .set_write_path_in_schema(false)
5122 .build(),
5123 );
5124
5125 let mut buf2 = Vec::with_capacity(1024);
5126 let mut writer =
5127 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5128 .unwrap();
5129 writer.write(&batch).unwrap();
5130 writer.close().unwrap();
5131
5132 assert!(buf.len() > buf2.len());
5134 }
5135
5136 #[test]
5137 fn mismatched_schemas() {
5138 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5139 let file_schema = Arc::new(Schema::new(vec![Field::new(
5140 "temperature",
5141 DataType::Float64,
5142 false,
5143 )]));
5144
5145 let batch = RecordBatch::try_new(
5146 Arc::new(batch_schema),
5147 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5148 )
5149 .unwrap();
5150
5151 let mut buf = Vec::with_capacity(1024);
5152 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5153
5154 let err = writer.write(&batch).unwrap_err().to_string();
5155 assert_eq!(
5156 err,
5157 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5158 );
5159 }
5160
5161 #[test]
5162 fn test_page_stats_not_written_by_default() {
5163 let string_field = Field::new("a", DataType::Utf8, false);
5164 let schema = Schema::new(vec![string_field]);
5165 let raw_string_values = vec!["Blart Versenwald III"];
5166 let string_values = StringArray::from(raw_string_values.clone());
5167 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5168
5169 let props = WriterProperties::builder()
5170 .set_statistics_enabled(EnabledStatistics::Page)
5171 .set_dictionary_enabled(false)
5172 .set_encoding(Encoding::PLAIN)
5173 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5174 .build();
5175
5176 let file = roundtrip_opts(&batch, props);
5177
5178 let first_page = &file[4..];
5183 let mut prot = ThriftSliceInputProtocol::new(first_page);
5184 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5185 let stats = hdr.data_page_header.unwrap().statistics;
5186
5187 assert!(stats.is_none());
5188 }
5189
5190 #[test]
5191 fn test_page_stats_when_enabled() {
5192 let string_field = Field::new("a", DataType::Utf8, false);
5193 let schema = Schema::new(vec![string_field]);
5194 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5195 let string_values = StringArray::from(raw_string_values.clone());
5196 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5197
5198 let props = WriterProperties::builder()
5199 .set_statistics_enabled(EnabledStatistics::Page)
5200 .set_dictionary_enabled(false)
5201 .set_encoding(Encoding::PLAIN)
5202 .set_write_page_header_statistics(true)
5203 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5204 .build();
5205
5206 let file = roundtrip_opts(&batch, props);
5207
5208 let first_page = &file[4..];
5213 let mut prot = ThriftSliceInputProtocol::new(first_page);
5214 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5215 let stats = hdr.data_page_header.unwrap().statistics;
5216
5217 let stats = stats.unwrap();
5218 assert!(stats.is_max_value_exact.unwrap());
5220 assert!(stats.is_min_value_exact.unwrap());
5221 assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5222 assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5223 }
5224
5225 #[test]
5226 fn test_page_stats_truncation() {
5227 let string_field = Field::new("a", DataType::Utf8, false);
5228 let binary_field = Field::new("b", DataType::Binary, false);
5229 let schema = Schema::new(vec![string_field, binary_field]);
5230
5231 let raw_string_values = vec!["Blart Versenwald III"];
5232 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5233 let raw_binary_value_refs = raw_binary_values
5234 .iter()
5235 .map(|x| x.as_slice())
5236 .collect::<Vec<_>>();
5237
5238 let string_values = StringArray::from(raw_string_values.clone());
5239 let binary_values = BinaryArray::from(raw_binary_value_refs);
5240 let batch = RecordBatch::try_new(
5241 Arc::new(schema),
5242 vec![Arc::new(string_values), Arc::new(binary_values)],
5243 )
5244 .unwrap();
5245
5246 let props = WriterProperties::builder()
5247 .set_statistics_truncate_length(Some(2))
5248 .set_dictionary_enabled(false)
5249 .set_encoding(Encoding::PLAIN)
5250 .set_write_page_header_statistics(true)
5251 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5252 .build();
5253
5254 let file = roundtrip_opts(&batch, props);
5255
5256 let first_page = &file[4..];
5261 let mut prot = ThriftSliceInputProtocol::new(first_page);
5262 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5263 let stats = hdr.data_page_header.unwrap().statistics;
5264 assert!(stats.is_some());
5265 let stats = stats.unwrap();
5266 assert!(!stats.is_max_value_exact.unwrap());
5268 assert!(!stats.is_min_value_exact.unwrap());
5269 assert_eq!(stats.max_value.unwrap(), b"Bm");
5270 assert_eq!(stats.min_value.unwrap(), b"Bl");
5271
5272 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5274 let mut prot = ThriftSliceInputProtocol::new(second_page);
5275 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5276 let stats = hdr.data_page_header.unwrap().statistics;
5277 assert!(stats.is_some());
5278 let stats = stats.unwrap();
5279 assert!(!stats.is_max_value_exact.unwrap());
5281 assert!(!stats.is_min_value_exact.unwrap());
5282 assert_eq!(stats.max_value.unwrap(), b"Bm");
5283 assert_eq!(stats.min_value.unwrap(), b"Bl");
5284 }
5285
5286 #[test]
5287 fn test_page_encoding_statistics_roundtrip() {
5288 let batch_schema = Schema::new(vec![Field::new(
5289 "int32",
5290 arrow_schema::DataType::Int32,
5291 false,
5292 )]);
5293
5294 let batch = RecordBatch::try_new(
5295 Arc::new(batch_schema.clone()),
5296 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5297 )
5298 .unwrap();
5299
5300 let mut file: File = tempfile::tempfile().unwrap();
5301 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5302 writer.write(&batch).unwrap();
5303 let file_metadata = writer.close().unwrap();
5304
5305 assert_eq!(file_metadata.num_row_groups(), 1);
5306 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5307 assert!(
5308 file_metadata
5309 .row_group(0)
5310 .column(0)
5311 .page_encoding_stats()
5312 .is_some()
5313 );
5314 let chunk_page_stats = file_metadata
5315 .row_group(0)
5316 .column(0)
5317 .page_encoding_stats()
5318 .unwrap();
5319
5320 let options = ReadOptionsBuilder::new()
5322 .with_page_index()
5323 .with_encoding_stats_as_mask(false)
5324 .build();
5325 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5326
5327 let rowgroup = reader.get_row_group(0).expect("row group missing");
5328 assert_eq!(rowgroup.num_columns(), 1);
5329 let column = rowgroup.metadata().column(0);
5330 assert!(column.page_encoding_stats().is_some());
5331 let file_page_stats = column.page_encoding_stats().unwrap();
5332 assert_eq!(chunk_page_stats, file_page_stats);
5333 }
5334
5335 #[test]
5336 #[cfg_attr(miri, ignore)] fn test_different_dict_page_size_limit() {
5338 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5339 let schema = Arc::new(Schema::new(vec![
5340 Field::new("col0", arrow_schema::DataType::Int64, false),
5341 Field::new("col1", arrow_schema::DataType::Int64, false),
5342 ]));
5343 let batch =
5344 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5345
5346 let props = WriterProperties::builder()
5347 .set_dictionary_page_size_limit(1024 * 1024)
5348 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5349 .build();
5350 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5351 writer.write(&batch).unwrap();
5352 let data = Bytes::from(writer.into_inner().unwrap());
5353
5354 let mut metadata = ParquetMetaDataReader::new();
5355 metadata.try_parse(&data).unwrap();
5356 let metadata = metadata.finish().unwrap();
5357 let col0_meta = metadata.row_group(0).column(0);
5358 let col1_meta = metadata.row_group(0).column(1);
5359
5360 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5361 let mut reader =
5362 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5363 let page = reader.get_next_page().unwrap().unwrap();
5364 match page {
5365 Page::DictionaryPage { buf, .. } => buf.len(),
5366 _ => panic!("expected DictionaryPage"),
5367 }
5368 };
5369
5370 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5371 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5372 }
5373
5374 #[test]
5375 #[cfg_attr(miri, ignore)] fn test_arrow_writer_granular_mode_roundtrip() {
5377 let small = "tiny".to_string();
5386 let big = "x".repeat(64 * 1024);
5387 let strings: Vec<String> = (0..256)
5388 .map(|i| {
5389 if i % 16 == 0 {
5390 big.clone()
5391 } else {
5392 small.clone()
5393 }
5394 })
5395 .collect();
5396
5397 let schema = Arc::new(Schema::new(vec![Field::new(
5398 "col",
5399 ArrowDataType::Utf8,
5400 false,
5401 )]));
5402 let batch = RecordBatch::try_new(
5403 schema.clone(),
5404 vec![Arc::new(StringArray::from(strings.clone())) as _],
5405 )
5406 .unwrap();
5407
5408 let props = WriterProperties::builder()
5409 .set_dictionary_enabled(false)
5410 .set_data_page_size_limit(16 * 1024)
5411 .build();
5412 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5413 writer.write(&batch).unwrap();
5414 let data = Bytes::from(writer.into_inner().unwrap());
5415
5416 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5417 let read = reader.next().unwrap().unwrap();
5418 assert!(reader.next().is_none(), "expected one batch");
5419 let col = read
5420 .column(0)
5421 .as_any()
5422 .downcast_ref::<StringArray>()
5423 .unwrap();
5424 assert_eq!(col.len(), strings.len());
5425 for (i, expected) in strings.iter().enumerate() {
5426 assert_eq!(
5427 col.value(i),
5428 expected.as_str(),
5429 "value mismatch at index {i}"
5430 );
5431 }
5432 }
5433
5434 #[test]
5435 fn test_arrow_writer_all_null_string_column() {
5436 let num_rows = 1024;
5441 let schema = Arc::new(Schema::new(vec![Field::new(
5442 "col",
5443 ArrowDataType::Utf8,
5444 true,
5445 )]));
5446 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5447 let batch = RecordBatch::try_new(
5448 schema.clone(),
5449 vec![Arc::new(StringArray::from(nulls)) as _],
5450 )
5451 .unwrap();
5452
5453 let props = WriterProperties::builder()
5454 .set_dictionary_enabled(false)
5455 .set_data_page_size_limit(16 * 1024)
5456 .build();
5457 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5458 writer.write(&batch).unwrap();
5459 let data = Bytes::from(writer.into_inner().unwrap());
5460
5461 let mut metadata = ParquetMetaDataReader::new();
5464 metadata.try_parse(&data).unwrap();
5465 let metadata = metadata.finish().unwrap();
5466 let row_group = metadata.row_group(0);
5467 let col_meta = row_group.column(0);
5468 assert_eq!(row_group.num_rows() as usize, num_rows);
5469 if let Some(stats) = col_meta.statistics() {
5472 assert_eq!(
5473 stats.null_count_opt().unwrap_or(0) as usize,
5474 num_rows,
5475 "expected all-null column to report null_count = num_rows"
5476 );
5477 }
5478
5479 let mut reader =
5480 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5481 let mut total_values = 0u32;
5482 while let Some(page) = reader.get_next_page().unwrap() {
5483 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5484 total_values += page.num_values();
5485 }
5486 }
5487 assert_eq!(
5488 total_values as usize, num_rows,
5489 "expected every level position to be represented in some page"
5490 );
5491 }
5492
5493 struct WriteBatchesShape {
5494 num_batches: usize,
5495 rows_per_batch: usize,
5496 row_size: usize,
5497 }
5498
5499 fn write_batches(
5501 WriteBatchesShape {
5502 num_batches,
5503 rows_per_batch,
5504 row_size,
5505 }: WriteBatchesShape,
5506 props: WriterProperties,
5507 ) -> ParquetRecordBatchReaderBuilder<File> {
5508 let schema = Arc::new(Schema::new(vec![Field::new(
5509 "str",
5510 ArrowDataType::Utf8,
5511 false,
5512 )]));
5513 let file = tempfile::tempfile().unwrap();
5514 let mut writer =
5515 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5516
5517 for batch_idx in 0..num_batches {
5518 let strings: Vec<String> = (0..rows_per_batch)
5519 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5520 .collect();
5521 let array = StringArray::from(strings);
5522 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5523 writer.write(&batch).unwrap();
5524 }
5525 writer.close().unwrap();
5526 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5527 }
5528
5529 #[test]
5530 fn test_row_group_limit_none_writes_single_row_group() {
5532 let props = WriterProperties::builder()
5533 .set_max_row_group_row_count(None)
5534 .set_max_row_group_bytes(None)
5535 .build();
5536
5537 let builder = write_batches(
5538 WriteBatchesShape {
5539 num_batches: 1,
5540 rows_per_batch: 1000,
5541 row_size: 4,
5542 },
5543 props,
5544 );
5545
5546 assert_eq!(
5547 &row_group_sizes(builder.metadata()),
5548 &[1000],
5549 "With no limits, all rows should be in a single row group"
5550 );
5551 }
5552
5553 #[test]
5554 fn test_row_group_limit_rows_only() {
5556 let props = WriterProperties::builder()
5557 .set_max_row_group_row_count(Some(300))
5558 .set_max_row_group_bytes(None)
5559 .build();
5560
5561 let builder = write_batches(
5562 WriteBatchesShape {
5563 num_batches: 1,
5564 rows_per_batch: 1000,
5565 row_size: 4,
5566 },
5567 props,
5568 );
5569
5570 assert_eq!(
5571 &row_group_sizes(builder.metadata()),
5572 &[300, 300, 300, 100],
5573 "Row groups should be split by row count"
5574 );
5575 }
5576
5577 #[test]
5578 #[cfg_attr(miri, ignore)] fn test_row_group_limit_rows_only_many_splits() {
5582 let props = WriterProperties::builder()
5583 .set_max_row_group_row_count(Some(1))
5584 .set_max_row_group_bytes(None)
5585 .build();
5586
5587 let rows = 50_000;
5588 let builder = write_batches(
5589 WriteBatchesShape {
5590 num_batches: 1,
5591 rows_per_batch: rows,
5592 row_size: 4,
5593 },
5594 props,
5595 );
5596
5597 let sizes = row_group_sizes(builder.metadata());
5598 assert_eq!(sizes.len(), rows, "Every row should get its own row group");
5599 assert_eq!(
5600 sizes.iter().sum::<i64>(),
5601 rows as i64,
5602 "Total rows should be preserved"
5603 );
5604 }
5605
5606 #[test]
5607 fn test_row_group_limit_bytes_only() {
5609 let props = WriterProperties::builder()
5610 .set_max_row_group_row_count(None)
5611 .set_max_row_group_bytes(Some(3500))
5613 .build();
5614
5615 let builder = write_batches(
5616 WriteBatchesShape {
5617 num_batches: 10,
5618 rows_per_batch: 10,
5619 row_size: 100,
5620 },
5621 props,
5622 );
5623
5624 let sizes = row_group_sizes(builder.metadata());
5625
5626 assert!(
5627 sizes.len() > 1,
5628 "Should have multiple row groups due to byte limit, got {sizes:?}",
5629 );
5630
5631 let total_rows: i64 = sizes.iter().sum();
5632 assert_eq!(total_rows, 100, "Total rows should be preserved");
5633 }
5634
5635 #[test]
5636 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5638 let schema = Arc::new(Schema::new(vec![Field::new(
5639 "str",
5640 ArrowDataType::Utf8,
5641 false,
5642 )]));
5643 let file = tempfile::tempfile().unwrap();
5644
5645 let props = WriterProperties::builder()
5647 .set_max_row_group_row_count(None)
5648 .set_max_row_group_bytes(None)
5649 .build();
5650 let mut writer =
5651 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5652
5653 let first_array = StringArray::from(
5654 (0..10)
5655 .map(|i| format!("{i:0>100}"))
5656 .collect::<Vec<String>>(),
5657 );
5658 let first_batch =
5659 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5660 writer.write(&first_batch).unwrap();
5661 assert_eq!(writer.in_progress_rows(), 10);
5662
5663 writer.max_row_group_bytes = Some(1);
5666
5667 let second_array = StringArray::from(vec!["x".to_string()]);
5668 let second_batch =
5669 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5670 writer.write(&second_batch).unwrap();
5671 writer.close().unwrap();
5672 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5673
5674 assert_eq!(
5675 &row_group_sizes(builder.metadata()),
5676 &[10, 1],
5677 "The second write should flush an oversized in-progress row group first",
5678 );
5679 }
5680
5681 #[test]
5682 fn test_row_group_limit_both_row_wins_single_batch() {
5684 let props = WriterProperties::builder()
5685 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
5688
5689 let builder = write_batches(
5690 WriteBatchesShape {
5691 num_batches: 1,
5692 row_size: 4,
5693 rows_per_batch: 1000,
5694 },
5695 props,
5696 );
5697
5698 assert_eq!(
5699 &row_group_sizes(builder.metadata()),
5700 &[200, 200, 200, 200, 200],
5701 "Row limit should trigger before byte limit"
5702 );
5703 }
5704
5705 #[test]
5706 fn test_row_group_limit_both_row_wins_multiple_batches() {
5708 let props = WriterProperties::builder()
5709 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
5712
5713 let builder = write_batches(
5714 WriteBatchesShape {
5715 num_batches: 10,
5716 rows_per_batch: 10,
5717 row_size: 100,
5718 },
5719 props,
5720 );
5721
5722 assert_eq!(
5723 &row_group_sizes(builder.metadata()),
5724 &[5; 20],
5725 "Row limit should trigger before byte limit"
5726 );
5727 }
5728
5729 #[test]
5730 fn test_row_group_limit_both_bytes_wins() {
5732 let props = WriterProperties::builder()
5733 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
5736
5737 let builder = write_batches(
5738 WriteBatchesShape {
5739 num_batches: 10,
5740 rows_per_batch: 10,
5741 row_size: 100,
5742 },
5743 props,
5744 );
5745
5746 let sizes = row_group_sizes(builder.metadata());
5747
5748 assert!(
5749 sizes.len() > 1,
5750 "Byte limit should trigger before row limit, got {sizes:?}",
5751 );
5752
5753 assert!(
5754 sizes.iter().all(|&s| s < 1000),
5755 "No row group should hit the row limit"
5756 );
5757
5758 let total_rows: i64 = sizes.iter().sum();
5759 assert_eq!(total_rows, 100, "Total rows should be preserved");
5760 }
5761
5762 #[test]
5763 fn test_row_group_limit_both_apply_to_same_batch() {
5766 let props = WriterProperties::builder()
5767 .set_max_row_group_row_count(Some(15))
5768 .set_max_row_group_bytes(Some(1500))
5769 .build();
5770
5771 let builder = write_batches(
5772 WriteBatchesShape {
5773 num_batches: 2,
5774 rows_per_batch: 10,
5775 row_size: 100,
5776 },
5777 props,
5778 );
5779
5780 assert_eq!(
5781 &row_group_sizes(builder.metadata()),
5782 &[14, 6],
5783 "Byte limit should still apply to a batch the row limit already split"
5784 );
5785 }
5786
5787 #[test]
5788 fn test_row_group_limit_bytes_splits_first_batch() {
5791 let props = WriterProperties::builder()
5792 .set_max_row_group_row_count(None)
5793 .set_max_row_group_bytes(Some(3500))
5794 .build();
5795
5796 let builder = write_batches(
5797 WriteBatchesShape {
5798 num_batches: 3,
5799 rows_per_batch: 100,
5800 row_size: 100,
5801 },
5802 props,
5803 );
5804
5805 let sizes = row_group_sizes(builder.metadata());
5806 assert!(
5807 sizes.iter().all(|&s| s < 100),
5808 "No row group should hold a whole batch, got {sizes:?}",
5809 );
5810 assert_eq!(
5811 sizes.iter().sum::<i64>(),
5812 300,
5813 "Total rows should be preserved"
5814 );
5815 for row_group in builder.metadata().row_groups() {
5816 assert!(
5817 row_group.total_byte_size() <= 3500 + 100,
5818 "Row groups should stay within one row of the byte limit, got {}",
5819 row_group.total_byte_size(),
5820 );
5821 }
5822 }
5823
5824 #[test]
5825 fn arrow_column_chunk_close_mut_drops_column_index() {
5826 use crate::arrow::ArrowSchemaConverter;
5827 use crate::file::writer::SerializedFileWriter;
5828
5829 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5830 let props = Arc::new(
5831 WriterProperties::builder()
5832 .set_statistics_enabled(EnabledStatistics::Page)
5833 .build(),
5834 );
5835 let parquet_schema = ArrowSchemaConverter::new()
5836 .with_coerce_types(props.coerce_types())
5837 .convert(&schema)
5838 .unwrap();
5839
5840 let mut buf = Vec::with_capacity(1024);
5841 let mut writer =
5842 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5843 .unwrap();
5844
5845 let factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap();
5846 let mut col_writers = factory.create_column_writers(0).unwrap();
5847 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5848 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5849 col_writers[0].write(&leaves).unwrap();
5850 }
5851 let mut chunk = col_writers.pop().unwrap().close().unwrap();
5852
5853 assert!(
5855 chunk.close().column_index.is_some(),
5856 "EnabledStatistics::Page should produce a column_index"
5857 );
5858
5859 chunk.close_mut().column_index = None;
5861 assert!(chunk.close().column_index.is_none());
5862
5863 let mut rg = writer.next_row_group().unwrap();
5864 chunk.append_to_row_group(&mut rg).unwrap();
5865 rg.close().unwrap();
5866 let file_meta = writer.close().unwrap();
5867
5868 let cc = file_meta.row_group(0).column(0);
5871 assert!(cc.column_index_range().is_none());
5872 }
5873
5874 #[test]
5875 #[cfg_attr(miri, ignore)] fn test_number_distinct_values_exact_count() {
5877 let cardinality = 50u32;
5880 let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
5881 if i % 7 == 0 {
5882 None
5883 } else {
5884 Some((i % cardinality) as i32)
5885 }
5886 })));
5887 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
5888 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
5889
5890 let props = WriterProperties::builder()
5891 .set_write_row_group_number_distinct_values(true)
5892 .build();
5893 let mut buf = Vec::new();
5894 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
5895 writer.write(&batch).unwrap();
5896 let metadata = writer.close().unwrap();
5897
5898 let count = metadata
5899 .row_group(0)
5900 .column(0)
5901 .statistics()
5902 .and_then(|s| s.distinct_count_opt())
5903 .expect("distinct_count should be set");
5904 assert_eq!(count, cardinality as u64);
5906 }
5907
5908 #[test]
5909 fn test_number_distinct_values_view_types() {
5910 let cardinality = 5u32;
5913 let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
5914
5915 let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
5916 if i % 4 == 0 {
5917 None
5918 } else {
5919 Some(distinct_strings[(i % cardinality) as usize])
5920 }
5921 })));
5922
5923 let schema = Arc::new(Schema::new(vec![Field::new(
5924 "string_view_col",
5925 DataType::Utf8View,
5926 true,
5927 )]));
5928 let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
5929
5930 let props = WriterProperties::builder()
5931 .set_write_row_group_number_distinct_values(true)
5932 .build();
5933 let mut parquet_bytes = Vec::new();
5934 let mut writer =
5935 ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
5936 writer.write(&batch).unwrap();
5937 let metadata = writer.close().unwrap();
5938
5939 let distinct_count = metadata
5940 .row_group(0)
5941 .column(0)
5942 .statistics()
5943 .and_then(|s| s.distinct_count_opt())
5944 .expect("distinct_count should be set for Utf8View column");
5945 assert_eq!(distinct_count, cardinality as u64);
5946 }
5947
5948 #[test]
5949 fn test_number_distinct_values_not_written_by_default() {
5950 let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
5951 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
5952 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
5953
5954 let mut buf = Vec::new();
5955 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5956 writer.write(&batch).unwrap();
5957 let metadata = writer.close().unwrap();
5958
5959 let count = metadata
5960 .row_group(0)
5961 .column(0)
5962 .statistics()
5963 .and_then(|s| s.distinct_count_opt());
5964 assert!(count.is_none());
5965 }
5966
5967 #[test]
5968 fn test_dictionary_ndv_single_batch() {
5969 let keys = Int32Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
5973 let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "bird"]));
5974 let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
5975
5976 let schema = Arc::new(Schema::new(vec![Field::new(
5977 "x",
5978 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
5979 false,
5980 )]));
5981 let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
5982
5983 let props = WriterProperties::builder()
5984 .set_write_row_group_number_distinct_values(true)
5985 .build();
5986 let mut buf = Vec::new();
5987 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
5988 writer.write(&batch).unwrap();
5989 let metadata = writer.close().unwrap();
5990
5991 let count = metadata
5992 .row_group(0)
5993 .column(0)
5994 .statistics()
5995 .and_then(|s| s.distinct_count_opt())
5996 .expect("distinct_count should be set");
5997 assert_eq!(count, 3);
5998 }
5999
6000 #[test]
6001 fn test_dictionary_ndv_excludes_unreferenced_values() {
6002 let keys = Int32Array::from(vec![0, 1, 0, 1]);
6005 let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "unreferenced"]));
6006 let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6007
6008 let schema = Arc::new(Schema::new(vec![Field::new(
6009 "x",
6010 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6011 false,
6012 )]));
6013 let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6014
6015 let props = WriterProperties::builder()
6016 .set_write_row_group_number_distinct_values(true)
6017 .build();
6018 let mut buf = Vec::new();
6019 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6020 writer.write(&batch).unwrap();
6021 let metadata = writer.close().unwrap();
6022
6023 let count = metadata
6024 .row_group(0)
6025 .column(0)
6026 .statistics()
6027 .and_then(|s| s.distinct_count_opt())
6028 .expect("distinct_count should be set");
6029 assert_eq!(
6030 count, 2,
6031 "unreferenced dictionary values must not count toward NDV"
6032 );
6033 }
6034
6035 #[test]
6036 fn test_dictionary_ndv_across_batches_regression() {
6037 let make_dict_batch = |a: &str, b: &str| -> RecordBatch {
6039 let keys = Int32Array::from(vec![0, 1, 0, 1]);
6040 let values: ArrayRef = Arc::new(StringArray::from(vec![a, b]));
6041 let dict: ArrayRef =
6042 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6043 let schema = Arc::new(Schema::new(vec![Field::new(
6044 "x",
6045 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6046 false,
6047 )]));
6048 RecordBatch::try_new(schema, vec![dict]).unwrap()
6049 };
6050
6051 let batch1 = make_dict_batch("cat", "dog");
6054 let batch2 = make_dict_batch("fish", "cat");
6055
6056 let props = WriterProperties::builder()
6057 .set_write_row_group_number_distinct_values(true)
6058 .build();
6059 let mut buf = Vec::new();
6060 let mut writer = ArrowWriter::try_new(&mut buf, batch1.schema(), Some(props)).unwrap();
6061 writer.write(&batch1).unwrap();
6062 writer.write(&batch2).unwrap();
6063 let metadata = writer.close().unwrap();
6064
6065 let count = metadata
6066 .row_group(0)
6067 .column(0)
6068 .statistics()
6069 .and_then(|s| s.distinct_count_opt())
6070 .expect("distinct_count should be set");
6071 assert_eq!(
6072 count, 3,
6073 "NDV should count distinct values, not distinct key indices"
6074 );
6075 }
6076
6077 #[test]
6078 fn ree_struct_with_ree_child() {
6079 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6082
6083 let col_a: ArrayRef = Arc::new(
6084 RunArray::try_new(
6085 &run_ends,
6086 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6087 )
6088 .unwrap(),
6089 );
6090 let col_b: ArrayRef = Arc::new(
6091 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6092 );
6093
6094 let struct_array: ArrayRef = Arc::new(StructArray::new(
6095 Fields::from(vec![
6096 Field::new("a", col_a.data_type().clone(), true),
6097 Field::new("b", col_b.data_type().clone(), true),
6098 ]),
6099 vec![col_a, col_b],
6100 None,
6101 ));
6102
6103 let schema = Arc::new(Schema::new(vec![Field::new(
6104 "row",
6105 struct_array.data_type().clone(),
6106 true,
6107 )]));
6108 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6109
6110 let mut buf = Vec::new();
6111 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6112 writer.write(&batch).unwrap();
6113 let metadata = writer.close().unwrap();
6114
6115 let parquet_schema = metadata.file_metadata().schema_descr();
6116 assert_eq!(parquet_schema.num_columns(), 2);
6117 assert_eq!(
6118 parquet_schema.column(0).physical_type(),
6119 crate::basic::Type::BYTE_ARRAY
6120 );
6121 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6122 assert_eq!(
6123 parquet_schema.column(1).physical_type(),
6124 crate::basic::Type::INT32
6125 );
6126 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6127 }
6128
6129 #[test]
6136 fn test_decimal_byte_array_statistics_signed_comparison() {
6137 use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
6138 use crate::data_type::{ByteArray, ByteArrayType};
6139 use crate::schema::types::Type as SchemaType;
6140
6141 const VALUES: [&[u8]; 3] = [&[0xffu8], &[0x00], &[0x01]];
6143
6144 fn new_file(physical: PhysicalType) -> SerializedFileWriter<Vec<u8>> {
6145 let mut field = SchemaType::primitive_type_builder("value", physical)
6146 .with_repetition(Repetition::REQUIRED)
6147 .with_logical_type(Some(LogicalType::decimal(0, 2)))
6148 .with_precision(2)
6149 .with_scale(0);
6150 if physical == PhysicalType::FIXED_LEN_BYTE_ARRAY {
6151 field = field.with_length(1);
6152 }
6153 let root = SchemaType::group_type_builder("root")
6154 .with_fields(vec![Arc::new(field.build().unwrap())])
6155 .build()
6156 .unwrap();
6157 SerializedFileWriter::new(Vec::new(), Arc::new(root), Default::default()).unwrap()
6158 }
6159
6160 fn read_minmax(file: Vec<u8>) -> (i8, i8) {
6161 let reader = SerializedFileReader::new(Bytes::from(file)).unwrap();
6162 let metadata = reader.metadata();
6163 assert_eq!(metadata.file_metadata().num_rows(), 3);
6164 let stats = metadata
6165 .row_group(0)
6166 .column(0)
6167 .statistics()
6168 .expect("missing statistics");
6169 (
6170 stats.min_bytes_opt().unwrap()[0] as i8,
6171 stats.max_bytes_opt().unwrap()[0] as i8,
6172 )
6173 }
6174
6175 fn direct_parquet() -> Vec<u8> {
6176 let mut file = new_file(PhysicalType::BYTE_ARRAY);
6177 let mut group = file.next_row_group().unwrap();
6178 let mut column = group.next_column().unwrap().unwrap();
6179 let values = VALUES
6180 .iter()
6181 .map(|v| ByteArray::from(v.to_vec()))
6182 .collect::<Vec<_>>();
6183 column
6184 .typed::<ByteArrayType>()
6185 .write_batch(&values, None, None)
6186 .unwrap();
6187 column.close().unwrap();
6188 group.close().unwrap();
6189 file.into_inner().unwrap()
6190 }
6191
6192 fn arrow_bounds(mut file: SerializedFileWriter<Vec<u8>>, array: ArrayRef) -> Vec<u8> {
6193 let field = Field::new("value", array.data_type().clone(), false);
6194 let schema = Arc::new(Schema::new(vec![field.clone()]));
6195 let factory = ArrowRowGroupWriterFactory::try_new(&file, schema).unwrap();
6196 let mut writer: ArrowColumnWriter = factory.create_column_writers(0).unwrap().remove(0);
6197 let leaf = compute_leaves(&field, &array).unwrap().remove(0);
6198 writer.write(&leaf).unwrap();
6199 let mut group = file.next_row_group().unwrap();
6200 writer
6201 .close()
6202 .unwrap()
6203 .append_to_row_group(&mut group)
6204 .unwrap();
6205 group.close().unwrap();
6206 file.into_inner().unwrap()
6207 }
6208
6209 fn arrow_binary() -> Vec<u8> {
6210 let file = new_file(PhysicalType::BYTE_ARRAY);
6211 let array: ArrayRef = Arc::new(BinaryArray::from_vec(VALUES.to_vec()));
6212 arrow_bounds(file, array)
6213 }
6214
6215 fn arrow_fixed_size_binary() -> Vec<u8> {
6216 let file = new_file(PhysicalType::FIXED_LEN_BYTE_ARRAY);
6217 let array: ArrayRef =
6218 Arc::new(FixedSizeBinaryArray::try_from_iter(VALUES.into_iter()).unwrap());
6219 arrow_bounds(file, array)
6220 }
6221
6222 let direct = read_minmax(direct_parquet());
6223 let fixed = read_minmax(arrow_fixed_size_binary());
6224 let binary = read_minmax(arrow_binary());
6225
6226 assert_eq!(direct, (-1, 1), "incorrect direct Parquet stats");
6227 assert_eq!(fixed, (-1, 1), "incorrect Arrow FixedSizeBinaryArray stats");
6228 assert_eq!(binary, (-1, 1), "incorrect Arrow BinaryArray stats");
6229 }
6230}