feat(parquet): project arrow array for parquet

Author: wgtmac

src/iceberg/parquet/parquet_data_util.cc

@@ -17,15 +17,266 @@
  * under the License.
  */
 
+#include <arrow/array.h>
+#include <arrow/array/builder_primitive.h>
+#include <arrow/compute/api.h>
+#include <arrow/record_batch.h>
+#include <arrow/type.h>
+
+#include "iceberg/arrow/arrow_error_transform_internal.h"
 #include "iceberg/parquet/parquet_data_util_internal.h"
+#include "iceberg/schema.h"
+#include "iceberg/schema_util.h"
+#include "iceberg/type.h"
+#include "iceberg/util/checked_cast.h"
+#include "iceberg/util/macros.h"
 
 namespace iceberg::parquet {
 
+namespace {
+
+// Forward declaration
+Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
+    const std::shared_ptr<::arrow::Array>& array,
+    const std::shared_ptr<::arrow::DataType>& output_arrow_type,
+    const NestedType& nested_type, std::span<const FieldProjection> projections,
+    ::arrow::MemoryPool* pool);
+
+/// \brief Create a null array of the given type and length.
+Result<std::shared_ptr<::arrow::Array>> MakeNullArray(
+    const std::shared_ptr<::arrow::DataType>& type, int64_t length,
+    ::arrow::MemoryPool* pool) {
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(auto builder, ::arrow::MakeBuilder(type, pool));
+  ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendNulls(length));
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(auto array, builder->Finish());
+  return array;
+}
+
+Result<std::shared_ptr<::arrow::Array>> ProjectPrimitiveArray(
+    const std::shared_ptr<::arrow::Array>& array,
+    const std::shared_ptr<::arrow::DataType>& output_arrow_type,
+    ::arrow::MemoryPool* pool) {
+  if (array->type()->Equals(output_arrow_type)) {
+    return array;
+  }
+
+  // TODO(gangwu): Check schema evolution rule.
+
+  // Use Arrow compute cast function for type conversions.
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(auto cast_result,
+                                 ::arrow::compute::Cast(array, output_arrow_type));
+  return cast_result.make_array();
+}
+
+Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
+    const std::shared_ptr<::arrow::StructArray>& struct_array,
+    const std::shared_ptr<::arrow::StructType>& output_struct_type,
+    const StructType& struct_type, std::span<const FieldProjection> projections,
+    ::arrow::MemoryPool* pool) {
+  if (struct_type.fields().size() != projections.size()) {
+    return InvalidSchema(
+        "Inconsistent number of fields ({}) and number of projections ({})",
+        struct_type.fields().size(), projections.size());
+  }
+  if (struct_type.fields().size() != output_struct_type->num_fields()) {
+    return InvalidSchema(
+        "Inconsistent number of fields ({}) and number of output fields ({})",
+        struct_type.fields().size(), output_struct_type->num_fields());
+  }
+
+  std::vector<std::shared_ptr<::arrow::Array>> projected_arrays;
+  projected_arrays.reserve(projections.size());
+
+  for (size_t i = 0; i < projections.size(); ++i) {
+    const auto& projected_field = struct_type.fields()[i];
+    const auto& field_projection = projections[i];
+    const auto& output_arrow_type = output_struct_type->fields()[i]->type();
+
+    std::shared_ptr<::arrow::Array> projected_array;
+
+    if (field_projection.kind == FieldProjection::Kind::kProjected) {
+      size_t parquet_field_index = std::get<size_t>(field_projection.from);
+      if (parquet_field_index >= struct_array->num_fields()) {
+        return InvalidArgument("Parquet field index {} out of bound {}",
+                               parquet_field_index, struct_array->num_fields());
+      }
+      const auto& parquet_array = struct_array->field(parquet_field_index);
+      if (projected_field.type()->is_nested()) {
+        const auto& nested_type =
+            internal::checked_cast<const NestedType&>(*projected_field.type());
+        ICEBERG_ASSIGN_OR_RAISE(
+            projected_array,
+            ProjectNestedArray(parquet_array, output_arrow_type, nested_type,
+                               field_projection.children, pool));
+      } else {
+        ICEBERG_ASSIGN_OR_RAISE(
+            projected_array,
+            ProjectPrimitiveArray(parquet_array, output_arrow_type, pool));
+      }
+    } else if (field_projection.kind == FieldProjection::Kind::kNull) {
+      ICEBERG_ASSIGN_OR_RAISE(
+          projected_array,
+          MakeNullArray(output_arrow_type, struct_array->length(), pool));
+    } else {
+      return NotImplemented("Unsupported field projection kind: {}",
+                            ToString(field_projection.kind));
+    }
+
+    projected_arrays.emplace_back(std::move(projected_array));
+  }
+
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(
+      auto output_array,
+      ::arrow::StructArray::Make(projected_arrays, output_struct_type->fields(),
+                                 struct_array->null_bitmap(), struct_array->null_count(),
+                                 struct_array->offset()));
+  return output_array;
+}
+
+/// FIXME: Support ::arrow::LargeListType.
+Result<std::shared_ptr<::arrow::Array>> ProjectListArray(
+    const std::shared_ptr<::arrow::ListArray>& list_array,
+    const std::shared_ptr<::arrow::ListType>& output_list_type, const ListType& list_type,
+    std::span<const FieldProjection> projections, ::arrow::MemoryPool* pool) {
+  if (projections.size() != 1) {
+    return InvalidArgument("Expected 1 projection for list, got: {}", projections.size());
+  }
+
+  const auto& element_field = list_type.fields().back();
+  const auto& element_projection = projections[0];
+  const auto& output_element_type = output_list_type->value_type();
+
+  std::shared_ptr<::arrow::Array> projected_values;
+  if (element_field.type()->is_nested()) {
+    const auto& nested_type =
+        internal::checked_cast<const NestedType&>(*element_field.type());
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_values,
+        ProjectNestedArray(list_array->values(), output_element_type, nested_type,
+                           element_projection.children, pool));
+  } else {
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_values,
+        ProjectPrimitiveArray(list_array->values(), output_element_type, pool));
+  }
+
+  return std::make_shared<::arrow::ListArray>(
+      output_list_type, list_array->length(), list_array->value_offsets(),
+      std::move(projected_values), list_array->null_bitmap(), list_array->null_count(),
+      list_array->offset());
+}
+
+/// FIXME: Support ::arrow::LargeMapType.
+Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
+    const std::shared_ptr<::arrow::MapArray>& map_array,
+    const std::shared_ptr<::arrow::MapType>& output_map_type, const MapType& map_type,
+    std::span<const FieldProjection> projections, ::arrow::MemoryPool* pool) {
+  if (projections.size() != 2) {
+    return InvalidArgument("Expected 2 projections for map, got: {}", projections.size());
+  }
+
+  const auto& key_projection = projections[0];
+  const auto& value_projection = projections[1];
+  const auto& key_type = map_type.key().type();
+  const auto& value_type = map_type.value().type();
+
+  // Project keys
+  std::shared_ptr<::arrow::Array> projected_keys;
+  if (key_type->is_nested()) {
+    const auto& nested_type = internal::checked_cast<const NestedType&>(*key_type);
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_keys, ProjectNestedArray(map_array->keys(), output_map_type->key_type(),
+                                           nested_type, key_projection.children, pool));
+  } else {
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_keys,
+        ProjectPrimitiveArray(map_array->keys(), output_map_type->key_type(), pool));
+  }
+
+  // Project values
+  std::shared_ptr<::arrow::Array> projected_items;
+  if (value_type->is_nested()) {
+    const auto& nested_type = internal::checked_cast<const NestedType&>(*value_type);
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_items,
+        ProjectNestedArray(map_array->items(), output_map_type->item_type(), nested_type,
+                           value_projection.children, pool));
+  } else {
+    ICEBERG_ASSIGN_OR_RAISE(
+        projected_items,
+        ProjectPrimitiveArray(map_array->items(), output_map_type->item_type(), pool));
+  }
+
+  return std::make_shared<::arrow::MapArray>(
+      output_map_type, map_array->length(), map_array->value_offsets(),
+      std::move(projected_keys), std::move(projected_items), map_array->null_bitmap(),
+      map_array->null_count(), map_array->offset());
+}
+
+Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
+    const std::shared_ptr<::arrow::Array>& array,
+    const std::shared_ptr<::arrow::DataType>& output_arrow_type,
+    const NestedType& nested_type, std::span<const FieldProjection> projections,
+    ::arrow::MemoryPool* pool) {
+  switch (nested_type.type_id()) {
+    case TypeId::kStruct: {
+      if (output_arrow_type->id() != ::arrow::Type::STRUCT) {
+        return InvalidSchema("Expected struct type, got: {}",
+                             output_arrow_type->ToString());
+      }
+      auto struct_array = internal::checked_pointer_cast<::arrow::StructArray>(array);
+      auto output_struct_type =
+          internal::checked_pointer_cast<::arrow::StructType>(output_arrow_type);
+      const auto& struct_type = internal::checked_cast<const StructType&>(nested_type);
+      return ProjectStructArray(struct_array, output_struct_type, struct_type,
+                                projections, pool);
+    }
+    case TypeId::kList: {
+      if (output_arrow_type->id() != ::arrow::Type::LIST) {
+        return InvalidSchema("Expected list type, got: {}",
+                             output_arrow_type->ToString());
+      }
+
+      auto list_array = internal::checked_pointer_cast<::arrow::ListArray>(array);
+      auto output_list_type =
+          internal::checked_pointer_cast<::arrow::ListType>(output_arrow_type);
+      const auto& list_type = internal::checked_cast<const ListType&>(nested_type);
+      return ProjectListArray(list_array, output_list_type, list_type, projections, pool);
+    }
+    case TypeId::kMap: {
+      if (output_arrow_type->id() != ::arrow::Type::MAP) {
+        return InvalidSchema("Expected map type, got: {}", output_arrow_type->ToString());
+      }
+
+      auto map_array = internal::checked_pointer_cast<::arrow::MapArray>(array);
+      auto output_map_type =
+          internal::checked_pointer_cast<::arrow::MapType>(output_arrow_type);
+      const auto& map_type = internal::checked_cast<const MapType&>(nested_type);
+      return ProjectMapArray(map_array, output_map_type, map_type, projections, pool);
+    }
+    default:
+      return InvalidSchema("Cannot project array of unsupported nested type: {}",
+                           nested_type.ToString());
+  }
+}
+
+}  // namespace
+
 Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
     std::shared_ptr<::arrow::RecordBatch> record_batch,
     const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
