GH-15151: [C++] ]Adding RecordBatchReaderSource to solve an issue in R API #15183

vibhatha · 2023-01-04T08:21:16Z

This PR includes the factory record_batch_reader_source for the Acero. This is a source node which takes in a RecordBatchReader as the data source along an executor which gives the freedom to choose the threadpool required for the execution. Also an example shows how this can be used in Acero.

Self-review

Closes: Should GetIOThreadPool() be accessible from installed headers? #15151

github-actions · 2023-01-04T08:21:43Z

Closes: Should GetIOThreadPool() be accessible from installed headers? #15151

github-actions · 2023-01-04T08:21:44Z

⚠️ GitHub issue #15151 has no components, please add labels for components.

github-actions · 2023-01-04T12:30:02Z

⚠️ GitHub issue #15151 has no components, please add labels for components.

vibhatha · 2023-01-04T12:32:08Z

cc @westonpace @paleolimbot could you please take a look.

paleolimbot

Awesome! I think this will fix my problem, but could you also update the R package code? That will tell you in a hurry whether or not this works. I tried:

// [[arrow::export]]
std::shared_ptr<compute::ExecNode> ExecNode_SourceNode(
    const std::shared_ptr<compute::ExecPlan>& plan,
    const std::shared_ptr<arrow::RecordBatchReader>& reader) {
  arrow::compute::RecordBatchReaderSourceNodeOptions options{reader->schema(), reader};
  return MakeExecNodeOrStop("record_batch_reader_source", plan.get(), {}, options);
}

Replacing

arrow/r/src/compute-exec.cpp

Lines 459 to 468 in 63b91cc

    
           std::shared_ptr<compute::ExecNode> ExecNode_SourceNode( 
        
               const std::shared_ptr<compute::ExecPlan>& plan, 
        
               const std::shared_ptr<arrow::RecordBatchReader>& reader) { 
        
             arrow::compute::SourceNodeOptions options{ 
        
                 /*output_schema=*/reader->schema(), 
        
                 /*generator=*/ValueOrStop( 
        
                     compute::MakeReaderGenerator(reader, arrow::io::internal::GetIOThreadPool()))}; 
        
             return MakeExecNodeOrStop("source", plan.get(), {}, options); 
        
           }

...and that seemed to work. Maybe also remove

arrow/r/src/compute-exec.cpp

Lines 33 to 40 in 63b91cc

    
           // GH-15151: Best path forward to make this available without a hack like this one 
        
           namespace arrow { 
        
           namespace io { 
        
           namespace internal { 
        
           arrow::internal::ThreadPool* GetIOThreadPool(); 
        
           } 
        
           }  // namespace io 
        
           }  // namespace arrow

while you're there?

westonpace

Thanks for doing this. I have a few suggestions.

cpp/src/arrow/compute/exec/options.h

westonpace · 2023-01-04T16:14:33Z

cpp/src/arrow/compute/exec/plan_test.cc

+void TestRecordBatchReaderSourceSink(
+    std::function<Result<std::shared_ptr<RecordBatchReader>>(const BatchesWithSchema&)>
+        to_reader) {
+  ASSERT_OK_AND_ASSIGN(auto executor, arrow::internal::ThreadPool::Make(1));
+  ExecContext exec_context(default_memory_pool(), executor.get());
+  ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(exec_context));
+  AsyncGenerator<std::optional<ExecBatch>> sink_gen;
+
+  auto exp_batches = MakeBasicBatches();
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<RecordBatchReader> reader, to_reader(exp_batches));
+
+  ASSERT_OK(Declaration::Sequence(
+                {
+                    {"record_batch_reader_source",
+                     RecordBatchReaderSourceNodeOptions{exp_batches.schema, reader}},
+                    {"sink", SinkNodeOptions{&sink_gen}},
+                })
+                .AddToPlan(plan.get()));
+
+  ASSERT_THAT(StartAndCollect(plan.get(), sink_gen),
+              Finishes(ResultWith(UnorderedElementsAreArray(exp_batches.batches))));
+}


You are specifying an executor but you aren't verifying that it is being used. Also, there is no test where we do not specify an executor.

Testing that the executor gets used would be a little tricky (maybe someday we can add a TotalTasksRun method to the executor). You would probably need to wrap the batch generator with something that verifies it is on the I/O executor (e.g. using OwnsThisThread or something similar).

I am probably being a bit harsh here as I suspect we test some of the other sources in a similar fashion. If we're only going to have one simple test I would prefer a test using the default instead of one specifying a custom executor as that is the more common case.

@westonpace I am working on a wrapper to address this. Let me evaluate options and update the PR. Also, I wouldn't say this is harsh but thorough as I see other sources are not properly tested like this. I think it would be better to test them too. I will create a ticket to cover this. Thanks for thinking in this angel.

By the way TotalTasksRun sounds fun, could you elaborate a bit what you expect here, I am happy to work on that task.

Probably external to this PR. Appreciate your thoughts.

If we're only going to have one simple test I would prefer a test using the default instead of one specifying a custom executor as that is the more common case.

