Function MakeCSFSparseTensorMessage

cpp/src/arrow/ipc/tensor_test.cc:513–555

Build a SparseTensor IPC message carrying a CSF index with caller-controlled buffer counts. This lets us exercise the reader's validation directly, since such inconsistent indices can only be produced by hand-crafted flatbuffers.

// buffer counts. This lets us exercise the reader's validation directly, since
// such inconsistent indices can only be produced by hand-crafted flatbuffers.
Result<std::shared_ptr<Message>> MakeCSFSparseTensorMessage(
    const std::vector<int64_t>& shape, int num_indptr_buffers, int num_indices_buffers,
    int axis_order_size, int64_t non_zero_length = 0) {
  flatbuffers::FlatBufferBuilder fbb;

  auto value_type = flatbuf::CreateInt(fbb, 64, /*is_signed=*/true);

  std::vector<flatbuffers::Offset<flatbuf::TensorDim>> dims;
  for (int64_t dim_size : shape) {
    dims.push_back(flatbuf::CreateTensorDim(fbb, dim_size, /*name=*/0));
  }
  auto fb_shape = fbb.CreateVector(dims);

  auto indptr_type = flatbuf::CreateInt(fbb, 64, /*is_signed=*/false);
  auto indices_type = flatbuf::CreateInt(fbb, 64, /*is_signed=*/false);

  std::vector<flatbuf::Buffer> indptr(num_indptr_buffers, flatbuf::Buffer(0, 0));
  std::vector<flatbuf::Buffer> indices(num_indices_buffers, flatbuf::Buffer(0, 0));
  auto fb_indptr = fbb.CreateVectorOfStructs(indptr);
  auto fb_indices = fbb.CreateVectorOfStructs(indices);

  std::vector<int32_t> axis_order(axis_order_size, 0);
  auto fb_axis_order = fbb.CreateVector(axis_order);

  auto csf = flatbuf::CreateSparseTensorIndexCSF(fbb, indptr_type, fb_indptr,
                                                 indices_type, fb_indices, fb_axis_order);

  flatbuf::Buffer data(0, 0);
  auto sparse_tensor = flatbuf::CreateSparseTensor(
      fbb, flatbuf::Type_Int, value_type.Union(), fb_shape, non_zero_length,
      flatbuf::SparseTensorIndex::SparseTensorIndex_SparseTensorIndexCSF, csf.Union(),
      &data);

  fbb.Finish(flatbuf::CreateMessage(fbb, internal::kCurrentMetadataVersion,
                                    flatbuf::MessageHeader::MessageHeader_SparseTensor,
                                    sparse_tensor.Union(),
                                    /*bodyLength=*/0));

  ARROW_ASSIGN_OR_RAISE(auto metadata, internal::WriteFlatbufferBuilder(fbb));
  auto body = Buffer::FromString(std::string(8, '\0'));
  ARROW_ASSIGN_OR_RAISE(auto message, Message::Open(metadata, body));
  return std::shared_ptr<Message>(std::move(message));
}

IpcPayload MakeSparseTensorPayload(const std::shared_ptr<Message>& message,
                                   int num_body_buffers) {