[C++] Potential memory leak in Parquet reading with Dataset

[icexelloss]

Describe the bug, including details regarding any error messages, version, and platform.

Version

Arrow 12.0

Platform

Debian 5.4.228

Description

I have been testing "What is the max rss needed to scan through ~100G of data in a parquet stored in gcs using Arrow C++".

The current answer is about ~6G of memory which seems a bit high so I looked into it. What I observed during the process led me to think that there are some potential cache/memory issues in the dataset/parquet cpp code.

Main observation:
(1) As I am scanning through the dataset, I printed out (a) memory allocated by the memory pool from ScanOptions (b) process rss. I found that while (a) stays pretty stable throughout the scan (stays < 1G), (b) keeps increasing during the scan (looks linear to the number of files scanned).
(2) I tested ScanNode in Arrow as well as an in-house library that implements its own "S3Dataset" similar to Arrow dataset, both showing similar rss usage. (Which led me to think the issue is more likely to be in the parquet cpp code instead of dataset code).
(3) Scan the same dataset twice in the same process doesn't increase the max rss.

Following suggestions from the mailing list, I also did a memory profiling of the test program and the results seem to indicate potential memory leak.

Test code

https://gist.github.com/icexelloss/88195de046962e1d043c99d96e1b8b43

Massif output

massif.out.3505660.txt

Component(s)

C++, Parquet

[lidavidm]

You could upload it in a gist and link that here, perhaps?

[icexelloss]

You could upload it in a gist and link that here, perhaps?

I ended up changing the file extension and that worked :)

[icexelloss]

Not too familiar with the async/future code (nor massif outputs), but does this look like the Future objects are not destructed somehow (which has a reference to the parquet metadata object?) (You need to expand the massif output in visualizer to see the Future object at the bottom of the call chain)

[icexelloss]

cc @mapleFU @wgtmac @pitrou @felipecrv (Folks that have been opining on the mailing list)

[mapleFU]

Future in arrow is an shared_future, which means if reference is hold, the object might not be dtor. I guess maybe not that problem, but I'll try to run the script and reproduce the problem later

[westonpace]

Note that FileFragment in the datasets API caches the parquet metadata (with no option to disable this at the moment). So if you are scanning many files you will see memory grow over the lifetime of the scan as more and more metadatas are cached. I would expect a second scan would not grow the memory.

[icexelloss]

Note that FileFragment in the datasets API caches the parquet metadata (with no option to disable this at the moment). So if you are scanning many files you will see memory grow over the lifetime of the scan as more and more metadatas are cached. I would expect a second scan would not grow the memory.

Thanks @westonpace, can you give a pointer to where that is happening?

[westonpace]

In file_parquet.h. The ParquetFileFragment class has these variables:

  /// Indices of row groups selected by this fragment,
  /// or std::nullopt if all row groups are selected.
  std::optional<std::vector<int>> row_groups_;

  std::vector<compute::Expression> statistics_expressions_;
  std::vector<bool> statistics_expressions_complete_;
  std::shared_ptr<parquet::FileMetaData> metadata_;
  std::shared_ptr<parquet::arrow::SchemaManifest> manifest_;

They initially start empty/null and are initialized during a call to EnsureCompleteMetadata.

[icexelloss]

In file_parquet.h. The ParquetFileFragment class has these variables:

  /// Indices of row groups selected by this fragment,
  /// or std::nullopt if all row groups are selected.
  std::optional<std::vector<int>> row_groups_;

  std::vector<compute::Expression> statistics_expressions_;
  std::vector<bool> statistics_expressions_complete_;
  std::shared_ptr<parquet::FileMetaData> metadata_;
  std::shared_ptr<parquet::arrow::SchemaManifest> manifest_;

They initially start empty/null and are initialized during a call to EnsureCompleteMetadata.

Hmm, what do you mean by they are cached? Doesn't ParquetFileFragment get dtored after we scanned the fragment?

[icexelloss]

Oh I see we keep a vector of Fragment, each Fragment has a shared_ptr to the schema of the file

[westonpace]

Yes, the fragments are part of the dataset. So they stay around until the dataset is destroyed which has to be after the scan.

[icexelloss]

