GH-52997: [Python][Parquet] Fix ambiguous dotted column projection

python/pyarrow/parquet/core.py

@@ -347,7 +347,8 @@ def __init__(self, source, *, metadata=None, common_metadata=None,
             arrow_extensions_enabled=arrow_extensions_enabled,
         )
         self.common_metadata = common_metadata
-        self._nested_paths_by_prefix = self._build_nested_paths()
+        (self._nested_paths_by_prefix,
+         self._top_level_paths_by_name) = self._build_nested_paths()
 
     def __enter__(self):
         return self
@@ -359,9 +360,11 @@ def _build_nested_paths(self):
         paths = self.reader.column_paths
 
         result = defaultdict(list)
+        top_level_paths = defaultdict(list)
 
         for i, path in enumerate(paths):
             key = path[0]
+            top_level_paths[key].append(i)
             rest = path[1:]
             while True:
                 result[key].append(i)
@@ -372,7 +375,7 @@ def _build_nested_paths(self):
                 key = '.'.join((key, rest[0]))
                 rest = rest[1:]
 
-        return result
+        return result, top_level_paths
 
     @property
     def metadata(self):
@@ -454,7 +457,8 @@ def read_row_group(self, i, columns=None, use_threads=True,
         columns : list
             If not None, only these columns will be read from the row group. A
             column name may be a prefix of a nested field, e.g. 'a' will select
-            'a.b', 'a.c', and 'a.d.e'.
+            'a.b', 'a.c', and 'a.d.e'. An exact top-level name takes precedence
+            over a nested field path with the same dotted name.
         use_threads : bool, default True
             Perform multi-threaded column reads.
         use_pandas_metadata : bool, default False
@@ -501,7 +505,8 @@ def read_row_groups(self, row_groups, columns=None, use_threads=True,
         columns : list
             If not None, only these columns will be read from the row group. A
             column name may be a prefix of a nested field, e.g. 'a' will select
-            'a.b', 'a.c', and 'a.d.e'.
+            'a.b', 'a.c', and 'a.d.e'. An exact top-level name takes precedence
+            over a nested field path with the same dotted name.
         use_threads : bool, default True
             Perform multi-threaded column reads.
         use_pandas_metadata : bool, default False
@@ -552,7 +557,8 @@ def iter_batches(self, batch_size=65536, row_groups=None, columns=None,
         columns : list
             If not None, only these columns will be read from the file. A
             column name may be a prefix of a nested field, e.g. 'a' will select
-            'a.b', 'a.c', and 'a.d.e'.
+            'a.b', 'a.c', and 'a.d.e'. An exact top-level name takes precedence
+            over a nested field path with the same dotted name.
         use_threads : boolean, default True
             Perform multi-threaded column reads.
         use_pandas_metadata : boolean, default False
@@ -611,7 +617,8 @@ def read(self, columns=None, use_threads=True, use_pandas_metadata=False):
         columns : list
             If not None, only these columns will be read from the file. A
             column name may be a prefix of a nested field, e.g. 'a' will select
-            'a.b', 'a.c', and 'a.d.e'.
+            'a.b', 'a.c', and 'a.d.e'. An exact top-level name takes precedence
+            over a nested field path with the same dotted name.
         use_threads : bool, default True
             Perform multi-threaded column reads.
         use_pandas_metadata : bool, default False
@@ -693,7 +700,9 @@ def _get_column_indices(self, column_names, use_pandas_metadata=False):
         indices = []
 
         for name in column_names:
-            if name in self._nested_paths_by_prefix:
+            if name in self._top_level_paths_by_name:
+                indices.extend(self._top_level_paths_by_name[name])
+            elif name in self._nested_paths_by_prefix:
                 indices.extend(self._nested_paths_by_prefix[name])
 
         if use_pandas_metadata:

python/pyarrow/tests/parquet/test_parquet_file.py

@@ -201,6 +201,81 @@ def test_read_column_invalid_index():
             f.reader.read_column(index)
 
 
+def test_dotted_top_level_column_takes_precedence():
+    table = pa.table({
+        'a.b': [10, 20],
+        'a': pa.array([{'b': 1}, {'b': 2}]),
+        'other': [3, 4],
+    })
+    sink = pa.BufferOutputStream()
+    pq.write_table(table, sink)
+    data = sink.getvalue()
+    file_ = pq.ParquetFile(data)
+    expected = table.select(['a.b'])
+
+    assert file_.read(columns=['a.b']).equals(expected)
+    assert file_.read_row_group(0, columns=['a.b']).equals(expected)
+    assert pa.Table.from_batches(
+        list(file_.iter_batches(columns=['a.b']))).equals(expected)
+    assert pq.read_table(
+        pa.BufferReader(data), columns=['a.b']).equals(expected)
+    assert file_.read(columns=['a.b', 'a.b']).equals(expected)
+    assert file_.read(columns=['a.b', 'a']).equals(
+        table.select(['a.b', 'a']))
+
+
+def test_nested_column_selection_without_collision():
+    table = pa.table({
+        'a': pa.array([
+            {'b': {'c': 1, 'd': 2}, 'x': 3},
+            {'b': {'c': 4, 'd': 5}, 'x': 6},
+        ]),
+        'other': [5, 6],
+    })
+    sink = pa.BufferOutputStream()
+    pq.write_table(table, sink)
+    file_ = pq.ParquetFile(sink.getvalue())
+
+    assert file_.read(columns=['a']).equals(table.select(['a']))
+    assert file_.read(columns=['a.b']).to_pydict() == {
+        'a': [{'b': {'c': 1, 'd': 2}}, {'b': {'c': 4, 'd': 5}}],
+    }
+    assert file_.read(columns=['a.b.c']).to_pydict() == {
+        'a': [{'b': {'c': 1}}, {'b': {'c': 4}}],
+    }
+
+
+def test_non_conflicting_dotted_column_names():
+    table = pa.table({
+        'metric.value': [10, 20],
+        'a': pa.array([{'b': 1}, {'b': 2}]),
+    })
+    sink = pa.BufferOutputStream()
+    pq.write_table(table, sink)
+    file_ = pq.ParquetFile(sink.getvalue())
+
+    assert file_.read(columns=['metric.value']).equals(
+        table.select(['metric.value']))
+    assert file_.read(columns=['a.b']).to_pydict() == {
+        'a': [{'b': 1}, {'b': 2}],
+    }
+
+
+@pytest.mark.pandas
+def test_dotted_top_level_column_with_pandas_metadata():
+    df = pd.DataFrame(
+        {'a.b': [10, 20], 'a': [{'b': 1}, {'b': 2}]},
+        index=pd.Index([7, 8], name='row_id'))
+    table = pa.Table.from_pandas(df, preserve_index=True)
+    sink = pa.BufferOutputStream()
+    pq.write_table(table, sink)
+    file_ = pq.ParquetFile(sink.getvalue())
+
+    result = file_.read(columns=['a.b'], use_pandas_metadata=True)
+    assert result.column_names == ['a.b', 'row_id']
+    tm.assert_frame_equal(result.to_pandas(), df[['a.b']])
+
+
 @pytest.mark.pandas
 @pytest.mark.parametrize('batch_size', [300, 1000, 1300])
 def test_iter_batches_columns_reader(tempdir, batch_size):