src/iceberg/arrow_c_data_util.cc

/*
 * 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.
 */

#include <cstdint>
#include <limits>
#include <memory>
#include <mutex>
#include <span>
#include <utility>
#include <vector>

#include <nanoarrow/nanoarrow.h>

#include "iceberg/arrow_c_data_guard_internal.h"
#include "iceberg/arrow_c_data_util_internal.h"
#include "iceberg/file_reader.h"
#include "iceberg/nanoarrow_status_internal.h"
#include "iceberg/result.h"
#include "iceberg/schema.h"
#include "iceberg/schema_field.h"
#include "iceberg/schema_internal.h"
#include "iceberg/util/macros.h"

namespace iceberg {

Result<ArrowArrayStream> MakeArrowArrayStream(std::unique_ptr<Reader> reader) {
  return MakeArrowArrayStream<Reader>(std::move(reader));
}

namespace {

Result<size_t> FindFieldIndexById(std::span<const SchemaField> fields, int32_t field_id) {
  for (size_t index = 0; index < fields.size(); ++index) {
    if (fields[index].field_id() == field_id) {
      return index;
    }
  }
  return InvalidArgument("Required schema does not contain projected field id {}",
                         field_id);
}

std::mutex g_project_batch_function_mutex;
ProjectionContext::ProjectBatchFunction g_project_batch_function = nullptr;

ProjectionContext::ProjectBatchFunction GetProjectBatchFunction() {
  std::lock_guard lock(g_project_batch_function_mutex);
  return g_project_batch_function;
}

Result<std::vector<int32_t>> BuildSelectedFieldIndices(
    std::span<const SchemaField> input_fields,
    std::span<const SchemaField> output_fields) {
  std::vector<int32_t> selected_field_indices;
  selected_field_indices.reserve(output_fields.size());

  for (const auto& output_field : output_fields) {
    ICEBERG_ASSIGN_OR_RAISE(auto input_index,
                            FindFieldIndexById(input_fields, output_field.field_id()));
    const auto& input_field = input_fields[input_index];
    if (*input_field.type() != *output_field.type()) {
      return InvalidArgument(
          "ProjectBatch only supports complete top-level fields, but field id "
          "{} changes type from {} to {}",
          output_field.field_id(), input_field.type()->ToString(),
          output_field.type()->ToString());
    }
    ICEBERG_PRECHECK(
        input_index <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
        "Input field index {} exceeds int32 range", input_index);
    selected_field_indices.push_back(static_cast<int32_t>(input_index));
  }

  return selected_field_indices;
}

Status AppendValue(const ArrowSchema& input_schema, const ArrowArray& input_array,
                   const ArrowArrayView& input_view, int64_t row_index,
                   ArrowArray* output_array);

Status AppendListValues(const ArrowSchema& input_schema, const ArrowArray& input_array,
                        const ArrowArrayView& input_view, int64_t row_index,
                        ArrowArray* output_array) {
  const int64_t begin = ArrowArrayViewListChildOffset(&input_view, row_index);
  const int64_t end = ArrowArrayViewListChildOffset(&input_view, row_index + 1);
  for (int64_t element_index = begin; element_index < end; ++element_index) {
    ICEBERG_RETURN_UNEXPECTED(
        AppendValue(*input_schema.children[0], *input_array.children[0],
                    *input_view.children[0], element_index, output_array->children[0]));
  }
  ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayFinishElement(output_array));
  return {};
}

Status AppendMapValues(const ArrowSchema& input_schema, const ArrowArray& input_array,
                       const ArrowArrayView& input_view, int64_t row_index,
                       ArrowArray* output_array) {
  const int64_t begin = ArrowArrayViewListChildOffset(&input_view, row_index);
  const int64_t end = ArrowArrayViewListChildOffset(&input_view, row_index + 1);
  for (int64_t entry_index = begin; entry_index < end; ++entry_index) {
    ICEBERG_RETURN_UNEXPECTED(
        AppendValue(*input_schema.children[0], *input_array.children[0],
                    *input_view.children[0], entry_index, output_array->children[0]));
  }
  ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayFinishElement(output_array));
  return {};
}

Status AppendDecimal(const ArrowSchema& input_schema, const ArrowArrayView& input_view,
                     int64_t row_index, ArrowArray* output_array) {
  ArrowError error;
  ArrowSchemaView schema_view;
  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowSchemaViewInit(&schema_view, &input_schema, &error), error);

  ArrowDecimal value;
  ArrowDecimalInit(&value, schema_view.decimal_bitwidth, schema_view.decimal_precision,
                   schema_view.decimal_scale);
  ArrowArrayViewGetDecimalUnsafe(&input_view, row_index, &value);
  ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendDecimal(output_array, &value));
  return {};
}

