adriangb commented on code in PR #9972:
URL: https://github.com/apache/arrow-rs/pull/9972#discussion_r3352027424
##########
parquet/tests/arrow_writer_layout.rs:
##########
Review Comment:
Done — added `test_large_string`, `test_large_string_view`,
`test_large_values_in_list`, `test_nullable_large_values`, and
`test_dictionary_spill_large_values`, porting the two arrow-writer page-size
tests over to `LayoutTest`. See d12657b040. Resolved.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4904,6 +4904,245 @@ mod tests {
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_strings() {
+ // End-to-end coverage for
`ByteArrayEncoder::estimated_value_bytes_gather`
+ // and the offsets-buffer fast path in `estimate_byte_size_offsets`.
+ //
+ // The standard column-writer regression test covers the same
+ // logic for the non-arrow path; this one drives writes through
+ // `ArrowWriter`, which is what real users hit and which uses
+ // `write_gather` with `non_null_indices` (always contiguous for
+ // a non-null column).
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_rows = 32;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let strings: Vec<String> = (0..num_rows).map(|_|
"x".repeat(value_size)).collect();
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings)) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
Review Comment:
Yes — the default is dictionary *enabled*, so this forces plain encoding.
This test was ported to a `LayoutTest` in d12657b040. Resolved.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4904,6 +4904,245 @@ mod tests {
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_strings() {
+ // End-to-end coverage for
`ByteArrayEncoder::estimated_value_bytes_gather`
+ // and the offsets-buffer fast path in `estimate_byte_size_offsets`.
+ //
+ // The standard column-writer regression test covers the same
+ // logic for the non-arrow path; this one drives writes through
+ // `ArrowWriter`, which is what real users hit and which uses
+ // `write_gather` with `non_null_indices` (always contiguous for
+ // a non-null column).
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_rows = 32;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let strings: Vec<String> = (0..num_rows).map(|_|
"x".repeat(value_size)).collect();
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings)) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(page_byte_limit)
+ .build();
+ let mut writer = ArrowWriter::try_new(Vec::new(), schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let data = Bytes::from(writer.into_inner().unwrap());
+
+ let mut metadata = ParquetMetaDataReader::new();
+ metadata.try_parse(&data).unwrap();
+ let metadata = metadata.finish().unwrap();
+ let col_meta = metadata.row_group(0).column(0);
+ let mut reader =
+ SerializedPageReader::new(Arc::new(data.clone()), col_meta,
num_rows, None).unwrap();
+
+ let mut data_pages = Vec::new();
+ while let Some(page) = reader.get_next_page().unwrap() {
+ if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. })
{
+ data_pages.push(page.buffer().len());
+ }
+ }
+
+ assert!(
+ data_pages.len() >= num_rows / 2,
+ "expected at least {} data pages for {num_rows} large strings via
ArrowWriter, got {} ({data_pages:?})",
+ num_rows / 2,
+ data_pages.len(),
+ );
+ for size in &data_pages {
+ assert!(
+ *size <= value_size + 64,
+ "page size {size} exceeds one-value bound; pages
{data_pages:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_arrow_writer_granular_mode_roundtrip() {
+ // Granular mode subdivides chunks and writes more pages than
+ // `main`. Make sure the data we write back is bit-identical to
+ // what went in — page-count assertions elsewhere only prove
+ // pages were cut, not that the encoded data is correct.
+ //
+ // Mix value sizes so that the cumulative-byte-budget cutoff
+ // lands mid-chunk, exercising both batched and granular paths
+ // within the same `write_batch_internal` call.
+ let small = "tiny".to_string();
+ let big = "x".repeat(64 * 1024);
+ let strings: Vec<String> = (0..256)
+ .map(|i| {
+ if i % 16 == 0 {
+ big.clone()
+ } else {
+ small.clone()
+ }
+ })
+ .collect();
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings.clone())) as _],
Review Comment:
This test was ported to a `LayoutTest` in d12657b040; the new version no
longer clones. Resolved.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4904,6 +4904,245 @@ mod tests {
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_strings() {
+ // End-to-end coverage for
`ByteArrayEncoder::estimated_value_bytes_gather`
+ // and the offsets-buffer fast path in `estimate_byte_size_offsets`.
+ //
+ // The standard column-writer regression test covers the same
+ // logic for the non-arrow path; this one drives writes through
+ // `ArrowWriter`, which is what real users hit and which uses
+ // `write_gather` with `non_null_indices` (always contiguous for
+ // a non-null column).
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_rows = 32;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let strings: Vec<String> = (0..num_rows).map(|_|
"x".repeat(value_size)).collect();
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings)) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(page_byte_limit)
+ .build();
+ let mut writer = ArrowWriter::try_new(Vec::new(), schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let data = Bytes::from(writer.into_inner().unwrap());
+
+ let mut metadata = ParquetMetaDataReader::new();
+ metadata.try_parse(&data).unwrap();
+ let metadata = metadata.finish().unwrap();
+ let col_meta = metadata.row_group(0).column(0);
+ let mut reader =
+ SerializedPageReader::new(Arc::new(data.clone()), col_meta,
num_rows, None).unwrap();
+
+ let mut data_pages = Vec::new();
+ while let Some(page) = reader.get_next_page().unwrap() {
+ if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. })
{
+ data_pages.push(page.buffer().len());
+ }
+ }
+
+ assert!(
+ data_pages.len() >= num_rows / 2,
+ "expected at least {} data pages for {num_rows} large strings via
ArrowWriter, got {} ({data_pages:?})",
+ num_rows / 2,
+ data_pages.len(),
+ );
+ for size in &data_pages {
+ assert!(
+ *size <= value_size + 64,
+ "page size {size} exceeds one-value bound; pages
{data_pages:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_arrow_writer_granular_mode_roundtrip() {
+ // Granular mode subdivides chunks and writes more pages than
+ // `main`. Make sure the data we write back is bit-identical to
+ // what went in — page-count assertions elsewhere only prove
+ // pages were cut, not that the encoded data is correct.
+ //
+ // Mix value sizes so that the cumulative-byte-budget cutoff
+ // lands mid-chunk, exercising both batched and granular paths
+ // within the same `write_batch_internal` call.
+ let small = "tiny".to_string();
+ let big = "x".repeat(64 * 1024);
+ let strings: Vec<String> = (0..256)
+ .map(|i| {
+ if i % 16 == 0 {
+ big.clone()
+ } else {
+ small.clone()
+ }
+ })
+ .collect();
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings.clone())) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(16 * 1024)
+ .build();
+ let mut writer = ArrowWriter::try_new(Vec::new(), schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let data = Bytes::from(writer.into_inner().unwrap());
+
+ let mut reader = ParquetRecordBatchReader::try_new(data,
1024).unwrap();
+ let read = reader.next().unwrap().unwrap();
+ assert!(reader.next().is_none(), "expected one batch");
+ let col = read
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(col.len(), strings.len());
+ for (i, expected) in strings.iter().enumerate() {
+ assert_eq!(
+ col.value(i),
+ expected.as_str(),
+ "value mismatch at index {i}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_string_view() {
+ // Coverage for the view/dict fallback in
+ // `ByteArrayEncoder::count_values_within_byte_budget_gather`.
+ // `Utf8View` arrays don't expose a single contiguous offsets
+ // buffer, so the gather method falls back to the per-value walk
+ // via `ArrayAccessor::value`.
+ use arrow_array::StringViewArray;
Review Comment:
The arrow-writer tests now use the declarative `LayoutTest` fixture
(d12657b040). The raw column-writer tests still carry the boilerplate you
flagged — tracking that in the thread on `column/writer/mod.rs` below. Resolved.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -4904,6 +4904,245 @@ mod tests {
assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
}
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_strings() {
+ // End-to-end coverage for
`ByteArrayEncoder::estimated_value_bytes_gather`
+ // and the offsets-buffer fast path in `estimate_byte_size_offsets`.
+ //
+ // The standard column-writer regression test covers the same
+ // logic for the non-arrow path; this one drives writes through
+ // `ArrowWriter`, which is what real users hit and which uses
+ // `write_gather` with `non_null_indices` (always contiguous for
+ // a non-null column).
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_rows = 32;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let strings: Vec<String> = (0..num_rows).map(|_|
"x".repeat(value_size)).collect();
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings)) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(page_byte_limit)
+ .build();
+ let mut writer = ArrowWriter::try_new(Vec::new(), schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let data = Bytes::from(writer.into_inner().unwrap());
+
+ let mut metadata = ParquetMetaDataReader::new();
+ metadata.try_parse(&data).unwrap();
+ let metadata = metadata.finish().unwrap();
+ let col_meta = metadata.row_group(0).column(0);
+ let mut reader =
+ SerializedPageReader::new(Arc::new(data.clone()), col_meta,
num_rows, None).unwrap();
+
+ let mut data_pages = Vec::new();
+ while let Some(page) = reader.get_next_page().unwrap() {
+ if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. })
{
+ data_pages.push(page.buffer().len());
+ }
+ }
+
+ assert!(
+ data_pages.len() >= num_rows / 2,
+ "expected at least {} data pages for {num_rows} large strings via
ArrowWriter, got {} ({data_pages:?})",
+ num_rows / 2,
+ data_pages.len(),
+ );
+ for size in &data_pages {
+ assert!(
+ *size <= value_size + 64,
+ "page size {size} exceeds one-value bound; pages
{data_pages:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_arrow_writer_granular_mode_roundtrip() {
+ // Granular mode subdivides chunks and writes more pages than
+ // `main`. Make sure the data we write back is bit-identical to
+ // what went in — page-count assertions elsewhere only prove
+ // pages were cut, not that the encoded data is correct.
+ //
+ // Mix value sizes so that the cumulative-byte-budget cutoff
+ // lands mid-chunk, exercising both batched and granular paths
+ // within the same `write_batch_internal` call.
+ let small = "tiny".to_string();
+ let big = "x".repeat(64 * 1024);
+ let strings: Vec<String> = (0..256)
+ .map(|i| {
+ if i % 16 == 0 {
+ big.clone()
+ } else {
+ small.clone()
+ }
+ })
+ .collect();
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8,
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(StringArray::from(strings.clone())) as _],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(16 * 1024)
+ .build();
+ let mut writer = ArrowWriter::try_new(Vec::new(), schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let data = Bytes::from(writer.into_inner().unwrap());
+
+ let mut reader = ParquetRecordBatchReader::try_new(data,
1024).unwrap();
+ let read = reader.next().unwrap().unwrap();
+ assert!(reader.next().is_none(), "expected one batch");
+ let col = read
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(col.len(), strings.len());
+ for (i, expected) in strings.iter().enumerate() {
+ assert_eq!(
+ col.value(i),
+ expected.as_str(),
+ "value mismatch at index {i}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_arrow_writer_caps_page_size_for_large_string_view() {
+ // Coverage for the view/dict fallback in
+ // `ByteArrayEncoder::count_values_within_byte_budget_gather`.
+ // `Utf8View` arrays don't expose a single contiguous offsets
+ // buffer, so the gather method falls back to the per-value walk
+ // via `ArrayAccessor::value`.
+ use arrow_array::StringViewArray;
+ let value_size = 64 * 1024;
+ let page_byte_limit = 16 * 1024;
+ let num_rows = 32usize;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ ArrowDataType::Utf8View,
+ false,
+ )]));
+ let strings: Vec<String> = (0..num_rows).map(|_|
"x".repeat(value_size)).collect();
Review Comment:
This test was ported to a `LayoutTest` (d12657b040), so it's much smaller
now (it still builds the 32 values up front, but that's a couple of lines).
Resolved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
