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> {
1806 type Output;
1807
1808 fn slice(self, values: &[T]) -> Result<Self::Output>;
1810
1811 fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<Self::Output>;
1813}
1814
1815pub(crate) fn visit_i32<V: PrimitiveVisitor<i32>>(
1817 column: &dyn arrow_array::Array,
1818 visitor: V,
1819) -> Result<V::Output> {
1820 match column.data_type() {
1822 ArrowDataType::Null => visitor.slice(Int32Array::new_null(column.len()).values()),
1823 ArrowDataType::Int8 => {
1824 visitor.map(column.as_primitive::<Int8Type>().values(), |x| x as i32)
1825 }
1826 ArrowDataType::Int16 => {
1827 visitor.map(column.as_primitive::<Int16Type>().values(), |x| x as i32)
1828 }
1829 ArrowDataType::Int32 => visitor.slice(column.as_primitive::<Int32Type>().values()),
1830 ArrowDataType::UInt8 => {
1831 visitor.map(column.as_primitive::<UInt8Type>().values(), |x| x as i32)
1832 }
1833 ArrowDataType::UInt16 => {
1834 visitor.map(column.as_primitive::<UInt16Type>().values(), |x| x as i32)
1835 }
1836 ArrowDataType::UInt32 => {
1837 let values = column.as_primitive::<UInt32Type>().values();
1840 visitor.slice(values.inner().typed_data())
1841 }
1842 ArrowDataType::Date32 => visitor.slice(column.as_primitive::<Date32Type>().values()),
1843 ArrowDataType::Time32(TimeUnit::Second) => {
1844 visitor.slice(column.as_primitive::<Time32SecondType>().values())
1845 }
1846 ArrowDataType::Time32(TimeUnit::Millisecond) => {
1847 visitor.slice(column.as_primitive::<Time32MillisecondType>().values())
1848 }
1849 ArrowDataType::Date64 => {
1850 visitor.map(column.as_primitive::<Date64Type>().values(), |x| {
1852 (x / 86_400_000) as i32
1853 })
1854 }
1855 ArrowDataType::Decimal32(_, _) => {
1856 visitor.slice(column.as_primitive::<Decimal32Type>().values())
1857 }
1858 ArrowDataType::Decimal64(_, _) => {
1859 visitor.map(column.as_primitive::<Decimal64Type>().values(), |v| {
1861 v as i32
1862 })
1863 }
1864 ArrowDataType::Decimal128(_, _) => {
1865 visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1867 v as i32
1868 })
1869 }
1870 ArrowDataType::Decimal256(_, _) => {
1871 visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1873 v.as_i128() as i32
1874 })
1875 }
1876 d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1877 }
1878}
1879
1880pub(crate) fn visit_i64<V: PrimitiveVisitor<i64>>(
1882 column: &dyn arrow_array::Array,
1883 visitor: V,
1884) -> Result<V::Output> {
1885 match column.data_type() {
1887 ArrowDataType::Date64 => visitor.slice(column.as_primitive::<Date64Type>().values()),
1888 ArrowDataType::Int64 => visitor.slice(column.as_primitive::<Int64Type>().values()),
1889 ArrowDataType::UInt64 => {
1890 let values = column.as_primitive::<UInt64Type>().values();
1893 visitor.slice(values.inner().typed_data())
1894 }
1895 ArrowDataType::Time64(TimeUnit::Microsecond) => {
1896 visitor.slice(column.as_primitive::<Time64MicrosecondType>().values())
1897 }
1898 ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1899 visitor.slice(column.as_primitive::<Time64NanosecondType>().values())
1900 }
1901 ArrowDataType::Timestamp(unit, _) => match unit {
1902 TimeUnit::Second => {
1903 visitor.slice(column.as_primitive::<TimestampSecondType>().values())
1904 }
1905 TimeUnit::Millisecond => {
1906 visitor.slice(column.as_primitive::<TimestampMillisecondType>().values())
1907 }
1908 TimeUnit::Microsecond => {
1909 visitor.slice(column.as_primitive::<TimestampMicrosecondType>().values())
1910 }
1911 TimeUnit::Nanosecond => {
1912 visitor.slice(column.as_primitive::<TimestampNanosecondType>().values())
1913 }
1914 },
1915 ArrowDataType::Duration(unit) => match unit {
1916 TimeUnit::Second => visitor.slice(column.as_primitive::<DurationSecondType>().values()),
1917 TimeUnit::Millisecond => {
1918 visitor.slice(column.as_primitive::<DurationMillisecondType>().values())
1919 }
1920 TimeUnit::Microsecond => {
1921 visitor.slice(column.as_primitive::<DurationMicrosecondType>().values())
1922 }
1923 TimeUnit::Nanosecond => {
1924 visitor.slice(column.as_primitive::<DurationNanosecondType>().values())
1925 }
1926 },
1927 ArrowDataType::Decimal32(_, _) => {
1928 visitor.map(column.as_primitive::<Decimal32Type>().values(), |v| {
1931 v as i64
1932 })
1933 }
1934 ArrowDataType::Decimal64(_, _) => {
1935 visitor.slice(column.as_primitive::<Decimal64Type>().values())
1936 }
1937 ArrowDataType::Decimal128(_, _) => {
1938 visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1940 v as i64
1941 })
1942 }
1943 ArrowDataType::Decimal256(_, _) => {
1944 visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1946 v.as_i128() as i64
1947 })
1948 }
1949 d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1950 }
1951}
1952
1953struct WritePrimitive<'a, 'b, E: ColumnValueEncoder> {
1955 writer: &'a mut GenericColumnWriter<'b, E>,
1956 levels: &'a ArrayLevels,
1957}
1958
1959impl<T, E> PrimitiveVisitor<T> for WritePrimitive<'_, '_, E>
1960where
1961 E: ColumnValueEncoder<Values = [T]>,
1962{
1963 type Output = usize;
1964
1965 fn slice(self, values: &[T]) -> Result<usize> {
1966 write_primitive(self.writer, values, self.levels)
1967 }
1968
1969 fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<usize> {
1970 let values: Vec<T> = values.iter().map(|v| convert(*v)).collect();
1971 write_primitive(self.writer, &values, self.levels)
1972 }
1973}
1974
1975pub(crate) trait FixedLenByteArrayVisitor {
1981 type Output;
1982
1983 fn decimal_length(&self, precision: u8, native_length: usize) -> usize;
1988
1989 fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<Self::Output>;
1994}
1995
1996pub(crate) fn visit_fixed_len_byte_array<V: FixedLenByteArrayVisitor>(
1999 column: &dyn arrow_array::Array,
2000 visitor: V,
2001) -> Result<V::Output> {
2002 match column.data_type() {
2003 ArrowDataType::Interval(interval_unit) => match interval_unit {
2004 IntervalUnit::YearMonth => {
2005 let array = column.as_primitive::<IntervalYearMonthType>();
2008 visitor.values(12, |i, value| {
2009 value[0..4].copy_from_slice(&array.value(i).to_le_bytes());
2010 })
2011 }
2012 IntervalUnit::DayTime => {
2013 let array = column.as_primitive::<IntervalDayTimeType>();
2017 visitor.values(12, |i, value| {
2018 let interval = array.value(i);
2019 value[4..8].copy_from_slice(&interval.days.to_le_bytes());
2020 value[8..12].copy_from_slice(&interval.milliseconds.to_le_bytes());
2021 })
2022 }
2023 IntervalUnit::MonthDayNano => Err(ParquetError::NYI(format!(
2024 "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
2025 ))),
2026 },
2027 ArrowDataType::FixedSizeBinary(_) => {
2028 let array = column.as_fixed_size_binary();
2029 visitor.values(array.value_size(), |i, value| {
2030 value.copy_from_slice(array.value(i));
2031 })
2032 }
2033 ArrowDataType::Decimal32(_, _) => {
2034 visit_decimal(column.as_primitive::<Decimal32Type>(), visitor)
2035 }
2036 ArrowDataType::Decimal64(_, _) => {
2037 visit_decimal(column.as_primitive::<Decimal64Type>(), visitor)
2038 }
2039 ArrowDataType::Decimal128(_, _) => {
2040 visit_decimal(column.as_primitive::<Decimal128Type>(), visitor)
2041 }
2042 ArrowDataType::Decimal256(_, _) => {
2043 visit_decimal(column.as_primitive::<Decimal256Type>(), visitor)
2044 }
2045 ArrowDataType::Float16 => {
2046 let array = column.as_primitive::<Float16Type>();
2047 visitor.values(2, |i, value| {
2048 value.copy_from_slice(&array.value(i).to_le_bytes());
2049 })
2050 }
2051 _ => Err(ParquetError::NYI(
2052 "Attempting to write an Arrow type that is not yet implemented".to_string(),
2053 )),
2054 }
2055}
2056
2057fn visit_decimal<T: NativeDecimalType, V: FixedLenByteArrayVisitor>(
2064 array: &PrimitiveArray<T>,
2065 visitor: V,
2066) -> Result<V::Output> {
2067 let length = visitor.decimal_length(array.precision(), T::BYTE_LENGTH);
2068 let be_bytes = |i| T::to_be_bytes(array.value(i));
2069 match length.checked_sub(T::BYTE_LENGTH) {
2071 Some(0) => visitor.values(length, |i, value| {
2073 value.copy_from_slice(be_bytes(i).as_ref());
2074 }),
2075 Some(extension) => visitor.values(length, |i, value| {
2076 let bytes = be_bytes(i);
2077 value[..extension].fill(sign_byte(bytes.as_ref()));
2078 value[extension..].copy_from_slice(bytes.as_ref());
2079 }),
2080 None => visitor.values(length, |i, value| {
2081 value.copy_from_slice(&be_bytes(i).as_ref()[T::BYTE_LENGTH - length..]);
2082 }),
2083 }
2084}
2085
2086fn sign_byte(bytes: &[u8]) -> u8 {
2088 if bytes.first().is_some_and(|b| (*b as i8) < 0) {
2089 0xFF
2090 } else {
2091 0
2092 }
2093}
2094
2095struct WriteFixedLenByteArray<'a, 'b> {
2097 writer: &'a mut ColumnWriterImpl<'b, FixedLenByteArrayType>,
2098 levels: &'a ArrayLevels,
2099}
2100
2101impl FixedLenByteArrayVisitor for WriteFixedLenByteArray<'_, '_> {
2102 type Output = usize;
2103
2104 fn decimal_length(&self, precision: u8, native_length: usize) -> usize {
2105 let length = decimal_length_from_precision(precision);
2106 assert!(length <= native_length);
2108 length
2109 }
2110
2111 fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<usize> {
2112 let indices = self.levels.non_null_indices().iter().copied();
2113 let values = chunk_array_slice(length, indices, fill);
2114 self.writer.write_batch_internal(
2115 values.as_slice(),
2116 None,
2117 self.levels.def_level_data().as_ref(),
2118 self.levels.rep_level_data().as_ref(),
2119 None,
2120 None,
2121 None,
2122 )
2123 }
2124}
2125
2126fn write_primitive<E: ColumnValueEncoder>(
2127 writer: &mut GenericColumnWriter<E>,
2128 values: &E::Values,
2129 levels: &ArrayLevels,
2130) -> Result<usize> {
2131 writer.write_batch_internal(
2132 values,
2133 Some(levels.non_null_indices()),
2134 levels.def_level_data().as_ref(),
2135 levels.rep_level_data().as_ref(),
2136 None,
2137 None,
2138 None,
2139 )
2140}
2141
2142fn get_bool_array_slice(
2143 array: &arrow_array::BooleanArray,
2144 indices: impl ExactSizeIterator<Item = usize>,
2145) -> Vec<bool> {
2146 let mut values = Vec::with_capacity(indices.len());
2147 for i in indices {
2148 values.push(array.value(i))
2149 }
2150 values
2151}
2152
2153trait NativeDecimalType: DecimalType {
2154 type NativeBytes: AsRef<[u8]>;
2155
2156 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
2157}
2158impl NativeDecimalType for Decimal32Type {
2159 type NativeBytes = [u8; Self::BYTE_LENGTH];
2160
2161 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2162 value.to_be_bytes()
2163 }
2164}
2165impl NativeDecimalType for Decimal64Type {
2166 type NativeBytes = [u8; Self::BYTE_LENGTH];
2167
2168 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2169 value.to_be_bytes()
2170 }
2171}
2172impl NativeDecimalType for Decimal128Type {
2173 type NativeBytes = [u8; Self::BYTE_LENGTH];
2174
2175 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2176 value.to_be_bytes()
2177 }
2178}
2179impl NativeDecimalType for Decimal256Type {
2180 type NativeBytes = [u8; Self::BYTE_LENGTH];
2181
2182 fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2183 value.to_be_bytes()
2184 }
2185}
2186
2187#[inline]
2188fn chunk_array_slice(
2189 chunk_size: usize,
2190 indices: impl ExactSizeIterator<Item = usize>,
2191 writer: impl Fn(usize, &mut [u8]),
2192) -> Vec<FixedLenByteArray> {
2193 let capacity = indices.len() * chunk_size;
2194 let mut arena = vec![0; capacity];
2197 for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
2198 writer(i, chunk);
2199 }
2200 chunk_contiguous_vec(arena, chunk_size)
2201}
2202
2203fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
2204 let mut values = Vec::with_capacity(arena.len() / chunk_size);
2205 let mut arena = Bytes::from(arena);
2206 while arena.len() >= chunk_size {
2207 let slice = arena.split_to(chunk_size);
2208 values.push(FixedLenByteArray::from(ByteArray::from(slice)));
2209 }
2210 values
2211}
2212
2213#[inline]
2215fn hash_bytes(bytes: &[u8]) -> u64 {
2216 twox_hash::XxHash64::oneshot(0, bytes)
2217}
2218
2219fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
2221 use ArrowDataType::*;
2222 match dt {
2223 Int8 | UInt8 => Some(1),
2224 Int16 | UInt16 | Float16 => Some(2),
2225 Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) => Some(4),
2226 Int64
2227 | UInt64
2228 | Float64
2229 | Date64
2230 | Time64(_)
2231 | Timestamp(_, _)
2232 | Duration(_)
2233 | Decimal64(_, _) => Some(8),
2234 Interval(IntervalUnit::YearMonth) => Some(4),
2235 Interval(IntervalUnit::DayTime) => Some(8),
2236 Interval(IntervalUnit::MonthDayNano) => Some(16),
2237 Decimal128(_, _) => Some(16),
2238 Decimal256(_, _) => Some(32),
2239 _ => None,
2240 }
2241}
2242
2243fn update_distinct_values_seen(
2249 array: &dyn arrow_array::Array,
2250 non_null_indices: &[usize],
2251 seen: &mut DistinctValuesSet,
2252) {
2253 let data = array.to_data();
2254 let offset = data.offset();
2255
2256 match array.data_type() {
2257 ArrowDataType::Boolean => {
2258 let arr = array
2259 .as_any()
2260 .downcast_ref::<arrow_array::BooleanArray>()
2261 .unwrap();
2262 for &row in non_null_indices {
2263 seen.insert(arr.value(row) as u64);
2264 }
2265 }
2266 ArrowDataType::Utf8 | ArrowDataType::Binary => {
2267 let offsets = data.buffers()[0].typed_data::<i32>();
2268 let values = data.buffers()[1].as_slice();
2269 for &row in non_null_indices {
2270 let start = offsets[offset + row] as usize;
2271 let end = offsets[offset + row + 1] as usize;
2272 seen.insert(hash_bytes(&values[start..end]));
2273 }
2274 }
2275 ArrowDataType::LargeUtf8 | ArrowDataType::LargeBinary => {
2276 let offsets = data.buffers()[0].typed_data::<i64>();
2277 let values = data.buffers()[1].as_slice();
2278 for &row in non_null_indices {
2279 let start = offsets[offset + row] as usize;
2280 let end = offsets[offset + row + 1] as usize;
2281 seen.insert(hash_bytes(&values[start..end]));
2282 }
2283 }
2284 ArrowDataType::FixedSizeBinary(byte_width) => {
2285 let byte_width = *byte_width as usize;
2286 let buffer = data.buffers()[0].as_slice();
2287 for &row in non_null_indices {
2288 let start = (offset + row) * byte_width;
2289 seen.insert(hash_bytes(&buffer[start..start + byte_width]));
2290 }
2291 }
2292 ArrowDataType::Utf8View => {
2293 let string_view_array = array.as_string_view();
2294 for &row in non_null_indices {
2295 seen.insert(hash_bytes(string_view_array.value(row).as_bytes()));
2296 }
2297 }
2298 ArrowDataType::BinaryView => {
2299 let binary_view_array = array.as_binary_view();
2300 for &row in non_null_indices {
2301 seen.insert(hash_bytes(binary_view_array.value(row)));
2302 }
2303 }
2304 data_type => {
2305 if let Some(width) = fixed_byte_width(data_type) {
2306 let buffer = data.buffers()[0].as_slice();
2307 for &row in non_null_indices {
2308 let pos = (offset + row) * width;
2309 seen.insert(hash_bytes(&buffer[pos..pos + width]));
2310 }
2311 }
2312 }
2314 }
2315}
2316
2317#[cfg(test)]
2319use crate as parquet_crate;
2320
2321#[cfg(test)]
2322#[path = "../../../tests/arrow_writer/roundtrip_helpers.rs"]
2323mod roundtrip_helpers;
2324
2325#[cfg(test)]
2326mod tests {
2327 use super::roundtrip_helpers::{
2328 RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip_opts,
2329 roundtrip_opts_with_array_validation,
2330 };
2331 use super::*;
2332 use std::cmp::Ordering;
2333 use std::collections::HashMap;
2334
2335 use std::fs::File;
2336
2337 use crate::arrow::ARROW_SCHEMA_META_KEY;
2338 use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
2339 use crate::column::page::{Page, PageReader};
2340 use crate::file::metadata::thrift::PageHeader;
2341 use crate::file::page_index::column_index::ColumnIndexMetaData;
2342 use crate::file::reader::SerializedPageReader;
2343 use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
2344 use crate::schema::types::ColumnPath;
2345 use arrow::datatypes::{DataType, Schema};
2346 use arrow::error::Result as ArrowResult;
2347 use arrow::util::data_gen::create_random_array;
2348 use arrow::util::pretty::pretty_format_batches;
2349 use arrow::{array::*, buffer::Buffer, compute::cast};
2350 use arrow_buffer::{IntervalMonthDayNano, NullBuffer, OffsetBuffer};
2351 use arrow_schema::Fields;
2352 use arrow_select::concat::concat;
2353
2354 use crate::basic::{Encoding, EncodingMask};
2355 use crate::data_type::AsBytes;
2356 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
2357 use crate::file::properties::{
2358 BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
2359 };
2360 use crate::file::serialized_reader::ReadOptionsBuilder;
2361 use crate::file::{
2362 reader::{FileReader, SerializedFileReader},
2363 statistics::Statistics,
2364 };
2365
2366 #[derive(Debug, Default)]
2371 struct RecordingPageStore {
2372 next: u64,
2373 blobs: HashMap<u64, Bytes>,
2374 puts: Arc<std::sync::atomic::AtomicUsize>,
2375 }
2376
2377 impl PageStore for RecordingPageStore {
2378 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2379 let id = 100 + self.next * 7;
2381 self.next += 1;
2382 self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2383 self.blobs.insert(id, value);
2384 Ok(PageKey::new(id))
2385 }
2386
2387 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2388 self.blobs
2389 .remove(&key.get())
2390 .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2391 }
2392 }
2393
2394 #[derive(Debug)]
2395 struct RecordingPageStoreFactory {
2396 puts: Arc<std::sync::atomic::AtomicUsize>,
2397 }
2398
2399 impl PageStoreFactory for RecordingPageStoreFactory {
2400 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2401 Ok(Box::new(RecordingPageStore {
2402 puts: self.puts.clone(),
2403 ..Default::default()
2404 }))
2405 }
2406 }
2407
2408 #[test]
2412 fn custom_page_store_is_byte_identical_to_default() {
2413 let schema = Arc::new(Schema::new(vec![
2414 Field::new("i", DataType::Int32, true),
2415 Field::new("s", DataType::Utf8, true),
2417 ]));
2418 let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2419 let s = StringArray::from(vec![
2420 Some("a"),
2421 Some("bb"),
2422 Some("a"),
2423 None,
2424 Some("bb"),
2425 Some("ccc"),
2426 ]);
2427 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2428
2429 let props = WriterProperties::builder()
2432 .set_max_row_group_row_count(Some(3))
2433 .build();
2434
2435 let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2436 let mut buffer = Vec::new();
2437 let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2438 if let Some(factory) = factory {
2439 opts = opts.with_page_store_factory(factory);
2440 }
2441 let mut writer =
2442 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2443 writer.write(&batch).unwrap();
2444 writer.close().unwrap();
2445 buffer
2446 };
2447
2448 let default_bytes = write(None);
2449
2450 let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2451 let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2452 puts: puts.clone(),
2453 })));
2454
2455 assert!(
2456 puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2457 "custom PageStore was never written to"
2458 );
2459 assert_eq!(
2460 default_bytes, custom_bytes,
2461 "a custom PageStore must produce byte-identical output to the default"
2462 );
2463 }
2464
2465 #[test]
2471 #[cfg_attr(miri, ignore)] fn dictionary_column_round_trips_with_offset_index_disabled() {
2473 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2474
2475 let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2478 let array = Int32Array::from(values.clone());
2479 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2480
2481 let props = WriterProperties::builder()
2482 .set_offset_index_disabled(true)
2483 .set_data_page_row_count_limit(4096)
2484 .build();
2485 let opts = ArrowWriterOptions::new().with_properties(props);
2486
2487 let mut buffer = Vec::new();
2488 let mut writer =
2489 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2490 writer.write(&batch).unwrap();
2491 writer.close().unwrap();
2492
2493 let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2494 let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2495 let read_values: Vec<Option<i32>> = read
2496 .iter()
2497 .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2498 .collect();
2499 assert_eq!(read_values, values);
2500 }
2501
2502 #[test]
2507 fn dictionary_page_is_routed_through_the_store() {
2508 #[derive(Debug, Default)]
2510 struct SizeRecordingPageStore {
2511 blobs: Vec<Bytes>,
2512 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2513 }
2514 impl PageStore for SizeRecordingPageStore {
2515 fn put(&mut self, value: Bytes) -> Result<PageKey> {
2516 self.bytes_put
2517 .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2518 let key = PageKey::new(self.blobs.len() as u64);
2519 self.blobs.push(value);
2520 Ok(key)
2521 }
2522 fn take(&mut self, key: PageKey) -> Result<Bytes> {
2523 Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2524 }
2525 }
2526 #[derive(Debug)]
2527 struct Factory {
2528 bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2529 }
2530 impl PageStoreFactory for Factory {
2531 fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2532 Ok(Box::new(SizeRecordingPageStore {
2533 bytes_put: self.bytes_put.clone(),
2534 ..Default::default()
2535 }))
2536 }
2537 }
2538
2539 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2540 let values: Vec<&str> = (0..2048)
2543 .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2544 .collect();
2545 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2546 .unwrap();
2547
2548 let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2549 let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2550 bytes_put: bytes_put.clone(),
2551 }));
2552
2553 let mut buffer = Vec::new();
2556 let mut writer =
2557 ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2558 writer.write(&batch).unwrap();
2559 writer.close().unwrap();
2560
2561 let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2562 let column = reader.metadata().row_group(0).column(0);
2563 assert!(
2564 column.dictionary_page_offset().is_some(),
2565 "expected the column to be dictionary-encoded"
2566 );
2567
2568 assert_eq!(
2572 bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2573 column.compressed_size(),
2574 "the dictionary page must pass through the store like any other page"
2575 );
2576 }
2577
2578 #[test]
2583 fn dictionary_page_encoding_stats_lists_dictionary_first() {
2584 use crate::basic::PageType;
2585 use crate::file::metadata::ParquetMetaDataOptions;
2586
2587 let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2590 let values: Vec<Option<i32>> = (0..4096).map(|i| Some(i % 4)).collect();
2591 let array = Int32Array::from(values);
2592 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2593
2594 let mut buffer = Vec::new();
2595 let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None).unwrap();
2596 writer.write(&batch).unwrap();
2597 writer.close().unwrap();
2598 let buffer = Bytes::from(buffer);
2599
2600 let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2604 let metadata = ParquetMetaDataReader::new()
2605 .with_metadata_options(Some(options))
2606 .parse_and_finish(&buffer)
2607 .unwrap();
2608
2609 let col = metadata.row_group(0).column(0);
2610 assert!(
2611 col.dictionary_page_offset().is_some(),
2612 "column should be dictionary encoded"
2613 );
2614 let stats = col
2615 .page_encoding_stats()
2616 .expect("full (non-mask) page encoding stats should be present");
2617
2618 assert_eq!(
2621 stats[0].page_type,
2622 PageType::DICTIONARY_PAGE,
2623 "the DICTIONARY_PAGE encoding stat must come first, got {stats:?}"
2624 );
2625 assert!(
2626 stats[1..]
2627 .iter()
2628 .all(|s| s.page_type == PageType::DATA_PAGE)
2629 );
2630 }
2631
2632 fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2633 let mut buffer = vec![];
2634
2635 let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2636 writer.write(expected_batch).unwrap();
2637 writer.close().unwrap();
2638
2639 buffer
2640 }
2641
2642 fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2643 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2644 writer.write(expected_batch).unwrap();
2645 writer.into_inner().unwrap()
2646 }
2647
2648 fn write_string_dictionaries(arrays: &[ArrayRef]) -> usize {
2651 let schema = Arc::new(Schema::new(vec![Field::new(
2652 "col",
2653 arrays[0].data_type().clone(),
2654 true,
2655 )]));
2656 let props = WriterProperties::builder().set_write_batch_size(2).build();
2657 let mut file = vec![];
2658 let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), Some(props)).unwrap();
2659 for array in arrays {
2660 let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();
2661 writer.write(&batch).unwrap();
2662 }
2663 writer.close().unwrap();
2664
2665 let reader = SerializedFileReader::new(Bytes::from(file.clone())).unwrap();
2666 let mut pages = reader
2667 .get_row_group(0)
2668 .unwrap()
2669 .get_column_page_reader(0)
2670 .unwrap();
2671 let dictionary_page = pages.get_next_page().unwrap().unwrap();
2672 assert!(matches!(dictionary_page, Page::DictionaryPage { .. }));
2673
2674 let expected: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
2675 let expected = concat(&expected).unwrap();
2676 let actual = read_column(file);
2677 assert_eq!(
2678 cast(&actual, &DataType::Utf8).unwrap().as_ref(),
2679 cast(&expected, &DataType::Utf8).unwrap().as_ref()
2680 );
2681
2682 dictionary_page.num_values() as usize
2683 }
2684
2685 #[test]
2686 fn arrow_writer_string_dictionary_cache_reused_for_same_values() {
2687 let shared: ArrayRef = Arc::new(StringArray::from(vec![
2689 "a",
2690 "bb",
2691 "longer than eight bytes",
2692 "dup",
2693 "dup",
2694 ]));
2695 let keys = Int32Array::from(vec![
2696 Some(0),
2697 Some(1),
2698 None,
2699 Some(2),
2700 Some(3),
2701 Some(4),
2702 Some(0),
2703 ]);
2704 let first = DictionaryArray::new(keys, shared.clone());
2705 let keys = Int32Array::from(vec![Some(4), Some(4), Some(1), None, Some(2), Some(0)]);
2708 let second = DictionaryArray::new(keys, shared).slice(1, 4);
2709
2710 let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second)];
2711 assert_eq!(write_string_dictionaries(&arrays), 4);
2713 }
2714
2715 #[test]
2716 fn arrow_writer_string_dictionary_cache_replaced_for_new_values() {
2717 let first_values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
2718 let first = DictionaryArray::new(Int32Array::from(vec![0, 1, 2, 0]), first_values.clone());
2719 let second_values: ArrayRef = Arc::new(StringArray::from(vec!["zz", "a", "new value!!"]));
2722 let second = DictionaryArray::new(
2723 Int32Array::from(vec![Some(2), Some(1), None, Some(0), Some(1)]),
2724 second_values,
2725 );
2726 let third = DictionaryArray::new(Int32Array::from(vec![2, 1, 0]), first_values);
2728
2729 let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second), Arc::new(third)];
2730 assert_eq!(write_string_dictionaries(&arrays), 5);
2732 }
2733
2734 fn dictionary_writer_memory_size(array: ArrayRef, batch_size: usize) -> usize {
2735 let schema = Arc::new(Schema::new(vec![Field::new(
2736 "col",
2737 array.data_type().clone(),
2738 true,
2739 )]));
2740 let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
2741 let props = WriterProperties::builder()
2742 .set_write_batch_size(batch_size)
2743 .build();
2744 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
2745 writer.write(&batch).unwrap();
2746 writer.memory_size()
2747 }
2748
2749 #[test]
2750 fn arrow_writer_dictionary_cache_ignores_null_rows() {
2751 let values: ArrayRef = Arc::new(StringArray::from(vec![
2752 Some("value"),
2753 None,
2754 Some("unused"),
2755 Some("also unused"),
2756 ]));
2757
2758 for non_null in [0, 2] {
2759 let valid = |i| (3..3 + non_null).contains(&i);
2760 let keys = Int32Array::from_iter((0..10).map(|i| valid(i).then_some(0)));
2761 let null_keys: ArrayRef =
2762 Arc::new(DictionaryArray::new(keys, values.clone()).slice(1, 8));
2763 let keys = Int32Array::from_iter_values((0..10).map(|i| i32::from(!valid(i))));
2764 let null_values: ArrayRef =
2765 Arc::new(DictionaryArray::new(keys, values.clone()).slice(1, 8));
2766 let child: ArrayRef = Arc::new(DictionaryArray::new(
2767 Int32Array::from(vec![0; 10]),
2768 values.clone(),
2769 ));
2770 let parent: ArrayRef = Arc::new(
2771 StructArray::new(
2772 vec![Arc::new(Field::new(
2773 "item",
2774 child.data_type().clone(),
2775 true,
2776 ))]
2777 .into(),
2778 vec![child],
2779 Some(NullBuffer::from_iter((0..10).map(valid))),
2780 )
2781 .slice(1, 8),
2782 );
2783
2784 for array in [null_keys, null_values, parent] {
2785 let plain_type = match array.data_type() {
2786 DataType::Struct(_) => DataType::Struct(
2787 vec![Arc::new(Field::new("item", DataType::Utf8, true))].into(),
2788 ),
2789 _ => DataType::Utf8,
2790 };
2791 let plain = cast(&array, &plain_type).unwrap();
2792 assert_eq!(
2794 dictionary_writer_memory_size(array, 2),
2795 dictionary_writer_memory_size(plain, 2)
2796 );
2797 }
2798 }
2799 }
2800
2801 #[test]
2802 fn arrow_writer_dictionary_cache_counts_rows_across_mini_batches() {
2803 let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
2804 let array: ArrayRef = Arc::new(DictionaryArray::new(
2805 Int32Array::from(vec![
2806 Some(0),
2807 Some(1),
2808 None,
2809 Some(2),
2810 Some(3),
2811 Some(0),
2812 Some(1),
2813 Some(2),
2814 Some(3),
2815 ]),
2816 values,
2817 ));
2818 let plain = cast(&array, &DataType::Utf8).unwrap();
2819 assert_eq!(
2821 dictionary_writer_memory_size(array, 2),
2822 dictionary_writer_memory_size(plain, 2) + 4 * std::mem::size_of::<u64>()
2823 );
2824 }
2825
2826 #[test]
2827 fn roundtrip_bytes() {
2828 let schema = Arc::new(Schema::new(vec![
2830 Field::new("a", DataType::Int32, false),
2831 Field::new("b", DataType::Int32, true),
2832 ]));
2833
2834 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2836 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2837
2838 let expected_batch =
2840 RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2841
2842 for buffer in [
2843 get_bytes_after_close(schema.clone(), &expected_batch),
2844 get_bytes_by_into_inner(schema, &expected_batch),
2845 ] {
2846 let cursor = Bytes::from(buffer);
2847 let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2848
2849 let actual_batch = record_batch_reader
2850 .next()
2851 .expect("No batch found")
2852 .expect("Unable to get batch");
2853
2854 assert_eq!(expected_batch.schema(), actual_batch.schema());
2855 assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2856 assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2857 for i in 0..expected_batch.num_columns() {
2858 let expected_data = expected_batch.column(i).to_data();
2859 let actual_data = actual_batch.column(i).to_data();
2860
2861 assert_eq!(expected_data, actual_data);
2862 }
2863 }
2864 }
2865
2866 fn read_column(file: Vec<u8>) -> ArrayRef {
2867 let reader = ParquetRecordBatchReader::try_new(Bytes::from(file), 4096).unwrap();
2868 let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2869 let arrays: Vec<&dyn Array> = batches.iter().map(|b| b.column(0).as_ref()).collect();
2870 concat(&arrays).unwrap()
2871 }
2872
2873 fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
2874 let writer_schema = Arc::new(Schema::new(vec![field]));
2875 let batch_schema = Arc::new(Schema::new(vec![Field::new(
2876 "c",
2877 col.data_type().clone(),
2878 col.logical_null_count() != 0,
2879 )]));
2880 let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
2881 let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
2884 let mut file = vec![];
2885 let mut writer =
2886 ArrowWriter::try_new_with_options(&mut file, writer_schema, options).unwrap();
2887 writer.write(&batch).unwrap();
2888 writer.close().unwrap();
2889 read_column(file)
2890 }
2891
2892 fn assert_compatible_roundtrip(field: Field, actual: ArrayRef, expected: &ArrayRef) {
2893 assert!(!compute_leaves(&field, &actual).unwrap().is_empty());
2894 assert_eq!(
2895 roundtrip_compatible_column(field, actual).as_ref(),
2896 expected.as_ref()
2897 );
2898 }
2899
2900 fn run_end_type(value_type: DataType, nullable_values: bool) -> DataType {
2901 DataType::RunEndEncoded(
2902 Arc::new(Field::new(
2903 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2904 DataType::Int32,
2905 false,
2906 )),
2907 Arc::new(Field::new(
2908 Field::REE_VALUES_FIELD_DEFAULT_NAME,
2909 value_type,
2910 nullable_values,
2911 )),
2912 )
2913 }
2914
2915 #[test]
2916 fn arrow_writer_required_ree_schema_rejects_nullable_dense_input() {
2917 let writer_schema = Arc::new(Schema::new(vec![Field::new(
2918 "c",
2919 run_end_type(DataType::Int32, false),
2920 false,
2921 )]));
2922 let batch_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2923 let batch = RecordBatch::try_new(
2924 batch_schema.clone(),
2925 vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
2926 )
2927 .unwrap();
2928
2929 for skip_metadata in [false, true] {
2930 for cdc in [false, true] {
2931 let props = WriterProperties::builder()
2932 .set_content_defined_chunking(cdc.then(Default::default))
2933 .build();
2934 let options = ArrowWriterOptions::new()
2935 .with_properties(props)
2936 .with_skip_arrow_metadata(skip_metadata);
2937 let mut writer = ArrowWriter::try_new_with_options(
2938 Vec::new(),
2939 writer_schema.clone(),
2940 options.clone(),
2941 )
2942 .unwrap();
2943 let error = writer.write(&batch).unwrap_err().to_string();
2944 assert!(
2945 error.contains("Found null at index 1 for required field 'c'"),
2946 "skip_metadata={skip_metadata}, cdc={cdc}: {error}"
2947 );
2948
2949 let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 8, 9]));
2951 let batch =
2952 RecordBatch::try_new(batch_schema.clone(), vec![values.clone()]).unwrap();
2953 let mut file = vec![];
2954 let mut writer =
2955 ArrowWriter::try_new_with_options(&mut file, writer_schema.clone(), options)
2956 .unwrap();
2957 writer.write(&batch).unwrap();
2958 writer.close().unwrap();
2959 assert_eq!(read_column(file).as_ref(), values.as_ref());
2960 }
2961 }
2962 }
2963
2964 #[test]
2965 fn ree_nulls_respect_dense_schema_nullability() {
2966 let run_ends = Int32Array::from(vec![2, 4, 5]);
2967 let values = Int32Array::from(vec![Some(1), None, Some(2)]);
2968 let ree: ArrayRef = Arc::new(Int32RunArray::try_new(&run_ends, &values).unwrap());
2969 let ree_batch = RecordBatch::try_new(
2970 Arc::new(Schema::new(vec![Field::new(
2971 "c",
2972 ree.data_type().clone(),
2973 true,
2974 )])),
2975 vec![ree],
2976 )
2977 .unwrap();
2978
2979 let nullable_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2980 let mut file = vec![];
2981 let mut writer = ArrowWriter::try_new(&mut file, nullable_schema, None).unwrap();
2982 writer.write(&ree_batch).unwrap();
2983 writer.close().unwrap();
2984 assert_eq!(
2985 read_column(file).as_primitive::<Int32Type>(),
2986 &Int32Array::from(vec![Some(1), Some(1), None, None, Some(2)])
2987 );
2988
2989 let required_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
2990 let mut writer = ArrowWriter::try_new(Vec::new(), required_schema, None).unwrap();
2991 let error = writer.write(&ree_batch).unwrap_err().to_string();
2992 assert!(error.contains("Found null") && error.contains("required field 'c'"));
2993 }
2994
2995 #[test]
2996 fn required_null_validation_respects_reachability_and_write_timing() {
2997 let required = Field::new("c", DataType::Int32, false);
2998 let nullable: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None]));
2999 let leaves = compute_leaves(&required, &nullable).unwrap();
3000 let schema = Arc::new(Schema::new(vec![required.clone()]));
3001 let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3002 let (_file, factory) = writer.into_serialized_writer().unwrap();
3003 let mut columns = factory.create_column_writers(0).unwrap();
3004 assert!(
3005 columns[0]
3006 .write(&leaves[0])
3007 .unwrap_err()
3008 .to_string()
3009 .contains("required field 'c'")
3010 );
3011
3012 let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None]));
3014 let dict: ArrayRef = Arc::new(DictionaryArray::new(
3015 Int8Array::from(vec![0, 0]),
3016 values.clone(),
3017 ));
3018 assert_eq!(
3019 roundtrip_compatible_column(required.clone(), dict).as_primitive::<Int32Type>(),
3020 &Int32Array::from(vec![7, 7])
3021 );
3022 let dict: ArrayRef = Arc::new(DictionaryArray::new(Int8Array::from(vec![0, 1]), values));
3023 let leaves = compute_leaves(&required, &dict).unwrap();
3024 assert!(leaves[0].0.validate().is_err());
3025 let dict: ArrayRef = Arc::new(DictionaryArray::new(
3026 Int8Array::from(vec![Some(0), None]),
3027 Arc::new(Int32Array::from(vec![7])),
3028 ));
3029 let leaves = compute_leaves(&required, &dict).unwrap();
3030 assert!(leaves[0].0.validate().is_err());
3031
3032 let actual_fields = Fields::from(vec![Field::new("source_child", DataType::Int32, true)]);
3035 let array: ArrayRef = Arc::new(StructArray::new(
3036 actual_fields.clone(),
3037 vec![nullable.clone()],
3038 Some(NullBuffer::from(vec![true, false])),
3039 ));
3040 let unmasked: ArrayRef = Arc::new(StructArray::new(
3041 actual_fields,
3042 vec![nullable.clone()],
3043 None,
3044 ));
3045 for child_type in [DataType::Int32, run_end_type(DataType::Int32, false)] {
3046 let target = Field::new(
3047 "c",
3048 DataType::Struct(Fields::from(vec![Field::new(
3049 "target_child",
3050 child_type,
3051 false,
3052 )])),
3053 true,
3054 );
3055 let leaves = compute_leaves(&target, &unmasked).unwrap();
3056 assert!(leaves[0].0.validate().is_err());
3057 let leaves = compute_leaves(&target, &array).unwrap();
3058 assert!(leaves[0].0.validate().is_ok());
3059 let result = roundtrip_compatible_column(target, array.clone());
3060 assert_eq!(result.as_struct().fields()[0].name(), "target_child");
3061 assert!(result.is_null(1));
3062 }
3063
3064 let values: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(9), None]));
3066 let list: ArrayRef = Arc::new(ListArray::new(
3067 Arc::new(Field::new_list_field(DataType::Int32, true)),
3068 OffsetBuffer::new(vec![1_i32, 2].into()),
3069 values,
3070 None,
3071 ));
3072 let target = Field::new(
3073 "c",
3074 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
3075 false,
3076 );
3077 assert!(
3078 compute_leaves(&target, &list).unwrap()[0]
3079 .0
3080 .validate()
3081 .is_ok()
3082 );
3083 assert_eq!(
3084 roundtrip_compatible_column(required, nullable.slice(0, 1)).as_primitive::<Int32Type>(),
3085 &Int32Array::from(vec![1])
3086 );
3087 }
3088
3089 fn check_required_list_validation(list_type: fn(Arc<Field>) -> DataType) {
3090 let make_array = |values: Vec<Option<i32>>,
3091 offsets: Vec<i32>,
3092 sizes: Vec<i32>,
3093 nulls: Option<NullBuffer>| {
3094 let child = Arc::new(Field::new("source_child", DataType::Int32, true));
3095 let values: ArrayRef = Arc::new(Int32Array::from(values));
3096 match list_type(child.clone()) {
3097 DataType::List(_) => Arc::new(ListArray::new(
3098 child,
3099 OffsetBuffer::new(offsets.into()),
3100 values,
3101 nulls,
3102 )) as ArrayRef,
3103 DataType::LargeList(_) => Arc::new(LargeListArray::new(
3104 child,
3105 OffsetBuffer::new(
3106 offsets
3107 .into_iter()
3108 .map(i64::from)
3109 .collect::<Vec<_>>()
3110 .into(),
3111 ),
3112 values,
3113 nulls,
3114 )) as ArrayRef,
3115 DataType::ListView(_) => Arc::new(ListViewArray::new(
3116 child,
3117 offsets.into(),
3118 sizes.into(),
3119 values,
3120 nulls,
3121 )) as ArrayRef,
3122 DataType::LargeListView(_) => Arc::new(LargeListViewArray::new(
3123 child,
3124 offsets
3125 .into_iter()
3126 .map(i64::from)
3127 .collect::<Vec<_>>()
3128 .into(),
3129 sizes.into_iter().map(i64::from).collect::<Vec<_>>().into(),
3130 values,
3131 nulls,
3132 )) as ArrayRef,
3133 _ => unreachable!(),
3134 }
3135 };
3136 let is_view = matches!(
3137 list_type(Arc::new(Field::new("item", DataType::Int32, true))),
3138 DataType::ListView(_) | DataType::LargeListView(_)
3139 );
3140 let offsets = |starts: Vec<i32>, end: i32| {
3141 let mut result = starts;
3142 if !is_view {
3143 result.push(end);
3144 }
3145 result
3146 };
3147 let target = Field::new(
3148 "c",
3149 list_type(Arc::new(Field::new("target_child", DataType::Int32, false))),
3150 true,
3151 );
3152 let invalid = make_array(vec![None, Some(7)], offsets(vec![0], 2), vec![2], None);
3153 let valid = make_array(vec![Some(7), Some(8)], offsets(vec![0], 2), vec![2], None);
3154 let masked = make_array(
3155 vec![None, Some(7)],
3156 offsets(vec![0, 1], 2),
3157 vec![1, 1],
3158 Some(NullBuffer::from(vec![false, true])),
3159 );
3160 let unused = make_array(
3161 vec![None, Some(7), None],
3162 offsets(vec![1], 2),
3163 vec![1],
3164 None,
3165 );
3166 let empty = make_array(vec![None], offsets(vec![0], 0), vec![0], None);
3167 let sliced = make_array(
3168 vec![None, Some(7)],
3169 offsets(vec![0, 1], 2),
3170 vec![1, 1],
3171 None,
3172 )
3173 .slice(1, 1);
3174 let schema = Arc::new(Schema::new(vec![target.clone()]));
3175
3176 for (name, array, invalid) in [
3178 ("reachable null", invalid, true),
3179 ("valid required child", valid.clone(), false),
3180 ("null ancestor", masked, false),
3181 ("unused child slots", unused, false),
3182 ("empty list", empty, false),
3183 ("out-of-slice null", sliced, false),
3184 ] {
3185 let leaves = compute_leaves(&target, &array).unwrap();
3186 assert_eq!(leaves.len(), 1);
3187 for cdc in [false, true] {
3188 let props = WriterProperties::builder()
3189 .set_content_defined_chunking(cdc.then(Default::default))
3190 .build();
3191 let mut writer =
3192 ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3193 let mut chunkers = writer.cdc_chunkers.take();
3194 let (_file, factory) = writer.into_serialized_writer().unwrap();
3195 let mut columns = factory.create_column_writers(0).unwrap();
3196 let before = (
3197 columns[0].memory_size(),
3198 columns[0].get_estimated_total_bytes(),
3199 );
3200 let result = match &mut chunkers {
3201 Some(chunkers) => columns[0].write_with_chunker(&leaves[0], &mut chunkers[0]),
3202 None => columns[0].write(&leaves[0]),
3203 };
3204 if invalid {
3205 let error = result.unwrap_err().to_string();
3206 assert!(
3207 error.contains("Found null at index 0 for required field 'target_child'"),
3208 "{name}, cdc={cdc}: {error}"
3209 );
3210 assert_eq!(
3211 before,
3212 (
3213 columns[0].memory_size(),
3214 columns[0].get_estimated_total_bytes()
3215 ),
3216 "invalid leaf must be rejected before encoding"
3217 );
3218 let valid_leaves = compute_leaves(&target, &valid).unwrap();
3220 match &mut chunkers {
3221 Some(chunkers) => columns[0]
3222 .write_with_chunker(&valid_leaves[0], &mut chunkers[0])
3223 .unwrap(),
3224 None => columns[0].write(&valid_leaves[0]).unwrap(),
3225 }
3226 assert_eq!(
3227 columns
3228 .pop()
3229 .unwrap()
3230 .close()
3231 .unwrap()
3232 .close
3233 .metadata
3234 .num_values(),
3235 2
3236 );
3237 } else {
3238 result.unwrap();
3239 columns.pop().unwrap().close().unwrap();
3240 }
3241
3242 let batch = RecordBatch::try_new(
3244 Arc::new(Schema::new(vec![Field::new(
3245 "c",
3246 array.data_type().clone(),
3247 true,
3248 )])),
3249 vec![array.clone()],
3250 )
3251 .unwrap();
3252 let props = WriterProperties::builder()
3253 .set_content_defined_chunking(cdc.then(Default::default))
3254 .build();
3255 let mut writer =
3256 ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3257 let result = writer.write(&batch);
3258 if invalid {
3259 assert!(
3260 result
3261 .unwrap_err()
3262 .to_string()
3263 .contains("required field 'target_child'")
3264 );
3265 } else {
3266 result.unwrap();
3267 writer.close().unwrap();
3268 }
3269 }
3270 }
3271 }
3272
3273 #[test]
3274 fn required_list_validation_preserves_level_generation() {
3275 check_required_list_validation(DataType::List);
3276 }
3277
3278 #[test]
3279 fn required_large_list_validation_preserves_level_generation() {
3280 check_required_list_validation(DataType::LargeList);
3281 }
3282
3283 #[test]
3284 fn required_list_view_validation_preserves_level_generation() {
3285 check_required_list_validation(DataType::ListView);
3286 }
3287
3288 #[test]
3289 fn required_large_list_view_validation_preserves_level_generation() {
3290 check_required_list_validation(DataType::LargeListView);
3291 }
3292
3293 fn nullable_int32_list() -> ArrayRef {
3294 Arc::new(ListArray::new(
3295 Arc::new(Field::new_list_field(DataType::Int32, true)),
3296 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3297 Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
3298 None,
3299 ))
3300 }
3301
3302 #[test]
3303 fn arrow_writer_dense_list_under_ree_schema() {
3304 let list = nullable_int32_list();
3305 let field = Field::new("c", run_end_type(list.data_type().clone(), false), false);
3307 assert_compatible_roundtrip(field, list.clone(), &list);
3308 }
3309
3310 #[test]
3311 fn arrow_writer_dense_list_under_dictionary_schema() {
3312 let list = nullable_int32_list();
3313 let field = Field::new(
3314 "c",
3315 DataType::Dictionary(Box::new(DataType::Int8), Box::new(list.data_type().clone())),
3316 false,
3317 );
3318 assert_compatible_roundtrip(field, list.clone(), &list);
3319 }
3320
3321 #[test]
3322 fn arrow_writer_dense_values_under_nested_ree_schema() {
3323 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3324 let field = Field::new(
3325 "c",
3326 run_end_type(run_end_type(DataType::Int32, true), false),
3327 false,
3328 );
3329 assert_compatible_roundtrip(field, dense.clone(), &dense);
3330 }
3331
3332 #[test]
3333 fn arrow_writer_dense_values_under_dictionary_ree_schema() {
3334 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3335 let field = Field::new(
3336 "c",
3337 DataType::Dictionary(
3338 Box::new(DataType::Int8),
3339 Box::new(run_end_type(DataType::Int32, true)),
3340 ),
3341 false,
3342 );
3343 assert_compatible_roundtrip(field, dense.clone(), &dense);
3344 }
3345
3346 #[test]
3347 fn arrow_writer_dense_values_under_ree_dictionary_schema() {
3348 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3349 let field = Field::new(
3352 "c",
3353 run_end_type(
3354 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3355 true,
3356 ),
3357 false,
3358 );
3359 assert_compatible_roundtrip(field, dense.clone(), &dense);
3360 }
3361
3362 #[test]
3363 fn arrow_writer_dense_struct_under_ree_child_schema() {
3364 let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3365 let struct_fields = Fields::from(vec![Field::new("value", DataType::Int32, true)]);
3366 let dense_struct: ArrayRef = Arc::new(StructArray::new(struct_fields, vec![dense], None));
3367 let field = Field::new(
3369 "c",
3370 DataType::Struct(Fields::from(vec![Field::new(
3371 "value",
3372 run_end_type(DataType::Int32, true),
3373 false,
3374 )])),
3375 false,
3376 );
3377 assert_compatible_roundtrip(field, dense_struct.clone(), &dense_struct);
3378 }
3379
3380 #[test]
3381 fn arrow_writer_dense_list_under_ree_child_schema() {
3382 let list = nullable_int32_list();
3383 let field = Field::new(
3384 "c",
3385 DataType::List(Arc::new(Field::new(
3386 "item",
3387 run_end_type(DataType::Int32, true),
3388 false,
3389 ))),
3390 false,
3391 );
3392 assert_compatible_roundtrip(field, list.clone(), &list);
3393 }
3394
3395 #[test]
3396 fn arrow_writer_dense_list_under_dictionary_child_schema() {
3397 let list = nullable_int32_list();
3398 let field = Field::new(
3399 "c",
3400 DataType::List(Arc::new(Field::new_list_field(
3401 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3402 true,
3403 ))),
3404 false,
3405 );
3406 assert_compatible_roundtrip(field, list.clone(), &list);
3407 }
3408
3409 #[test]
3410 fn arrow_writer_dictionary_list_under_dense_child_schema() {
3411 let dict: ArrayRef = Arc::new(DictionaryArray::new(
3412 Int8Array::from(vec![Some(0), None, Some(1)]),
3413 Arc::new(Int32Array::from(vec![1, 3])),
3414 ));
3415 let dict_list: ArrayRef = Arc::new(ListArray::new(
3416 Arc::new(Field::new_list_field(dict.data_type().clone(), true)),
3417 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3418 dict,
3419 None,
3420 ));
3421 let expected = nullable_int32_list();
3422 assert_compatible_roundtrip(
3423 Field::new("c", expected.data_type().clone(), false),
3424 dict_list,
3425 &expected,
3426 );
3427 }
3428
3429 #[test]
3430 fn arrow_writer_dense_map_under_ree_value_schema() {
3431 let key_field = Arc::new(Field::new("keys", DataType::Utf8, false));
3432 let entries = StructArray::new(
3433 Fields::from(vec![
3434 key_field.clone(),
3435 Arc::new(Field::new("values", DataType::Int32, true)),
3436 ]),
3437 vec![
3438 Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
3439 Arc::new(Int32Array::from(vec![Some(7), None, Some(9)])),
3440 ],
3441 None,
3442 );
3443 let map: ArrayRef = Arc::new(MapArray::new(
3444 Arc::new(Field::new("entries", entries.data_type().clone(), false)),
3445 OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3446 entries,
3447 None,
3448 false,
3449 ));
3450 let field = Field::new(
3452 "c",
3453 DataType::Map(
3454 Arc::new(Field::new(
3455 "entries",
3456 DataType::Struct(Fields::from(vec![
3457 key_field,
3458 Arc::new(Field::new(
3459 "values",
3460 run_end_type(DataType::Int32, true),
3461 false,
3462 )),
3463 ])),
3464 false,
3465 )),
3466 false,
3467 ),
3468 false,
3469 );
3470 assert_compatible_roundtrip(field, map.clone(), &map);
3471 }
3472
3473 #[test]
3474 fn arrow_writer_nested_dictionary_metadata_roundtrips_with_default_reader() {
3475 let dense: ArrayRef = Arc::new(StructArray::new(
3476 Fields::from(vec![Field::new("item", DataType::Int32, true)]),
3477 vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
3478 None,
3479 ));
3480 let field = Field::new(
3481 "value",
3482 DataType::Dictionary(
3483 Box::new(DataType::Int8),
3484 Box::new(dense.data_type().clone()),
3485 ),
3486 false,
3487 )
3488 .with_metadata(HashMap::from([(
3489 "description".to_string(),
3490 "nested dictionary".to_string(),
3491 )]));
3492 let writer_schema = Arc::new(Schema::new(vec![field.clone()]));
3493 let expected_schema = Arc::new(Schema::new(vec![
3494 field.with_data_type(dense.data_type().clone()),
3495 ]));
3496 let batch = RecordBatch::try_new(expected_schema.clone(), vec![dense]).unwrap();
3497
3498 let mut file = vec![];
3501 let mut writer = ArrowWriter::try_new(&mut file, writer_schema, None).unwrap();
3502 writer.write(&batch).unwrap();
3503 writer.close().unwrap();
3504
3505 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file)).unwrap();
3506 assert!(
3507 builder
3508 .metadata()
3509 .file_metadata()
3510 .key_value_metadata()
3511 .unwrap()
3512 .iter()
3513 .any(|kv| kv.key == ARROW_SCHEMA_META_KEY)
3514 );
3515 assert_eq!(builder.schema(), &expected_schema);
3516 let decoded = builder
3517 .build()
3518 .unwrap()
3519 .collect::<ArrowResult<Vec<_>>>()
3520 .unwrap();
3521 assert_eq!(decoded, vec![batch]);
3522 }
3523
3524 #[test]
3525 fn arrow_writer_rejects_missing_or_extra_top_level_columns_cleanly() {
3526 let one_field = Arc::new(Schema::new(vec![Field::new(
3527 "a",
3528 ArrowDataType::Int32,
3529 false,
3530 )]));
3531 let two_fields = Arc::new(Schema::new(vec![
3532 Field::new("a", ArrowDataType::Int32, false),
3533 Field::new("b", ArrowDataType::Int32, false),
3534 ]));
3535 let one_column = RecordBatch::try_new(
3536 Arc::clone(&one_field),
3537 vec![Arc::new(Int32Array::from(vec![1]))],
3538 )
3539 .unwrap();
3540 let two_columns = RecordBatch::try_new(
3541 Arc::clone(&two_fields),
3542 vec![
3543 Arc::new(Int32Array::from(vec![1])),
3544 Arc::new(Int32Array::from(vec![2])),
3545 ],
3546 )
3547 .unwrap();
3548
3549 let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&two_fields), None).unwrap();
3550 let err = writer.write(&one_column).unwrap_err();
3551 assert!(
3552 err.to_string()
3553 .contains("writer has 2 top-level fields but batch has 1 columns"),
3554 "{err}"
3555 );
3556 assert_eq!(writer.in_progress_rows(), 0);
3557
3558 let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&one_field), None).unwrap();
3559 let err = writer.write(&two_columns).unwrap_err();
3560 assert!(
3561 err.to_string()
3562 .contains("writer has 1 top-level fields but batch has 2 columns"),
3563 "{err}"
3564 );
3565 assert_eq!(writer.in_progress_rows(), 0);
3566 }
3567
3568 #[test]
3569 fn arrow_row_group_factory_rejects_parquet_leaf_count_mismatch_cleanly() {
3570 let one_field = Arc::new(Schema::new(vec![Field::new(
3571 "a",
3572 ArrowDataType::Int32,
3573 false,
3574 )]));
3575 let two_fields = Arc::new(Schema::new(vec![
3576 Field::new("a", ArrowDataType::Int32, false),
3577 Field::new("b", ArrowDataType::Int32, false),
3578 ]));
3579 let parquet = ArrowSchemaConverter::new().convert(&one_field).unwrap();
3580 let props = Arc::new(WriterProperties::default());
3581 let file_writer =
3582 SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), Arc::clone(&props))
3583 .unwrap();
3584
3585 let eager_error =
3586 ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&two_fields)).unwrap_err();
3587 assert!(matches!(&eager_error, ParquetError::General(_)));
3588 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&two_fields));
3590 let Err(err) = factory.create_column_writers(0) else {
3591 panic!("mismatched schemas should not create column writers")
3592 };
3593 assert_eq!(err.to_string(), eager_error.to_string());
3594 assert!(
3595 err.to_string()
3596 .contains("Arrow schema maps to 2 leaf columns but Parquet schema has 1"),
3597 "{err}"
3598 );
3599
3600 let parquet = ArrowSchemaConverter::new().convert(&two_fields).unwrap();
3601 let file_writer =
3602 SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), props).unwrap();
3603 let eager_error =
3604 ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&one_field)).unwrap_err();
3605 assert!(matches!(&eager_error, ParquetError::General(_)));
3606 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file_writer, one_field);
3608 let Err(err) = factory.create_column_writers(0) else {
3609 panic!("mismatched schemas should not create column writers")
3610 };
3611 assert_eq!(err.to_string(), eager_error.to_string());
3612 assert!(
3613 err.to_string()
3614 .contains("Arrow schema maps to 1 leaf columns but Parquet schema has 2"),
3615 "{err}"
3616 );
3617 }
3618
3619 #[test]
3620 fn arrow_writer_required_map_keys_reject_only_reachable_nulls() {
3621 let keys: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
3622 Int32Array::from(vec![1, 0, 1]),
3623 Arc::new(Int32Array::from(vec![None, Some(7)])),
3624 ));
3625 let values = Int32Array::from(vec![Some(10), Some(20), None]);
3626 let fields = Fields::from(vec![
3627 Field::new("key", keys.data_type().clone(), false),
3628 Field::new("value", DataType::Int32, true),
3629 ]);
3630 assert!(
3631 StructArray::try_new(
3632 fields.clone(),
3633 vec![keys.clone(), Arc::new(values.clone())],
3634 None
3635 )
3636 .is_err()
3637 );
3638
3639 let data = ArrayData::builder(DataType::Struct(fields))
3642 .len(3)
3643 .add_child_data(keys.to_data())
3644 .add_child_data(values.to_data())
3645 .build()
3646 .unwrap();
3647 data.validate_full().unwrap();
3648 let entries = StructArray::from(data);
3649 let entries_field = Arc::new(Field::new("entries", entries.data_type().clone(), false));
3650 let offsets = OffsetBuffer::new(vec![0i32, 1, 2, 3].into());
3651 let map = MapArray::try_new(
3652 entries_field.clone(),
3653 offsets.clone(),
3654 entries.clone(),
3655 None,
3656 false,
3657 )
3658 .unwrap();
3659 assert_eq!(map.keys().logical_null_count(), 1);
3660 let masked = MapArray::try_new(
3661 entries_field,
3662 offsets,
3663 entries,
3664 Some(NullBuffer::from(vec![true, false, true])),
3665 false,
3666 )
3667 .unwrap();
3668 let writer_schema = Arc::new(Schema::new(vec![Field::new(
3669 "map",
3670 DataType::Map(
3671 Arc::new(Field::new(
3672 "entries",
3673 DataType::Struct(
3674 vec![
3675 Field::new("key", DataType::Int32, false),
3676 Field::new("value", DataType::Int32, true),
3677 ]
3678 .into(),
3679 ),
3680 false,
3681 )),
3682 false,
3683 ),
3684 true,
3685 )]));
3686 for cdc in [false, true] {
3687 let props = WriterProperties::builder()
3688 .set_content_defined_chunking(cdc.then(Default::default))
3689 .build();
3690 for (array, expected_keys, expected_null_maps) in [
3691 (map.clone(), None, 0),
3692 (map.slice(0, 1), Some(1), 0),
3693 (map.slice(2, 1), Some(1), 0),
3694 (masked.clone(), Some(2), 1),
3695 ] {
3696 let schema = Arc::new(Schema::new(vec![Field::new(
3697 "map",
3698 array.data_type().clone(),
3699 true,
3700 )]));
3701 let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
3702 let mut bytes = Vec::new();
3703 let mut writer =
3704 ArrowWriter::try_new(&mut bytes, writer_schema.clone(), Some(props.clone()))
3705 .unwrap();
3706 let result = writer.write(&batch);
3707 let Some(expected_keys) = expected_keys else {
3708 let err = result.unwrap_err();
3709 assert!(
3710 err.to_string().contains("Found null")
3711 && err.to_string().contains("required field 'key'"),
3712 "{err}"
3713 );
3714 continue;
3715 };
3716 result.unwrap();
3717 writer.close().unwrap();
3718 let actual = ParquetRecordBatchReader::try_new(Bytes::from(bytes), 1024)
3719 .unwrap()
3720 .next()
3721 .unwrap()
3722 .unwrap();
3723 let actual = actual.column(0).as_map();
3724 assert_eq!(actual.null_count(), expected_null_maps);
3725 assert_eq!(
3726 actual.keys().as_primitive::<Int32Type>().values().as_ref(),
3727 vec![7; expected_keys]
3728 );
3729 }
3730 }
3731 }
3732
3733 #[test]
3734 fn arrow_writer_rejects_nullable_map_key_schemas_early() {
3735 let schema_with_key = |key| {
3736 Arc::new(Schema::new(vec![Field::new(
3737 "map",
3738 DataType::Map(
3739 Arc::new(Field::new(
3740 "entries",
3741 DataType::Struct(
3742 vec![key, Field::new("value", DataType::Int32, true)].into(),
3743 ),
3744 false,
3745 )),
3746 false,
3747 ),
3748 true,
3749 )]))
3750 };
3751 let required = schema_with_key(Field::new("key", DataType::Int32, false));
3752 let parquet = ArrowSchemaConverter::new().convert(&required).unwrap();
3753 let ree = DataType::RunEndEncoded(
3754 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3755 Arc::new(Field::new("values", DataType::Int32, true)),
3756 );
3757 for key in [
3758 Field::new("key", DataType::Int32, true),
3759 Field::new("key", ree.clone(), false),
3760 Field::new(
3761 "key",
3762 DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
3763 false,
3764 ),
3765 ] {
3766 let schema = schema_with_key(key);
3767 for explicit_schema in [false, true] {
3768 for skip_metadata in [false, true] {
3769 let mut options =
3770 ArrowWriterOptions::new().with_skip_arrow_metadata(skip_metadata);
3771 if explicit_schema {
3772 options = options.with_parquet_schema(parquet.clone());
3773 }
3774 let mut bytes = Vec::new();
3775 let err =
3776 ArrowWriter::try_new_with_options(&mut bytes, schema.clone(), options)
3777 .unwrap_err();
3778 assert!(
3779 err.to_string()
3780 .contains("Map key field 'key' must be non-nullable"),
3781 "{err}"
3782 );
3783 assert!(
3784 bytes.is_empty(),
3785 "schema rejection must precede file initialization"
3786 );
3787 }
3788 }
3789 let file = SerializedFileWriter::new(
3790 Vec::new(),
3791 parquet.root_schema_ptr(),
3792 Arc::new(WriterProperties::default()),
3793 )
3794 .unwrap();
3795 let eager_error =
3796 ArrowRowGroupWriterFactory::try_new(&file, Arc::clone(&schema)).unwrap_err();
3797 assert!(matches!(&eager_error, ParquetError::ArrowError(_)));
3798 assert_eq!(
3799 eager_error.to_string(),
3800 "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3801 );
3802 #[expect(deprecated)] let factory = ArrowRowGroupWriterFactory::new(&file, schema);
3804 for row_group_index in 0..2 {
3805 let Err(err) = factory.create_column_writers(row_group_index) else {
3806 panic!("nullable map-key schema should not create column writers")
3807 };
3808 assert!(matches!(&err, ParquetError::ArrowError(_)), "{err:?}");
3809 assert_eq!(
3810 err.to_string(),
3811 "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3812 );
3813 }
3814 }
3815 }
3816
3817 #[test]
3818 fn eager_row_group_validation_rejects_writer_and_chunker_count_mismatches() {
3819 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
3820 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))])
3821 .unwrap();
3822 let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3823 let (_file, factory) = writer.into_serialized_writer().unwrap();
3824 let mut group = factory.create_row_group_writer(0).unwrap();
3825 assert!(
3826 group
3827 .write_with_chunkers(&batch, &mut [])
3828 .unwrap_err()
3829 .to_string()
3830 .contains("chunkers")
3831 );
3832 assert_eq!(group.buffered_rows, 0);
3833 group.writers.clear();
3834 assert!(
3835 group
3836 .write(&batch)
3837 .unwrap_err()
3838 .to_string()
3839 .contains("column writers")
3840 );
3841 assert_eq!(group.buffered_rows, 0);
3842 }
3843
3844 #[test]
3845 fn arrow_write_schema_plan_preserves_error_variants() {
3846 for error in [
3847 ParquetError::General("schema mismatch".to_owned()),
3848 ParquetError::NYI("unsupported type".to_owned()),
3849 ParquetError::ArrowError("invalid schema".to_owned()),
3850 ] {
3851 let expected_variant = std::mem::discriminant(&error);
3852 let expected_message = error.to_string();
3853 let cached = ArrowWriteSchemaPlanError::from_parquet_error(error);
3854 for _ in 0..2 {
3855 let error = cached.to_parquet_error();
3856 assert_eq!(std::mem::discriminant(&error), expected_variant);
3857 assert_eq!(error.to_string(), expected_message);
3858 }
3859 }
3860 }
3861
3862 #[test]
3863 fn arrow_write_schema_plan_caches_nested_leaf_ranges() {
3864 let nested = ArrowDataType::Struct(Fields::from(vec![
3865 Field::new("i", ArrowDataType::Int32, false),
3866 Field::new("s", ArrowDataType::Utf8, true),
3867 ]));
3868 let list = ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Boolean, true)));
3869 let ree = ArrowDataType::RunEndEncoded(
3870 Arc::new(Field::new("run_ends", ArrowDataType::Int32, false)),
3871 Arc::new(Field::new("values", ArrowDataType::Int64, false)),
3872 );
3873 let dictionary_fsb = ArrowDataType::Dictionary(
3874 Box::new(ArrowDataType::Int8),
3875 Box::new(ArrowDataType::FixedSizeBinary(2)),
3876 );
3877 let schema = Arc::new(Schema::new(vec![
3878 Field::new("nested", nested, true),
3879 Field::new("list", list, true),
3880 Field::new("ree", ree, false),
3881 Field::new("dictionary_fsb", dictionary_fsb, true),
3882 ]));
3883 let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
3884
3885 let plan = ArrowWriteSchemaPlan::try_new(&parquet, &schema).unwrap();
3886 let ranges = plan
3887 .fields
3888 .iter()
3889 .map(|field| field.leaf_range.clone())
3890 .collect::<Vec<_>>();
3891 assert_eq!(ranges, vec![0..2, 2..3, 3..4, 4..5]);
3892 assert_eq!(plan.leaves.len(), parquet.num_columns());
3893 assert_eq!(
3894 plan.leaves[4].physical_type(),
3895 crate::basic::Type::FIXED_LEN_BYTE_ARRAY
3896 );
3897 }
3898
3899 #[test]
3900 fn cached_schema_plan_preserves_compatible_physical_layout_alternation() {
3901 let writer_schema = Arc::new(Schema::new(vec![
3902 Field::new("number", ArrowDataType::Int32, false),
3903 Field::new("text", ArrowDataType::Utf8, false),
3904 Field::new("bytes", ArrowDataType::Binary, false),
3905 ]));
3906
3907 let dense = RecordBatch::try_new(
3908 Arc::clone(&writer_schema),
3909 vec![
3910 Arc::new(Int32Array::from(vec![1])),
3911 Arc::new(StringArray::from(vec!["a"])),
3912 Arc::new(BinaryArray::from_iter_values([b"a".as_slice()])),
3913 ],
3914 )
3915 .unwrap();
3916
3917 let number_dictionary = DictionaryArray::<Int8Type>::try_new(
3918 Int8Array::from(vec![0]),
3919 Arc::new(Int32Array::from(vec![2])),
3920 )
3921 .unwrap();
3922 let text_dictionary = DictionaryArray::<Int8Type>::try_new(
3923 Int8Array::from(vec![0]),
3924 Arc::new(StringArray::from(vec!["b"])),
3925 )
3926 .unwrap();
3927 let bytes_dictionary = DictionaryArray::<Int8Type>::try_new(
3928 Int8Array::from(vec![0]),
3929 Arc::new(BinaryArray::from_iter_values([b"b".as_slice()])),
3930 )
3931 .unwrap();
3932 let dictionary = RecordBatch::try_from_iter(vec![
3933 ("number", Arc::new(number_dictionary) as ArrayRef),
3934 ("text", Arc::new(text_dictionary) as ArrayRef),
3935 ("bytes", Arc::new(bytes_dictionary) as ArrayRef),
3936 ])
3937 .unwrap();
3938
3939 let run_ends = Int32Array::from(vec![1]);
3940 let number_ree: ArrayRef =
3941 Arc::new(Int32RunArray::try_new(&run_ends, &Int32Array::from(vec![3])).unwrap());
3942 let text_ree: ArrayRef =
3943 Arc::new(Int32RunArray::try_new(&run_ends, &StringArray::from(vec!["c"])).unwrap());
3944 let bytes_ree: ArrayRef = Arc::new(
3945 Int32RunArray::try_new(&run_ends, &BinaryArray::from_iter_values([b"c".as_slice()]))
3946 .unwrap(),
3947 );
3948 let ree = RecordBatch::try_from_iter(vec![
3949 ("number", number_ree),
3950 ("text", text_ree),
3951 ("bytes", bytes_ree),
3952 ])
3953 .unwrap();
3954
3955 let alternate = RecordBatch::try_from_iter(vec![
3956 ("number", Arc::new(Int32Array::from(vec![4])) as ArrayRef),
3957 (
3958 "text",
3959 Arc::new(LargeStringArray::from(vec!["d"])) as ArrayRef,
3960 ),
3961 (
3962 "bytes",
3963 Arc::new(BinaryViewArray::from_iter_values([b"d".as_slice()])) as ArrayRef,
3964 ),
3965 ])
3966 .unwrap();
3967
3968 let props = WriterProperties::builder()
3969 .set_max_row_group_row_count(Some(1))
3970 .build();
3971 let mut out = Vec::new();
3972 let mut writer =
3973 ArrowWriter::try_new(&mut out, Arc::clone(&writer_schema), Some(props)).unwrap();
3974 for batch in [&dense, &dictionary, &ree, &alternate] {
3975 writer.write(batch).unwrap();
3976 }
3977 writer.close().unwrap();
3978
3979 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(out)).unwrap();
3980 assert_eq!(builder.metadata().num_row_groups(), 4);
3981 let mut reader = builder.build().unwrap();
3982 let actual = reader.next().unwrap().unwrap();
3983 assert_eq!(
3984 actual.column(0).as_primitive::<Int32Type>(),
3985 &Int32Array::from(vec![1, 2, 3, 4])
3986 );
3987 assert_eq!(
3988 actual.column(1).as_string::<i32>(),
3989 &StringArray::from(vec!["a", "b", "c", "d"])
3990 );
3991 assert_eq!(
3992 actual.column(2).as_binary::<i32>(),
3993 &BinaryArray::from_iter_values([
3994 b"a".as_slice(),
3995 b"b".as_slice(),
3996 b"c".as_slice(),
3997 b"d".as_slice(),
3998 ])
3999 );
4000 }
4001
4002 #[test]
4003 fn cached_schema_plan_does_not_retain_batch_arrays() {
4004 let schema = Arc::new(Schema::new(vec![Field::new(
4005 "a",
4006 ArrowDataType::Int32,
4007 false,
4008 )]));
4009 let array = Arc::new(Int32Array::from(vec![1, 2, 3]));
4010 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap();
4011 let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
4012
4013 writer.write(&batch).unwrap();
4014 drop(batch);
4015 assert_eq!(
4016 Arc::strong_count(&array),
4017 1,
4018 "schema and row-group caches must not retain input ArrayRefs"
4019 );
4020 writer.close().unwrap();
4021 }
4022
4023 #[test]
4024 fn arrow_writer_page_size() {
4025 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
4026
4027 let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
4028
4029 for i in 0..10 {
4031 let value = i
4032 .to_string()
4033 .repeat(10)
4034 .chars()
4035 .take(10)
4036 .collect::<String>();
4037
4038 builder.append_value(value);
4039 }
4040
4041 let array = Arc::new(builder.finish());
4042
4043 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
4044
4045 let file = tempfile::tempfile().unwrap();
4046
4047 let props = WriterProperties::builder()
4049 .set_data_page_size_limit(1)
4050 .set_dictionary_page_size_limit(1)
4051 .set_write_batch_size(1)
4052 .build();
4053
4054 let mut writer =
4055 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
4056 .expect("Unable to write file");
4057 writer.write(&batch).unwrap();
4058 writer.close().unwrap();
4059
4060 let options = ReadOptionsBuilder::new().with_page_index().build();
4061 let reader =
4062 SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
4063
4064 let column = reader.metadata().row_group(0).columns();
4065
4066 assert_eq!(column.len(), 1);
4067
4068 assert!(
4071 column[0].dictionary_page_offset().is_some(),
4072 "Expected a dictionary page"
4073 );
4074
4075 let page_index = reader
4076 .metadata()
4077 .page_index()
4078 .expect("page index should be present");
4079 let page_locations = page_index
4080 .page_locations(0, 0)
4081 .expect("page locations should exist");
4082
4083 assert_eq!(
4086 page_locations.len(),
4087 10,
4088 "Expected 10 pages but got {page_locations:#?}"
4089 );
4090 }
4091
4092 const MEDIUM_SIZE: usize = 63;
4093
4094 fn check_bloom_filter<T: AsBytes>(
4095 files: Vec<Bytes>,
4096 file_column: String,
4097 positive_values: Vec<T>,
4098 negative_values: Vec<T>,
4099 ) {
4100 files.into_iter().take(1).for_each(|file| {
4101 let file_reader = SerializedFileReader::new_with_options(
4102 file,
4103 ReadOptionsBuilder::new()
4104 .with_reader_properties(
4105 ReaderProperties::builder()
4106 .set_read_bloom_filter(true)
4107 .build(),
4108 )
4109 .build(),
4110 )
4111 .expect("Unable to open file as Parquet");
4112 let metadata = file_reader.metadata();
4113
4114 let mut bloom_filters: Vec<_> = vec![];
4116 for (ri, row_group) in metadata.row_groups().iter().enumerate() {
4117 if let Some((column_index, _)) = row_group
4118 .columns()
4119 .iter()
4120 .enumerate()
4121 .find(|(_, column)| column.column_path().string() == file_column)
4122 {
4123 let row_group_reader = file_reader
4124 .get_row_group(ri)
4125 .expect("Unable to read row group");
4126 if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
4127 bloom_filters.push(sbbf.clone());
4128 } else {
4129 panic!("No bloom filter for column named {file_column} found");
4130 }
4131 } else {
4132 panic!("No column named {file_column} found");
4133 }
4134 }
4135
4136 positive_values.iter().for_each(|value| {
4137 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4138 assert!(
4139 found.is_some(),
4140 "{}",
4141 format!("Value {:?} should be in bloom filter", value.as_bytes())
4142 );
4143 });
4144
4145 negative_values.iter().for_each(|value| {
4146 let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4147 assert!(
4148 found.is_none(),
4149 "{}",
4150 format!("Value {:?} should not be in bloom filter", value.as_bytes())
4151 );
4152 });
4153 });
4154 }
4155
4156 #[test]
4157 #[cfg_attr(miri, ignore)] fn bool_large_single_column() {
4159 let values = Arc::new(
4160 [None, Some(true), Some(false)]
4161 .iter()
4162 .cycle()
4163 .copied()
4164 .take(200_000)
4165 .collect::<BooleanArray>(),
4166 );
4167 let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
4168 let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4169 let file = tempfile::tempfile().unwrap();
4170
4171 let mut writer =
4172 ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
4173 .expect("Unable to write file");
4174 writer.write(&expected_batch).unwrap();
4175 writer.close().unwrap();
4176 }
4177
4178 #[test]
4179 fn check_page_offset_index_with_nan() {
4180 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4181 let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
4182 let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4183
4184 let mut out = Vec::with_capacity(1024);
4185 let mut writer =
4186 ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
4187 writer.write(&batch).unwrap();
4188 let file_meta_data = writer.close().unwrap();
4189 for row_group in file_meta_data.row_groups() {
4190 for column in row_group.columns() {
4191 assert!(column.offset_index_offset().is_some());
4192 assert!(column.offset_index_length().is_some());
4193 assert!(column.column_index_offset().is_some());
4194 assert!(column.column_index_length().is_some());
4195 }
4196 }
4197 if let Some(page_index) = file_meta_data.page_index() {
4198 for rg in 0..file_meta_data.num_row_groups() {
4199 for col in 0..file_meta_data.row_group(rg).num_columns() {
4200 let idx = page_index
4201 .column_index(rg, col)
4202 .expect("column index should exist");
4203 assert!(idx.nan_counts().is_some());
4204 let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
4205 panic!("expected double statistics")
4206 };
4207 for i in 0..idx.num_pages() as usize {
4208 assert_eq!(float_idx.nan_count(i), Some(10));
4209 assert_eq!(
4210 f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
4211 Ordering::Equal
4212 );
4213 assert_eq!(
4214 f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
4215 Ordering::Equal
4216 );
4217 }
4218 }
4219 }
4220 } else {
4221 panic!("page index should be present");
4222 }
4223 }
4224
4225 #[test]
4226 fn check_page_offset_index_with_mixed_nan() {
4227 let schema = Arc::new(Schema::new(vec![Field::new(
4228 "col",
4229 DataType::Float64,
4230 true,
4231 )]));
4232
4233 let mut out = Vec::with_capacity(1024);
4234 let props = WriterProperties::builder()
4235 .set_data_page_row_count_limit(10)
4236 .build();
4237 let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
4238 .expect("Unable to write file");
4239
4240 let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4242 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4243 writer.write(&batch).unwrap();
4244
4245 let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
4247 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4248 writer.write(&batch).unwrap();
4249
4250 let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
4252 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4253 writer.write(&batch).unwrap();
4254
4255 let values = Arc::new(Float64Array::from(vec![
4257 -1.0,
4258 0.0,
4259 f64::NAN,
4260 -f64::NAN,
4261 1.0,
4262 ]));
4263 let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4264 writer.write(&batch).unwrap();
4265
4266 let file_meta_data = writer.close().unwrap();
4267
4268 let col_stats = file_meta_data
4270 .row_group(0)
4271 .column(0)
4272 .statistics()
4273 .expect("missing column chunk statistics");
4274
4275 assert_eq!(col_stats.nan_count_opt(), Some(22));
4276 assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
4277 assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
4278
4279 assert!(file_meta_data.page_index().is_some());
4280 let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
4281 assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
4282
4283 let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
4285 panic!("expected double statistics")
4286 };
4287
4288 assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
4289 assert_eq!(
4290 f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
4291 Ordering::Equal
4292 );
4293 assert_eq!(
4294 f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
4295 Ordering::Equal
4296 );
4297 assert_eq!(
4298 (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
4299 Ordering::Equal
4300 );
4301 assert_eq!(
4302 (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
4303 Ordering::Equal
4304 );
4305 assert_eq!(float_idx.min_value(2), Some(&0.0));
4306 assert_eq!(float_idx.max_value(2), Some(&0.0));
4307 assert_eq!(float_idx.min_value(3), Some(&-1.0));
4308 assert_eq!(float_idx.max_value(3), Some(&1.0));
4309 }
4310
4311 #[test]
4312 #[should_panic(
4313 expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
4314 )]
4315 fn interval_month_day_nano_single_column() {
4316 required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
4317 IntervalMonthDayNano::new(0, 1, 5),
4318 IntervalMonthDayNano::new(0, 3, 2),
4319 IntervalMonthDayNano::new(3, -2, -5),
4320 IntervalMonthDayNano::new(-200, 4, -1),
4321 ]);
4322 }
4323
4324 #[test]
4325 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_at_end() {
4327 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4328 let files = RoundTripTest::new(array)
4329 .with_nullable(false)
4330 .with_bloom_filter(true)
4331 .with_bloom_filter_position(BloomFilterPosition::End)
4332 .run();
4333
4334 check_bloom_filter(
4335 files,
4336 "col".to_string(),
4337 (0..SMALL_SIZE as i32).collect(),
4338 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4339 );
4340 }
4341
4342 #[test]
4343 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter() {
4345 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4346 let files = RoundTripTest::new(array)
4347 .with_nullable(false)
4348 .with_bloom_filter(true)
4349 .run();
4350
4351 check_bloom_filter(
4352 files,
4353 "col".to_string(),
4354 (0..SMALL_SIZE as i32).collect(),
4355 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4356 );
4357 }
4358
4359 fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
4360 let schema = Arc::new(Schema::new(vec![Field::new(
4361 "col",
4362 array.data_type().clone(),
4363 false,
4364 )]));
4365 let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
4366 let props = WriterProperties::builder()
4367 .set_dictionary_enabled(true)
4368 .set_dictionary_page_size_limit(dictionary_page_size_limit)
4369 .set_write_batch_size(256)
4370 .set_bloom_filter_enabled(true)
4371 .build();
4372 let mut buf = Vec::new();
4373 let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
4374 writer.write(&batch).unwrap();
4375 writer.close().unwrap();
4376 Bytes::from(buf)
4377 }
4378
4379 fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4380 let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4381 *metadata
4382 .row_group(0)
4383 .column(0)
4384 .page_encoding_stats_mask()
4385 .unwrap()
4386 }
4387
4388 #[test]
4391 fn string_column_bloom_filter_populated_from_dictionary() {
4392 let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4393 let array = Arc::new(StringArray::from_iter_values(&values));
4394 let file = write_with_bloom_filter(array, 1024 * 1024);
4395 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4396
4397 check_bloom_filter(
4398 vec![file],
4399 "col".to_string(),
4400 (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4401 (10..20)
4402 .map(|i| format!("value-{i}").into_bytes())
4403 .collect(),
4404 );
4405 }
4406
4407 #[test]
4410 fn string_column_bloom_filter_across_dictionary_fallback() {
4411 let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4412 let array = Arc::new(StringArray::from_iter_values(&values));
4413 let file = write_with_bloom_filter(array, 1024);
4414 let encodings = data_page_encoding_mask(&file);
4415 assert!(
4416 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4417 "expected dictionary and plain data pages, got {encodings:?}"
4418 );
4419
4420 check_bloom_filter(
4421 vec![file],
4422 "col".to_string(),
4423 values.into_iter().map(String::into_bytes).collect(),
4424 (2000..2010)
4425 .map(|i| format!("value-{i}").into_bytes())
4426 .collect(),
4427 );
4428 }
4429
4430 #[test]
4431 fn i64_column_bloom_filter_populated_from_dictionary() {
4432 let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4433 let file = write_with_bloom_filter(array, 1024 * 1024);
4434 assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4435
4436 check_bloom_filter(
4437 vec![file],
4438 "col".to_string(),
4439 (0..10i64).collect(),
4440 (10..20i64).collect(),
4441 );
4442 }
4443
4444 #[test]
4445 fn i64_column_bloom_filter_across_dictionary_fallback() {
4446 let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4447 let file = write_with_bloom_filter(array, 1024);
4448 let encodings = data_page_encoding_mask(&file);
4449 assert!(
4450 encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4451 "expected dictionary and plain data pages, got {encodings:?}"
4452 );
4453
4454 check_bloom_filter(
4455 vec![file],
4456 "col".to_string(),
4457 (0..2000i64).collect(),
4458 (2000..2010i64).collect(),
4459 );
4460 }
4461
4462 #[test]
4467 #[cfg_attr(miri, ignore)] fn i32_column_bloom_filter_fixed_ndv() {
4469 let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4470
4471 let files = RoundTripTest::new(array.clone())
4473 .with_nullable(false)
4474 .with_bloom_filter(true)
4475 .with_bloom_filter_ndv(1_000_000)
4476 .run();
4477
4478 check_bloom_filter(
4479 files,
4480 "col".to_string(),
4481 (0..SMALL_SIZE as i32).collect(),
4482 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4483 );
4484
4485 let files = RoundTripTest::new(array)
4487 .with_nullable(false)
4488 .with_bloom_filter(true)
4489 .with_bloom_filter_ndv(3)
4490 .run();
4491
4492 check_bloom_filter(
4493 files,
4494 "col".to_string(),
4495 (0..SMALL_SIZE as i32).collect(),
4496 (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4497 );
4498 }
4499
4500 #[test]
4501 #[cfg_attr(miri, ignore)] fn binary_column_bloom_filter() {
4503 let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4504 let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4505 let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4506
4507 let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4508 let files = RoundTripTest::new(array)
4509 .with_nullable(false)
4510 .with_bloom_filter(true)
4511 .run();
4512
4513 check_bloom_filter(
4514 files,
4515 "col".to_string(),
4516 many_vecs,
4517 vec![vec![(SMALL_SIZE + 1) as u8]],
4518 );
4519 }
4520
4521 #[test]
4522 #[cfg_attr(miri, ignore)] fn empty_string_null_column_bloom_filter() {
4524 let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4525 let raw_strs = raw_values.iter().map(|s| s.as_str());
4526
4527 let array = Arc::new(StringArray::from_iter_values(raw_strs));
4528 let files = RoundTripTest::new(array)
4529 .with_nullable(false)
4530 .with_bloom_filter(true)
4531 .run();
4532
4533 let optional_raw_values: Vec<_> = raw_values
4534 .iter()
4535 .enumerate()
4536 .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4537 .collect();
4538 check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4540 }
4541
4542 #[test]
4543 fn list_and_map_coerced_names() {
4544 let list_field =
4546 Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4547 let map_field = Field::new_map(
4548 "my_map",
4549 "my_entries",
4550 Field::new("my_keys", DataType::Int32, false),
4551 Field::new("my_values", DataType::Int32, true),
4552 false,
4553 true,
4554 );
4555
4556 let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4557 let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4558
4559 let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4560
4561 let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4563 let file = tempfile::tempfile().unwrap();
4564 let mut writer =
4565 ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4566
4567 let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4568 writer.write(&batch).unwrap();
4569 let file_metadata = writer.close().unwrap();
4570
4571 let schema = file_metadata.file_metadata().schema();
4572 let list_field = &schema.get_fields()[0].get_fields()[0];
4574 assert_eq!(list_field.get_fields()[0].name(), "element");
4575
4576 let map_field = &schema.get_fields()[1].get_fields()[0];
4577 assert_eq!(map_field.name(), "key_value");
4579 assert_eq!(map_field.get_fields()[0].name(), "key");
4581 assert_eq!(map_field.get_fields()[1].name(), "value");
4583
4584 let reader = SerializedFileReader::new(file).unwrap();
4586 let file_schema = reader.metadata().file_metadata().schema();
4587 let fields = file_schema.get_fields();
4588 let list_field = &fields[0].get_fields()[0];
4589 assert_eq!(list_field.get_fields()[0].name(), "element");
4590 let map_field = &fields[1].get_fields()[0];
4591 assert_eq!(map_field.name(), "key_value");
4592 assert_eq!(map_field.get_fields()[0].name(), "key");
4593 assert_eq!(map_field.get_fields()[1].name(), "value");
4594 }
4595
4596 #[test]
4597 #[cfg_attr(miri, ignore)] fn fallback_flush_data_page() {
4599 let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4601 let values = Arc::new(StringArray::from(raw_values));
4602 let encodings = vec![
4603 Encoding::DELTA_BYTE_ARRAY,
4604 Encoding::DELTA_LENGTH_BYTE_ARRAY,
4605 ];
4606 let data_type = values.data_type().clone();
4607 let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4608 let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4609
4610 let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4611 let data_page_size_limit: usize = 32;
4612 let write_batch_size: usize = 16;
4613
4614 for encoding in &encodings {
4615 for row_group_size in row_group_sizes {
4616 let props = WriterProperties::builder()
4617 .set_writer_version(WriterVersion::PARQUET_2_0)
4618 .set_max_row_group_row_count(Some(row_group_size))
4619 .set_dictionary_enabled(false)
4620 .set_encoding(*encoding)
4621 .set_data_page_size_limit(data_page_size_limit)
4622 .set_write_batch_size(write_batch_size)
4623 .build();
4624
4625 roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4626 let string_array_a = StringArray::from(a.clone());
4627 let string_array_b = StringArray::from(b.clone());
4628 let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4629 let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4630 assert_eq!(
4631 vec_a, vec_b,
4632 "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4633 );
4634 });
4635 }
4636 }
4637 }
4638
4639 #[test]
4640 #[cfg_attr(miri, ignore)] fn u32_min_max() {
4642 let src = [
4644 u32::MIN,
4645 1,
4646 (i32::MAX as u32) - 1,
4647 i32::MAX as u32,
4648 (i32::MAX as u32) + 1,
4649 u32::MAX - 1,
4650 u32::MAX,
4651 ];
4652 let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4653 let files = RoundTripTest::new(values).with_nullable(false).run();
4654
4655 for file in files {
4656 let reader = SerializedFileReader::new(file).unwrap();
4658 let metadata = reader.metadata();
4659
4660 let mut row_offset = 0;
4661 for row_group in metadata.row_groups() {
4662 assert_eq!(row_group.num_columns(), 1);
4663 let column = row_group.column(0);
4664
4665 let num_values = column.num_values() as usize;
4666 let src_slice = &src[row_offset..row_offset + num_values];
4667 row_offset += column.num_values() as usize;
4668
4669 let stats = column.statistics().unwrap();
4670 if let Statistics::Int32(stats) = stats {
4671 assert_eq!(
4672 *stats.min_opt().unwrap() as u32,
4673 *src_slice.iter().min().unwrap()
4674 );
4675 assert_eq!(
4676 *stats.max_opt().unwrap() as u32,
4677 *src_slice.iter().max().unwrap()
4678 );
4679 } else {
4680 panic!("Statistics::Int32 missing")
4681 }
4682 }
4683 }
4684 }
4685
4686 #[test]
4687 #[cfg_attr(miri, ignore)] fn u64_min_max() {
4689 let src = [
4691 u64::MIN,
4692 1,
4693 (i64::MAX as u64) - 1,
4694 i64::MAX as u64,
4695 (i64::MAX as u64) + 1,
4696 u64::MAX - 1,
4697 u64::MAX,
4698 ];
4699 let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4700 let files = RoundTripTest::new(values).with_nullable(false).run();
4701
4702 for file in files {
4703 let reader = SerializedFileReader::new(file).unwrap();
4705 let metadata = reader.metadata();
4706
4707 let mut row_offset = 0;
4708 for row_group in metadata.row_groups() {
4709 assert_eq!(row_group.num_columns(), 1);
4710 let column = row_group.column(0);
4711
4712 let num_values = column.num_values() as usize;
4713 let src_slice = &src[row_offset..row_offset + num_values];
4714 row_offset += column.num_values() as usize;
4715
4716 let stats = column.statistics().unwrap();
4717 if let Statistics::Int64(stats) = stats {
4718 assert_eq!(
4719 *stats.min_opt().unwrap() as u64,
4720 *src_slice.iter().min().unwrap()
4721 );
4722 assert_eq!(
4723 *stats.max_opt().unwrap() as u64,
4724 *src_slice.iter().max().unwrap()
4725 );
4726 } else {
4727 panic!("Statistics::Int64 missing")
4728 }
4729 }
4730 }
4731 }
4732
4733 #[test]
4734 #[cfg_attr(miri, ignore)] fn statistics_null_counts_only_nulls() {
4736 let values = Arc::new(UInt64Array::from(vec![None, None]));
4738 let files = RoundTripTest::new(values).run();
4739
4740 for file in files {
4741 let reader = SerializedFileReader::new(file).unwrap();
4743 let metadata = reader.metadata();
4744 assert_eq!(metadata.num_row_groups(), 1);
4745 let row_group = metadata.row_group(0);
4746 assert_eq!(row_group.num_columns(), 1);
4747 let column = row_group.column(0);
4748 let stats = column.statistics().unwrap();
4749 assert_eq!(stats.null_count_opt(), Some(2));
4750 }
4751 }
4752
4753 fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4754 metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4755 }
4756
4757 #[test]
4758 fn test_aggregates_records() {
4759 let arrays = [
4760 Int32Array::from((0..100).collect::<Vec<_>>()),
4761 Int32Array::from((0..50).collect::<Vec<_>>()),
4762 Int32Array::from((200..500).collect::<Vec<_>>()),
4763 ];
4764
4765 let schema = Arc::new(Schema::new(vec![Field::new(
4766 "int",
4767 ArrowDataType::Int32,
4768 false,
4769 )]));
4770
4771 let file = tempfile::tempfile().unwrap();
4772
4773 let props = WriterProperties::builder()
4774 .set_max_row_group_row_count(Some(200))
4775 .build();
4776
4777 let mut writer =
4778 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4779
4780 for array in arrays {
4781 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4782 writer.write(&batch).unwrap();
4783 }
4784
4785 writer.close().unwrap();
4786
4787 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4788 assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4789
4790 let batches = builder
4791 .with_batch_size(100)
4792 .build()
4793 .unwrap()
4794 .collect::<ArrowResult<Vec<_>>>()
4795 .unwrap();
4796
4797 assert_eq!(batches.len(), 5);
4798 assert!(batches.iter().all(|x| x.num_columns() == 1));
4799
4800 let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4801
4802 assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4803
4804 let values: Vec<_> = batches
4805 .iter()
4806 .flat_map(|x| {
4807 x.column(0)
4808 .as_any()
4809 .downcast_ref::<Int32Array>()
4810 .unwrap()
4811 .values()
4812 .iter()
4813 .copied()
4814 })
4815 .collect();
4816
4817 let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4818 assert_eq!(&values, &expected_values)
4819 }
4820
4821 #[test]
4822 fn complex_aggregate() {
4823 let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4825 let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4826 let struct_a = Arc::new(Field::new(
4827 "struct_a",
4828 DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4829 true,
4830 ));
4831
4832 let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4833 let struct_b = Arc::new(Field::new(
4834 "struct_b",
4835 DataType::Struct(vec![list_a.clone()].into()),
4836 false,
4837 ));
4838
4839 let schema = Arc::new(Schema::new(vec![struct_b]));
4840
4841 let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4843 let field_b_array =
4844 Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4845
4846 let struct_a_array = StructArray::from(vec![
4847 (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4848 (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4849 ]);
4850
4851 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4852 .len(5)
4853 .add_buffer(Buffer::from_iter(vec![
4854 0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4855 ]))
4856 .null_bit_buffer(Some(Buffer::from_iter(vec![
4857 true, false, true, false, true,
4858 ])))
4859 .child_data(vec![struct_a_array.into_data()])
4860 .build()
4861 .unwrap();
4862
4863 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4864 let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4865
4866 let batch1 =
4867 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4868 .unwrap();
4869
4870 let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4871 let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4872
4873 let struct_a_array = StructArray::from(vec![
4874 (field_a, Arc::new(field_a_array) as ArrayRef),
4875 (field_b, Arc::new(field_b_array) as ArrayRef),
4876 ]);
4877
4878 let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4879 .len(2)
4880 .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4881 .child_data(vec![struct_a_array.into_data()])
4882 .build()
4883 .unwrap();
4884
4885 let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4886 let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4887
4888 let batch2 =
4889 RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4890 .unwrap();
4891
4892 let batches = &[batch1, batch2];
4893
4894 let expected = r"
4897 +-------------------------------------------------------------------------------------------------------+
4898 | struct_b |
4899 +-------------------------------------------------------------------------------------------------------+
4900 | {list: [{leaf_a: 1, leaf_b: 1}]} |
4901 | {list: } |
4902 | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]} |
4903 | {list: } |
4904 | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]} |
4905 | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4906 | {list: [{leaf_a: 10, leaf_b: }]} |
4907 +-------------------------------------------------------------------------------------------------------+
4908 ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4909
4910 let actual = pretty_format_batches(batches).unwrap().to_string();
4911 assert_eq!(actual, expected);
4912
4913 let file = tempfile::tempfile().unwrap();
4915 let props = WriterProperties::builder()
4916 .set_max_row_group_row_count(Some(6))
4917 .build();
4918
4919 let mut writer =
4920 ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4921
4922 for batch in batches {
4923 writer.write(batch).unwrap();
4924 }
4925 writer.close().unwrap();
4926
4927 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4932 assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4933
4934 let batches = builder
4935 .with_batch_size(2)
4936 .build()
4937 .unwrap()
4938 .collect::<ArrowResult<Vec<_>>>()
4939 .unwrap();
4940
4941 assert_eq!(batches.len(), 4);
4942 let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4943 assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4944
4945 let actual = pretty_format_batches(&batches).unwrap().to_string();
4946 assert_eq!(actual, expected);
4947 }
4948
4949 #[test]
4950 fn test_arrow_writer_metadata() {
4951 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4952 let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4953
4954 let batch = RecordBatch::try_new(
4955 Arc::new(batch_schema),
4956 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4957 )
4958 .unwrap();
4959
4960 let mut buf = Vec::with_capacity(1024);
4961 let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4962 writer.write(&batch).unwrap();
4963 writer.close().unwrap();
4964 }
4965
4966 #[test]
4967 fn in_progress_accounting() {
4968 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4970
4971 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4973
4974 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4976
4977 let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4978
4979 assert_eq!(writer.in_progress_size(), 0);
4981 assert_eq!(writer.in_progress_rows(), 0);
4982 assert_eq!(writer.memory_size(), 0);
4983 assert_eq!(writer.bytes_written(), 4); writer.write(&batch).unwrap();
4985
4986 let initial_size = writer.in_progress_size();
4988 assert!(initial_size > 0);
4989 assert_eq!(writer.in_progress_rows(), 5);
4990 let initial_memory = writer.memory_size();
4991 assert!(initial_memory > 0);
4992 assert!(
4994 initial_size <= initial_memory,
4995 "{initial_size} <= {initial_memory}"
4996 );
4997
4998 writer.write(&batch).unwrap();
5000 assert!(writer.in_progress_size() > initial_size);
5001 assert_eq!(writer.in_progress_rows(), 10);
5002 assert!(writer.memory_size() > initial_memory);
5003 assert!(
5004 writer.in_progress_size() <= writer.memory_size(),
5005 "in_progress_size {} <= memory_size {}",
5006 writer.in_progress_size(),
5007 writer.memory_size()
5008 );
5009
5010 let pre_flush_bytes_written = writer.bytes_written();
5012 writer.flush().unwrap();
5013 assert_eq!(writer.in_progress_size(), 0);
5014 assert_eq!(writer.memory_size(), 0);
5015 assert!(writer.bytes_written() > pre_flush_bytes_written);
5016
5017 writer.close().unwrap();
5018 }
5019
5020 #[test]
5021 fn test_writer_all_null() {
5022 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5023 let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
5024 let batch = RecordBatch::try_from_iter(vec![
5025 ("a", Arc::new(a) as ArrayRef),
5026 ("b", Arc::new(b) as ArrayRef),
5027 ])
5028 .unwrap();
5029
5030 let mut buf = Vec::with_capacity(1024);
5031 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5032 writer.write(&batch).unwrap();
5033 writer.close().unwrap();
5034
5035 let bytes = Bytes::from(buf);
5036 let options = ReadOptionsBuilder::new().with_page_index().build();
5037 let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
5038 let index = reader.metadata().page_index().unwrap();
5039
5040 assert_eq!(index.num_data_pages(0, 0), Some(1)); assert_eq!(index.num_data_pages(0, 1), Some(1)); }
5043
5044 #[test]
5045 fn test_disabled_statistics_with_page() {
5046 let file_schema = Schema::new(vec![
5047 Field::new("a", DataType::Utf8, true),
5048 Field::new("b", DataType::Utf8, true),
5049 ]);
5050 let file_schema = Arc::new(file_schema);
5051
5052 let batch = RecordBatch::try_new(
5053 file_schema.clone(),
5054 vec![
5055 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5056 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5057 ],
5058 )
5059 .unwrap();
5060
5061 let props = WriterProperties::builder()
5062 .set_statistics_enabled(EnabledStatistics::None)
5063 .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5064 .build();
5065
5066 let mut buf = Vec::with_capacity(1024);
5067 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5068 writer.write(&batch).unwrap();
5069
5070 let metadata = writer.close().unwrap();
5071 assert_eq!(metadata.num_row_groups(), 1);
5072 let row_group = metadata.row_group(0);
5073 assert_eq!(row_group.num_columns(), 2);
5074 assert!(row_group.column(0).offset_index_offset().is_some());
5076 assert!(row_group.column(0).column_index_offset().is_some());
5077 assert!(row_group.column(1).offset_index_offset().is_some());
5079 assert!(row_group.column(1).column_index_offset().is_none());
5080
5081 let options = ReadOptionsBuilder::new().with_page_index().build();
5082 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5083
5084 let row_group = reader.get_row_group(0).unwrap();
5085 let a_col = row_group.metadata().column(0);
5086 let b_col = row_group.metadata().column(1);
5087
5088 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5090 let min = byte_array_stats.min_opt().unwrap();
5091 let max = byte_array_stats.max_opt().unwrap();
5092
5093 assert_eq!(min.as_bytes(), b"a");
5094 assert_eq!(max.as_bytes(), b"d");
5095 } else {
5096 panic!("expecting Statistics::ByteArray");
5097 }
5098
5099 assert!(b_col.statistics().is_none());
5101
5102 let page_index = reader.metadata().page_index().unwrap();
5103
5104 let a_idx = page_index.column_index(0, 0);
5105 assert!(
5106 matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
5107 "{a_idx:?}"
5108 );
5109 let b_idx = page_index.column_index(0, 1);
5110 assert!(b_idx.is_none(), "{b_idx:?}");
5111 }
5112
5113 #[test]
5114 fn test_disabled_statistics_with_chunk() {
5115 let file_schema = Schema::new(vec![
5116 Field::new("a", DataType::Utf8, true),
5117 Field::new("b", DataType::Utf8, true),
5118 ]);
5119 let file_schema = Arc::new(file_schema);
5120
5121 let batch = RecordBatch::try_new(
5122 file_schema.clone(),
5123 vec![
5124 Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5125 Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5126 ],
5127 )
5128 .unwrap();
5129
5130 let props = WriterProperties::builder()
5131 .set_statistics_enabled(EnabledStatistics::None)
5132 .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5133 .build();
5134
5135 let mut buf = Vec::with_capacity(1024);
5136 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5137 writer.write(&batch).unwrap();
5138
5139 let metadata = writer.close().unwrap();
5140 assert_eq!(metadata.num_row_groups(), 1);
5141 let row_group = metadata.row_group(0);
5142 assert_eq!(row_group.num_columns(), 2);
5143 assert!(row_group.column(0).offset_index_offset().is_some());
5145 assert!(row_group.column(0).column_index_offset().is_none());
5146 assert!(row_group.column(1).offset_index_offset().is_some());
5148 assert!(row_group.column(1).column_index_offset().is_none());
5149
5150 let options = ReadOptionsBuilder::new().with_page_index().build();
5151 let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5152
5153 let row_group = reader.get_row_group(0).unwrap();
5154 let a_col = row_group.metadata().column(0);
5155 let b_col = row_group.metadata().column(1);
5156
5157 if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5159 let min = byte_array_stats.min_opt().unwrap();
5160 let max = byte_array_stats.max_opt().unwrap();
5161
5162 assert_eq!(min.as_bytes(), b"a");
5163 assert_eq!(max.as_bytes(), b"d");
5164 } else {
5165 panic!("expecting Statistics::ByteArray");
5166 }
5167
5168 assert!(b_col.statistics().is_none());
5170
5171 let page_index = reader.metadata().page_index().unwrap();
5172
5173 let a_idx = page_index.column_index(0, 0);
5174 assert!(a_idx.is_none(), "{a_idx:?}");
5175 let b_idx = page_index.column_index(0, 1);
5176 assert!(b_idx.is_none(), "{b_idx:?}");
5177 }
5178
5179 #[test]
5180 fn test_arrow_writer_skip_metadata() {
5181 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5182 let file_schema = Arc::new(batch_schema.clone());
5183
5184 let batch = RecordBatch::try_new(
5185 Arc::new(batch_schema),
5186 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5187 )
5188 .unwrap();
5189 let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5190
5191 let mut buf = Vec::with_capacity(1024);
5192 let mut writer =
5193 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5194 writer.write(&batch).unwrap();
5195 writer.close().unwrap();
5196
5197 let bytes = Bytes::from(buf);
5198 let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5199 assert_eq!(file_schema, *reader_builder.schema());
5200 if let Some(key_value_metadata) = reader_builder
5201 .metadata()
5202 .file_metadata()
5203 .key_value_metadata()
5204 {
5205 assert!(
5206 !key_value_metadata
5207 .iter()
5208 .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5209 );
5210 }
5211 }
5212
5213 #[test]
5214 fn test_arrow_writer_skip_path_in_schema() {
5215 let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5216 let file_schema = Arc::new(batch_schema.clone());
5217
5218 let batch = RecordBatch::try_new(
5219 Arc::new(batch_schema),
5220 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5221 )
5222 .unwrap();
5223
5224 let skip_options = ArrowWriterOptions::new();
5226
5227 let mut buf = Vec::with_capacity(1024);
5228 let mut writer =
5229 ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5230 writer.write(&batch).unwrap();
5231 writer.close().unwrap();
5232
5233 let skip_options = ArrowWriterOptions::new().with_properties(
5235 WriterProperties::builder()
5236 .set_write_path_in_schema(false)
5237 .build(),
5238 );
5239
5240 let mut buf2 = Vec::with_capacity(1024);
5241 let mut writer =
5242 ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5243 .unwrap();
5244 writer.write(&batch).unwrap();
5245 writer.close().unwrap();
5246
5247 assert!(buf.len() > buf2.len());
5249 }
5250
5251 #[test]
5252 fn mismatched_schemas() {
5253 let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5254 let file_schema = Arc::new(Schema::new(vec![Field::new(
5255 "temperature",
5256 DataType::Float64,
5257 false,
5258 )]));
5259
5260 let batch = RecordBatch::try_new(
5261 Arc::new(batch_schema),
5262 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5263 )
5264 .unwrap();
5265
5266 let mut buf = Vec::with_capacity(1024);
5267 let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5268
5269 let err = writer.write(&batch).unwrap_err().to_string();
5270 assert_eq!(
5271 err,
5272 "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5273 );
5274 }
5275
5276 #[test]
5277 fn test_page_stats_not_written_by_default() {
5278 let string_field = Field::new("a", DataType::Utf8, false);
5279 let schema = Schema::new(vec![string_field]);
5280 let raw_string_values = vec!["Blart Versenwald III"];
5281 let string_values = StringArray::from(raw_string_values.clone());
5282 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5283
5284 let props = WriterProperties::builder()
5285 .set_statistics_enabled(EnabledStatistics::Page)
5286 .set_dictionary_enabled(false)
5287 .set_encoding(Encoding::PLAIN)
5288 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5289 .build();
5290
5291 let file = roundtrip_opts(&batch, props);
5292
5293 let first_page = &file[4..];
5298 let mut prot = ThriftSliceInputProtocol::new(first_page);
5299 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5300 let stats = hdr.data_page_header.unwrap().statistics;
5301
5302 assert!(stats.is_none());
5303 }
5304
5305 #[test]
5306 fn test_page_stats_when_enabled() {
5307 let string_field = Field::new("a", DataType::Utf8, false);
5308 let schema = Schema::new(vec![string_field]);
5309 let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5310 let string_values = StringArray::from(raw_string_values.clone());
5311 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5312
5313 let props = WriterProperties::builder()
5314 .set_statistics_enabled(EnabledStatistics::Page)
5315 .set_dictionary_enabled(false)
5316 .set_encoding(Encoding::PLAIN)
5317 .set_write_page_header_statistics(true)
5318 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5319 .build();
5320
5321 let file = roundtrip_opts(&batch, props);
5322
5323 let first_page = &file[4..];
5328 let mut prot = ThriftSliceInputProtocol::new(first_page);
5329 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5330 let stats = hdr.data_page_header.unwrap().statistics;
5331
5332 let stats = stats.unwrap();
5333 assert!(stats.is_max_value_exact.unwrap());
5335 assert!(stats.is_min_value_exact.unwrap());
5336 assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5337 assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5338 }
5339
5340 #[test]
5341 fn test_page_stats_truncation() {
5342 let string_field = Field::new("a", DataType::Utf8, false);
5343 let binary_field = Field::new("b", DataType::Binary, false);
5344 let schema = Schema::new(vec![string_field, binary_field]);
5345
5346 let raw_string_values = vec!["Blart Versenwald III"];
5347 let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5348 let raw_binary_value_refs = raw_binary_values
5349 .iter()
5350 .map(|x| x.as_slice())
5351 .collect::<Vec<_>>();
5352
5353 let string_values = StringArray::from(raw_string_values.clone());
5354 let binary_values = BinaryArray::from(raw_binary_value_refs);
5355 let batch = RecordBatch::try_new(
5356 Arc::new(schema),
5357 vec![Arc::new(string_values), Arc::new(binary_values)],
5358 )
5359 .unwrap();
5360
5361 let props = WriterProperties::builder()
5362 .set_statistics_truncate_length(Some(2))
5363 .set_dictionary_enabled(false)
5364 .set_encoding(Encoding::PLAIN)
5365 .set_write_page_header_statistics(true)
5366 .set_compression(crate::basic::Compression::UNCOMPRESSED)
5367 .build();
5368
5369 let file = roundtrip_opts(&batch, props);
5370
5371 let first_page = &file[4..];
5376 let mut prot = ThriftSliceInputProtocol::new(first_page);
5377 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5378 let stats = hdr.data_page_header.unwrap().statistics;
5379 assert!(stats.is_some());
5380 let stats = stats.unwrap();
5381 assert!(!stats.is_max_value_exact.unwrap());
5383 assert!(!stats.is_min_value_exact.unwrap());
5384 assert_eq!(stats.max_value.unwrap(), b"Bm");
5385 assert_eq!(stats.min_value.unwrap(), b"Bl");
5386
5387 let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5389 let mut prot = ThriftSliceInputProtocol::new(second_page);
5390 let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5391 let stats = hdr.data_page_header.unwrap().statistics;
5392 assert!(stats.is_some());
5393 let stats = stats.unwrap();
5394 assert!(!stats.is_max_value_exact.unwrap());
5396 assert!(!stats.is_min_value_exact.unwrap());
5397 assert_eq!(stats.max_value.unwrap(), b"Bm");
5398 assert_eq!(stats.min_value.unwrap(), b"Bl");
5399 }
5400
5401 #[test]
5402 fn test_page_encoding_statistics_roundtrip() {
5403 let batch_schema = Schema::new(vec![Field::new(
5404 "int32",
5405 arrow_schema::DataType::Int32,
5406 false,
5407 )]);
5408
5409 let batch = RecordBatch::try_new(
5410 Arc::new(batch_schema.clone()),
5411 vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5412 )
5413 .unwrap();
5414
5415 let mut file: File = tempfile::tempfile().unwrap();
5416 let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5417 writer.write(&batch).unwrap();
5418 let file_metadata = writer.close().unwrap();
5419
5420 assert_eq!(file_metadata.num_row_groups(), 1);
5421 assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5422 assert!(
5423 file_metadata
5424 .row_group(0)
5425 .column(0)
5426 .page_encoding_stats()
5427 .is_some()
5428 );
5429 let chunk_page_stats = file_metadata
5430 .row_group(0)
5431 .column(0)
5432 .page_encoding_stats()
5433 .unwrap();
5434
5435 let options = ReadOptionsBuilder::new()
5437 .with_page_index()
5438 .with_encoding_stats_as_mask(false)
5439 .build();
5440 let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5441
5442 let rowgroup = reader.get_row_group(0).expect("row group missing");
5443 assert_eq!(rowgroup.num_columns(), 1);
5444 let column = rowgroup.metadata().column(0);
5445 assert!(column.page_encoding_stats().is_some());
5446 let file_page_stats = column.page_encoding_stats().unwrap();
5447 assert_eq!(chunk_page_stats, file_page_stats);
5448 }
5449
5450 #[test]
5451 #[cfg_attr(miri, ignore)] fn test_different_dict_page_size_limit() {
5453 let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5454 let schema = Arc::new(Schema::new(vec![
5455 Field::new("col0", arrow_schema::DataType::Int64, false),
5456 Field::new("col1", arrow_schema::DataType::Int64, false),
5457 ]));
5458 let batch =
5459 arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5460
5461 let props = WriterProperties::builder()
5462 .set_dictionary_page_size_limit(1024 * 1024)
5463 .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5464 .build();
5465 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5466 writer.write(&batch).unwrap();
5467 let data = Bytes::from(writer.into_inner().unwrap());
5468
5469 let mut metadata = ParquetMetaDataReader::new();
5470 metadata.try_parse(&data).unwrap();
5471 let metadata = metadata.finish().unwrap();
5472 let col0_meta = metadata.row_group(0).column(0);
5473 let col1_meta = metadata.row_group(0).column(1);
5474
5475 let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5476 let mut reader =
5477 SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5478 let page = reader.get_next_page().unwrap().unwrap();
5479 match page {
5480 Page::DictionaryPage { buf, .. } => buf.len(),
5481 _ => panic!("expected DictionaryPage"),
5482 }
5483 };
5484
5485 assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5486 assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5487 }
5488
5489 #[test]
5490 #[cfg_attr(miri, ignore)] fn test_arrow_writer_granular_mode_roundtrip() {
5492 let small = "tiny".to_string();
5501 let big = "x".repeat(64 * 1024);
5502 let strings: Vec<String> = (0..256)
5503 .map(|i| {
5504 if i % 16 == 0 {
5505 big.clone()
5506 } else {
5507 small.clone()
5508 }
5509 })
5510 .collect();
5511
5512 let schema = Arc::new(Schema::new(vec![Field::new(
5513 "col",
5514 ArrowDataType::Utf8,
5515 false,
5516 )]));
5517 let batch = RecordBatch::try_new(
5518 schema.clone(),
5519 vec![Arc::new(StringArray::from(strings.clone())) as _],
5520 )
5521 .unwrap();
5522
5523 let props = WriterProperties::builder()
5524 .set_dictionary_enabled(false)
5525 .set_data_page_size_limit(16 * 1024)
5526 .build();
5527 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5528 writer.write(&batch).unwrap();
5529 let data = Bytes::from(writer.into_inner().unwrap());
5530
5531 let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5532 let read = reader.next().unwrap().unwrap();
5533 assert!(reader.next().is_none(), "expected one batch");
5534 let col = read
5535 .column(0)
5536 .as_any()
5537 .downcast_ref::<StringArray>()
5538 .unwrap();
5539 assert_eq!(col.len(), strings.len());
5540 for (i, expected) in strings.iter().enumerate() {
5541 assert_eq!(
5542 col.value(i),
5543 expected.as_str(),
5544 "value mismatch at index {i}"
5545 );
5546 }
5547 }
5548
5549 #[test]
5550 fn test_arrow_writer_all_null_string_column() {
5551 let num_rows = 1024;
5556 let schema = Arc::new(Schema::new(vec![Field::new(
5557 "col",
5558 ArrowDataType::Utf8,
5559 true,
5560 )]));
5561 let nulls: Vec<Option<&str>> = vec![None; num_rows];
5562 let batch = RecordBatch::try_new(
5563 schema.clone(),
5564 vec![Arc::new(StringArray::from(nulls)) as _],
5565 )
5566 .unwrap();
5567
5568 let props = WriterProperties::builder()
5569 .set_dictionary_enabled(false)
5570 .set_data_page_size_limit(16 * 1024)
5571 .build();
5572 let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5573 writer.write(&batch).unwrap();
5574 let data = Bytes::from(writer.into_inner().unwrap());
5575
5576 let mut metadata = ParquetMetaDataReader::new();
5579 metadata.try_parse(&data).unwrap();
5580 let metadata = metadata.finish().unwrap();
5581 let row_group = metadata.row_group(0);
5582 let col_meta = row_group.column(0);
5583 assert_eq!(row_group.num_rows() as usize, num_rows);
5584 if let Some(stats) = col_meta.statistics() {
5587 assert_eq!(
5588 stats.null_count_opt().unwrap_or(0) as usize,
5589 num_rows,
5590 "expected all-null column to report null_count = num_rows"
5591 );
5592 }
5593
5594 let mut reader =
5595 SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5596 let mut total_values = 0u32;
5597 while let Some(page) = reader.get_next_page().unwrap() {
5598 if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5599 total_values += page.num_values();
5600 }
5601 }
5602 assert_eq!(
5603 total_values as usize, num_rows,
5604 "expected every level position to be represented in some page"
5605 );
5606 }
5607
5608 struct WriteBatchesShape {
5609 num_batches: usize,
5610 rows_per_batch: usize,
5611 row_size: usize,
5612 }
5613
5614 fn write_batches(
5616 WriteBatchesShape {
5617 num_batches,
5618 rows_per_batch,
5619 row_size,
5620 }: WriteBatchesShape,
5621 props: WriterProperties,
5622 ) -> ParquetRecordBatchReaderBuilder<File> {
5623 let schema = Arc::new(Schema::new(vec![Field::new(
5624 "str",
5625 ArrowDataType::Utf8,
5626 false,
5627 )]));
5628 let file = tempfile::tempfile().unwrap();
5629 let mut writer =
5630 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5631
5632 for batch_idx in 0..num_batches {
5633 let strings: Vec<String> = (0..rows_per_batch)
5634 .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5635 .collect();
5636 let array = StringArray::from(strings);
5637 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5638 writer.write(&batch).unwrap();
5639 }
5640 writer.close().unwrap();
5641 ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5642 }
5643
5644 #[test]
5645 fn test_row_group_limit_none_writes_single_row_group() {
5647 let props = WriterProperties::builder()
5648 .set_max_row_group_row_count(None)
5649 .set_max_row_group_bytes(None)
5650 .build();
5651
5652 let builder = write_batches(
5653 WriteBatchesShape {
5654 num_batches: 1,
5655 rows_per_batch: 1000,
5656 row_size: 4,
5657 },
5658 props,
5659 );
5660
5661 assert_eq!(
5662 &row_group_sizes(builder.metadata()),
5663 &[1000],
5664 "With no limits, all rows should be in a single row group"
5665 );
5666 }
5667
5668 #[test]
5669 fn test_row_group_limit_rows_only() {
5671 let props = WriterProperties::builder()
5672 .set_max_row_group_row_count(Some(300))
5673 .set_max_row_group_bytes(None)
5674 .build();
5675
5676 let builder = write_batches(
5677 WriteBatchesShape {
5678 num_batches: 1,
5679 rows_per_batch: 1000,
5680 row_size: 4,
5681 },
5682 props,
5683 );
5684
5685 assert_eq!(
5686 &row_group_sizes(builder.metadata()),
5687 &[300, 300, 300, 100],
5688 "Row groups should be split by row count"
5689 );
5690 }
5691
5692 #[test]
5693 #[cfg_attr(miri, ignore)] fn test_row_group_limit_rows_only_many_splits() {
5697 let props = WriterProperties::builder()
5698 .set_max_row_group_row_count(Some(1))
5699 .set_max_row_group_bytes(None)
5700 .build();
5701
5702 let rows = 50_000;
5703 let builder = write_batches(
5704 WriteBatchesShape {
5705 num_batches: 1,
5706 rows_per_batch: rows,
5707 row_size: 4,
5708 },
5709 props,
5710 );
5711
5712 let sizes = row_group_sizes(builder.metadata());
5713 assert_eq!(sizes.len(), rows, "Every row should get its own row group");
5714 assert_eq!(
5715 sizes.iter().sum::<i64>(),
5716 rows as i64,
5717 "Total rows should be preserved"
5718 );
5719 }
5720
5721 #[test]
5722 fn test_row_group_limit_bytes_only() {
5724 let props = WriterProperties::builder()
5725 .set_max_row_group_row_count(None)
5726 .set_max_row_group_bytes(Some(3500))
5728 .build();
5729
5730 let builder = write_batches(
5731 WriteBatchesShape {
5732 num_batches: 10,
5733 rows_per_batch: 10,
5734 row_size: 100,
5735 },
5736 props,
5737 );
5738
5739 let sizes = row_group_sizes(builder.metadata());
5740
5741 assert!(
5742 sizes.len() > 1,
5743 "Should have multiple row groups due to byte limit, got {sizes:?}",
5744 );
5745
5746 let total_rows: i64 = sizes.iter().sum();
5747 assert_eq!(total_rows, 100, "Total rows should be preserved");
5748 }
5749
5750 #[test]
5751 fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5753 let schema = Arc::new(Schema::new(vec![Field::new(
5754 "str",
5755 ArrowDataType::Utf8,
5756 false,
5757 )]));
5758 let file = tempfile::tempfile().unwrap();
5759
5760 let props = WriterProperties::builder()
5762 .set_max_row_group_row_count(None)
5763 .set_max_row_group_bytes(None)
5764 .build();
5765 let mut writer =
5766 ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5767
5768 let first_array = StringArray::from(
5769 (0..10)
5770 .map(|i| format!("{i:0>100}"))
5771 .collect::<Vec<String>>(),
5772 );
5773 let first_batch =
5774 RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5775 writer.write(&first_batch).unwrap();
5776 assert_eq!(writer.in_progress_rows(), 10);
5777
5778 writer.max_row_group_bytes = Some(1);
5781
5782 let second_array = StringArray::from(vec!["x".to_string()]);
5783 let second_batch =
5784 RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5785 writer.write(&second_batch).unwrap();
5786 writer.close().unwrap();
5787 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5788
5789 assert_eq!(
5790 &row_group_sizes(builder.metadata()),
5791 &[10, 1],
5792 "The second write should flush an oversized in-progress row group first",
5793 );
5794 }
5795
5796 #[test]
5797 fn test_row_group_limit_both_row_wins_single_batch() {
5799 let props = WriterProperties::builder()
5800 .set_max_row_group_row_count(Some(200)) .set_max_row_group_bytes(Some(1024 * 1024)) .build();
5803
5804 let builder = write_batches(
5805 WriteBatchesShape {
5806 num_batches: 1,
5807 row_size: 4,
5808 rows_per_batch: 1000,
5809 },
5810 props,
5811 );
5812
5813 assert_eq!(
5814 &row_group_sizes(builder.metadata()),
5815 &[200, 200, 200, 200, 200],
5816 "Row limit should trigger before byte limit"
5817 );
5818 }
5819
5820 #[test]
5821 fn test_row_group_limit_both_row_wins_multiple_batches() {
5823 let props = WriterProperties::builder()
5824 .set_max_row_group_row_count(Some(5)) .set_max_row_group_bytes(Some(9999)) .build();
5827
5828 let builder = write_batches(
5829 WriteBatchesShape {
5830 num_batches: 10,
5831 rows_per_batch: 10,
5832 row_size: 100,
5833 },
5834 props,
5835 );
5836
5837 assert_eq!(
5838 &row_group_sizes(builder.metadata()),
5839 &[5; 20],
5840 "Row limit should trigger before byte limit"
5841 );
5842 }
5843
5844 #[test]
5845 fn test_row_group_limit_both_bytes_wins() {
5847 let props = WriterProperties::builder()
5848 .set_max_row_group_row_count(Some(1000)) .set_max_row_group_bytes(Some(3500)) .build();
5851
5852 let builder = write_batches(
5853 WriteBatchesShape {
5854 num_batches: 10,
5855 rows_per_batch: 10,
5856 row_size: 100,
5857 },
5858 props,
5859 );
5860
5861 let sizes = row_group_sizes(builder.metadata());
5862
5863 assert!(
5864 sizes.len() > 1,
5865 "Byte limit should trigger before row limit, got {sizes:?}",
5866 );
5867
5868 assert!(
5869 sizes.iter().all(|&s| s < 1000),
5870 "No row group should hit the row limit"
5871 );
5872
5873 let total_rows: i64 = sizes.iter().sum();
5874 assert_eq!(total_rows, 100, "Total rows should be preserved");
5875 }
5876
5877 #[test]
5878 fn test_row_group_limit_both_apply_to_same_batch() {
5881 let props = WriterProperties::builder()
5882 .set_max_row_group_row_count(Some(15))
5883 .set_max_row_group_bytes(Some(1500))
5884 .build();
5885
5886 let builder = write_batches(
5887 WriteBatchesShape {
5888 num_batches: 2,
5889 rows_per_batch: 10,
5890 row_size: 100,
5891 },
5892 props,
5893 );
5894
5895 assert_eq!(
5896 &row_group_sizes(builder.metadata()),
5897 &[14, 6],
5898 "Byte limit should still apply to a batch the row limit already split"
5899 );
5900 }
5901
5902 #[test]
5903 fn test_row_group_limit_bytes_splits_first_batch() {
5906 let props = WriterProperties::builder()
5907 .set_max_row_group_row_count(None)
5908 .set_max_row_group_bytes(Some(3500))
5909 .build();
5910
5911 let builder = write_batches(
5912 WriteBatchesShape {
5913 num_batches: 3,
5914 rows_per_batch: 100,
5915 row_size: 100,
5916 },
5917 props,
5918 );
5919
5920 let sizes = row_group_sizes(builder.metadata());
5921 assert!(
5922 sizes.iter().all(|&s| s < 100),
5923 "No row group should hold a whole batch, got {sizes:?}",
5924 );
5925 assert_eq!(
5926 sizes.iter().sum::<i64>(),
5927 300,
5928 "Total rows should be preserved"
5929 );
5930 for row_group in builder.metadata().row_groups() {
5931 assert!(
5932 row_group.total_byte_size() <= 3500 + 100,
5933 "Row groups should stay within one row of the byte limit, got {}",
5934 row_group.total_byte_size(),
5935 );
5936 }
5937 }
5938
5939 #[test]
5940 fn arrow_column_chunk_close_mut_drops_column_index() {
5941 use crate::arrow::ArrowSchemaConverter;
5942 use crate::file::writer::SerializedFileWriter;
5943
5944 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5945 let props = Arc::new(
5946 WriterProperties::builder()
5947 .set_statistics_enabled(EnabledStatistics::Page)
5948 .build(),
5949 );
5950 let parquet_schema = ArrowSchemaConverter::new()
5951 .with_coerce_types(props.coerce_types())
5952 .convert(&schema)
5953 .unwrap();
5954
5955 let mut buf = Vec::with_capacity(1024);
5956 let mut writer =
5957 SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5958 .unwrap();
5959
5960 let factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap();
5961 let mut col_writers = factory.create_column_writers(0).unwrap();
5962 let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5963 for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5964 col_writers[0].write(&leaves).unwrap();
5965 }
5966 let mut chunk = col_writers.pop().unwrap().close().unwrap();
5967
5968 assert!(
5970 chunk.close().column_index.is_some(),
5971 "EnabledStatistics::Page should produce a column_index"
5972 );
5973
5974 chunk.close_mut().column_index = None;
5976 assert!(chunk.close().column_index.is_none());
5977
5978 let mut rg = writer.next_row_group().unwrap();
5979 chunk.append_to_row_group(&mut rg).unwrap();
5980 rg.close().unwrap();
5981 let file_meta = writer.close().unwrap();
5982
5983 let cc = file_meta.row_group(0).column(0);
5986 assert!(cc.column_index_range().is_none());
5987 }
5988
5989 #[test]
5990 #[cfg_attr(miri, ignore)] fn test_number_distinct_values_exact_count() {
5992 let cardinality = 50u32;
5995 let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
5996 if i % 7 == 0 {
5997 None
5998 } else {
5999 Some((i % cardinality) as i32)
6000 }
6001 })));
6002 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
6003 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6004
6005 let props = WriterProperties::builder()
6006 .set_write_row_group_number_distinct_values(true)
6007 .build();
6008 let mut buf = Vec::new();
6009 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6010 writer.write(&batch).unwrap();
6011 let metadata = writer.close().unwrap();
6012
6013 let count = metadata
6014 .row_group(0)
6015 .column(0)
6016 .statistics()
6017 .and_then(|s| s.distinct_count_opt())
6018 .expect("distinct_count should be set");
6019 assert_eq!(count, cardinality as u64);
6021 }
6022
6023 #[test]
6024 fn test_number_distinct_values_view_types() {
6025 let cardinality = 5u32;
6028 let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
6029
6030 let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
6031 if i % 4 == 0 {
6032 None
6033 } else {
6034 Some(distinct_strings[(i % cardinality) as usize])
6035 }
6036 })));
6037
6038 let schema = Arc::new(Schema::new(vec![Field::new(
6039 "string_view_col",
6040 DataType::Utf8View,
6041 true,
6042 )]));
6043 let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
6044
6045 let props = WriterProperties::builder()
6046 .set_write_row_group_number_distinct_values(true)
6047 .build();
6048 let mut parquet_bytes = Vec::new();
6049 let mut writer =
6050 ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
6051 writer.write(&batch).unwrap();
6052 let metadata = writer.close().unwrap();
6053
6054 let distinct_count = metadata
6055 .row_group(0)
6056 .column(0)
6057 .statistics()
6058 .and_then(|s| s.distinct_count_opt())
6059 .expect("distinct_count should be set for Utf8View column");
6060 assert_eq!(distinct_count, cardinality as u64);
6061 }
6062
6063 #[test]
6064 fn test_number_distinct_values_not_written_by_default() {
6065 let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
6066 let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
6067 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6068
6069 let mut buf = Vec::new();
6070 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
6071 writer.write(&batch).unwrap();
6072 let metadata = writer.close().unwrap();
6073
6074 let count = metadata
6075 .row_group(0)
6076 .column(0)
6077 .statistics()
6078 .and_then(|s| s.distinct_count_opt());
6079 assert!(count.is_none());
6080 }
6081
6082 #[test]
6083 fn test_dictionary_ndv_single_batch() {
6084 let keys = Int32Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
6088 let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "bird"]));
6089 let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6090
6091 let schema = Arc::new(Schema::new(vec![Field::new(
6092 "x",
6093 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6094 false,
6095 )]));
6096 let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6097
6098 let props = WriterProperties::builder()
6099 .set_write_row_group_number_distinct_values(true)
6100 .build();
6101 let mut buf = Vec::new();
6102 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6103 writer.write(&batch).unwrap();
6104 let metadata = writer.close().unwrap();
6105
6106 let count = metadata
6107 .row_group(0)
6108 .column(0)
6109 .statistics()
6110 .and_then(|s| s.distinct_count_opt())
6111 .expect("distinct_count should be set");
6112 assert_eq!(count, 3);
6113 }
6114
6115 #[test]
6116 fn test_dictionary_ndv_excludes_unreferenced_values() {
6117 let keys = Int32Array::from(vec![0, 1, 0, 1]);
6120 let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "unreferenced"]));
6121 let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6122
6123 let schema = Arc::new(Schema::new(vec![Field::new(
6124 "x",
6125 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6126 false,
6127 )]));
6128 let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6129
6130 let props = WriterProperties::builder()
6131 .set_write_row_group_number_distinct_values(true)
6132 .build();
6133 let mut buf = Vec::new();
6134 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6135 writer.write(&batch).unwrap();
6136 let metadata = writer.close().unwrap();
6137
6138 let count = metadata
6139 .row_group(0)
6140 .column(0)
6141 .statistics()
6142 .and_then(|s| s.distinct_count_opt())
6143 .expect("distinct_count should be set");
6144 assert_eq!(
6145 count, 2,
6146 "unreferenced dictionary values must not count toward NDV"
6147 );
6148 }
6149
6150 #[test]
6151 fn test_dictionary_ndv_across_batches_regression() {
6152 let make_dict_batch = |a: &str, b: &str| -> RecordBatch {
6154 let keys = Int32Array::from(vec![0, 1, 0, 1]);
6155 let values: ArrayRef = Arc::new(StringArray::from(vec![a, b]));
6156 let dict: ArrayRef =
6157 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6158 let schema = Arc::new(Schema::new(vec![Field::new(
6159 "x",
6160 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6161 false,
6162 )]));
6163 RecordBatch::try_new(schema, vec![dict]).unwrap()
6164 };
6165
6166 let batch1 = make_dict_batch("cat", "dog");
6169 let batch2 = make_dict_batch("fish", "cat");
6170
6171 let props = WriterProperties::builder()
6172 .set_write_row_group_number_distinct_values(true)
6173 .build();
6174 let mut buf = Vec::new();
6175 let mut writer = ArrowWriter::try_new(&mut buf, batch1.schema(), Some(props)).unwrap();
6176 writer.write(&batch1).unwrap();
6177 writer.write(&batch2).unwrap();
6178 let metadata = writer.close().unwrap();
6179
6180 let count = metadata
6181 .row_group(0)
6182 .column(0)
6183 .statistics()
6184 .and_then(|s| s.distinct_count_opt())
6185 .expect("distinct_count should be set");
6186 assert_eq!(
6187 count, 3,
6188 "NDV should count distinct values, not distinct key indices"
6189 );
6190 }
6191
6192 #[test]
6193 fn ree_struct_with_ree_child() {
6194 let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6197
6198 let col_a: ArrayRef = Arc::new(
6199 RunArray::try_new(
6200 &run_ends,
6201 &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6202 )
6203 .unwrap(),
6204 );
6205 let col_b: ArrayRef = Arc::new(
6206 RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6207 );
6208
6209 let struct_array: ArrayRef = Arc::new(StructArray::new(
6210 Fields::from(vec![
6211 Field::new("a", col_a.data_type().clone(), true),
6212 Field::new("b", col_b.data_type().clone(), true),
6213 ]),
6214 vec![col_a, col_b],
6215 None,
6216 ));
6217
6218 let schema = Arc::new(Schema::new(vec![Field::new(
6219 "row",
6220 struct_array.data_type().clone(),
6221 true,
6222 )]));
6223 let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6224
6225 let mut buf = Vec::new();
6226 let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6227 writer.write(&batch).unwrap();
6228 let metadata = writer.close().unwrap();
6229
6230 let parquet_schema = metadata.file_metadata().schema_descr();
6231 assert_eq!(parquet_schema.num_columns(), 2);
6232 assert_eq!(
6233 parquet_schema.column(0).physical_type(),
6234 crate::basic::Type::BYTE_ARRAY
6235 );
6236 assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6237 assert_eq!(
6238 parquet_schema.column(1).physical_type(),
6239 crate::basic::Type::INT32
6240 );
6241 assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6242 }
6243
6244 #[test]
6251 fn test_decimal_byte_array_statistics_signed_comparison() {
6252 use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
6253 use crate::data_type::{ByteArray, ByteArrayType};
6254 use crate::schema::types::Type as SchemaType;
6255
6256 const VALUES: [&[u8]; 3] = [&[0xffu8], &[0x00], &[0x01]];
6258
6259 fn new_file(physical: PhysicalType) -> SerializedFileWriter<Vec<u8>> {
6260 let mut field = SchemaType::primitive_type_builder("value", physical)
6261 .with_repetition(Repetition::REQUIRED)
6262 .with_logical_type(Some(LogicalType::decimal(0, 2)))
6263 .with_precision(2)
6264 .with_scale(0);
6265 if physical == PhysicalType::FIXED_LEN_BYTE_ARRAY {
6266 field = field.with_length(1);
6267 }
6268 let root = SchemaType::group_type_builder("root")
6269 .with_fields(vec![Arc::new(field.build().unwrap())])
6270 .build()
6271 .unwrap();
6272 SerializedFileWriter::new(Vec::new(), Arc::new(root), Default::default()).unwrap()
6273 }
6274
6275 fn read_minmax(file: Vec<u8>) -> (i8, i8) {
6276 let reader = SerializedFileReader::new(Bytes::from(file)).unwrap();
6277 let metadata = reader.metadata();
6278 assert_eq!(metadata.file_metadata().num_rows(), 3);
6279 let stats = metadata
6280 .row_group(0)
6281 .column(0)
6282 .statistics()
6283 .expect("missing statistics");
6284 (
6285 stats.min_bytes_opt().unwrap()[0] as i8,
6286 stats.max_bytes_opt().unwrap()[0] as i8,
6287 )
6288 }
6289
6290 fn direct_parquet() -> Vec<u8> {
6291 let mut file = new_file(PhysicalType::BYTE_ARRAY);
6292 let mut group = file.next_row_group().unwrap();
6293 let mut column = group.next_column().unwrap().unwrap();
6294 let values = VALUES
6295 .iter()
6296 .map(|v| ByteArray::from(v.to_vec()))
6297 .collect::<Vec<_>>();
6298 column
6299 .typed::<ByteArrayType>()
6300 .write_batch(&values, None, None)
6301 .unwrap();
6302 column.close().unwrap();
6303 group.close().unwrap();
6304 file.into_inner().unwrap()
6305 }
6306
6307 fn arrow_bounds(mut file: SerializedFileWriter<Vec<u8>>, array: ArrayRef) -> Vec<u8> {
6308 let field = Field::new("value", array.data_type().clone(), false);
6309 let schema = Arc::new(Schema::new(vec![field.clone()]));
6310 let factory = ArrowRowGroupWriterFactory::try_new(&file, schema).unwrap();
6311 let mut writer: ArrowColumnWriter = factory.create_column_writers(0).unwrap().remove(0);
6312 let leaf = compute_leaves(&field, &array).unwrap().remove(0);
6313 writer.write(&leaf).unwrap();
6314 let mut group = file.next_row_group().unwrap();
6315 writer
6316 .close()
6317 .unwrap()
6318 .append_to_row_group(&mut group)
6319 .unwrap();
6320 group.close().unwrap();
6321 file.into_inner().unwrap()
6322 }
6323
6324 fn arrow_binary() -> Vec<u8> {
6325 let file = new_file(PhysicalType::BYTE_ARRAY);
6326 let array: ArrayRef = Arc::new(BinaryArray::from_vec(VALUES.to_vec()));
6327 arrow_bounds(file, array)
6328 }
6329
6330 fn arrow_fixed_size_binary() -> Vec<u8> {
6331 let file = new_file(PhysicalType::FIXED_LEN_BYTE_ARRAY);
6332 let array: ArrayRef =
6333 Arc::new(FixedSizeBinaryArray::try_from_iter(VALUES.into_iter()).unwrap());
6334 arrow_bounds(file, array)
6335 }
6336
6337 let direct = read_minmax(direct_parquet());
6338 let fixed = read_minmax(arrow_fixed_size_binary());
6339 let binary = read_minmax(arrow_binary());
6340
6341 assert_eq!(direct, (-1, 1), "incorrect direct Parquet stats");
6342 assert_eq!(fixed, (-1, 1), "incorrect Arrow FixedSizeBinaryArray stats");
6343 assert_eq!(binary, (-1, 1), "incorrect Arrow BinaryArray stats");
6344 }
6345}