In a local build I cleared

  std::shared_ptr<parquet::FileMetaData> metadata_;
  std::shared_ptr<parquet::arrow::SchemaManifest> manifest_;

after scanning a fragment and things appear to be in better shape. Memory usage goes from 4G -> 700+Mb.

There still seem to be some leaking that I couldn't figure out but at least things seem to be in a much better shape after clearing those two fields.

massif.out.593871.txt

[mapleFU]

I was so busy these two weeks, sorry for late reply. Have you find out the reason? Seems that you're suffering from too many arrow::field? @icexelloss

[IamJeffG]

I believe I might be be experiencing this same problem through the Python API. Having to maintain a local build of Arrow doesn't sound like the right solution, so I wonder if there are ideas of how to achieve the same result as @icexelloss 's last comment, but through Arrow's normal APIs?

[mapleFU]

Result<std::vector<std::shared_ptr<FileFragment>>>
ParquetDatasetFactory::CollectParquetFragments(const Partitioning& partitioning) {
  std::vector<std::shared_ptr<FileFragment>> fragments(paths_with_row_group_ids_.size());

  size_t i = 0;
  for (const auto& e : paths_with_row_group_ids_) {
    const auto& path = e.first;
    auto metadata_subset = metadata_->Subset(e.second);

    auto row_groups = Iota(metadata_subset->num_row_groups());

    auto partition_expression =
        partitioning.Parse(StripPrefix(path, options_.partition_base_dir))
            .ValueOr(compute::literal(true));

    ARROW_ASSIGN_OR_RAISE(
        auto fragment,
        format_->MakeFragment({path, filesystem_}, std::move(partition_expression),
                              physical_schema_, std::move(row_groups)));

    RETURN_NOT_OK(fragment->SetMetadata(metadata_subset, manifest_));
    fragments[i++] = std::move(fragment);
  }

  return fragments;
}

I noticed that these metadata and manifest is shared_ptr and share among fragments, maybe when users have too many columns, this would be a huge cost?

[IamJeffG]

maybe when users have too many columns

Maybe but I think there must be more to it. In my example I have a partitioned parquet dataset on local disk, 8.6GB in total, with 13 columns and 38,747 fragments. Writing this dataset to a new location on disk (i.e. to compact the fragments) consumes all 8GB of RAM on my machine and then swaps to disk. I can't imagine that 13 columns or even 13×38,747 takes upwards of 8GB of memory.

[icexelloss]

I was so busy these two weeks, sorry for late reply. Have you find out the reason? Seems that you're suffering from too many arrow::field? @icexelloss

@mapleFU Sorry I have not found out the reason for the seemly leak of "arrow::Field" in #37630 (comment) post.

[icexelloss]

I believe I might be be experiencing this same problem through the Python API. Having to maintain a local build of Arrow doesn't sound like the right solution, so I wonder if there are ideas of how to achieve the same result as @icexelloss 's last comment, but through Arrow's normal APIs?

I don't think you can at the moment. Because there is no way to tell Dataset to not to cache these metadata through normal API. I had a local fix but didn't find time the push it out (the fix is actually quite simple - clear the metadata from segment after done reading a segment)

[icexelloss]

I do have plan to push the patch upstream proper when I have some spare time

FWIW, This is my local patch:

diff --git a/c/src/arrow/dataset/dataset.h b/c/src/arrow/dataset/dataset.h
index 1db230b16e9c2e52ad58c02255f87375307271d0..37c2ddca4884d32bb1721dd3a2d937f9faf6b93f 100644
--- a/c/src/arrow/dataset/dataset.h
+++ b/c/src/arrow/dataset/dataset.h
@@ -199,6 +199,9 @@ class ARROW_DS_EXPORT Fragment : public std::enable_shared_from_this<Fragment> {
     return partition_expression_;
   }
 
+  /// \brief Clear metadata in the fragment. By default this is an noop.
+  virtual void ClearMetadata() { return; }
+
   virtual ~Fragment() = default;
 
  protected:
diff --git a/c/src/arrow/dataset/file_parquet.h b/c/src/arrow/dataset/file_parquet.h
index f33190bd93347f781df91a6cda612031f09caf75..632a555894a7b91ee6ac722a6820a9eaff61f232 100644
--- a/c/src/arrow/dataset/file_parquet.h
+++ b/c/src/arrow/dataset/file_parquet.h
@@ -174,6 +174,11 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
   static std::optional<compute::Expression> EvaluateStatisticsAsExpression(
       const Field& field, const parquet::Statistics& statistics);
 
+  void ClearMetadata() override {
+    metadata_ = NULLPTR;
+    manifest_ = NULLPTR;
+  }
+
  private:
   ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
                       compute::Expression partition_expression,
diff --git a/c/src/arrow/dataset/scanner.cc b/c/src/arrow/dataset/scanner.cc
index 18981d1451980f4b775fb87e42dcaede89a1fc7c..8f83a7c2dc52b8f5c728b91be2f93442168002d7 100644
--- a/c/src/arrow/dataset/scanner.cc
+++ b/c/src/arrow/dataset/scanner.cc
@@ -1037,6 +1037,9 @@ Result<acero::ExecNode*> MakeScanNode(acero::ExecPlan* plan,
         // unnecessarily materialized columns in batch. We could drop them now instead of
         // letting them coast through the rest of the plan.
         auto guarantee = partial.fragment.value->partition_expression();
+        // This is safe because by the time we reach here we have already decoded
+        // the fragment into a batch so the metadata is no longer needed.
+        partial.fragment.value->ClearMetadata();
 
         ARROW_ASSIGN_OR_RAISE(
             std::optional<compute::ExecBatch> batch,

[felipecrv]

You can do .reset() on shared_ptrs to achieve the same effect. [1] I would call the method ResetMetadata() because "clear" in C++ usually means semantically clearing data, but not really deallocating the buffers (e.g. std::string::clear() and std::vector<T>::clear()).

[1] https://en.cppreference.com/w/cpp/memory/shared_ptr/reset

[westonpace]

I believe I might be be experiencing this same problem through the Python API. Having to maintain a local build of Arrow doesn't sound like the right solution, so I wonder if there are ideas of how to achieve the same result as @icexelloss 's last comment, but through Arrow's normal APIs?

We should definitely make metadata caching an optional feature of the scanner and/or dataset. I think the API could be as simple as...

my dataset = pyarrow.dataset.dataset(..., cache_metadata=False)

Any place that is using a dataset "temporarily" should also set this to false (e.g. when we run pyarrow.parquet.read_table it creates a dataset and scans it. That temporary dataset should NOT cache metadata)

[mapleFU]

Sorry I have not found out the reason for the seemly leak of "arrow::Field" in #37630 (comment) post.

Oh, in your graph the arrow::Field and schema occupied lots of memory. Release Manifest would release them. Previously I wonder if a schema cache would be better. But I think clear the pointer in scanner api would be much better.

[pitrou]

I've submitted a draft PR that adds a cache_metadata option when scanning.

Here are some memory consumption measurements on a synthetic dataset similar to the one in #45287:

• with cache_metadata=True (the default):

$ /usr/bin/time -f "\n%E real\n%U user\n%S sys\n%M kB peak RSS" python  -c 'import pyarrow.dataset as pd; pd.dataset("/home/antoine/arrow/data/bamboo-streaming-parquet-test-data/1000col-dataset/").to_table()'

0:11.01 real
107.66 user
4.37 sys
2069284 kB peak RSS

• with cache_metadata=False:

$ /usr/bin/time -f "\n%E real\n%U user\n%S sys\n%M kB peak RSS" python  -c 'import pyarrow.dataset as pd; pd.dataset("/home/antoine/arrow/data/bamboo-streaming-parquet-test-data/1000col-dataset/").to_table(cache_metadata=False)'

0:10.97 real
109.11 user
3.87 sys
1381696 kB peak RSS

[pitrou]

@icexelloss It would be useful it you could test the above PR on a more real-world dataset.

[icexelloss]

Thanks @pitrou! We actually already have an internal patch similar to what you have now but still observed higher than expected memory usage when scanning a parquet dataset - my colleague @timothydijamco is working a repro and should have it soon

I also noticed that you also cleared original_metadata_.reset(); - we will test that as well