Status AppendValue(const ArrowSchema& input_schema, const ArrowArray& input_array,
                   const ArrowArrayView& input_view, int64_t row_index,
                   ArrowArray* output_array) {
  if (ArrowArrayViewIsNull(&input_view, row_index)) {
    ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendNull(output_array, 1));
    return {};
  }

  switch (input_view.storage_type) {
    case NANOARROW_TYPE_NA:
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendNull(output_array, 1));
      return {};
    case NANOARROW_TYPE_BOOL:
    case NANOARROW_TYPE_INT8:
    case NANOARROW_TYPE_INT16:
    case NANOARROW_TYPE_INT32:
    case NANOARROW_TYPE_INT64:
    case NANOARROW_TYPE_DATE32:
    case NANOARROW_TYPE_DATE64:
    case NANOARROW_TYPE_TIME32:
    case NANOARROW_TYPE_TIME64:
    case NANOARROW_TYPE_TIMESTAMP:
    case NANOARROW_TYPE_DURATION:
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendInt(
          output_array, ArrowArrayViewGetIntUnsafe(&input_view, row_index)));
      return {};
    case NANOARROW_TYPE_UINT8:
    case NANOARROW_TYPE_UINT16:
    case NANOARROW_TYPE_UINT32:
    case NANOARROW_TYPE_UINT64:
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendUInt(
          output_array, ArrowArrayViewGetUIntUnsafe(&input_view, row_index)));
      return {};
    case NANOARROW_TYPE_HALF_FLOAT:
    case NANOARROW_TYPE_FLOAT:
    case NANOARROW_TYPE_DOUBLE:
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendDouble(
          output_array, ArrowArrayViewGetDoubleUnsafe(&input_view, row_index)));
      return {};
    case NANOARROW_TYPE_STRING:
    case NANOARROW_TYPE_LARGE_STRING:
    case NANOARROW_TYPE_STRING_VIEW: {
      auto value = ArrowArrayViewGetStringUnsafe(&input_view, row_index);
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendString(output_array, value));
      return {};
    }
    case NANOARROW_TYPE_BINARY:
    case NANOARROW_TYPE_LARGE_BINARY:
    case NANOARROW_TYPE_FIXED_SIZE_BINARY:
    case NANOARROW_TYPE_BINARY_VIEW: {
      auto value = ArrowArrayViewGetBytesUnsafe(&input_view, row_index);
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayAppendBytes(output_array, value));
      return {};
    }
    case NANOARROW_TYPE_DECIMAL128:
    case NANOARROW_TYPE_DECIMAL256:
      return AppendDecimal(input_schema, input_view, row_index, output_array);
    case NANOARROW_TYPE_STRUCT: {
      for (int64_t child_index = 0; child_index < input_schema.n_children;
           ++child_index) {
        ICEBERG_RETURN_UNEXPECTED(AppendValue(
            *input_schema.children[child_index], *input_array.children[child_index],
            *input_view.children[child_index], row_index,
            output_array->children[child_index]));
      }
      ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayFinishElement(output_array));
      return {};
    }
    case NANOARROW_TYPE_LIST:
    case NANOARROW_TYPE_LARGE_LIST:
    case NANOARROW_TYPE_FIXED_SIZE_LIST:
      return AppendListValues(input_schema, input_array, input_view, row_index,
                              output_array);
    case NANOARROW_TYPE_MAP:
      return AppendMapValues(input_schema, input_array, input_view, row_index,
                             output_array);
    default:
      return NotImplemented("Unsupported Arrow type for merge-on-read projection: {}",
                            static_cast<int>(input_view.storage_type));
  }
}

}  // namespace

ProjectionContext::ProjectionContext(ProjectionContext&& other) noexcept
    : input_schema_(std::exchange(other.input_schema_, nullptr)),
      output_schema_(std::exchange(other.output_schema_, nullptr)),
      selected_field_indices_(std::move(other.selected_field_indices_)),
      input_arrow_schema_(other.input_arrow_schema_),
      output_arrow_schema_(other.output_arrow_schema_),
      project_batch_function_(std::exchange(other.project_batch_function_, nullptr)),
      project_batch_state_(std::move(other.project_batch_state_)) {
  other.input_arrow_schema_.release = nullptr;
  other.output_arrow_schema_.release = nullptr;
}

ProjectionContext& ProjectionContext::operator=(ProjectionContext&& other) noexcept {
  if (this == &other) {
    return *this;
  }

  if (input_arrow_schema_.release != nullptr) {
    input_arrow_schema_.release(&input_arrow_schema_);
  }
  if (output_arrow_schema_.release != nullptr) {
    output_arrow_schema_.release(&output_arrow_schema_);
  }

  input_schema_ = std::exchange(other.input_schema_, nullptr);
  output_schema_ = std::exchange(other.output_schema_, nullptr);
  selected_field_indices_ = std::move(other.selected_field_indices_);
  input_arrow_schema_ = other.input_arrow_schema_;
  other.input_arrow_schema_.release = nullptr;
  output_arrow_schema_ = other.output_arrow_schema_;
  other.output_arrow_schema_.release = nullptr;
  project_batch_function_ = std::exchange(other.project_batch_function_, nullptr);
  project_batch_state_ = std::move(other.project_batch_state_);
  return *this;
}

