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"
 
@@ -881,6 +883,8 @@ class StreamingReaderImpl : public ReaderMixin,
  }
 
  Future<std::shared_ptr<RecordBatch>> ReadNextAsync() override {
+ util::tracing::Span span;
+ START_SPAN(span, "arrow::csv::ReadNextAsync");
  return record_batch_gen_();
  }
 
@@ -892,6 +896,11 @@ class StreamingReaderImpl : public ReaderMixin,
  return Status::Invalid("Empty CSV file");
  }
 
+ // 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 +920,12 @@ 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) {
+ auto init_finished = rb_gen().Then([self, rb_gen, max_readahead,
+ read_span = std::move(read_span)
+ ](const DecodedBlock& first_block) {
  return self->InitFromBlock(first_block, std::move(rb_gen), max_readahead, 0);
  });
+ 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});
  }

cpp/src/arrow/dataset/dataset_writer.cc

@@ -27,6 +27,7 @@
 #include "arrow/record_batch.h"
 #include "arrow/result.h"
 #include "arrow/table.h"
+#include "arrow/util/byte_size.h"
 #include "arrow/util/future.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/map.h"
@@ -191,9 +192,13 @@ class DatasetWriterFileQueue {
  }
 
  Result<int64_t> PopAndDeliverStagedBatch() {
+ util::tracing::Span span;
+ START_SPAN(span, "DatasetWriter::Pop");
  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> next_batch, PopStagedBatch());
  int64_t rows_popped = next_batch->num_rows();
  rows_currently_staged_ -= next_batch->num_rows();
+ ATTRIBUTE_ON_CURRENT_SPAN("batch.size_rows", next_batch->num_rows());
+ ATTRIBUTE_ON_CURRENT_SPAN("rows_currently_staged", rows_currently_staged_);
  ScheduleBatch(std::move(next_batch));
  return rows_popped;
  }
