alamb commented on code in PR #10545:
URL: https://github.com/apache/arrow-rs/pull/10545#discussion_r3713943663
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -3175,101 +3179,138 @@ mod tests {
})
}
- struct RoundTripOptions {
+ /// Round trip testing fixture:
+ ///
+ /// Tests based on this fixture write data to parquet and then read it
back.
+ struct RoundTripTest {
Review Comment:
I renamed and moved a few parameters to fields
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -3175,101 +3179,138 @@ mod tests {
})
}
- struct RoundTripOptions {
+ /// Round trip testing fixture:
+ ///
+ /// Tests based on this fixture write data to parquet and then read it
back.
+ struct RoundTripTest {
values: ArrayRef,
- schema: SchemaRef,
+ /// Optionally supplied schema
+ schema: Option<SchemaRef>,
+ /// If the created schema should be nullable. Defaults to true. Ignored
+ /// if schema is set to Some.
+ nullable: bool,
bloom_filter: bool,
bloom_filter_ndv: Option<u64>,
bloom_filter_position: BloomFilterPosition,
}
- impl RoundTripOptions {
- fn new(values: ArrayRef, nullable: bool) -> Self {
- let data_type = values.data_type().clone();
- let schema = Schema::new(vec![Field::new("col", data_type,
nullable)]);
+ impl RoundTripTest {
+ /// Create a test for round tripping values with a nullable schema
+ fn new(values: ArrayRef) -> Self {
Self {
values,
- schema: Arc::new(schema),
+ schema: None,
+ nullable: true,
bloom_filter: false,
bloom_filter_ndv: None,
bloom_filter_position: BloomFilterPosition::AfterRowGroup,
}
}
- }
- fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec<Bytes> {
- one_column_roundtrip_with_options(RoundTripOptions::new(values,
nullable))
- }
+ /// Set the schema
+ pub fn with_schema(mut self, schema: SchemaRef) -> Self {
+ self.schema = Some(schema);
+ self
+ }
- fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef)
-> Vec<Bytes> {
- let mut options = RoundTripOptions::new(values, false);
- options.schema = schema;
- one_column_roundtrip_with_options(options)
- }
+ /// Set the nullable flag
+ pub fn with_nullable(mut self, nullable: bool) -> Self {
+ self.nullable = nullable;
+ self
+ }
- fn one_column_roundtrip_with_options(options: RoundTripOptions) ->
Vec<Bytes> {
- let RoundTripOptions {
- values,
- schema,
- bloom_filter,
- bloom_filter_ndv,
- bloom_filter_position,
- } = options;
-
- let encodings = match values.data_type() {
- DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary |
DataType::LargeBinary => {
- vec![
- Encoding::PLAIN,
- Encoding::DELTA_BYTE_ARRAY,
- Encoding::DELTA_LENGTH_BYTE_ARRAY,
- ]
- }
- DataType::Int64
- | DataType::Int32
- | DataType::Int16
- | DataType::Int8
- | DataType::UInt64
- | DataType::UInt32
- | DataType::UInt16
- | DataType::UInt8 => vec![
- Encoding::PLAIN,
- Encoding::DELTA_BINARY_PACKED,
- Encoding::BYTE_STREAM_SPLIT,
- ],
- DataType::Float32 | DataType::Float64 => {
- vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
- }
- _ => vec![Encoding::PLAIN],
- };
+ /// Set bloom filter
+ pub fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
+ self.bloom_filter = bloom_filter;
+ self
+ }
- let expected_batch = RecordBatch::try_new(schema,
vec![values]).unwrap();
+ /// Set bloom filter max ndv
+ pub fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: Option<u64>)
-> Self {
+ self.bloom_filter_ndv = bloom_filter_ndv;
+ self
+ }
- let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE /
2 + 1, 10];
+ /// Set bloom filter position
+ pub fn with_bloom_filter_position(
+ mut self,
+ bloom_filter_position: BloomFilterPosition,
+ ) -> Self {
+ self.bloom_filter_position = bloom_filter_position;
+ self
+ }
- let mut files = vec![];
- for dictionary_size in [0, 1, 1024] {
- for encoding in &encodings {
- for version in [WriterVersion::PARQUET_1_0,
WriterVersion::PARQUET_2_0] {
- for row_group_size in row_group_sizes {
- let mut builder = WriterProperties::builder()
- .set_writer_version(version)
- .set_max_row_group_row_count(Some(row_group_size))
- .set_dictionary_enabled(dictionary_size != 0)
-
.set_dictionary_page_size_limit(dictionary_size.max(1))
- .set_encoding(*encoding)
- .set_bloom_filter_enabled(bloom_filter)
- .set_bloom_filter_position(bloom_filter_position);
- if let Some(ndv) = bloom_filter_ndv {
- builder = builder.set_bloom_filter_max_ndv(ndv);
- }
- let props = builder.build();
+ /// Run the test specified by the options, returning the encoded
Parquet bytes
Review Comment:
Viewing this diff without whitespace makes this change clearer I think:
https://github.com/apache/arrow-rs/pull/10545/changes?w=1
The method is indented but the only new code is this
```rust
let schema = schema.unwrap_or_else(|| {
let data_type = values.data_type().clone();
Arc::new(Schema::new(vec![Field::new("col", data_type,
nullable)]))
});
```
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -2575,8 +2575,12 @@ mod tests {
let binary_view_values: ArrayRef =
Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
- one_column_roundtrip(Arc::clone(&string_view_values), false);
- one_column_roundtrip(Arc::clone(&binary_view_values), false);
+ RoundTripTest::new(Arc::clone(&string_view_values))
Review Comment:
I think this is much easier to understand now -- the nullable_flag is now
self documenting
--
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]
