diff --git a/cpp/src/arrow/dataset/file_parquet.cc b/cpp/src/arrow/dataset/file_parquet.cc index 8aa96926e493..7373a7035301 100644 --- a/cpp/src/arrow/dataset/file_parquet.cc +++ b/cpp/src/arrow/dataset/file_parquet.cc @@ -681,6 +681,37 @@ Result> ParquetFileFormat::MakeFragment( std::move(physical_schema), std::nullopt)); } +Result> ParquetFileFormat::MakeFragment( + FileSource source, compute::Expression partition_expression, + std::shared_ptr physical_schema, std::optional> row_groups, + std::shared_ptr metadata) { + if (metadata == nullptr) { + return Status::Invalid("ParquetFileFormat::MakeFragment: metadata must not be null"); + } + + auto properties = MakeArrowReaderProperties(*this, *metadata); + std::shared_ptr metadata_schema; + RETURN_NOT_OK(parquet::arrow::FromParquetSchema( + metadata->schema(), properties, metadata->key_value_metadata(), &metadata_schema)); + if (physical_schema && !physical_schema->Equals(*metadata_schema)) { + return Status::Invalid("Fragment initialized with physical schema ", *physical_schema, + " but the provided metadata for ", source.path(), + " has schema ", *metadata_schema); + } + ARROW_ASSIGN_OR_RAISE(auto manifest, GetSchemaManifest(*metadata, properties)); + + if (!row_groups) { + row_groups = Iota(metadata->num_row_groups()); + } + + std::shared_ptr fragment(new ParquetFileFragment( + std::move(source), shared_from_this(), std::move(partition_expression), + std::move(metadata_schema), std::move(row_groups))); + // `metadata` owns the SchemaDescriptor referenced by `manifest`. + RETURN_NOT_OK(fragment->SetMetadata(std::move(metadata), std::move(manifest))); + return fragment; +} + // // ParquetFileWriter, ParquetFileWriteOptions // @@ -834,7 +865,7 @@ Status ParquetFileFragment::SetMetadata( for (int row_group : *row_groups_) { // Ensure RowGroups are indexing valid RowGroups before augmenting. - if (row_group < metadata_->num_row_groups()) continue; + if (row_group >= 0 && row_group < metadata_->num_row_groups()) continue; return Status::IndexError("ParquetFileFragment references row group ", row_group, " but ", source_.path(), " only has ", diff --git a/cpp/src/arrow/dataset/file_parquet.h b/cpp/src/arrow/dataset/file_parquet.h index 1811a96bf986..145f27af9c05 100644 --- a/cpp/src/arrow/dataset/file_parquet.h +++ b/cpp/src/arrow/dataset/file_parquet.h @@ -120,6 +120,26 @@ class ARROW_DS_EXPORT ParquetFileFormat : public FileFormat { FileSource source, compute::Expression partition_expression, std::shared_ptr physical_schema, std::vector row_groups); + /// \brief Create a Fragment initialized with already loaded FileMetaData. + /// + /// No I/O is performed. The fragment uses `metadata` instead of reading the + /// footer of `source`, so row groups can be pruned using statistics and scanned + /// without a redundant footer read. `metadata` must describe the file referenced + /// by `source`, e.g. a per-file subset of a `_metadata` file obtained with + /// parquet::FileMetaData::Subset(). It is kept alive by the fragment. + /// + /// \param[in] source the file to scan + /// \param[in] partition_expression an expression guaranteed true for all rows + /// \param[in] physical_schema if null, inferred from `metadata`; otherwise it must + /// equal the schema inferred from `metadata` + /// \param[in] row_groups indices of the selected row groups, or std::nullopt to + /// select all row groups described by `metadata` + /// \param[in] metadata the FileMetaData of `source`, must not be null + Result> MakeFragment( + FileSource source, compute::Expression partition_expression, + std::shared_ptr physical_schema, std::optional> row_groups, + std::shared_ptr metadata); + /// \brief Return a FileReader on the given source. Result> GetReader( const FileSource& source, const std::shared_ptr& options) const; diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 0d86f5eaa0d3..42db2f18bdf9 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -433,6 +433,233 @@ TEST_F(TestParquetFileFormat, CachedMetadata) { ASSERT_NE(nullptr, pq_fragment->metadata()); } +namespace { + +// The last 8 bytes of a Parquet file hold the footer length and the magic +// bytes, which are always read when the footer is parsed. +bool ReadsFooter(const io::TrackedRandomAccessFile& file, int64_t file_size) { + for (const auto& range : file.get_read_ranges()) { + if (range.offset + range.length > file_size - 8) return true; + } + return false; +} + +// A FileSource over an in-memory Parquet file counting opens and reads. +struct TrackedParquetSource { + explicit TrackedParquetSource(std::shared_ptr buffer) + : buffer(std::move(buffer)), + buffer_reader(std::make_shared(this->buffer)), + tracked(io::TrackedRandomAccessFile::Make(buffer_reader.get())) {} + + FileSource source() { + return FileSource( + [this]() -> Result> { + ++num_opens; + return tracked; + }, + buffer->size()); + } + + std::shared_ptr buffer; + std::shared_ptr buffer_reader; + std::shared_ptr tracked; + int num_opens = 0; +}; + +} // namespace + +TEST_F(TestParquetFileFormat, MakeFragmentWithMetadata) { + constexpr int kNumRowGroups = 16; + constexpr int64_t kTotalNumRows = kNumRowGroups * (kNumRowGroups + 1) / 2; + + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(kNumRowGroups); + ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get())); + auto metadata = parquet::ReadMetaData(std::make_shared(buffer)); + ASSERT_EQ(metadata->num_row_groups(), kNumRowGroups); + + TrackedParquetSource file(buffer); + ASSERT_OK_AND_ASSIGN(auto fragment, + format_->MakeFragment(file.source(), literal(true), + /*physical_schema=*/nullptr, + /*row_groups=*/std::nullopt, metadata)); + + // Construction, schema and metadata access perform no I/O. + ASSERT_EQ(fragment->metadata(), metadata); + ASSERT_EQ(fragment->row_groups(), ::arrow::internal::Iota(kNumRowGroups)); + ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema()); + ASSERT_OK(fragment->EnsureCompleteMetadata()); + ASSERT_EQ(file.num_opens, 0); + ASSERT_EQ(file.tracked->num_reads(), 0); + + // The schema matches the one read from the file by a regular fragment. + auto regular = MakeFragment(FileSource(buffer)); + ASSERT_OK_AND_ASSIGN(auto expected_schema, regular->ReadPhysicalSchema()); + AssertSchemaEqual(*expected_schema, *physical_schema); + + // Scanning reads data pages but not the footer. + SetSchema(reader->schema()->fields()); + SetFilter(literal(true)); + CountRowsAndBatchesInScan(fragment, kTotalNumRows, kNumRowGroups); + ASSERT_EQ(file.num_opens, 1); + ASSERT_GT(file.tracked->num_reads(), 0); + ASSERT_FALSE(ReadsFooter(*file.tracked, buffer->size())); + ASSERT_EQ(fragment->metadata(), metadata); + + // Contrast: a regular fragment over the same file reads the footer on scan. + TrackedParquetSource regular_file(buffer); + ASSERT_OK_AND_ASSIGN(auto regular_fragment, + format_->MakeFragment(regular_file.source(), literal(true))); + CountRowsAndBatchesInScan(regular_fragment, kTotalNumRows, kNumRowGroups); + ASSERT_TRUE(ReadsFooter(*regular_file.tracked, buffer->size())); +} + +TEST_F(TestParquetFileFormat, MakeFragmentWithMetadataRowGroups) { + constexpr int kNumRowGroups = 16; + + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(kNumRowGroups); + ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get())); + auto metadata = parquet::ReadMetaData(std::make_shared(buffer)); + SetSchema(reader->schema()->fields()); + + // Row group i holds i + 1 rows. + ASSERT_OK_AND_ASSIGN(auto fragment, + format_->MakeFragment(FileSource(buffer), literal(true), + /*physical_schema=*/nullptr, + std::vector{1, 3, 10}, metadata)); + ASSERT_EQ(fragment->row_groups(), std::vector({1, 3, 10})); + + SetFilter(literal(true)); + CountRowsAndBatchesInScan(fragment, 2 + 4 + 11, 3); + // Pruning applies to the selected row groups only. + SetFilter(less(field_ref("i64"), literal(5))); + CountRowsAndBatchesInScan(fragment, 2 + 4, 2); + CountRowGroupsInFragment(fragment, {1, 3}, less(field_ref("i64"), literal(5))); + SetFilter(equal(field_ref("i64"), literal(1))); + CountRowsAndBatchesInScan(fragment, 0, 0); + + ASSERT_OK_AND_ASSIGN(fragment, format_->MakeFragment(FileSource(buffer), literal(true), + /*physical_schema=*/nullptr, + std::vector{}, metadata)); + SetFilter(literal(true)); + CountRowsAndBatchesInScan(fragment, 0, 0); +} + +TEST_F(TestParquetFileFormat, MakeFragmentWithMetadataPredicatePushdown) { + constexpr int kNumRowGroups = 16; + constexpr int64_t kTotalNumRows = kNumRowGroups * (kNumRowGroups + 1) / 2; + + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(kNumRowGroups); + ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get())); + auto metadata = parquet::ReadMetaData(std::make_shared(buffer)); + SetSchema(reader->schema()->fields()); + + TrackedParquetSource file(buffer); + ASSERT_OK_AND_ASSIGN(auto fragment, + format_->MakeFragment(file.source(), literal(true), + /*physical_schema=*/nullptr, + /*row_groups=*/std::nullopt, metadata)); + + // Filters matching no row group are resolved from the metadata alone. + SetFilter(literal(false)); + CountRowsAndBatchesInScan(fragment, 0, 0); + SetFilter(equal(field_ref("i64"), literal(kNumRowGroups + 1))); + CountRowsAndBatchesInScan(fragment, 0, 0); + SetFilter(and_(equal(field_ref("i64"), literal(1)), + equal(field_ref("u8"), literal(2)))); + CountRowsAndBatchesInScan(fragment, 0, 0); + ASSERT_EQ(file.num_opens, 0); + + for (int64_t i = 1; i <= kNumRowGroups; i++) { + SetFilter(equal(field_ref("i64"), literal(i))); + CountRowsAndBatchesInScan(fragment, i, 1); + } + SetFilter(less(field_ref("i64"), literal(6))); + CountRowsAndBatchesInScan(fragment, 5 * (5 + 1) / 2, 5); + SetFilter(greater_equal(field_ref("i64"), literal(6))); + CountRowsAndBatchesInScan(fragment, kTotalNumRows - (5 * (5 + 1) / 2), + kNumRowGroups - 5); + ASSERT_FALSE(ReadsFooter(*file.tracked, buffer->size())); + + CountRowGroupsInFragment(fragment, {5, 6}, + and_(greater_equal(field_ref("i64"), literal(6)), + less(field_ref("i64"), literal(8)))); + + auto options = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto predicate, + less_equal(field_ref("i64"), literal(4)).Bind(*reader->schema())); + ASSERT_FINISHES_OK_AND_EQ(std::make_optional(4 * (4 + 1) / 2), + fragment->CountRows(predicate, options)); +} + +TEST_F(TestParquetFileFormat, MakeFragmentWithMetadataInvalid) { + constexpr int kNumRowGroups = 4; + + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(kNumRowGroups); + ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get())); + auto metadata = parquet::ReadMetaData(std::make_shared(buffer)); + + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, + ::testing::HasSubstr("metadata must not be null"), + format_->MakeFragment(FileSource(buffer), literal(true), + nullptr, std::nullopt, nullptr)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + IndexError, ::testing::HasSubstr("references row group 4"), + format_->MakeFragment(FileSource(buffer), literal(true), nullptr, + std::vector{0, kNumRowGroups}, metadata)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + IndexError, ::testing::HasSubstr("references row group -1"), + format_->MakeFragment(FileSource(buffer), literal(true), nullptr, + std::vector{-1}, metadata)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Fragment initialized with physical schema"), + format_->MakeFragment(FileSource(buffer), literal(true), + schema({field("x", int32())}), std::nullopt, metadata)); +} + +TEST_F(TestParquetFileFormat, MakeFragmentWithMetadataSubset) { + // Emulate a `_metadata` file: the row groups of two files concatenated in + // a single FileMetaData. Each file's fragment gets its own subset. + std::vector> buffers; + std::shared_ptr combined; + std::vector> row_groups_per_file; + int offset = 0; + for (int num_row_groups : {3, 5}) { + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(num_row_groups); + ASSERT_OK_AND_ASSIGN(auto buffer, ParquetFormatHelper::Write(reader.get())); + auto metadata = parquet::ReadMetaData(std::make_shared(buffer)); + if (combined == nullptr) { + combined = metadata->Subset(::arrow::internal::Iota(metadata->num_row_groups())); + } else { + combined->AppendRowGroups(*metadata); + } + row_groups_per_file.push_back( + ::arrow::internal::Iota(offset, offset + num_row_groups)); + offset += num_row_groups; + buffers.push_back(std::move(buffer)); + } + ASSERT_EQ(combined->num_row_groups(), 8); + SetSchema(ArithmeticDatasetFixture::schema()->fields()); + + std::vector> fragments; + for (size_t i = 0; i < buffers.size(); i++) { + auto subset = combined->Subset(row_groups_per_file[i]); + ASSERT_EQ(subset->num_row_groups(), static_cast(row_groups_per_file[i].size())); + ASSERT_OK_AND_ASSIGN(auto fragment, + format_->MakeFragment(FileSource(buffers[i]), literal(true), + nullptr, std::nullopt, std::move(subset))); + fragments.push_back(std::move(fragment)); + } + // The fragments own their metadata and schema descriptors. + combined.reset(); + + SetFilter(literal(true)); + CountRowsAndBatchesInScan(fragments[0], 1 + 2 + 3, 3); + CountRowsAndBatchesInScan(fragments[1], 1 + 2 + 3 + 4 + 5, 5); + SetFilter(greater(field_ref("i64"), literal(3))); + CountRowsAndBatchesInScan(fragments[0], 0, 0); + CountRowsAndBatchesInScan(fragments[1], 4 + 5, 2); +} + TEST_F(TestParquetFileFormat, MultithreadedScan) { constexpr int64_t kNumRowGroups = 16; diff --git a/cpp/src/parquet/metadata.cc b/cpp/src/parquet/metadata.cc index 2bc871a8c865..43929a08fd79 100644 --- a/cpp/src/parquet/metadata.cc +++ b/cpp/src/parquet/metadata.cc @@ -954,7 +954,7 @@ class FileMetaData::FileMetaDataImpl { std::shared_ptr Subset(const std::vector& row_groups) { for (int i : row_groups) { - if (i < num_row_groups()) continue; + if (i >= 0 && i < num_row_groups()) continue; throw ParquetException( "The file only has ", num_row_groups(), diff --git a/cpp/src/parquet/metadata_test.cc b/cpp/src/parquet/metadata_test.cc index ee2f0619ebf2..362bbc90f4f5 100644 --- a/cpp/src/parquet/metadata_test.cc +++ b/cpp/src/parquet/metadata_test.cc @@ -275,6 +275,10 @@ TEST(Metadata, TestBuildAccess) { f_accessor_1 = f_accessor_2->Subset({0}); f_accessor_1->AppendRowGroups(*f_accessor->Subset({0})); ASSERT_TRUE(f_accessor_1->Equals(*f_accessor->Subset({2, 0}))); + + // Subset rejects out-of-range row group indices + ASSERT_THROW(f_accessor->Subset({8}), ParquetException); + ASSERT_THROW(f_accessor->Subset({0, -1}), ParquetException); } namespace { diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx index b3553708fb6a..d20262918e0a 100644 --- a/python/pyarrow/_dataset_parquet.pyx +++ b/python/pyarrow/_dataset_parquet.pyx @@ -242,7 +242,8 @@ cdef class ParquetFileFormat(FileFormat): return f"" def make_fragment(self, file, filesystem=None, - Expression partition_expression=None, row_groups=None, *, file_size=None): + Expression partition_expression=None, row_groups=None, *, + file_size=None, FileMetaData metadata=None): """ Make a FileFragment from a given file. @@ -261,6 +262,14 @@ cdef class ParquetFileFormat(FileFormat): file_size : int, optional The size of the file in bytes. Can improve performance with high-latency filesystems when file size needs to be known before reading. + metadata : pyarrow.parquet.FileMetaData, optional + The already loaded metadata of `file`. The fragment is then created + without any I/O, and scanning it uses this metadata instead of + reading the file footer again. Row group statistics are still used + to skip row groups. To get the metadata of one file from a combined + ``_metadata`` file, use :meth:`pyarrow.parquet.FileMetaData.subset` + with the indices of that file's row groups. If `row_groups` is also + given, its indices refer to the row groups of `metadata`. Returns ------- @@ -269,8 +278,21 @@ cdef class ParquetFileFormat(FileFormat): """ cdef: vector[int] c_row_groups + optional[vector[int]] c_opt_row_groups if partition_expression is None: partition_expression = _true + if metadata is not None: + c_source = _make_file_source(file, filesystem, file_size) + if row_groups is not None: + c_row_groups = [ row_group for row_group in set(row_groups)] + c_opt_row_groups = c_row_groups + c_fragment = GetResultValue( + self.parquet_format.MakeFragment(move(c_source), + partition_expression.unwrap(), + nullptr, + move(c_opt_row_groups), + metadata.sp_metadata)) + return Fragment.wrap(move(c_fragment)) if row_groups is None: return super().make_fragment(file, filesystem, partition_expression, file_size=file_size) diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 40c504ef3b04..bf64128fa625 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -1216,6 +1216,49 @@ cdef class FileMetaData(_Weakrefable): c_metadata = other.sp_metadata self._metadata.AppendRowGroups(deref(c_metadata)) + def subset(self, row_group_indices): + """ + Return a FileMetaData containing a subset of the row groups. + + This can be used to extract the metadata of a single data file from + a combined ``_metadata`` file (see :func:`pyarrow.parquet.write_metadata`), + to pass it to :meth:`pyarrow.dataset.ParquetFileFormat.make_fragment`. + + Parameters + ---------- + row_group_indices : list of int + Indices of the row groups to include, in the order they should + appear in the returned metadata. + + Returns + ------- + metadata : FileMetaData + A new FileMetaData with the same schema and the selected row groups. + Row groups are renumbered from 0 in the order given. + + Raises + ------ + IndexError + If an index is negative or not lower than ``num_row_groups``. + """ + cdef: + vector[int] c_row_groups + shared_ptr[CFileMetaData] c_subset + int num_row_groups = self._metadata.num_row_groups() + int index + FileMetaData result + + for index in row_group_indices: + if index < 0 or index >= num_row_groups: + raise IndexError('{0} out of bounds'.format(index)) + c_row_groups.push_back(index) + + with nogil: + c_subset = self._metadata.Subset(c_row_groups) + result = FileMetaData.__new__(FileMetaData) + result.init(c_subset) + return result + def write_metadata_file(self, where): """ Write the metadata to a metadata-only Parquet file. diff --git a/python/pyarrow/includes/libarrow_dataset_parquet.pxd b/python/pyarrow/includes/libarrow_dataset_parquet.pxd index 8f4917f6c0c8..99f89641ecca 100644 --- a/python/pyarrow/includes/libarrow_dataset_parquet.pxd +++ b/python/pyarrow/includes/libarrow_dataset_parquet.pxd @@ -74,6 +74,12 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil: CExpression partition_expression, shared_ptr[CSchema] physical_schema, vector[int] row_groups) + CResult[shared_ptr[CParquetFileFragment]] MakeFragment( + CFileSource source, + CExpression partition_expression, + shared_ptr[CSchema] physical_schema, + optional[vector[int]] row_groups, + shared_ptr[CFileMetaData] metadata) cdef cppclass CParquetFragmentScanOptions \ "arrow::dataset::ParquetFragmentScanOptions"(CFragmentScanOptions): diff --git a/python/pyarrow/includes/libparquet.pxd b/python/pyarrow/includes/libparquet.pxd index a7ad239feac6..ee3435af4b9b 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -403,6 +403,7 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil: void set_file_path(const c_string& path) void AppendRowGroups(const CFileMetaData& other) except + + shared_ptr[CFileMetaData] Subset(const vector[int]& row_groups) except + unique_ptr[CRowGroupMetaData] RowGroup(int i) const SchemaDescriptor* schema() diff --git a/python/pyarrow/tests/parquet/test_metadata.py b/python/pyarrow/tests/parquet/test_metadata.py index 9eee70b125e5..e642a69b12f2 100644 --- a/python/pyarrow/tests/parquet/test_metadata.py +++ b/python/pyarrow/tests/parquet/test_metadata.py @@ -754,6 +754,54 @@ def test_metadata_append_row_groups_diff(t1, t2, expected_error): meta1.append_row_groups(meta2) +def test_metadata_subset(tempdir, pickle_module): + # Row group i holds the values [3 * i, 3 * i + 2] in column 'a'. + table = pa.table({'a': list(range(15)), + 'b': [str(i) for i in range(15)]}) + buf = io.BytesIO() + pq.write_table(table, buf, row_group_size=3) + metadata = pq.ParquetFile(io.BytesIO(buf.getvalue())).metadata + assert metadata.num_row_groups == 5 + + subset = metadata.subset([3, 1]) + assert subset.num_row_groups == 2 + assert subset.num_rows == 6 + assert subset.num_columns == metadata.num_columns + assert subset.schema.equals(metadata.schema) + assert subset.created_by == metadata.created_by + assert subset.metadata == metadata.metadata + # Row groups are renumbered in the order given, with their statistics + assert subset.row_group(0).equals(metadata.row_group(3)) + assert subset.row_group(1).equals(metadata.row_group(1)) + stats = subset.row_group(0).column(0).statistics + assert (stats.min, stats.max) == (9, 11) + + assert metadata.subset(range(5)).equals(metadata) + assert metadata.subset([]).num_row_groups == 0 + assert metadata.num_row_groups == 5 + + # Serialization + assert pickle_module.loads(pickle_module.dumps(subset)).equals(subset) + path = str(tempdir / '_metadata') + subset.write_metadata_file(path) + assert pq.read_metadata(path).equals(subset) + + # The subset owns its schema and row groups + del metadata + assert subset.schema.to_arrow_schema().equals(table.schema) + assert subset.row_group(1).column(0).statistics.min == 3 + + # Invalid indices + metadata = subset + for indices in ([2], [0, -1], [2 ** 31 - 1]): + with pytest.raises(IndexError, match="out of bounds"): + metadata.subset(indices) + with pytest.raises(TypeError): + metadata.subset(['a']) + with pytest.raises(TypeError): + metadata.subset(0) + + @pytest.mark.s3 def test_write_metadata_fs_file_combinations(tempdir, s3_example_s3fs): s3_fs, s3_path = s3_example_s3fs diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index b50441a1ae05..554f27bbe2cf 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -1822,6 +1822,239 @@ def test_fragments_parquet_row_groups_reconstruct(tempdir, dataset_reader, dataset_reader.to_table(new_fragment) +class _ReadRecordingFile: + """A local binary file recording the (offset, length) of every read.""" + + def __init__(self, path, reads): + self._f = open(path, 'rb') + self._reads = reads + + def read(self, nbytes=-1): + offset = self._f.tell() + data = self._f.read(nbytes) + self._reads.append((offset, len(data))) + return data + + def seek(self, offset, whence=0): + return self._f.seek(offset, whence) + + def tell(self): + return self._f.tell() + + def close(self): + self._f.close() + + @property + def closed(self): + return self._f.closed + + def readable(self): + return True + + def seekable(self): + return True + + +class _ReadRecordingHandler(ProxyHandler): + """Local filesystem handler recording file opens and read ranges.""" + + def __init__(self): + super().__init__(fs.LocalFileSystem()) + self.opens = [] + self.reads = {} + + def open_input_file(self, path): + self.opens.append(path) + return pa.PythonFile( + _ReadRecordingFile(path, self.reads.setdefault(path, [])), mode='r') + + def reads_footer(self, path): + # The trailing 8 bytes (footer length and magic) are always read + # when the footer is parsed. + size = os.path.getsize(path) + return any(offset + length > size - 8 + for offset, length in self.reads.get(path, [])) + + +def _write_row_groups_file(path, num_row_groups=4): + # Row group i holds the values [2 * i, 2 * i + 1] in column 'f1'. + table = pa.table({'f1': list(range(2 * num_row_groups)), + 'f2': [str(i) for i in range(2 * num_row_groups)]}) + pq.write_table(table, path, row_group_size=2) + return table + + +@pytest.mark.parquet +def test_make_fragment_with_metadata(tempdir): + path = str(tempdir / 'data.parquet') + table = _write_row_groups_file(path) + metadata = pq.read_metadata(path) + assert metadata.num_row_groups == 4 + + handler = _ReadRecordingHandler() + filesystem = fs.PyFileSystem(handler) + parquet_format = ds.ParquetFileFormat() + + fragment = parquet_format.make_fragment(path, filesystem, metadata=metadata) + assert isinstance(fragment, ds.ParquetFileFragment) + # Construction and metadata access perform no I/O + assert fragment.row_groups == [0, 1, 2, 3] + assert fragment.num_row_groups == 4 + assert fragment.metadata.equals(metadata) + assert fragment.physical_schema.equals(table.schema) + assert [rg.statistics['f1'] for rg in fragment.row_groups] == [ + {'min': 2 * i, 'max': 2 * i + 1} for i in range(4)] + fragment.ensure_complete_metadata() + assert handler.opens == [] + + # Scanning reads the data, but not the footer again + assert fragment.to_table().equals(table) + assert handler.opens == [path] + assert handler.reads[path] + assert not handler.reads_footer(path) + + # file_size can be combined with metadata + fragment = parquet_format.make_fragment( + path, filesystem, file_size=os.path.getsize(path), metadata=metadata) + assert fragment.to_table().equals(table) + assert not handler.reads_footer(path) + + # Without metadata, the fragment reads the footer (unchanged behavior) + handler = _ReadRecordingHandler() + filesystem = fs.PyFileSystem(handler) + fragment = parquet_format.make_fragment(path, filesystem) + assert handler.opens == [] + assert fragment.to_table().equals(table) + assert handler.reads_footer(path) + assert fragment.metadata.equals(metadata) + + +@pytest.mark.parquet +def test_make_fragment_with_metadata_row_groups(tempdir, dataset_reader): + path = str(tempdir / 'data.parquet') + table = _write_row_groups_file(path) + metadata = pq.read_metadata(path) + + handler = _ReadRecordingHandler() + filesystem = fs.PyFileSystem(handler) + parquet_format = ds.ParquetFileFormat() + + fragment = parquet_format.make_fragment( + path, filesystem, row_groups=[1, 3], metadata=metadata) + assert fragment.row_groups == [1, 3] + assert fragment.num_row_groups == 2 + expected = pa.concat_tables([table.slice(2, 2), table.slice(6, 2)]) + assert dataset_reader.to_table(fragment).equals(expected) + + # Statistics are used to prune the selected row groups + result = dataset_reader.to_table(fragment, filter=ds.field('f1') >= 6) + assert result.equals(table.slice(6, 2)) + sub_fragments = list(fragment.split_by_row_group(ds.field('f1') < 4)) + assert [f.row_groups for f in sub_fragments] == [[1]] + + # A filter matching no row group is answered from the metadata alone + handler.opens.clear() + result = dataset_reader.to_table(fragment, filter=ds.field('f1') < 2) + assert result.num_rows == 0 + assert result.schema.equals(table.schema) + assert handler.opens == [] + assert fragment.count_rows(filter=ds.field('f1') > 100) == 0 + assert handler.opens == [] + + # All row groups, filters matching some of them + fragment = parquet_format.make_fragment(path, filesystem, metadata=metadata) + result = dataset_reader.to_table(fragment, filter=ds.field('f1') == 5) + assert result.equals(table.slice(5, 1)) + result = dataset_reader.to_table( + fragment, filter=(ds.field('f1') > 1) & (ds.field('f1') < 6)) + assert result.equals(table.slice(2, 4)) + assert not handler.reads_footer(path) + + # Fragments are usable in a dataset + dataset = ds.FileSystemDataset([fragment], table.schema, parquet_format, + filesystem) + assert dataset.to_table(filter=ds.field('f1') >= 4).equals(table.slice(4)) + + +@pytest.mark.parquet +def test_make_fragment_with_metadata_invalid(tempdir): + path = str(tempdir / 'data.parquet') + _write_row_groups_file(path) + metadata = pq.read_metadata(path) + parquet_format = ds.ParquetFileFormat() + localfs = fs.LocalFileSystem() + + with pytest.raises(TypeError): + parquet_format.make_fragment(path, localfs, metadata="not metadata") + with pytest.raises(TypeError): + parquet_format.make_fragment(path, localfs, metadata=metadata.schema) + with pytest.raises(TypeError): + # metadata is keyword-only + parquet_format.make_fragment(path, localfs, None, None, metadata) + with pytest.raises(IndexError, match="references row group 4"): + parquet_format.make_fragment(path, localfs, row_groups=[0, 4], + metadata=metadata) + with pytest.raises(IndexError, match="references row group -1"): + parquet_format.make_fragment(path, localfs, row_groups=[-1], + metadata=metadata) + + # metadata=None keeps the existing behavior + fragment = parquet_format.make_fragment(path, localfs, metadata=None) + assert fragment.row_groups == [0, 1, 2, 3] + + +@pytest.mark.parquet +def test_make_fragment_with_metadata_from_metadata_file(tempdir): + # Write a _metadata file combining the row groups of several files, then + # build one fragment per file from a subset of it. + root = tempdir / 'dataset' + root.mkdir() + metadata_collector = [] + tables = [] + for i in range(3): + table = pa.table({'f1': list(range(10 * i, 10 * i + 6)), + 'f2': [float(i)] * 6}) + relative_path = f'part-{i}.parquet' + pq.write_table(table, str(root / relative_path), row_group_size=2, + metadata_collector=metadata_collector) + metadata_collector[-1].set_file_path(relative_path) + tables.append(table) + metadata_path = str(root / '_metadata') + pq.write_metadata(tables[0].schema, metadata_path, + metadata_collector=metadata_collector) + expected = pa.concat_tables(tables) + + combined = pq.read_metadata(metadata_path) + assert combined.num_row_groups == 9 + row_groups_per_file = {} + for i in range(combined.num_row_groups): + file_path = combined.row_group(i).column(0).file_path + row_groups_per_file.setdefault(file_path, []).append(i) + + handler = _ReadRecordingHandler() + filesystem = fs.PyFileSystem(handler) + parquet_format = ds.ParquetFileFormat() + fragments = [] + for file_path, row_groups in row_groups_per_file.items(): + file_metadata = combined.subset(row_groups) + assert file_metadata.num_row_groups == len(row_groups) + fragments.append(parquet_format.make_fragment( + str(root / file_path), filesystem, metadata=file_metadata)) + del combined, file_metadata + assert handler.opens == [] + + dataset = ds.FileSystemDataset(fragments, expected.schema, parquet_format, + filesystem) + assert dataset.to_table().equals(expected) + reference = ds.parquet_dataset(metadata_path) + assert dataset.to_table().equals(reference.to_table()) + for expr in [ds.field('f1') == 13, ds.field('f1') >= 22, + ds.field('f2') == 1.0, ds.field('f1') > 100]: + assert dataset.to_table(filter=expr).equals( + reference.to_table(filter=expr)) + assert not any(handler.reads_footer(path) for path in handler.reads) + + @pytest.mark.parquet def test_fragments_parquet_subset_ids(tempdir, open_logging_fs, dataset_reader):