Rich-T-kid commented on code in PR #11278:
URL: https://github.com/apache/arrow-rs/pull/11278#discussion_r4170341756
##########
parquet/src/arrow/arrow_writer/levels.rs:
##########
@@ -208,16 +230,19 @@ impl LevelInfoBuilder {
let levels = ArrayLevels::new(parent_ctx, is_nullable,
array.clone());
Ok(Self::Primitive(levels))
}
- DataType::RunEndEncoded(_, value_field) => {
+ DataType::RunEndEncoded(_, _) => {
let flat = expand_ree_array(array)?;
- let flat_field = Field::new(
- field.name(),
- value_field.data_type().clone(),
- field.is_nullable(),
- );
- Self::try_new(&flat_field, parent_ctx, &flat)
+ Self::try_new(field, parent_ctx, &flat)
}
- DataType::Struct(children) => {
+ DataType::Dictionary(_, _) => {
Review Comment:
for correctness reasons we need to do this, would be interested to see the
performance ramifications of this though.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
Review Comment:
nice to remove this unwrap as well 👍
##########
parquet/src/arrow/arrow_writer/levels.rs:
##########
@@ -933,42 +967,36 @@ impl LevelInfoBuilder {
/// and the other is a native array, the dictionary values must have the
same type as the
/// native array
fn types_compatible(a: &DataType, b: &DataType) -> bool {
- // if the Arrow data types are equal, the types are deemed compatible
- if a.equals_datatype(b) {
- return true;
- }
-
- // get the values out of the dictionaries
- let (a, b) = match (a, b) {
- (DataType::Dictionary(_, va), DataType::Dictionary(_, vb)) => {
- (va.as_ref(), vb.as_ref())
+ let (a, _) = logical_type(a);
+ let (b, _) = logical_type(b);
+ match (a, b) {
+ (DataType::Struct(a), DataType::Struct(b)) => {
+ a.len() == b.len()
+ && a.iter()
+ .zip(b)
+ .all(|(a, b)| Self::types_compatible(a.data_type(),
b.data_type()))
}
- (DataType::Dictionary(_, v), b) => (v.as_ref(), b),
- (a, DataType::Dictionary(_, v)) => (a, v.as_ref()),
- _ => (a, b),
- };
-
- // now that we've got the values from one/both dictionaries, if the
values
- // have the same Arrow data type, they're compatible
- if a == b {
- return true;
- }
-
- // here we have different Arrow data types, but if the array contains
the same type of data
- // then we consider the type compatible
- match a {
- // String, StringView and LargeString are compatible
- DataType::Utf8 => matches!(b, DataType::LargeUtf8 |
DataType::Utf8View),
- DataType::Utf8View => matches!(b, DataType::LargeUtf8 |
DataType::Utf8),
- DataType::LargeUtf8 => matches!(b, DataType::Utf8 |
DataType::Utf8View),
-
- // Binary, BinaryView and LargeBinary are compatible
- DataType::Binary => matches!(b, DataType::LargeBinary |
DataType::BinaryView),
- DataType::BinaryView => matches!(b, DataType::LargeBinary |
DataType::Binary),
- DataType::LargeBinary => matches!(b, DataType::Binary |
DataType::BinaryView),
-
- // otherwise we have incompatible types
- _ => false,
+ (DataType::List(a), DataType::List(b))
Review Comment:
nice, I find the refactored code easier to reason about 👍
##########
parquet/src/arrow/arrow_writer/levels.rs:
##########
@@ -933,42 +967,36 @@ impl LevelInfoBuilder {
/// and the other is a native array, the dictionary values must have the
same type as the
/// native array
fn types_compatible(a: &DataType, b: &DataType) -> bool {
- // if the Arrow data types are equal, the types are deemed compatible
- if a.equals_datatype(b) {
- return true;
- }
-
- // get the values out of the dictionaries
- let (a, b) = match (a, b) {
- (DataType::Dictionary(_, va), DataType::Dictionary(_, vb)) => {
- (va.as_ref(), vb.as_ref())
+ let (a, _) = logical_type(a);
+ let (b, _) = logical_type(b);
+ match (a, b) {
+ (DataType::Struct(a), DataType::Struct(b)) => {
+ a.len() == b.len()
+ && a.iter()
+ .zip(b)
+ .all(|(a, b)| Self::types_compatible(a.data_type(),
b.data_type()))
}
- (DataType::Dictionary(_, v), b) => (v.as_ref(), b),
- (a, DataType::Dictionary(_, v)) => (a, v.as_ref()),
- _ => (a, b),
- };
-
- // now that we've got the values from one/both dictionaries, if the
values
- // have the same Arrow data type, they're compatible
- if a == b {
- return true;
- }
-
- // here we have different Arrow data types, but if the array contains
the same type of data
- // then we consider the type compatible
- match a {
- // String, StringView and LargeString are compatible
- DataType::Utf8 => matches!(b, DataType::LargeUtf8 |
DataType::Utf8View),
- DataType::Utf8View => matches!(b, DataType::LargeUtf8 |
DataType::Utf8),
- DataType::LargeUtf8 => matches!(b, DataType::Utf8 |
DataType::Utf8View),
-
- // Binary, BinaryView and LargeBinary are compatible
- DataType::Binary => matches!(b, DataType::LargeBinary |
DataType::BinaryView),
- DataType::BinaryView => matches!(b, DataType::LargeBinary |
DataType::Binary),
- DataType::LargeBinary => matches!(b, DataType::Binary |
DataType::BinaryView),
-
- // otherwise we have incompatible types
- _ => false,
+ (DataType::List(a), DataType::List(b))
+ | (DataType::LargeList(a), DataType::LargeList(b))
+ | (DataType::ListView(a), DataType::ListView(b))
+ | (DataType::LargeListView(a), DataType::LargeListView(b)) => {
+ Self::types_compatible(a.data_type(), b.data_type())
+ }
+ (DataType::FixedSizeList(a, an), DataType::FixedSizeList(b, bn))
if an == bn => {
Review Comment:
super nit: I dislike the short varible names, something a bit more
descritive would nice.
```suggestion
(DataType::FixedSizeList(left_field, left_len),
DataType::FixedSizeList(right_field, right_len)) if left_len == right_len => {
```
##########
parquet/src/arrow/schema/mod.rs:
##########
@@ -313,12 +313,132 @@ pub fn encode_arrow_schema(schema: &Schema) -> String {
BASE64_STANDARD.encode(&len_prefix_schema)
}
-fn flatten_ree_field(field: &Field) -> Field {
+/// Reject map keys that are nullable after resolving their encoded layout.
+/// This validates the writer's logical contract, not a batch's physical
layout.
+pub(crate) fn validate_map_key_type(data_type: &DataType) -> Result<()> {
+ match data_type {
+ DataType::Map(entries, _) => {
+ let DataType::Struct(fields) = entries.data_type() else {
+ return Err(arrow_err!(
+ "DataType::Map should contain a struct field child"
+ ));
+ };
+ if fields.len() != 2 {
+ return Err(arrow_err!("DataType::Map entries must contain two
fields"));
+ }
+ let key = &fields[0];
+ let mut key_type = key.data_type();
+ let mut nullable = key.is_nullable();
+ loop {
+ match key_type {
+ DataType::Dictionary(_, value) => key_type = value,
+ DataType::RunEndEncoded(_, value) => {
+ nullable |= value.is_nullable();
+ key_type = value.data_type();
+ }
+ _ => break,
+ }
+ }
+ if nullable {
+ return Err(arrow_err!(
+ "Map key field '{}' must be non-nullable, including its
encoded values",
+ key.name()
+ ));
+ }
+ validate_map_key_type(entries.data_type())?;
+ }
+ DataType::Struct(fields) => {
+ for field in fields {
+ validate_map_key_type(field.data_type())?;
+ }
+ }
+ DataType::List(field)
+ | DataType::LargeList(field)
+ | DataType::FixedSizeList(field, _)
+ | DataType::ListView(field)
+ | DataType::LargeListView(field)
+ | DataType::RunEndEncoded(_, field) =>
validate_map_key_type(field.data_type())?,
+ DataType::Dictionary(_, value) => validate_map_key_type(value)?,
+ _ => {}
+ }
+ Ok(())
+}
+
+/// Normalize `field` for Arrow schema metadata.
+///
+/// This removes wrappers the Parquet reader does not reconstruct: REE at any
+/// nesting depth and dictionaries around nested values. Scalar dictionaries,
+/// including `Utf8View` and `BinaryView`, are retained in the metadata.
+fn normalize_metadata_field(field: &Field) -> Field {
Review Comment:
when arrow-rs's parquet reader supports reading ree fields directly this
will need to be updated right?
in that case it would just like the dictionary path where it only strips it
down to scalar value types?
##########
parquet/src/arrow/schema/mod.rs:
##########
@@ -313,12 +313,132 @@ pub fn encode_arrow_schema(schema: &Schema) -> String {
BASE64_STANDARD.encode(&len_prefix_schema)
}
-fn flatten_ree_field(field: &Field) -> Field {
+/// Reject map keys that are nullable after resolving their encoded layout.
+/// This validates the writer's logical contract, not a batch's physical
layout.
+pub(crate) fn validate_map_key_type(data_type: &DataType) -> Result<()> {
+ match data_type {
+ DataType::Map(entries, _) => {
+ let DataType::Struct(fields) = entries.data_type() else {
+ return Err(arrow_err!(
+ "DataType::Map should contain a struct field child"
+ ));
+ };
+ if fields.len() != 2 {
+ return Err(arrow_err!("DataType::Map entries must contain two
fields"));
+ }
+ let key = &fields[0];
+ let mut key_type = key.data_type();
+ let mut nullable = key.is_nullable();
+ loop {
+ match key_type {
+ DataType::Dictionary(_, value) => key_type = value,
+ DataType::RunEndEncoded(_, value) => {
+ nullable |= value.is_nullable();
+ key_type = value.data_type();
+ }
+ _ => break,
+ }
+ }
+ if nullable {
+ return Err(arrow_err!(
+ "Map key field '{}' must be non-nullable, including its
encoded values",
+ key.name()
+ ));
+ }
+ validate_map_key_type(entries.data_type())?;
+ }
+ DataType::Struct(fields) => {
+ for field in fields {
+ validate_map_key_type(field.data_type())?;
+ }
+ }
+ DataType::List(field)
+ | DataType::LargeList(field)
+ | DataType::FixedSizeList(field, _)
+ | DataType::ListView(field)
+ | DataType::LargeListView(field)
+ | DataType::RunEndEncoded(_, field) =>
validate_map_key_type(field.data_type())?,
+ DataType::Dictionary(_, value) => validate_map_key_type(value)?,
+ _ => {}
+ }
+ Ok(())
+}
+
+/// Normalize `field` for Arrow schema metadata.
+///
+/// This removes wrappers the Parquet reader does not reconstruct: REE at any
+/// nesting depth and dictionaries around nested values. Scalar dictionaries,
+/// including `Utf8View` and `BinaryView`, are retained in the metadata.
+fn normalize_metadata_field(field: &Field) -> Field {
match field.data_type() {
- DataType::RunEndEncoded(_, value_field) => field
+ DataType::RunEndEncoded(_, value_field) => {
+ let value = normalize_metadata_field(value_field);
+ field
+ .clone()
+ .with_data_type(value.data_type().clone())
+ .with_nullable(field.is_nullable() || value.is_nullable())
+ }
+ DataType::Dictionary(key, value_type) => {
+ // Dictionary values have no Field, so use a synthetic one to
retain
+ // the nullability of any REE value while normalizing recursively.
+ let value_field = Field::new("", value_type.as_ref().clone(),
false);
Review Comment:
a bit unrelated but https://github.com/apache/arrow-rs/pull/10840 introduced
constants for ree field names I.E `run_ends`
[link](https://github.com/apache/arrow-rs/blob/705e6405b56cc01efb1c267203d240d0d85804ab/arrow-schema/src/field.rs#L164)
& `values`
[link](https://github.com/apache/arrow-rs/blob/705e6405b56cc01efb1c267203d240d0d85804ab/arrow-schema/src/field.rs#L168).
I know the [dictionary section of the
arrow-spec](https://github.com/apache/arrow/blob/main/docs/source/format/Columnar.rst#dictionary-encoded-layout)
doesnt outright name the keys and value fields but this may be a nice thing
to have.
cc @Jefffrey
##########
parquet/src/arrow/schema/mod.rs:
##########
@@ -1990,6 +2118,266 @@ mod tests {
converted_arrow_schema.unwrap();
}
+ #[test]
Review Comment:
Overall these test look very good. 👍
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -2513,6 +2517,186 @@ mod tests {
roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
}
+ fn read_column(file: Vec<u8>) -> ArrayRef {
+ let reader = ParquetRecordBatchReader::try_new(Bytes::from(file),
4096).unwrap();
+ let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
+ let arrays: Vec<&dyn Array> = batches.iter().map(|b|
b.column(0).as_ref()).collect();
+ arrow_select::concat::concat(&arrays).unwrap()
Review Comment:
```suggestion
use arrow_select::conact::conact;
... (in the test mod)
arrow_select::concat::concat(&arrays).unwrap()
```
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -2513,6 +2517,186 @@ mod tests {
roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
}
+ fn read_column(file: Vec<u8>) -> ArrayRef {
+ let reader = ParquetRecordBatchReader::try_new(Bytes::from(file),
4096).unwrap();
+ let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
+ let arrays: Vec<&dyn Array> = batches.iter().map(|b|
b.column(0).as_ref()).collect();
+ arrow_select::concat::concat(&arrays).unwrap()
+ }
+
+ fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
+ let writer_schema = Arc::new(Schema::new(vec![field]));
+ let batch_schema = Arc::new(Schema::new(vec![Field::new(
+ "c",
+ col.data_type().clone(),
+ col.logical_null_count() != 0,
+ )]));
+ let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
+ let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
+ let mut file = vec![];
+ let mut writer =
+ ArrowWriter::try_new_with_options(&mut file, writer_schema,
options).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ read_column(file)
+ }
+
+ #[test]
+ fn arrow_writer_dense_batches_under_nested_wrapper_schemas() {
+ let run_ends = || Arc::new(Field::new("run_ends", DataType::Int32,
false));
+ let ree = |value: DataType, nullable| {
+ DataType::RunEndEncoded(run_ends(), Arc::new(Field::new("values",
value, nullable)))
+ };
+ let assert_unified = |field: Field, actual: ArrayRef, expected:
&ArrayRef| {
+ assert!(!compute_leaves(&field, &actual).unwrap().is_empty());
+ assert_eq!(
+ roundtrip_compatible_column(field, actual).as_ref(),
+ expected.as_ref()
+ );
+ };
+
+ let item = Arc::new(Field::new_list_field(DataType::Int32, true));
+ let list: ArrayRef = Arc::new(ListArray::new(
+ item.clone(),
+ OffsetBuffer::new(vec![0_i32, 2, 3].into()),
+ Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
+ None,
+ ));
+
+ // A dense list under an REE<List> schema uses range traversal; the
+ // schema wrapper is physical, not a second logical list node.
+ let ree_list_field = Field::new("c", ree(list.data_type().clone(),
false), false);
+ assert_unified(ree_list_field, list.clone(), &list);
+
+ // Dictionary<List> has the same dense logical shape and follows the
+ // same list traversal after schema normalization.
+ let dictionary_list = Field::new(
+ "c",
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(list.data_type().clone())),
+ false,
+ );
+ assert_unified(dictionary_list, list.clone(), &list);
+
+ let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None,
Some(9)]));
+ let nested_ree = Field::new("c", ree(ree(DataType::Int32, true),
false), false);
+ assert_unified(nested_ree, dense.clone(), &dense);
+
+ let dictionary_ree = Field::new(
+ "c",
+ DataType::Dictionary(
+ Box::new(DataType::Int8),
+ Box::new(ree(DataType::Int32, true)),
+ ),
+ false,
+ );
+ assert_unified(dictionary_ree, dense.clone(), &dense);
+
+ // The inverse wrapper order must hoist REE value nullability before
+ // peeling the dictionary; the dense batch contains an actual null.
Review Comment:
this may be a good split point. half the test can cover the inverse wrapper
order and the inverse
##########
parquet/src/arrow/arrow_writer/levels.rs:
##########
@@ -208,16 +230,19 @@ impl LevelInfoBuilder {
let levels = ArrayLevels::new(parent_ctx, is_nullable,
array.clone());
Ok(Self::Primitive(levels))
}
- DataType::RunEndEncoded(_, value_field) => {
+ DataType::RunEndEncoded(_, _) => {
let flat = expand_ree_array(array)?;
- let flat_field = Field::new(
- field.name(),
- value_field.data_type().clone(),
- field.is_nullable(),
- );
- Self::try_new(&flat_field, parent_ctx, &flat)
+ Self::try_new(field, parent_ctx, &flat)
}
- DataType::Struct(children) => {
+ DataType::Dictionary(_, _) => {
+ let dictionary = array.as_any_dictionary();
+ let flat = arrow_select::take::take(dictionary.values(),
dictionary.keys(), None)?;
Review Comment:
```suggestion
use arrow_select::take::take;
...
let flat = take(dictionary.values(), dictionary.keys(),
None)?;
```
##########
parquet/src/arrow/arrow_writer/levels.rs:
##########
@@ -933,42 +967,36 @@ impl LevelInfoBuilder {
/// and the other is a native array, the dictionary values must have the
same type as the
/// native array
fn types_compatible(a: &DataType, b: &DataType) -> bool {
Review Comment:
this relates back to what I was refering to with scope creep. we are
updating many datatypes; `dictionarys`, `list`,`largelist`, `listview`
,`largeListView` ,`fixedsizedlist` , and `maps`.
this isnt too much of a change as they will just route to
`left.equals_datatype(right)` as before but its unrelated to the issues. I
think for this case its fine but in general I think its best to avoid this.
##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -2513,6 +2517,186 @@ mod tests {
roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
}
+ fn read_column(file: Vec<u8>) -> ArrayRef {
+ let reader = ParquetRecordBatchReader::try_new(Bytes::from(file),
4096).unwrap();
+ let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
+ let arrays: Vec<&dyn Array> = batches.iter().map(|b|
b.column(0).as_ref()).collect();
+ arrow_select::concat::concat(&arrays).unwrap()
+ }
+
+ fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
+ let writer_schema = Arc::new(Schema::new(vec![field]));
+ let batch_schema = Arc::new(Schema::new(vec![Field::new(
+ "c",
+ col.data_type().clone(),
+ col.logical_null_count() != 0,
+ )]));
+ let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
+ let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
+ let mut file = vec![];
+ let mut writer =
+ ArrowWriter::try_new_with_options(&mut file, writer_schema,
options).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ read_column(file)
+ }
+
+ #[test]
+ fn arrow_writer_dense_batches_under_nested_wrapper_schemas() {
Review Comment:
this test is a bit hard to follow. can it it broken up into seperate focused
test?
##########
parquet/src/arrow/schema/mod.rs:
##########
@@ -1990,6 +2118,266 @@ mod tests {
converted_arrow_schema.unwrap();
}
+ #[test]
+ fn test_run_end_schema_conversion_recurses_into_value_type() {
+ fn run_end_type(value_type: DataType, nullable_values: bool) ->
DataType {
+ DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", value_type, nullable_values)),
+ )
+ }
+
+ // A leaf REE contributes the value's physical type and combined
+ // nullability to the Parquet schema.
+ let leaf_schema = Schema::new(vec![Field::new(
+ "ree",
+ run_end_type(DataType::Int32, true),
+ false,
+ )]);
+ let converted =
ArrowSchemaConverter::new().convert(&leaf_schema).unwrap();
+ assert_eq!(converted.columns().len(), 1);
+ assert_eq!(converted.column(0).name(), "ree");
+ assert_eq!(converted.column(0).physical_type(), PhysicalType::INT32);
+ assert_eq!(converted.column(0).max_def_level(), 1);
+
+ // Non-leaf value types recurse into the value type rather than
+ // exposing Arrow's physical run-end wrapper in the Parquet schema.
+ let list_schema = Schema::new(vec![Field::new(
+ "ree_list",
+ run_end_type(
+ DataType::List(Arc::new(Field::new("element", DataType::Int32,
true))),
+ true,
+ ),
+ true,
+ )]);
+ let converted =
ArrowSchemaConverter::new().convert(&list_schema).unwrap();
+ assert_eq!(converted.columns().len(), 1);
+ assert_eq!(converted.column(0).physical_type(), PhysicalType::INT32);
+ assert!(
+ converted.column(0).max_rep_level() >= 1,
+ "list element should be repeated"
+ );
+
+ let dictionary_schema = Schema::new(vec![Field::new(
+ "ree_dictionary",
+ run_end_type(
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
+ true,
+ ),
+ true,
+ )]);
+ let converted = ArrowSchemaConverter::new()
+ .convert(&dictionary_schema)
+ .unwrap();
+ assert_eq!(converted.columns().len(), 1);
+ assert_eq!(
+ converted.column(0).physical_type(),
+ PhysicalType::BYTE_ARRAY
+ );
+ }
+
+ fn map_type_with_key(key: Field) -> DataType {
+ DataType::Map(
+ Arc::new(Field::new(
+ "entries",
+ DataType::Struct(vec![key, Field::new("value",
DataType::Int32, true)].into()),
+ false,
+ )),
+ false,
+ )
+ }
+
+ #[test]
+ fn test_map_key_schema_rejects_logically_nullable_keys() {
+ let ree = DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", DataType::Int32, true)),
+ );
+ let dictionary_ree = DataType::Dictionary(Box::new(DataType::Int8),
Box::new(ree.clone()));
+ let nested_ree = DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", dictionary_ree.clone(), false)),
+ );
+ for key in [
+ Field::new("key", DataType::Int32, true),
+ Field::new("key", ree, false),
+ Field::new("key", dictionary_ree, false),
+ Field::new("key", nested_ree, false),
+ ] {
+ let map = map_type_with_key(key);
+ for data_type in [
+ map.clone(),
+ DataType::Struct(vec![Field::new("nested", map.clone(),
true)].into()),
+ DataType::List(Arc::new(Field::new("item", map.clone(),
true))),
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(map.clone())),
+ DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", map, true)),
+ ),
+ ] {
+ let schema = Schema::new(vec![Field::new("map", data_type,
true)]);
+ for coerce_types in [false, true] {
+ let err = ArrowSchemaConverter::new()
+ .with_coerce_types(coerce_types)
+ .convert(&schema)
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("Map key field 'key' must be
non-nullable"),
+ "{err}"
+ );
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_map_key_schema_preserves_required_keys_and_nullable_members() {
+ for key_type in [
+ DataType::Int32,
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Int32)),
+ DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", DataType::Int32, false)),
+ ),
+ DataType::Struct(vec![Field::new("member", DataType::Int32,
true)].into()),
+ ] {
+ let schema = Schema::new(vec![Field::new(
+ "map",
+ map_type_with_key(Field::new("key", key_type, false)),
+ true,
+ )]);
+ let parquet =
ArrowSchemaConverter::new().convert(&schema).unwrap();
+ let key =
&parquet.root_schema().get_fields()[0].get_fields()[0].get_fields()[0];
+ assert_eq!(key.get_basic_info().repetition(),
Repetition::REQUIRED);
+
+ let mut props = WriterProperties::builder().build();
+ add_encoded_arrow_schema_to_metadata(&schema, &mut props);
+ let encoded = props
+ .key_value_metadata()
+ .unwrap()
+ .iter()
+ .find(|kv| kv.key == crate::arrow::ARROW_SCHEMA_META_KEY)
+ .unwrap()
+ .value
+ .as_deref()
+ .unwrap();
+ let decoded = get_arrow_schema_from_metadata(encoded).unwrap();
+ let DataType::Map(entries, _) = decoded.field(0).data_type() else {
+ unreachable!()
+ };
+ let DataType::Struct(fields) = entries.data_type() else {
+ unreachable!()
+ };
+ assert!(!fields[0].is_nullable());
+ assert!(fields[1].is_nullable());
+ if let DataType::Struct(members) = fields[0].data_type() {
+ assert!(members[0].is_nullable());
+ }
+ }
+ }
+
+ #[test]
+ fn test_arrow_metadata_normalization_merges_ree_value_nullability() {
+ let run_ends = Arc::new(Field::new("run_ends", DataType::Int32,
false));
+ let inner = Arc::new(Field::new("inner", DataType::Int32, true));
+ let nested = Arc::new(Field::new(
+ "nested",
+ DataType::RunEndEncoded(Arc::clone(&run_ends), inner),
+ false,
+ ));
+ let field = Field::new("outer", DataType::RunEndEncoded(run_ends,
nested), false);
+
+ let normalized = normalize_metadata_field(&field);
+ assert_eq!(normalized.data_type(), &DataType::Int32);
+ assert!(normalized.is_nullable());
+ }
+
+ #[test]
+ fn test_arrow_metadata_normalization_descends_dictionary_values() {
+ let ree = DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", DataType::Int64, true)),
+ );
+ let field = Field::new(
+ "dictionary_ree",
+ DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
+ false,
+ );
+
+ assert!(needs_metadata_normalization(field.data_type()));
+ let normalized = normalize_metadata_field(&field);
+ assert_eq!(
+ normalized.data_type(),
+ &DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Int64))
+ );
+ assert!(normalized.is_nullable());
+ assert!(!needs_metadata_normalization(normalized.data_type()));
+
+ let mut props = WriterProperties::builder().build();
+ add_encoded_arrow_schema_to_metadata(&Schema::new(vec![field]), &mut
props);
+ let encoded = props
+ .key_value_metadata()
+ .unwrap()
+ .iter()
+ .find(|kv| kv.key == crate::arrow::ARROW_SCHEMA_META_KEY)
+ .unwrap()
+ .value
+ .as_deref()
+ .unwrap();
+ let decoded = get_arrow_schema_from_metadata(encoded).unwrap();
+ assert_eq!(decoded.field(0), &normalized);
+ }
+
+ #[test]
+ fn metadata_normalization_nested_and_view_dictionary_contract() {
Review Comment:
is there any reason why `dict<_,utf8/binary view` has a dedicated test?
--
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]