Are you referring to IOThreadPool being used as the default instead of NULLPTR?

Are you referring to IOThreadPool being used as the default instead of NULLPTR?

What you had before might actually have been correct (I thought you were using executor in the options but on closer examination it was only being used in the exec_ctx). I was saying the most common case is RecordBatchReaderSourceNodeOptions options(reader) where the executor is not specified and is left to be the default. We should have our test verify that this case works correctly. At the very least we should make sure we can run such a plan without crashing. Long term we should check, after we ran the plan, to ensure that there were tasks submitted to the I/O thread pool.

This kind of test should not need to create a thread pool at all. We should only need to create a thread pool when testing a custom executor in the node.

I updated the code, is it okay?

cpp/src/arrow/compute/exec/test_util.h

cpp/src/arrow/compute/exec/test_util.cc

cpp/src/arrow/compute/exec/source_node.cc

westonpace · 2023-01-04T16:22:31Z

cpp/src/arrow/compute/exec/plan_test.cc

+  ASSERT_THAT(StartAndCollect(plan.get(), sink_gen),
+              Finishes(ResultWith(UnorderedElementsAreArray(exp_batches.batches))));


Let's use DeclarationToExecBatches for new tests (feel free to fix existing tests too).

Would you mind, if I do this in another PR for the rest of the items?

I don't mind if you fix the other tests in a separate PR. However, any newly added tests should follow the new format. TEST(ExecPlanExecution, SourceGroupedSum) is a good example from this file that you can follow. This test, for example, would become something like:

auto exp_batches = MakeBasicBatches(); ASSERT_OK_AND_ASSIGN(std::shared_ptr<RecordBatchReader> reader, to_reader(exp_batches)); Declaration plan("record_batch_reader_source", RecordBatchReaderSourceNodeOptions{exp_batches.schema, reader}); ASSERT_OK_AND_ASSIGN(auto out_batches, DeclarationToExecBatches(std::move(plan))); AssertExecBatchesEqualIgnoringOrder(exp_batches.schema, exp_batches.batches, out_batches);

The follow up PR is WIP here: #15288

paleolimbot

From the R side of things...looks great! Thank you!

westonpace

A few more suggestions

cpp/examples/arrow/execution_plan_documentation_examples.cc

cpp/src/arrow/compute/exec/plan_test.cc

cpp/src/arrow/compute/exec/source_node.cc

westonpace

A final few cleanup nits and then we can merge.

cpp/examples/arrow/execution_plan_documentation_examples.cc

cpp/src/arrow/compute/exec/plan_test.cc

cpp/src/arrow/compute/exec/test_util.cc

vibhatha · 2023-01-12T01:07:50Z

@westonpace btw, I just searched for places where I could have missed move and found out the following in other sections of the code base. I haven't included it in here yet, but here is the diff (I tested locally and works fine).

diff --git a/cpp/src/arrow/compute/exec/test_util.cc b/cpp/src/arrow/compute/exec/test_util.cc
index 038171c84..912428acf 100644
--- a/cpp/src/arrow/compute/exec/test_util.cc
+++ b/cpp/src/arrow/compute/exec/test_util.cc
@@ -273,8 +273,8 @@ Result<std::vector<std::shared_ptr<ExecBatch>>> ToExecBatches(
     const BatchesWithSchema& batches_with_schema) {
   std::vector<std::shared_ptr<ExecBatch>> exec_batches;
   for (auto batch : batches_with_schema.batches) {
-    auto exec_batch = std::make_shared<ExecBatch>(batch);
-    exec_batches.push_back(exec_batch);
+    exec_batches.push_back(std::make_shared<ExecBatch>(batch));
   }
   return exec_batches;
 }
@@ -285,7 +285,7 @@ Result<std::vector<std::shared_ptr<RecordBatch>>> ToRecordBatches(
   for (auto batch : batches_with_schema.batches) {
     ARROW_ASSIGN_OR_RAISE(auto record_batch,
                           batch.ToRecordBatch(batches_with_schema.schema));
-    record_batches.push_back(record_batch);
+    record_batches.push_back(std::move(record_batch));
   }
   return record_batches;
 }

Is this acceptable?

westonpace · 2023-01-12T01:24:28Z

Is this acceptable?

Yes, both of those changes look correct. Do you want to apply them real quick?

vibhatha · 2023-01-12T01:25:02Z

Is this acceptable?

Yes, both of those changes look correct. Do you want to apply them real quick?

Yes I will. One minute.

vibhatha · 2023-01-12T01:34:11Z

@westonpace Done. Ready to review.

westonpace

Thanks for looking into this. I'll merge when CI passes.

vibhatha · 2023-01-12T02:27:24Z

@westonpace CIs passed.

github-actions bot added the Component: C++ label Jan 4, 2023

vibhatha marked this pull request as ready for review January 4, 2023 12:29

paleolimbot reviewed Jan 4, 2023

View reviewed changes

westonpace requested changes Jan 4, 2023

View reviewed changes

github-actions bot added Component: C# Component: Documentation Component: FlightRPC Component: GLib Component: Go Component: Java Component: Parquet Component: Python Component: R Component: Ruby labels Jan 9, 2023

