diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 4de88fe22997..41f4ef24346f 100644 --- a/python/pyarrow/parquet/core.py +++ b/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: diff --git a/python/pyarrow/tests/parquet/test_parquet_file.py b/python/pyarrow/tests/parquet/test_parquet_file.py index a62b5c3298c9..28ac93388f86 100644 --- a/python/pyarrow/tests/parquet/test_parquet_file.py +++ b/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):