-    const Schema& projected_schema, const SchemaProjection& projection) {
-  return NotImplemented("NYI");
+    const Schema& projected_schema, const SchemaProjection& projection,
+    ::arrow::MemoryPool* pool) {
+  auto array = std::make_shared<::arrow::StructArray>(
+      ::arrow::struct_(record_batch->schema()->fields()), record_batch->num_rows(),
+      record_batch->columns());
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto output_array,
+      ProjectNestedArray(array, ::arrow::struct_(output_arrow_schema->fields()),
+                         projected_schema, projection.fields, pool));
+  auto struct_array = internal::checked_pointer_cast<::arrow::StructArray>(output_array);
+  return ::arrow::RecordBatch::Make(output_arrow_schema, record_batch->num_rows(),
+                                    struct_array->fields());
 }
 
 }  // namespace iceberg::parquet

src/iceberg/parquet/parquet_data_util_internal.h

@@ -22,6 +22,7 @@
 #include "iceberg/schema_util.h"
 
 namespace arrow {
+class MemoryPool;
 class RecordBatch;
 class Schema;
 }  // namespace arrow
@@ -34,10 +35,12 @@ namespace iceberg::parquet {
 /// \param output_arrow_schema The Arrow schema to convert to.
 /// \param projected_schema The projected Iceberg schema.
 /// \param projection The projection from projected Iceberg schema to the record batch.
+/// \param pool The arrow memory pool.
 /// \return The converted record batch.
 Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
     std::shared_ptr<::arrow::RecordBatch> record_batch,
     const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
-    const Schema& projected_schema, const SchemaProjection& projection);
+    const Schema& projected_schema, const SchemaProjection& projection,
+    ::arrow::MemoryPool* pool = nullptr);
 
 }  // namespace iceberg::parquet