vibhatha force-pushed the gh-15151 branch 2 times, most recently from c000ed1 to 9792e7c Compare January 9, 2023 15:48

github-actions bot removed Component: R Component: Go Component: Ruby Component: Python Component: FlightRPC Component: GLib Component: Documentation Component: Java Component: C# Component: Parquet labels Jan 9, 2023

fix(reviews)

a638e41

vibhatha force-pushed the gh-15151 branch from 9792e7c to a638e41 Compare January 9, 2023 15:48

fix(r): adding record_batch_reader_source for R API

f105511

github-actions bot added the Component: R label Jan 9, 2023

fix(reviews): reformat test case to use test abstraction

dae7076

vibhatha requested review from westonpace and paleolimbot and removed request for westonpace and paleolimbot January 9, 2023 16:13

paleolimbot approved these changes Jan 9, 2023

View reviewed changes

westonpace requested changes Jan 10, 2023

View reviewed changes

vibhatha added 2 commits January 11, 2023 08:09

fix(reviews)

10610eb

fix(reviews)

746d7c2

vibhatha requested a review from westonpace January 11, 2023 04:06

westonpace requested changes Jan 11, 2023

View reviewed changes

fix(reviews): minor changes

ff495b6

vibhatha requested a review from westonpace January 12, 2023 01:11

fix(move)

76b3674

westonpace approved these changes Jan 12, 2023

View reviewed changes

westonpace merged commit fc53ff8 into apache:master Jan 12, 2023

vibhatha deleted the gh-15151 branch January 14, 2023 02:24

raulcd mentioned this pull request Jan 16, 2023

[C++][CI] arrow-compute-plan-test timeouts on nightly job for ubuntu 18.04 #33695

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-15151: [C++] ]Adding RecordBatchReaderSource to solve an issue in R API #15183

GH-15151: [C++] ]Adding RecordBatchReaderSource to solve an issue in R API #15183

vibhatha commented Jan 4, 2023 •

edited

Loading

github-actions bot commented Jan 4, 2023

github-actions bot commented Jan 4, 2023

github-actions bot commented Jan 4, 2023

vibhatha commented Jan 4, 2023

paleolimbot left a comment

westonpace left a comment

westonpace Jan 4, 2023

vibhatha Jan 5, 2023

vibhatha Jan 5, 2023

vibhatha Jan 5, 2023

vibhatha Jan 5, 2023

westonpace Jan 6, 2023

vibhatha Jan 9, 2023

westonpace Jan 4, 2023

vibhatha Jan 5, 2023

westonpace Jan 6, 2023

vibhatha Jan 12, 2023

paleolimbot left a comment

westonpace left a comment

westonpace left a comment

vibhatha commented Jan 12, 2023 •

edited

Loading

westonpace commented Jan 12, 2023

vibhatha commented Jan 12, 2023

vibhatha commented Jan 12, 2023

westonpace left a comment

vibhatha commented Jan 12, 2023

	std::shared_ptr<compute::ExecNode> ExecNode_SourceNode(
	const std::shared_ptr<compute::ExecPlan>& plan,
	const std::shared_ptr<arrow::RecordBatchReader>& reader) {
	arrow::compute::SourceNodeOptions options{
	/output_schema=/reader->schema(),
	/generator=/ValueOrStop(
	compute::MakeReaderGenerator(reader, arrow::io::internal::GetIOThreadPool()))};

	return MakeExecNodeOrStop("source", plan.get(), {}, options);
	}

	// GH-15151: Best path forward to make this available without a hack like this one
	namespace arrow {
	namespace io {
	namespace internal {
	arrow::internal::ThreadPool* GetIOThreadPool();
	}
	} // namespace io
	} // namespace arrow

		ASSERT_THAT(StartAndCollect(plan.get(), sink_gen),
		Finishes(ResultWith(UnorderedElementsAreArray(exp_batches.batches))));

GH-15151: [C++] ]Adding RecordBatchReaderSource to solve an issue in R API #15183

GH-15151: [C++] ]Adding RecordBatchReaderSource to solve an issue in R API #15183

Conversation

vibhatha commented Jan 4, 2023 • edited Loading

github-actions bot commented Jan 4, 2023

github-actions bot commented Jan 4, 2023

github-actions bot commented Jan 4, 2023

vibhatha commented Jan 4, 2023

paleolimbot left a comment

Choose a reason for hiding this comment

westonpace left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

paleolimbot left a comment

Choose a reason for hiding this comment

westonpace left a comment

Choose a reason for hiding this comment

westonpace left a comment

Choose a reason for hiding this comment

vibhatha commented Jan 12, 2023 • edited Loading

westonpace commented Jan 12, 2023

vibhatha commented Jan 12, 2023

vibhatha commented Jan 12, 2023

westonpace left a comment

Choose a reason for hiding this comment

vibhatha commented Jan 12, 2023

vibhatha commented Jan 4, 2023 •

edited

Loading

vibhatha commented Jan 12, 2023 •

edited

Loading