@@ -202,7 +207,15 @@ class DatasetWriterFileQueue {
  Status Push(std::shared_ptr<RecordBatch> batch) {
  uint64_t delta_staged = batch->num_rows();
  rows_currently_staged_ += delta_staged;
- staged_batches_.push_back(std::move(batch));
+ {
+ util::tracing::Span span;
+ START_SPAN(span, "DatasetWriter::Push",
+ {{"batch.size_rows", batch->num_rows()},
+ {"rows_currently_staged", rows_currently_staged_},
+ {"options_.min_rows_per_group", options_.min_rows_per_group},
+ {"max_rows_staged", writer_state_->max_rows_staged}});
+ staged_batches_.push_back(std::move(batch));
+ }
  while (!staged_batches_.empty() &&
  (writer_state_->StagingFull() ||
  rows_currently_staged_ >= options_.min_rows_per_group)) {
@@ -233,6 +246,18 @@ class DatasetWriterFileQueue {
  return DeferNotOk(options_.filesystem->io_context().executor()->Submit(
  [self = this, batch = std::move(next)]() {
  int64_t rows_to_release = batch->num_rows();
+#ifdef ARROW_WITH_OPENTELEMETRY
+ uint64_t size_bytes = util::TotalBufferSize(*batch);
+ uint64_t num_buffers = 0;
+ for (auto column : batch->columns()) {
+ num_buffers += column->data()->buffers.size();
+ }
+ util::tracing::Span span;
+ START_SPAN(span, "DatasetWriter::WriteNext",
+ {{"threadpool", "IO"},
+ {"batch.size_bytes", size_bytes},
+ {"batch.num_buffers", num_buffers}});
+#endif
  Status status = self->writer_->Write(batch);
  self->writer_state_->rows_in_flight_throttle.Release(rows_to_release);
  return status;
@@ -261,11 +286,6 @@ class DatasetWriterFileQueue {
  util::AsyncTaskScheduler* file_tasks_ = nullptr;
 };
 
-struct WriteTask {
- std::string filename;
- uint64_t num_rows;
-};
-
 class DatasetWriterDirectoryQueue {
  public:
  DatasetWriterDirectoryQueue(util::AsyncTaskScheduler* scheduler, std::string directory,
@@ -301,7 +321,6 @@ class DatasetWriterDirectoryQueue {
 
  Status StartWrite(const std::shared_ptr<RecordBatch>& batch) {
  rows_written_ += batch->num_rows();
- WriteTask task{current_filename_, static_cast<uint64_t>(batch->num_rows())};
  if (!latest_open_file_) {
  ARROW_RETURN_NOT_OK(OpenFileQueue(current_filename_));
  }
@@ -351,6 +370,8 @@ class DatasetWriterDirectoryQueue {
  latest_open_file_tasks_ = util::MakeThrottledAsyncTaskGroup(
  scheduler_, 1, /*queue=*/nullptr, std::move(file_finish_task));
  if (init_future_.is_valid()) {
+ util::tracing::Span span;
+ START_SPAN(span, "arrow::dataset::WaitForDirectoryInit");
  latest_open_file_tasks_->AddSimpleTask(
  [init_future = init_future_]() { return init_future; },
  "DatasetWriter::WaitForDirectoryInit"sv);
@@ -362,6 +383,8 @@ class DatasetWriterDirectoryQueue {
  uint64_t rows_written() const { return rows_written_; }
 
  void PrepareDirectory() {
+ util::tracing::Span span;
+ START_SPAN(span, "arrow::dataset::SubmitPrepareDirectoryTask");
  if (directory_.empty() || !write_options_.create_dir) {
  return;
  }
@@ -383,6 +406,8 @@ class DatasetWriterDirectoryQueue {
  if (write_options_.existing_data_behavior ==
  ExistingDataBehavior::kDeleteMatchingPartitions) {
  init_task = [this, create_dir_cb, notify_waiters_cb, notify_waiters_on_err_cb] {
+ util::tracing::Span span;
+ START_SPAN(span, "arrow::dataset::PrepareDirectory");
  return write_options_.filesystem
  ->DeleteDirContentsAsync(directory_,
  /*missing_dir_ok=*/true)
@@ -614,12 +639,14 @@ class DatasetWriter::DatasetWriterImpl {
  backpressure =
  writer_state_.rows_in_flight_throttle.Acquire(next_chunk->num_rows());
  if (!backpressure.is_finished()) {
+ EVENT(scheduler_->span(), "DatasetWriter::Backpressure::TooManyRowsQueued");
  EVENT_ON_CURRENT_SPAN("DatasetWriter::Backpressure::TooManyRowsQueued");
  break;
  }
  if (will_open_file) {
  backpressure = writer_state_.open_files_throttle.Acquire(1);
  if (!backpressure.is_finished()) {
+ EVENT(scheduler_->span(), "DatasetWriter::Backpressure::TooManyOpenFiles");
  EVENT_ON_CURRENT_SPAN("DatasetWriter::Backpressure::TooManyOpenFiles");
  RETURN_NOT_OK(TryCloseLargestFile());
  break;

cpp/src/arrow/dataset/file_csv.cc

@@ -278,10 +278,6 @@ static inline Result<csv::ReadOptions> GetReadOptions(
 static inline Future<std::shared_ptr<csv::StreamingReader>> OpenReaderAsync(
  const FileSource& source, const CsvFileFormat& format,
  const std::shared_ptr<ScanOptions>& scan_options, Executor* cpu_executor) {
-#ifdef ARROW_WITH_OPENTELEMETRY
- auto tracer = arrow::internal::tracing::GetTracer();
- auto span = tracer->StartSpan("arrow::dataset::CsvFileFormat::OpenReaderAsync");
-#endif
  ARROW_ASSIGN_OR_RAISE(
  auto fragment_scan_options,
  GetFragmentScanOptions<CsvFragmentScanOptions>(
@@ -300,31 +296,24 @@ static inline Future<std::shared_ptr<csv::StreamingReader>> OpenReaderAsync(
  // input->Peek call blocks so we run the whole thing on the I/O thread pool.
  auto reader_fut = DeferNotOk(input->io_context().executor()->Submit(
  [=]() -> Future<std::shared_ptr<csv::StreamingReader>> {
+
  ARROW_ASSIGN_OR_RAISE(auto first_block, input->Peek(reader_options.block_size));
  const auto& parse_options = format.parse_options;
  ARROW_ASSIGN_OR_RAISE(
  auto convert_options,
  GetConvertOptions(format, scan_options ? scan_options.get() : nullptr,
  first_block));
- return csv::StreamingReader::MakeAsync(io::default_io_context(), std::move(input),
- cpu_executor, reader_options,
-parse_options, convert_options);
+ return csv::StreamingReader::MakeAsync(
+io::default_io_context(), std::move(input), cpu_executor, reader_options,
+ parse_options, convert_options);
  }));
  return reader_fut.Then(
  // Adds the filename to the error
  [=](const std::shared_ptr<csv::StreamingReader>& reader)
  -> Result<std::shared_ptr<csv::StreamingReader>> {
-#ifdef ARROW_WITH_OPENTELEMETRY
- span->SetStatus(opentelemetry::trace::StatusCode::kOk);
- span->End();
-#endif
  return reader;
  },
  [=](const Status& err) -> Result<std::shared_ptr<csv::StreamingReader>> {
-#ifdef ARROW_WITH_OPENTELEMETRY
- arrow::internal::tracing::MarkSpan(err, span.get());
- span->End();
-#endif
  return err.WithMessage("Could not open CSV input source '", path, "': ", err);
  });
 }
@@ -384,8 +373,6 @@ Result<RecordBatchGenerator> CsvFileFormat::ScanBatchesAsync(
  auto reader_fut =
  OpenReaderAsync(source, *this, scan_options, ::arrow::internal::GetCpuThreadPool());
  auto generator = GeneratorFromReader(std::move(reader_fut), scan_options->batch_size);
- WRAP_ASYNC_GENERATOR_WITH_CHILD_SPAN(
- generator, "arrow::dataset::CsvFileFormat::ScanBatchesAsync::Next");
  return generator;
 }
 

cpp/src/arrow/ipc/reader.cc

@@ -57,6 +57,7 @@
 #include "arrow/util/parallel.h"
 #include "arrow/util/string.h"
 #include "arrow/util/thread_pool.h"
+#include "arrow/util/tracing_internal.h"
 #include "arrow/util/ubsan.h"
 #include "arrow/util/vector.h"
 #include "arrow/visit_type_inline.h"
@@ -523,6 +524,12 @@ Status DecompressBuffers(Compression::type compression, const IpcReadOptions& op
 
  return ::arrow::internal::OptionalParallelFor(
  options.use_threads, static_cast<int>(buffers.size()), [&](int i) {
+ util::tracing::Span span;
+ START_SPAN(span, "arrow::ipc::DecompressBuffer",
+ {{"buffer_index", i},
+ {"ipc.compression.codec", codec.get()->name().c_str()},
+ {"ipc.options.use_threads", options.use_threads},
+ {"size.uncompressed", (*buffers[i])->size() - sizeof(int64_t)}});
  ARROW_ASSIGN_OR_RAISE(*buffers[i],
  DecompressBuffer(*buffers[i], options, codec.get()));
  return Status::OK();