HippoBaro commented on code in PR #11278:
URL: https://github.com/apache/arrow-rs/pull/11278#discussion_r4174740261
##########
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:
Yeah, these are preservation controls: the reader supports scalar `Utf8View`
and `BinaryView` dictionaries, so stripping unsupported nested dictionary
wrappers must not strip those supported hints too. The test makes sure we don't
regress there.
--
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]
