alamb commented on code in PR #6871:
URL: https://github.com/apache/arrow-rs/pull/6871#discussion_r1887587225
##########
parquet/src/arrow/mod.rs:
##########
@@ -371,4 +442,109 @@ mod test {
.unwrap();
Bytes::from(buf)
}
+
+ #[test]
+ fn test_mask_from_column_names() {
+ let message_type = "
+ message test_schema {
+ OPTIONAL group a (MAP) {
+ REPEATED group key_value {
+ REQUIRED BYTE_ARRAY key (UTF8);
+ OPTIONAL group value (MAP) {
+ REPEATED group key_value {
+ REQUIRED INT32 key;
+ REQUIRED BOOLEAN value;
+ }
+ }
+ }
+ }
+ REQUIRED INT32 b;
+ REQUIRED DOUBLE c;
+ }
+ ";
+ let parquet_group_type = parse_message_type(message_type).unwrap();
+ let schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
+
+ let mask = ProjectionMask::columns(&schema, ["foo", "bar"]);
+ assert_eq!(mask.mask.unwrap(), vec![false; 5]);
+
+ let mask = ProjectionMask::columns(&schema, []);
+ assert_eq!(mask.mask.unwrap(), vec![false; 5]);
+
+ let mask = ProjectionMask::columns(&schema, ["a", "c"]);
+ assert_eq!(mask.mask.unwrap(), [true, true, true, false, true]);
+
+ let mask = ProjectionMask::columns(&schema, ["a.key_value.key", "c"]);
+ assert_eq!(mask.mask.unwrap(), [true, false, false, false, true]);
+
+ let mask = ProjectionMask::columns(&schema, ["a.key_value.value",
"b"]);
+ assert_eq!(mask.mask.unwrap(), [false, true, true, true, false]);
+
+ let message_type = "
+ message test_schema {
+ OPTIONAL group a (LIST) {
+ REPEATED group list {
+ OPTIONAL group element (LIST) {
+ REPEATED group list {
+ OPTIONAL group element (LIST) {
+ REPEATED group list {
+ OPTIONAL BYTE_ARRAY element (UTF8);
+ }
+ }
+ }
+ }
+ }
+ }
+ REQUIRED INT32 b;
+ }
+ ";
+ let parquet_group_type = parse_message_type(message_type).unwrap();
+ let schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
+
+ let mask = ProjectionMask::columns(&schema, ["a", "b"]);
+ assert_eq!(mask.mask.unwrap(), [true, true]);
+
+ let mask = ProjectionMask::columns(&schema, ["a.list.element", "b"]);
+ assert_eq!(mask.mask.unwrap(), [true, true]);
+
+ let mask =
+ ProjectionMask::columns(&schema,
["a.list.element.list.element.list.element", "b"]);
+ assert_eq!(mask.mask.unwrap(), [true, true]);
+
+ let mask = ProjectionMask::columns(&schema, ["b"]);
+ assert_eq!(mask.mask.unwrap(), [false, true]);
+
+ let message_type = "
+ message test_schema {
+ OPTIONAL INT32 a;
+ OPTIONAL INT32 b;
+ OPTIONAL INT32 c;
+ OPTIONAL INT32 d;
+ OPTIONAL INT32 e;
+ }
+ ";
+ let parquet_group_type = parse_message_type(message_type).unwrap();
+ let schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
+
+ let mask = ProjectionMask::columns(&schema, ["a", "b"]);
+ assert_eq!(mask.mask.unwrap(), [true, true, false, false, false]);
+
+ let mask = ProjectionMask::columns(&schema, ["d", "b", "d"]);
+ assert_eq!(mask.mask.unwrap(), [false, true, false, true, false]);
+
+ let message_type = "
+ message test_schema {
+ OPTIONAL INT32 a;
+ OPTIONAL INT32 b;
+ OPTIONAL INT32 a;
Review Comment:
This seems a nasty thing to do (repeat the name of a field in the parquet
file) but it seems to be allowed and your code handles it
--
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]