ProjectionContext::~ProjectionContext() {
  if (input_arrow_schema_.release != nullptr) {
    input_arrow_schema_.release(&input_arrow_schema_);
  }
  if (output_arrow_schema_.release != nullptr) {
    output_arrow_schema_.release(&output_arrow_schema_);
  }
}

Result<ProjectionContext> ProjectionContext::Make(
    const Schema& input_schema, const Schema& output_schema,
    ProjectionContext::ProjectBatchFunction project_batch_function) {
  ICEBERG_ASSIGN_OR_RAISE(
      auto selected_field_indices,
      BuildSelectedFieldIndices(input_schema.fields(), output_schema.fields()));

  ProjectionContext context;
  context.input_schema_ = &input_schema;
  context.output_schema_ = &output_schema;
  context.selected_field_indices_ = std::move(selected_field_indices);
  ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(input_schema, &context.input_arrow_schema_));
  ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_));
  context.project_batch_function_ = project_batch_function;

  return context;
}

const Schema& ProjectionContext::input_schema() const { return *input_schema_; }

const Schema& ProjectionContext::output_schema() const { return *output_schema_; }

const ArrowSchema& ProjectionContext::input_arrow_schema() const {
  return input_arrow_schema_;
}

const ArrowSchema& ProjectionContext::output_arrow_schema() const {
  return output_arrow_schema_;
}

std::span<const int32_t> ProjectionContext::selected_field_indices() const {
  return selected_field_indices_;
}

ProjectionContext::ProjectBatchFunction ProjectionContext::project_batch_function()
    const {
  return project_batch_function_;
}

ProjectionContext::ProjectBatchState& ProjectionContext::project_batch_state() {
  return project_batch_state_;
}

void ProjectionContext::RegisterProjectBatchFunction(
    ProjectionContext::ProjectBatchFunction project_batch_function) {
  ICEBERG_DCHECK(project_batch_function != nullptr,
                 "ProjectBatch implementation must not be null");
  if (project_batch_function == nullptr) {
    return;
  }
  std::lock_guard lock(g_project_batch_function_mutex);
  g_project_batch_function = project_batch_function;
}

bool ProjectionContext::HasProjectBatchFunction() {
  return GetProjectBatchFunction() != nullptr;
}

auto ProjectionContext::ResolveProjectBatchFunction()
    -> ProjectionContext::ProjectBatchFunction {
  return GetProjectBatchFunction();
}

namespace {

Result<ArrowArray> ProjectBatchNanoarrow(
    const ArrowSchema& input_arrow_schema, const ArrowArray& input_batch,
    std::span<const int32_t> row_indices, const ArrowSchema& output_arrow_schema,
    std::span<const int32_t> selected_field_indices) {
  ArrowArrayView input_view;
  ArrowError error;
  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowArrayViewInitFromSchema(&input_view, &input_arrow_schema, &error), error);
  internal::ArrowArrayViewGuard input_view_guard(&input_view);
  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowArrayViewSetArray(&input_view, &input_batch, &error), error);
  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowArrayViewValidate(&input_view, NANOARROW_VALIDATION_LEVEL_DEFAULT, &error),
      error);

  ArrowArray output_array;
  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowArrayInitFromSchema(&output_array, &output_arrow_schema, &error), error);
  internal::ArrowArrayGuard output_array_guard(&output_array);
  ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayStartAppending(&output_array));
  ICEBERG_NANOARROW_RETURN_UNEXPECTED(
      ArrowArrayReserve(&output_array, static_cast<int64_t>(row_indices.size())));

  for (int64_t row_index : row_indices) {
    ICEBERG_PRECHECK(row_index >= 0 && row_index < input_batch.length,
                     "Row index {} out of range for batch length {}", row_index,
                     input_batch.length);
    for (size_t output_index = 0; output_index < selected_field_indices.size();
         ++output_index) {
      const int32_t input_index = selected_field_indices[output_index];
      ICEBERG_RETURN_UNEXPECTED(AppendValue(*input_arrow_schema.children[input_index],
                                            *input_batch.children[input_index],
                                            *input_view.children[input_index], row_index,
                                            output_array.children[output_index]));
    }
    ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayFinishElement(&output_array));
  }

  ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
      ArrowArrayFinishBuildingDefault(&output_array, &error), error);

  return std::exchange(output_array, ArrowArray{});
}

}  // namespace

Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
                                std::span<const int32_t> row_indices,
                                ProjectionContext& projection) {
  ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null");
  internal::ArrowArrayGuard input_batch_guard(input_batch);

  auto project_batch_function = projection.project_batch_function();
  if (project_batch_function != nullptr) {
    // ProjectBatch owns input_batch. Arrow-backed implementations import it and clear
    // release, so input_batch_guard becomes a no-op instead of double-releasing.
    return project_batch_function(input_batch, row_indices, projection);
  }

  return ProjectBatchNanoarrow(projection.input_arrow_schema(), *input_batch, row_indices,
                               projection.output_arrow_schema(),
                               projection.selected_field_indices());
}

}  // namespace iceberg