Fix for column name based projection mask creation (#8447)

# Which issue does this PR close?

- Closes #8443.

# Rationale for this change

See issue.

# What changes are included in this PR?

Replaces `starts_with` on concatenated paths with an element-by-element
search of path components.

# Are these changes tested?

Yes, added a new test based on the issue

# Are there any user-facing changes?
No, just a bug fix

Author: etseidl

parquet/src/arrow/mod.rs

@@ -193,14 +193,12 @@ pub mod async_writer;
 mod record_reader;
 experimental!(mod schema);
 
-use std::sync::Arc;
-
 pub use self::arrow_writer::ArrowWriter;
 #[cfg(feature = "async")]
 pub use self::async_reader::ParquetRecordBatchStreamBuilder;
 #[cfg(feature = "async")]
 pub use self::async_writer::AsyncArrowWriter;
-use crate::schema::types::{SchemaDescriptor, Type};
+use crate::schema::types::SchemaDescriptor;
 use arrow_schema::{FieldRef, Schema};
 
 pub use self::schema::{
@@ -318,21 +316,6 @@ impl ProjectionMask {
         Self { mask: Some(mask) }
     }
 
-    // Given a starting point in the schema, do a DFS for that node adding leaf paths to `paths`.
-    fn find_leaves(root: &Arc<Type>, parent: Option<&String>, paths: &mut Vec<String>) {
-        let path = parent
-            .map(|p| [p, root.name()].join("."))
-            .unwrap_or(root.name().to_string());
-        if root.is_group() {
-            for child in root.get_fields() {
-                Self::find_leaves(child, Some(&path), paths);
-            }
-        } else {
-            // Reached a leaf, add to paths
-            paths.push(path);
-        }
-    }
-
     /// Create a [`ProjectionMask`] which selects only the named columns
     ///
     /// All leaf columns that fall below a given name will be selected. For example, given
@@ -360,21 +343,24 @@ impl ProjectionMask {
     /// Note: repeated or out of order indices will not impact the final mask.
     ///
     /// i.e. `["b", "c"]` will construct the same mask as `["c", "b", "c"]`.
+    ///
+    /// Also, this will not produce the desired results if a column contains a '.' in its name.
+    /// Use [`Self::leaves`] or [`Self::roots`] in that case.
     pub fn columns<'a>(
         schema: &SchemaDescriptor,
         names: impl IntoIterator<Item = &'a str>,
     ) -> Self {
-        // first make vector of paths for leaf columns
-        let mut paths: Vec<String> = vec![];
-        for root in schema.root_schema().get_fields() {
-            Self::find_leaves(root, None, &mut paths);
-        }
-        assert_eq!(paths.len(), schema.num_columns());
-
         let mut mask = vec![false; schema.num_columns()];
         for name in names {
-            for idx in 0..schema.num_columns() {
-                if paths[idx].starts_with(name) {
+            let name_path: Vec<&str> = name.split('.').collect();
+            for (idx, col) in schema.columns().iter().enumerate() {
+                let path = col.path().parts();
+                // searching for "a.b.c" cannot match "a.b"
+                if name_path.len() > path.len() {
+                    continue;
+                }
+                // now path >= name_path, so check that each element in name_path matches
+                if name_path.iter().zip(path.iter()).all(|(a, b)| a == b) {
                     mask[idx] = true;
                 }
             }
@@ -698,6 +684,18 @@ mod test {
 
         let mask = ProjectionMask::columns(&schema, ["a", "e"]);
         assert_eq!(mask.mask.unwrap(), [true, false, true, false, true]);
+
+        let message_type = "
+            message test_schema {
+                OPTIONAL INT32 a;
+                OPTIONAL INT32 aa;
+            }
+            ";
+        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"]);
+        assert_eq!(mask.mask.unwrap(), [true, false]);
     }
 
     #[test]