GH-45522: [Parquet][C++] Parquet GEOMETRY and GEOGRAPHY logical type implementations

cpp/src/parquet/CMakeLists.txt

@@ -172,6 +172,9 @@ set(PARQUET_SRCS
     exception.cc
     file_reader.cc
     file_writer.cc
+    geospatial/statistics.cc
+    geospatial/util_internal.cc
+    geospatial/util_json_internal.cc
     level_comparison.cc
     level_conversion.cc
     metadata.cc
@@ -260,6 +263,14 @@ endif()
 if(NOT PARQUET_MINIMAL_DEPENDENCY)
   list(APPEND PARQUET_SHARED_LINK_LIBS arrow_shared)
 
+  # Add RapidJSON libraries
+  if(NOT ARROW_JSON)
+    message(FATAL_ERROR "Building Parquet with Arrow support requires ARROW_JSON=ON")
+  endif()
+
+  list(APPEND PARQUET_SHARED_PRIVATE_LINK_LIBS RapidJSON)
+  list(APPEND PARQUET_STATIC_LINK_LIBS RapidJSON)
+
   # These are libraries that we will link privately with parquet_shared (as they
   # do not need to be linked transitively by other linkers)
   list(APPEND PARQUET_SHARED_PRIVATE_LINK_LIBS thrift::thrift)
@@ -357,6 +368,7 @@ endif()
 add_subdirectory(api)
 add_subdirectory(arrow)
 add_subdirectory(encryption)
+add_subdirectory(geospatial)
 
 arrow_install_all_headers("parquet")
 
@@ -367,15 +379,17 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/parquet_version.h"
 
 add_parquet_test(internals-test
                  SOURCES
-                 bloom_filter_test.cc
                  bloom_filter_reader_test.cc
-                 properties_test.cc
-                 statistics_test.cc
+                 bloom_filter_test.cc
                  encoding_test.cc
+                 geospatial/statistics_test.cc
+                 geospatial/util_internal_test.cc
                  metadata_test.cc
                  page_index_test.cc
+                 properties_test.cc
                  public_api_test.cc
                  size_statistics_test.cc
+                 statistics_test.cc
                  types_test.cc)
 
 set_source_files_properties(public_api_test.cc PROPERTIES SKIP_PRECOMPILE_HEADERS ON

cpp/src/parquet/api/reader.h

@@ -22,6 +22,7 @@
 #include "parquet/column_scanner.h"
 #include "parquet/exception.h"
 #include "parquet/file_reader.h"
+#include "parquet/geospatial/statistics.h"
 #include "parquet/metadata.h"
 #include "parquet/platform.h"
 #include "parquet/printer.h"

cpp/src/parquet/arrow/arrow_reader_writer_test.cc

@@ -43,6 +43,7 @@
 #include "arrow/scalar.h"
 #include "arrow/table.h"
 #include "arrow/testing/builder.h"
+#include "arrow/testing/extension_type.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/random.h"
 #include "arrow/testing/util.h"
@@ -1481,6 +1482,61 @@ TEST_F(TestJsonParquetIO, JsonExtension) {
   this->RoundTripSingleColumn(json_large_array, json_large_array, writer_properties);
 }
 
+using TestGeoArrowParquetIO = TestParquetIO<test::GeoArrowWkbExtensionType>;
+
+TEST_F(TestGeoArrowParquetIO, GeoArrowExtension) {
+  ::arrow::ExtensionTypeGuard guard(test::geoarrow_wkb());
+
+  // Build a binary WKB array with at least one null value
+  ::arrow::BinaryBuilder builder;
+
+  for (int k = 0; k < 10; k++) {
+    std::string item = test::MakeWKBPoint(
+        {static_cast<double>(k), static_cast<double>(k + 1)}, false, false);
+    ASSERT_OK(builder.Append(item));
+  }
+  ASSERT_OK(builder.AppendNull());
+  for (int k = 0; k < 5; k++) {
+    std::string item = test::MakeWKBPoint(
+        {static_cast<double>(k), static_cast<double>(k + 1)}, false, false);
+    ASSERT_OK(builder.Append(item));
+  }
+
+  ASSERT_OK_AND_ASSIGN(const auto binary_array, builder.Finish());
+  const auto wkb_type = test::geoarrow_wkb_lonlat();
+  const auto wkb_array = ::arrow::ExtensionType::WrapArray(wkb_type, binary_array);
+
+  const auto large_wkb_type = test::geoarrow_wkb_lonlat(::arrow::large_binary());
+  ASSERT_OK_AND_ASSIGN(const auto large_binary_array,
+                       ::arrow::compute::Cast(binary_array, ::arrow::large_binary()));
+  const auto large_wkb_array =
+      ::arrow::ExtensionType::WrapArray(large_wkb_type, large_binary_array.make_array());
+
+  // When the original Arrow schema isn't stored and Arrow extensions are disabled,
+  // LogicalType::GEOMETRY is read as utf8.
+  auto writer_properties = default_arrow_writer_properties();
+  ASSERT_NO_FATAL_FAILURE(
+      this->RoundTripSingleColumn(wkb_array, binary_array, writer_properties));
+  ASSERT_NO_FATAL_FAILURE(
+      this->RoundTripSingleColumn(large_wkb_array, binary_array, writer_properties));
+
+  // When the original Arrow schema isn't stored and Arrow extensions are enabled,
+  // LogicalType::GEOMETRY is read as geoarrow.wkb with binary storage.
+  ::parquet::ArrowReaderProperties reader_properties;
+  reader_properties.set_arrow_extensions_enabled(true);
+  ASSERT_NO_FATAL_FAILURE(this->RoundTripSingleColumn(
+      wkb_array, wkb_array, writer_properties, reader_properties));
+  ASSERT_NO_FATAL_FAILURE(this->RoundTripSingleColumn(
+      large_wkb_array, wkb_array, writer_properties, reader_properties));
+
+  // When the original Arrow schema is stored, the stored Arrow type is respected.
+  writer_properties = ::parquet::ArrowWriterProperties::Builder().store_schema()->build();
+  ASSERT_NO_FATAL_FAILURE(
+      this->RoundTripSingleColumn(wkb_array, wkb_array, writer_properties));
+  ASSERT_NO_FATAL_FAILURE(
+      this->RoundTripSingleColumn(large_wkb_array, large_wkb_array, writer_properties));
+}
+
 using TestNullParquetIO = TestParquetIO<::arrow::NullType>;
 
 TEST_F(TestNullParquetIO, NullColumn) {

cpp/src/parquet/arrow/arrow_schema_test.cc

@@ -19,7 +19,6 @@
 #include <vector>
 
 #include "gmock/gmock-matchers.h"
-#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include "parquet/arrow/reader.h"
@@ -35,6 +34,7 @@
 #include "arrow/extension/json.h"
 #include "arrow/extension/uuid.h"
 #include "arrow/ipc/writer.h"
+#include "arrow/testing/extension_type.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/type.h"
 #include "arrow/util/base64.h"
@@ -237,6 +237,10 @@ TEST_F(TestConvertParquetSchema, ParquetAnnotatedFields) {
        ::arrow::int64()},
       {"json", LogicalType::JSON(), ParquetType::BYTE_ARRAY, -1, ::arrow::utf8()},
       {"bson", LogicalType::BSON(), ParquetType::BYTE_ARRAY, -1, ::arrow::binary()},
+      {"geometry", LogicalType::Geometry(), ParquetType::BYTE_ARRAY, -1,
+       ::arrow::binary()},
+      {"geography", LogicalType::Geography(), ParquetType::BYTE_ARRAY, -1,
+       ::arrow::binary()},
       {"interval", LogicalType::Interval(), ParquetType::FIXED_LEN_BYTE_ARRAY, 12,
        ::arrow::fixed_size_binary(12)},
       {"uuid", LogicalType::UUID(), ParquetType::FIXED_LEN_BYTE_ARRAY, 16,
@@ -1090,6 +1094,49 @@ TEST_F(TestConvertParquetSchema, ParquetSchemaArrowUuidExtension) {
   }
 }
 
+TEST_F(TestConvertParquetSchema, ParquetSchemaGeoArrowExtensions) {
+  std::vector<NodePtr> parquet_fields;
+  parquet_fields.push_back(PrimitiveNode::Make("geometry", Repetition::OPTIONAL,
+                                               LogicalType::Geometry(),
+                                               ParquetType::BYTE_ARRAY));
+  parquet_fields.push_back(PrimitiveNode::Make("geography", Repetition::OPTIONAL,
+                                               LogicalType::Geography(),
+                                               ParquetType::BYTE_ARRAY));
+
+  {
+    // Parquet file does not contain Arrow schema.
+    // By default, both fields should be treated as binary() fields in Arrow.
+    auto arrow_schema = ::arrow::schema({::arrow::field("geometry", BINARY, true),
+                                         ::arrow::field("geography", BINARY, true)});
+    std::shared_ptr<KeyValueMetadata> metadata{};
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata));
+    CheckFlatSchema(arrow_schema);
+  }
+
+  {
+    // Parquet file does not contain Arrow schema.
+    // If Arrow extensions are enabled and extensions are registered,
+    // fields will be interpreted as geoarrow_wkb(binary()) extension fields.
+    ::arrow::ExtensionTypeGuard guard(test::geoarrow_wkb());
+
+    ArrowReaderProperties props;
+    props.set_arrow_extensions_enabled(true);
+    auto arrow_schema = ::arrow::schema(
+        {::arrow::field(
+             "geometry",
+             test::geoarrow_wkb(R"({"crs": "OGC:CRS84", "crs_type": "authority_code"})"),
+             true),
+         ::arrow::field(
+             "geography",
+             test::geoarrow_wkb(
+                 R"({"crs": "OGC:CRS84", "crs_type": "authority_code", "edges": "spherical"})"),
+             true)});
+    std::shared_ptr<KeyValueMetadata> metadata{};
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata, props));
+    CheckFlatSchema(arrow_schema);
+  }
+}
+
 class TestConvertArrowSchema : public ::testing::Test {
  public:
   virtual void SetUp() {}
@@ -1343,6 +1390,96 @@ TEST_F(TestConvertArrowSchema, ParquetFlatPrimitivesAsDictionaries) {
   ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields));
 }
 
