arrow_c_data_util_internal.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#pragma once
 
#include <cerrno>
#include <concepts>
#include <cstdint>
#include <cstring>
#include <memory>
#include <optional>
#include <span>
#include <string>
#include <utility>
#include <vector>
 
#include "iceberg/arrow_c_data.h"
#include "iceberg/iceberg_export.h"
#include "iceberg/result.h"
#include "iceberg/type_fwd.h"
 
namespace iceberg {
 
class ICEBERG_EXPORT ProjectionContext {
 public:
  using ProjectBatchState = std::shared_ptr<void>;
 
  using ProjectBatchFunction = auto (*)(ArrowArray* input_batch,
                                        std::span<const int32_t> row_indices,
                                        ProjectionContext& projection)
      -> Result<ArrowArray>;
 
  static void RegisterProjectBatchFunction(ProjectBatchFunction project_batch_function);
 
  static bool HasProjectBatchFunction();
 
  static auto ResolveProjectBatchFunction() -> ProjectBatchFunction;
 
  static Result<ProjectionContext> Make(const Schema& input_schema,
                                        const Schema& output_schema,
                                        ProjectBatchFunction project_batch_function);
 
  ProjectionContext(ProjectionContext&&) noexcept;
  ProjectionContext& operator=(ProjectionContext&&) noexcept;
  ~ProjectionContext();
 
  ProjectionContext(const ProjectionContext&) = delete;
  ProjectionContext& operator=(const ProjectionContext&) = delete;
 
  const Schema& input_schema() const;
 
  const Schema& output_schema() const;
 
  const ArrowSchema& input_arrow_schema() const;
 
  const ArrowSchema& output_arrow_schema() const;
 
  std::span<const int32_t> selected_field_indices() const;
 
  ProjectBatchFunction project_batch_function() const;
 
  ProjectBatchState& project_batch_state();
 
 private:
  ProjectionContext() = default;
 
  // Raw schema pointers are borrowed from caller-owned schemas. FileScanTaskReader
  // keeps those schema objects alive in the same stream source that owns this context.
  const Schema* input_schema_ = nullptr;
  const Schema* output_schema_ = nullptr;
  std::vector<int32_t> selected_field_indices_;
  ArrowSchema input_arrow_schema_{};
  ArrowSchema output_arrow_schema_{};
  ProjectBatchFunction project_batch_function_ = nullptr;
  ProjectBatchState project_batch_state_;
};
 
template <typename Source>
concept ArrowArrayStreamProvider = requires(Source& source) {
  { source.Close() } -> std::same_as<Status>;
  { source.Next() } -> std::same_as<Result<std::optional<ArrowArray>>>;
  { source.Schema() } -> std::same_as<Result<ArrowSchema>>;
};
 
namespace detail {
 
template <ArrowArrayStreamProvider Source>
struct ArrowArrayStreamPrivateData {
  std::unique_ptr<Source> source;
  std::string last_error;
 
  explicit ArrowArrayStreamPrivateData(std::unique_ptr<Source> src)
      : source(std::move(src)) {}
 
  ~ArrowArrayStreamPrivateData() {
    if (source != nullptr) {
      std::ignore = source->Close();
    }
  }
};
 
template <ArrowArrayStreamProvider Source>
int GetSchema(struct ArrowArrayStream* stream, struct ArrowSchema* out) {
  if (stream == nullptr || stream->private_data == nullptr) {
    return EINVAL;
  }
 
  auto* private_data =
      static_cast<ArrowArrayStreamPrivateData<Source>*>(stream->private_data);
  auto schema_result = private_data->source->Schema();
  if (!schema_result.has_value()) {
    private_data->last_error = schema_result.error().message;
    std::memset(out, 0, sizeof(ArrowSchema));
    return EIO;
  }
 
  *out = std::move(schema_result.value());
  return 0;
}
 
template <ArrowArrayStreamProvider Source>
int GetNext(struct ArrowArrayStream* stream, struct ArrowArray* out) {
  if (stream == nullptr || stream->private_data == nullptr) {
    return EINVAL;
  }
 
  auto* private_data =
      static_cast<ArrowArrayStreamPrivateData<Source>*>(stream->private_data);
  auto next_result = private_data->source->Next();
  if (!next_result.has_value()) {
    private_data->last_error = next_result.error().message;
    std::memset(out, 0, sizeof(ArrowArray));
    return EIO;
  }
 
  auto& optional_array = next_result.value();
  if (optional_array.has_value()) {
    *out = std::move(optional_array.value());
  } else {
    std::memset(out, 0, sizeof(ArrowArray));
  }
 
  return 0;
}
 
template <ArrowArrayStreamProvider Source>
const char* GetLastError(struct ArrowArrayStream* stream) {
  if (stream == nullptr || stream->private_data == nullptr) {
    return nullptr;
  }
 
  auto* private_data =
      static_cast<ArrowArrayStreamPrivateData<Source>*>(stream->private_data);
  return private_data->last_error.empty() ? nullptr : private_data->last_error.c_str();
}
 
template <ArrowArrayStreamProvider Source>
void Release(struct ArrowArrayStream* stream) {
  if (stream == nullptr || stream->private_data == nullptr) {
    return;
  }
 
  delete static_cast<ArrowArrayStreamPrivateData<Source>*>(stream->private_data);
  stream->private_data = nullptr;
  stream->release = nullptr;
}
 
}  // namespace detail
 
template <ArrowArrayStreamProvider Source>
Result<ArrowArrayStream> MakeArrowArrayStream(std::unique_ptr<Source> source) {
  if (source == nullptr) {
    return InvalidArgument("Cannot make ArrowArrayStream from null source");
  }
 
  auto private_data =
      std::make_unique<detail::ArrowArrayStreamPrivateData<Source>>(std::move(source));
  ArrowArrayStream stream{.get_schema = detail::GetSchema<Source>,
                          .get_next = detail::GetNext<Source>,
                          .get_last_error = detail::GetLastError<Source>,
                          .release = detail::Release<Source>,
                          .private_data = private_data.release()};
  return stream;
}
 
ICEBERG_EXPORT Result<ArrowArrayStream> MakeArrowArrayStream(
    std::unique_ptr<Reader> reader);
 
ICEBERG_EXPORT Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
                                               std::span<const int32_t> row_indices,
                                               ProjectionContext& projection);
 
}  // namespace iceberg