[GitHub] [arrow] edponce commented on a change in pull request #12089: ARROW-9285: [C++] Detect unauthorized memory allocations in function kernels

[GitBox <gi...@apache.org>]

edponce commented on a change in pull request #12089:
URL: https://github.com/apache/arrow/pull/12089#discussion_r779611886



##########
File path: cpp/src/arrow/compute/exec.cc
##########
@@ -110,6 +112,78 @@ int64_t ExecBatch::TotalBufferSize() const {
   return sum;
 }
 
+bool AddBuffersToSet(const ArrayData& array_data,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  bool insertion_occured = false;
+  for (const auto& buffer : array_data.buffers) {
+    insertion_occured = (buffer && seen_buffers->insert(buffer).second);
+  }
+  for (const auto& child : array_data.child_data) {
+    insertion_occured |= AddBuffersToSet(*child, seen_buffers);
+  }
+  if (array_data.dictionary) {
+    insertion_occured |= AddBuffersToSet(*array_data.dictionary, seen_buffers);
+  }
+  return insertion_occured;
+}
+
+bool AddBuffersToSet(const Array& array,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  return AddBuffersToSet(*array.data(), seen_buffers);
+}
+
+bool AddBuffersToSet(const ChunkedArray& chunked_array,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  bool insertion_occured = false;
+  for (const auto& chunk : chunked_array.chunks()) {
+    insertion_occured |= AddBuffersToSet(*chunk, seen_buffers);
+  }
+  return insertion_occured;
+}
+
+bool AddBuffersToSet(const RecordBatch& record_batch,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  bool insertion_occured = false;
+  for (const auto& column : record_batch.columns()) {
+    insertion_occured |= AddBuffersToSet(*column, seen_buffers);
+  }
+  return insertion_occured;
+}
+
+bool AddBuffersToSet(const Table& table,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  bool insertion_occured = false;
+  for (const auto& column : table.columns()) {
+    insertion_occured |= AddBuffersToSet(*column, seen_buffers);
+  }
+  return insertion_occured;
+}
+
+// Add all Buffers to a given set, return true if anything was actually added.
+// If all the buffers in the datum were already in the set, this will return false.
+bool AddBuffersToSet(Datum datum,
+                     std::unordered_set<std::shared_ptr<Buffer>>* seen_buffers) {
+  switch (datum.kind()) {
+    case Datum::ARRAY:
+      return AddBuffersToSet(*util::get<std::shared_ptr<ArrayData>>(datum.value),
+                             seen_buffers);
+    case Datum::CHUNKED_ARRAY:
+      return AddBuffersToSet(*util::get<std::shared_ptr<ChunkedArray>>(datum.value),
+                             seen_buffers);
+    case Datum::RECORD_BATCH:
+      return AddBuffersToSet(*util::get<std::shared_ptr<RecordBatch>>(datum.value),
+                             seen_buffers);
+    case Datum::TABLE:
+      return AddBuffersToSet(*util::get<std::shared_ptr<Table>>(datum.value),
+                             seen_buffers);

Review comment:
       I am curious if hashing all the values-buffers pointers would suffice instead of storing them in a set. We would still need to traverse nested data structures to capture all the child buffers.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org