GH-33880: [C++] Improve I/O tracing

cpp/src/arrow/acero/filter_node.cc

@@ -73,12 +73,12 @@ class FilterNode : public MapNode {
   Result<ExecBatch> ProcessBatch(ExecBatch batch) override {
     ARROW_ASSIGN_OR_RAISE(Expression simplified_filter,
                           SimplifyWithGuarantee(filter_, batch.guarantee));
-
     arrow::util::tracing::Span span;
     START_COMPUTE_SPAN(span, "Filter",
                        {{"filter.expression", ToStringExtra()},
                         {"filter.expression.simplified", simplified_filter.ToString()},
-                        {"filter.length", batch.length}});
+                        {"filter.length", batch.length},
+                        {"input_batch.size_bytes", batch.TotalBufferSize()}});
 
     ARROW_ASSIGN_OR_RAISE(
         Datum mask, ExecuteScalarExpression(simplified_filter, batch,
@@ -87,8 +87,10 @@ class FilterNode : public MapNode {
     if (mask.is_scalar()) {
       const auto& mask_scalar = mask.scalar_as<BooleanScalar>();
       if (mask_scalar.is_valid && mask_scalar.value) {
+        ATTRIBUTE_ON_CURRENT_SPAN("output_batch.size_bytes", batch.TotalBufferSize());
         return batch;
       }
+      ATTRIBUTE_ON_CURRENT_SPAN("output_batch.size_bytes", 0);
       return batch.Slice(0, 0);
     }
 
@@ -101,7 +103,10 @@ class FilterNode : public MapNode {
       if (value.is_scalar()) continue;
       ARROW_ASSIGN_OR_RAISE(value, Filter(value, mask, FilterOptions::Defaults()));
     }
-    return ExecBatch::Make(std::move(values));
+    auto filtered_batch = ExecBatch::Make(std::move(values));
+    ATTRIBUTE_ON_CURRENT_SPAN("output_batch.size_bytes",
+                              filtered_batch->TotalBufferSize());
+    return filtered_batch;
   }
 
  protected:

cpp/src/arrow/acero/project_node.cc

@@ -79,19 +79,22 @@ class ProjectNode : public MapNode {
 
   Result<ExecBatch> ProcessBatch(ExecBatch batch) override {
     std::vector<Datum> values{exprs_.size()};
+    arrow::util::tracing::Span span;
+    START_COMPUTE_SPAN(span, "Project",
+                       {{"project.length", batch.length},
+                        {"input_batch.size_bytes", batch.TotalBufferSize()}});
     for (size_t i = 0; i < exprs_.size(); ++i) {
-      arrow::util::tracing::Span span;
-      START_COMPUTE_SPAN(span, "Project",
-                         {{"project.type", exprs_[i].type()->ToString()},
-                          {"project.length", batch.length},
-                          {"project.expression", exprs_[i].ToString()}});
+      std::string project_name = "project[" + std::to_string(i) + "]";
+      ATTRIBUTE_ON_CURRENT_SPAN(project_name + ".type", exprs_[i].type()->ToString());
+      ATTRIBUTE_ON_CURRENT_SPAN(project_name + ".expression", exprs_[i].ToString());
       ARROW_ASSIGN_OR_RAISE(Expression simplified_expr,
                             SimplifyWithGuarantee(exprs_[i], batch.guarantee));
 
       ARROW_ASSIGN_OR_RAISE(
           values[i], ExecuteScalarExpression(simplified_expr, batch,
                                              plan()->query_context()->exec_context()));
     }
+    ATTRIBUTE_ON_CURRENT_SPAN("output_batch.size_bytes", batch.TotalBufferSize());
     return ExecBatch{std::move(values), batch.length};
   }
 

cpp/src/arrow/acero/sink_node.cc

@@ -539,6 +539,7 @@ struct OrderBySinkNode final : public SinkNode {
 
   Status Finish() override {
     arrow::util::tracing::Span span;
+    START_SPAN(span, std::string(kind_name()) + "::Finish");
     ARROW_RETURN_NOT_OK(DoFinish());
     return SinkNode::Finish();
   }

cpp/src/arrow/acero/source_node.cc

@@ -133,6 +133,8 @@ struct SourceNode : ExecNode, public TracedNode {
     plan_->query_context()->ScheduleTask(
         [this, morsel_length, use_legacy_batching, initial_batch_index, morsel,
          has_ordering = !ordering_.is_unordered()]() {
+          arrow::util::tracing::Span span;
+          START_SPAN(span, "SourceNode::ProcessMorsel");
           int64_t offset = 0;
           int batch_index = initial_batch_index;
           do {
@@ -163,6 +165,7 @@ struct SourceNode : ExecNode, public TracedNode {
 
   Status StartProducing() override {
     NoteStartProducing(ToStringExtra());
+
     {
       // If another exec node encountered an error during its StartProducing call
       // it might have already called StopProducing on all of its inputs (including this
@@ -184,6 +187,9 @@ struct SourceNode : ExecNode, public TracedNode {
     options.should_schedule = ShouldSchedule::IfDifferentExecutor;
     ARROW_ASSIGN_OR_RAISE(Future<> scan_task, plan_->query_context()->BeginExternalTask(
                                                   "SourceNode::DatasetScan"));
+    arrow::util::tracing::Span span;
+    START_SPAN(span, "SourceNode::DatasetScan");
+
     if (!scan_task.is_valid()) {
       // Plan has already been aborted, no need to start scanning
       return Status::OK();
@@ -195,9 +201,6 @@ struct SourceNode : ExecNode, public TracedNode {
       }
       lock.unlock();
 
-      arrow::util::tracing::Span fetch_batch_span;
-      auto fetch_batch_scope =
-          START_SCOPED_SPAN(fetch_batch_span, "SourceNode::ReadBatch");
       return generator_().Then(
           [this](
               const std::optional<ExecBatch>& morsel_or_end) -> Future<ControlFlow<int>> {

cpp/src/arrow/compute/function.cc

@@ -219,7 +219,7 @@ struct FunctionExecutorImpl : public FunctionExecutor {
   }
 
   Result<Datum> Execute(const std::vector<Datum>& args, int64_t passed_length) override {
-    util::tracing::Span span;
+    arrow::util::tracing::Span span;
 
     auto func_kind = func.kind();
     const auto& func_name = func.name();

cpp/src/arrow/csv/reader.cc

@@ -43,12 +43,14 @@
 #include "arrow/type.h"
 #include "arrow/type_fwd.h"
 #include "arrow/util/async_generator.h"
+#include "arrow/util/byte_size.h"
 #include "arrow/util/future.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/task_group.h"
 #include "arrow/util/thread_pool.h"
+#include "arrow/util/tracing_internal.h"
 #include "arrow/util/utf8_internal.h"
 #include "arrow/util/vector.h"
 
@@ -401,6 +403,8 @@ class BlockParsingOperator {
         num_rows_seen_(first_row) {}
 
   Result<ParsedBlock> operator()(const CSVBlock& block) {
+    util::tracing::Span span;
+    START_SPAN(span, "arrow::csv::BlockParsingOperator");
     constexpr int32_t max_num_rows = std::numeric_limits<int32_t>::max();
     auto parser = std::make_shared<BlockParser>(
         io_context_.pool(), parse_options_, num_csv_cols_, num_rows_seen_, max_num_rows);
@@ -431,6 +435,7 @@ class BlockParsingOperator {
       num_rows_seen_ += parser->total_num_rows();
     }
     RETURN_NOT_OK(block.consume_bytes(parsed_size));
+    ATTRIBUTE_ON_CURRENT_SPAN("parsed_size", parsed_size);
     return ParsedBlock{std::move(parser), block.block_index,
                        static_cast<int64_t>(parsed_size) + block.bytes_skipped};
   }
@@ -448,6 +453,8 @@ class BlockParsingOperator {
 class BlockDecodingOperator {
  public:
   Future<DecodedBlock> operator()(const ParsedBlock& block) {
+    util::tracing::Span span;
+    START_SPAN(span, "arrow::csv::BlockDecodingOperator");
     DCHECK(!state_->column_decoders.empty());
     std::vector<Future<std::shared_ptr<Array>>> decoded_array_futs;
     for (auto& decoder : state_->column_decoders) {
@@ -456,6 +463,23 @@ class BlockDecodingOperator {
     auto bytes_parsed_or_skipped = block.bytes_parsed_or_skipped;
     auto decoded_arrays_fut = All(std::move(decoded_array_futs));
     auto state = state_;
+#ifdef ARROW_WITH_OPENTELEMETRY
+    opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> raw_span =
+        ::arrow::internal::tracing::UnwrapSpan(span.details.get());
+    return decoded_arrays_fut.Then(
+        [state, bytes_parsed_or_skipped, raw_span](
+            const std::vector<Result<std::shared_ptr<Array>>>& maybe_decoded_arrays)
+            -> Result<DecodedBlock> {
+          ARROW_ASSIGN_OR_RAISE(auto decoded_arrays,
+                                arrow::internal::UnwrapOrRaise(maybe_decoded_arrays));
+
+          ARROW_ASSIGN_OR_RAISE(auto batch,
+                                state->DecodedArraysToBatch(std::move(decoded_arrays)));
+          raw_span->SetAttribute("arrow.csv.output_batch_size_bytes",
+                                 util::TotalBufferSize(*batch));
+          return DecodedBlock{std::move(batch), bytes_parsed_or_skipped};
+        });
+#else
     return decoded_arrays_fut.Then(
         [state, bytes_parsed_or_skipped](
             const std::vector<Result<std::shared_ptr<Array>>>& maybe_decoded_arrays)
@@ -467,6 +491,7 @@ class BlockDecodingOperator {
                                 state->DecodedArraysToBatch(std::move(decoded_arrays)));
           return DecodedBlock{std::move(batch), bytes_parsed_or_skipped};
         });
+#endif
   }
 
   static Result<BlockDecodingOperator> Make(io::IOContext io_context,
@@ -881,7 +906,27 @@ class StreamingReaderImpl : public ReaderMixin,
   }
 
   Future<std::shared_ptr<RecordBatch>> ReadNextAsync() override {
-    return record_batch_gen_();
+    util::tracing::Span span;
+    START_SPAN(span, "arrow::csv::ReadNextAsync");
+    auto future = record_batch_gen_();
+#ifdef ARROW_WITH_OPENTELEMETRY
+    auto longer_living_span = std::make_unique<util::tracing::Span>(std::move(span));
+    future.AddCallback(
+        [span = std::move(longer_living_span)](
+            const arrow::Result<std::shared_ptr<arrow::RecordBatch>>& result) {
+          opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> raw_span =
+              ::arrow::internal::tracing::UnwrapSpan(span->details.get());
+          if (result.ok()) {
+            auto result_batch = result.ValueOrDie();
+            if (result_batch) {
+              raw_span->SetAttribute("batch.size_bytes",
+                                     ::arrow::util::TotalBufferSize(*result_batch));
+            }
+          }
+          END_SPAN((*span));
+        });
+#endif
+    return future;
   }
 
  protected:
@@ -892,6 +937,14 @@ class StreamingReaderImpl : public ReaderMixin,
       return Status::Invalid("Empty CSV file");
     }
 
+    util::tracing::Span init_span;
+    START_SPAN(init_span, "arrow::csv::InitAfterFirstBuffer");
+
+    // Create a arrow::csv::ReadNextAsync span so that grouping by that name does not
+    // ignore the work performed for this first block.
+    util::tracing::Span read_span;
+    auto scope = START_SCOPED_SPAN(read_span, "arrow::csv::ReadNextAsync");
+
     std::shared_ptr<Buffer> after_header;
     ARROW_ASSIGN_OR_RAISE(auto header_bytes_consumed,
                           ProcessHeader(first_buffer, &after_header));
@@ -911,9 +964,25 @@ class StreamingReaderImpl : public ReaderMixin,
     auto rb_gen = MakeMappedGenerator(std::move(parsed_block_gen), std::move(decoder_op));
 
     auto self = shared_from_this();
-    return rb_gen().Then([self, rb_gen, max_readahead](const DecodedBlock& first_block) {
-      return self->InitFromBlock(first_block, std::move(rb_gen), max_readahead, 0);
+    auto init_finished = rb_gen().Then([self, rb_gen, max_readahead
+#ifdef ARROW_WITH_OPENTELEMETRY
+                                        ,
+                                        init_span = std::move(init_span),
+                                        read_span = std::move(read_span)
+#endif
+    ](const DecodedBlock& first_block) {
+      auto fut = self->InitFromBlock(first_block, std::move(rb_gen), max_readahead, 0);
+#ifdef ARROW_WITH_OPENTELEMETRY
+      opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> raw_span =
+          ::arrow::internal::tracing::UnwrapSpan(read_span.details.get());
+      raw_span->SetAttribute("batch.size_bytes",
+                             util::TotalBufferSize(*first_block.record_batch));
+#endif
+      END_SPAN(read_span);
+      END_SPAN(init_span);
+      return fut;
     });
+    return init_finished;
   }
 
   Future<> InitFromBlock(const DecodedBlock& block,

cpp/src/arrow/dataset/dataset_internal.h

@@ -29,8 +29,10 @@
 #include "arrow/scalar.h"
 #include "arrow/type.h"
 #include "arrow/util/async_generator.h"
+#include "arrow/util/byte_size.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/iterator.h"
+#include "arrow/util/tracing_internal.h"
 
 namespace arrow {
 namespace dataset {
@@ -144,6 +146,12 @@ inline RecordBatchGenerator MakeChunkedBatchGenerator(RecordBatchGenerator gen,
       [batch_size](const std::shared_ptr<RecordBatch>& batch)
           -> ::arrow::AsyncGenerator<std::shared_ptr<::arrow::RecordBatch>> {
         const int64_t rows = batch->num_rows();
+        util::tracing::Span span;
+        START_SPAN(span, "MakeChunkedBatchGenerator",
+                   {{"target_batch_size_rows", batch_size},
+                    {"batch.size_rows", rows},
+                    {"batch.size_bytes", util::TotalBufferSize(*batch)},
+                    {"output_batches", rows / batch_size + (rows % batch_size != 0)}});
         if (rows <= batch_size) {
           return ::arrow::MakeVectorGenerator<std::shared_ptr<RecordBatch>>({batch});
         }