[GitHub] [arrow] lidavidm commented on a change in pull request #12275: ARROW-15262: [C++] Create a ToTable sink node

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

lidavidm commented on a change in pull request #12275:
URL: https://github.com/apache/arrow/pull/12275#discussion_r796610755



##########
File path: cpp/src/arrow/compute/exec/sink_node.cc
##########
@@ -232,6 +232,51 @@ class ConsumingSinkNode : public ExecNode {
   std::shared_ptr<SinkNodeConsumer> consumer_;
 };
 
+/**
+ * @brief This node is an extension on ConsumingSinkNode
+ * to facilitate to get the output from an execution plan
+ * as a table. We define a custom SinkNodeConsumer to
+ * enable this functionality.
+ */
+
+struct TableSinkNodeConsumer : public arrow::compute::SinkNodeConsumer {
+ public:
+  TableSinkNodeConsumer(std::shared_ptr<Table>* out,
+                        std::shared_ptr<Schema> output_schema, MemoryPool* pool)
+      : out_(out), output_schema_(std::move(output_schema)), pool_(pool) {}
+
+  Status Consume(ExecBatch batch) override {
+    std::lock_guard<std::mutex> guard(consum_mutex);
+    ARROW_ASSIGN_OR_RAISE(auto rb, batch.ToRecordBatch(output_schema_, pool_));
+    batch_vector.push_back(rb);
+    return Status::OK();
+  }
+
+  Future<> Finish() override {
+    ARROW_ASSIGN_OR_RAISE(*out_, Table::FromRecordBatches(batch_vector));
+    return Status::OK();
+  }
+
+ private:
+  std::shared_ptr<Table>* out_;
+  std::shared_ptr<Schema> output_schema_;
+  MemoryPool* pool_;
+  std::vector<std::shared_ptr<RecordBatch>> batch_vector;
+  std::mutex consum_mutex;

Review comment:
       Typo? `consume_mutex_`? Or just `mutex_`

##########
File path: cpp/src/arrow/compute/exec/sink_node.cc
##########
@@ -232,6 +232,51 @@ class ConsumingSinkNode : public ExecNode {
   std::shared_ptr<SinkNodeConsumer> consumer_;
 };
 
+/**
+ * @brief This node is an extension on ConsumingSinkNode
+ * to facilitate to get the output from an execution plan
+ * as a table. We define a custom SinkNodeConsumer to
+ * enable this functionality.
+ */
+
+struct TableSinkNodeConsumer : public arrow::compute::SinkNodeConsumer {
+ public:
+  TableSinkNodeConsumer(std::shared_ptr<Table>* out,
+                        std::shared_ptr<Schema> output_schema, MemoryPool* pool)
+      : out_(out), output_schema_(std::move(output_schema)), pool_(pool) {}
+
+  Status Consume(ExecBatch batch) override {
+    std::lock_guard<std::mutex> guard(consum_mutex);
+    ARROW_ASSIGN_OR_RAISE(auto rb, batch.ToRecordBatch(output_schema_, pool_));
+    batch_vector.push_back(rb);
+    return Status::OK();
+  }
+
+  Future<> Finish() override {
+    ARROW_ASSIGN_OR_RAISE(*out_, Table::FromRecordBatches(batch_vector));
+    return Status::OK();
+  }
+
+ private:
+  std::shared_ptr<Table>* out_;
+  std::shared_ptr<Schema> output_schema_;
+  MemoryPool* pool_;
+  std::vector<std::shared_ptr<RecordBatch>> batch_vector;

Review comment:
       nit, but just `batches_`? Generally for private variables we suffix them with `_`




-- 
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