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