+TEST_F(TestConvertArrowSchema, ParquetGeoArrowCrsLonLat) {
+  // All the Arrow Schemas below should convert to the type defaults for GEOMETRY
+  // and GEOGRAPHY when GeoArrow extension types are registered and the appropriate
+  // writer option is set.
+  ::arrow::ExtensionTypeGuard guard(test::geoarrow_wkb());
+
+  std::vector<NodePtr> parquet_fields;
+  parquet_fields.push_back(PrimitiveNode::Make("geometry", Repetition::OPTIONAL,
+                                               LogicalType::Geometry(),
+                                               ParquetType::BYTE_ARRAY));
+  parquet_fields.push_back(PrimitiveNode::Make("geography", Repetition::OPTIONAL,
+                                               LogicalType::Geography(),
+                                               ParquetType::BYTE_ARRAY));
+
+  // There are several ways that longitude/latitude could be specified when coming from
+  // GeoArrow, which allows null, missing, arbitrary strings (e.g., Authority:Code), and
+  // PROJJSON.
+  std::vector<std::string> geoarrow_lonlat = {
+      "null", R"("OGC:CRS84")", R"("EPSG:4326")",
+      // Purely the parts of the PROJJSON that we inspect to check the lon/lat case
+      R"({"id": {"authority": "OGC", "code": "CRS84"}})",
+      R"({"id": {"authority": "EPSG", "code": 4326}})"};
+
+  std::string geoarrow_lonlatish_crs = geoarrow_lonlat[0];
+  for (const auto& geoarrow_lonlatish_crs : geoarrow_lonlat) {
+    SCOPED_TRACE(geoarrow_lonlatish_crs);
+    std::vector<std::shared_ptr<Field>> arrow_fields = {
+        ::arrow::field("geometry",
+                       test::geoarrow_wkb(R"({"crs": )" + geoarrow_lonlatish_crs + "}"),
+                       true),
+        ::arrow::field("geography",
+                       test::geoarrow_wkb(R"({"crs": )" + geoarrow_lonlatish_crs +
+                                          R"(, "edges": "spherical"})"),
+                       true)};
+
+    ASSERT_OK(ConvertSchema(arrow_fields));
+    ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields));
+  }
+}
+
+TEST_F(TestConvertArrowSchema, ParquetGeoArrowCrsSrid) {
+  // Checks that it is possible to write the srid:xxxx reccomendation from GeoArrow
+  ::arrow::ExtensionTypeGuard guard(test::geoarrow_wkb());
+
+  std::vector<NodePtr> parquet_fields;
+  parquet_fields.push_back(PrimitiveNode::Make("geometry", Repetition::OPTIONAL,
+                                               LogicalType::Geometry("srid:1234"),
+                                               ParquetType::BYTE_ARRAY));
+  parquet_fields.push_back(PrimitiveNode::Make("geography", Repetition::OPTIONAL,
+                                               LogicalType::Geography("srid:5678"),
+                                               ParquetType::BYTE_ARRAY));
+
+  std::vector<std::shared_ptr<Field>> arrow_fields = {
+      ::arrow::field("geometry", test::geoarrow_wkb(R"({"crs": "srid:1234"})"), true),
+      ::arrow::field("geography",
+                     test::geoarrow_wkb(R"({"crs": "srid:5678", "edges": "spherical"})"),
+                     true)};
+
+  ASSERT_OK(ConvertSchema(arrow_fields));
+  ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields));
+}
+
+TEST_F(TestConvertArrowSchema, ParquetGeoArrowCrsProjjson) {
+  // Checks the conversion between GeoArrow that contains non-lon/lat PROJJSON
+  // to Parquet. Almost all GeoArrow types that arrive at the Parquet reader
+  // will have their CRS expressed in this way.
+  ::arrow::ExtensionTypeGuard guard(test::geoarrow_wkb());
+
+  std::vector<std::shared_ptr<Field>> arrow_fields = {
+      ::arrow::field("geometry", test::geoarrow_wkb(R"({"crs": {"key0": "value0"}})"),
+                     true),
+      ::arrow::field(
+          "geography",
+          test::geoarrow_wkb(R"({"crs": {"key1": "value1"}, "edges": "spherical"})"),
+          true)};
+
+  auto arrow_properties = default_arrow_writer_properties();
+  ASSERT_OK(ConvertSchema(arrow_fields, arrow_properties));
+
+  std::vector<NodePtr> parquet_fields;
+  parquet_fields.push_back(PrimitiveNode::Make(
+      "geometry", Repetition::OPTIONAL, LogicalType::Geometry(R"({"key0":"value0"})"),
+      ParquetType::BYTE_ARRAY));
+  parquet_fields.push_back(PrimitiveNode::Make(
+      "geography", Repetition::OPTIONAL, LogicalType::Geography(R"({"key1":"value1"})"),
+      ParquetType::BYTE_ARRAY));
+
+  ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields));
+}
+
 TEST_F(TestConvertArrowSchema, ParquetLists) {
   std::vector<NodePtr> parquet_fields;
   std::vector<std::shared_ptr<Field>> arrow_fields;

cpp/src/parquet/arrow/schema.cc

@@ -38,6 +38,7 @@
 #include "parquet/arrow/schema_internal.h"
 #include "parquet/arrow/variant_internal.h"
 #include "parquet/exception.h"
+#include "parquet/geospatial/util_json_internal.h"
 #include "parquet/metadata.h"
 #include "parquet/properties.h"
 #include "parquet/types.h"
@@ -288,8 +289,8 @@ int FieldIdFromMetadata(
   if (::arrow::internal::ParseValue<::arrow::Int32Type>(
           field_id_str.c_str(), field_id_str.length(), &field_id)) {
     if (field_id < 0) {
-      // Thrift should convert any negative value to null but normalize to -1 here in case
-      // we later check this in logic.
+      // Thrift should convert any negative value to null but normalize to -1 here in
+      // case we later check this in logic.
       return -1;
     }
     return field_id;
@@ -449,9 +450,9 @@ Status FieldToNode(const std::string& name, const std::shared_ptr<Field>& field,
     }
     case ArrowTypeId::EXTENSION: {
       auto ext_type = std::static_pointer_cast<::arrow::ExtensionType>(field->type());
-      // Built-in JSON extension is handled differently.
-      if (ext_type->extension_name() == std::string("arrow.json")) {
-        // Set physical and logical types and instantiate primitive node.
+      // Set physical and logical types and instantiate primitive node
+      // for extension types
+      if (ext_type->extension_name() == std::string_view("arrow.json")) {
         type = ParquetType::BYTE_ARRAY;
         logical_type = LogicalType::JSON();
         break;
@@ -460,6 +461,11 @@ Status FieldToNode(const std::string& name, const std::shared_ptr<Field>& field,
         logical_type = LogicalType::UUID();
         length = 16;
         break;
+      } else if (ext_type->extension_name() == std::string_view("geoarrow.wkb")) {
+        type = ParquetType::BYTE_ARRAY;
+        ARROW_ASSIGN_OR_RAISE(logical_type,
+                              LogicalTypeFromGeoArrowMetadata(ext_type->Serialize()));
+        break;
       } else if (ext_type->extension_name() == std::string("parquet.variant")) {
         auto variant_type = std::static_pointer_cast<VariantExtensionType>(field->type());
 
@@ -495,6 +501,7 @@ struct SchemaTreeContext {
   SchemaManifest* manifest;
   ArrowReaderProperties properties;
   const SchemaDescriptor* schema;
+  std::shared_ptr<const KeyValueMetadata> metadata;
 
   void LinkParent(const SchemaField* child, const SchemaField* parent) {
     manifest->child_to_parent[child] = parent;
@@ -517,7 +524,7 @@ ::arrow::Result<std::shared_ptr<ArrowType>> GetTypeForNode(
     int column_index, const schema::PrimitiveNode& primitive_node,
     SchemaTreeContext* ctx) {
   ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrowType> storage_type,
-                        GetArrowType(primitive_node, ctx->properties));
+                        GetArrowType(primitive_node, ctx->properties, ctx->metadata));
   if (ctx->properties.read_dictionary(column_index) &&
       IsDictionaryReadSupported(*storage_type)) {
     return ::arrow::dictionary(::arrow::int32(), storage_type);
@@ -610,8 +617,8 @@ Status MapToSchemaField(const GroupNode& group, LevelInfo current_levels,
     return Status::Invalid("Map keys must be annotated as required.");
   }
   // Arrow doesn't support 1 column maps (i.e. Sets).  The options are to either
-  // make the values column nullable, or process the map as a list.  We choose the latter
-  // as it is simpler.
+  // make the values column nullable, or process the map as a list.  We choose the
+  // latter as it is simpler.
   if (key_value.field_count() == 1) {
     return ListToSchemaField(group, current_levels, ctx, parent, out);
   }
@@ -1197,6 +1204,7 @@ Status SchemaManifest::Make(const SchemaDescriptor* schema,
   ctx.manifest = manifest;
   ctx.properties = properties;
   ctx.schema = schema;
+  ctx.metadata = metadata;
   const GroupNode& schema_node = *schema->group_node();
   manifest->descr = schema;
   manifest->schema_fields.resize(schema_node.field_count());