GH-55325: [C++][Python] Accept known metadata in Parquet fragments

cpp/src/arrow/dataset/file_parquet.cc

@@ -681,6 +681,37 @@ Result<std::shared_ptr<FileFragment>> ParquetFileFormat::MakeFragment(
       std::move(physical_schema), std::nullopt));
 }
 
+Result<std::shared_ptr<ParquetFileFragment>> ParquetFileFormat::MakeFragment(
+    FileSource source, compute::Expression partition_expression,
+    std::shared_ptr<Schema> physical_schema, std::optional<std::vector<int>> row_groups,
+    std::shared_ptr<parquet::FileMetaData> metadata) {
+  if (metadata == nullptr) {
+    return Status::Invalid("ParquetFileFormat::MakeFragment: metadata must not be null");
+  }
+
+  auto properties = MakeArrowReaderProperties(*this, *metadata);
+  std::shared_ptr<Schema> 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<ParquetFileFragment> 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 ",

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<Schema> physical_schema, std::vector<int> 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<std::shared_ptr<ParquetFileFragment>> MakeFragment(
+      FileSource source, compute::Expression partition_expression,
+      std::shared_ptr<Schema> physical_schema, std::optional<std::vector<int>> row_groups,
+      std::shared_ptr<parquet::FileMetaData> metadata);
+
   /// \brief Return a FileReader on the given source.
   Result<std::shared_ptr<parquet::arrow::FileReader>> GetReader(
       const FileSource& source, const std::shared_ptr<ScanOptions>& options) const;

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)
+      : buffer(std::move(buffer)),
+        buffer_reader(std::make_shared<io::BufferReader>(this->buffer)),
+        tracked(io::TrackedRandomAccessFile::Make(buffer_reader.get())) {}
+
+  FileSource source() {
+    return FileSource(
+        [this]() -> Result<std::shared_ptr<io::RandomAccessFile>> {
+          ++num_opens;
+          return tracked;
+        },
+        buffer->size());
+  }
+
+  std::shared_ptr<Buffer> buffer;
+  std::shared_ptr<io::BufferReader> buffer_reader;
+  std::shared_ptr<io::TrackedRandomAccessFile> 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<io::BufferReader>(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<io::BufferReader>(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<int>{1, 3, 10}, metadata));
+  ASSERT_EQ(fragment->row_groups(), std::vector<int>({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<int64_t>(5)));
+  CountRowsAndBatchesInScan(fragment, 2 + 4, 2);
+  CountRowGroupsInFragment(fragment, {1, 3}, less(field_ref("i64"), literal(5)));
+  SetFilter(equal(field_ref("i64"), literal<int64_t>(1)));
+  CountRowsAndBatchesInScan(fragment, 0, 0);
+
+  ASSERT_OK_AND_ASSIGN(fragment, format_->MakeFragment(FileSource(buffer), literal(true),
+                                                       /*physical_schema=*/nullptr,
+                                                       std::vector<int>{}, 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<io::BufferReader>(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<int64_t>(kNumRowGroups + 1)));
+  CountRowsAndBatchesInScan(fragment, 0, 0);
+  SetFilter(and_(equal(field_ref("i64"), literal<int64_t>(1)),
+                 equal(field_ref("u8"), literal<uint8_t>(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<int64_t>(6)));
+  CountRowsAndBatchesInScan(fragment, 5 * (5 + 1) / 2, 5);
+  SetFilter(greater_equal(field_ref("i64"), literal<int64_t>(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<ScanOptions>();
+  ASSERT_OK_AND_ASSIGN(auto predicate,
+                       less_equal(field_ref("i64"), literal(4)).Bind(*reader->schema()));
+  ASSERT_FINISHES_OK_AND_EQ(std::make_optional<int64_t>(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<io::BufferReader>(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<int>{0, kNumRowGroups}, metadata));
+  EXPECT_RAISES_WITH_MESSAGE_THAT(
+      IndexError, ::testing::HasSubstr("references row group -1"),
+      format_->MakeFragment(FileSource(buffer), literal(true), nullptr,
+                            std::vector<int>{-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<std::shared_ptr<Buffer>> buffers;
+  std::shared_ptr<parquet::FileMetaData> combined;
+  std::vector<std::vector<int>> 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<io::BufferReader>(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<std::shared_ptr<ParquetFileFragment>> 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<int>(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<int64_t>(3)));
+  CountRowsAndBatchesInScan(fragments[0], 0, 0);
+  CountRowsAndBatchesInScan(fragments[1], 4 + 5, 2);
+}
+
 TEST_F(TestParquetFileFormat, MultithreadedScan) {
   constexpr int64_t kNumRowGroups = 16;
 

cpp/src/parquet/metadata.cc

@@ -954,7 +954,7 @@ class FileMetaData::FileMetaDataImpl {
 
   std::shared_ptr<FileMetaData> Subset(const std::vector<int>& 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(),

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 {

python/pyarrow/_dataset_parquet.pyx

@@ -242,7 +242,8 @@ cdef class ParquetFileFormat(FileFormat):
         return f"<ParquetFileFormat read_options={self.read_options}>"
 
     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 = [<int> row_group for row_group in set(row_groups)]
+                c_opt_row_groups = c_row_groups
+            c_fragment = <shared_ptr[CFragment]> GetResultValue(
+                self.parquet_format.MakeFragment(move(c_source),
+                                                 partition_expression.unwrap(),
+                                                 <shared_ptr[CSchema]>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)