Proof-of-concept Parquet GEOMETRY logical type implementation

cpp/src/parquet/CMakeLists.txt

@@ -375,6 +375,7 @@ add_parquet_test(internals-test
                  statistics_test.cc
                  encoding_test.cc
                  metadata_test.cc
+                 geometry_util_internal_test.cc
                  page_index_test.cc
                  public_api_test.cc
                  types_test.cc)

cpp/src/parquet/column_reader.cc

@@ -243,6 +243,10 @@ EncodedStatistics ExtractStatsFromHeader(const H& header) {
   if (stats.__isset.distinct_count) {
     page_statistics.set_distinct_count(stats.distinct_count);
   }
+  if (stats.__isset.geometry_stats) {
+    page_statistics.set_geometry(
+        FromThrift(stats.geometry_stats, stats.__isset.geometry_stats));
+  }
   return page_statistics;
 }
 

cpp/src/parquet/column_writer.cc

@@ -1224,8 +1224,11 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
     // Will be null if not using dictionary, but that's ok
     current_dict_encoder_ = dynamic_cast<DictEncoder<DType>*>(current_encoder_.get());
 
+    bool is_geometry =
+        (descr_->logical_type() != nullptr && descr_->logical_type()->is_geometry());
+    bool has_sort_order = SortOrder::UNKNOWN != descr_->sort_order();
     if (properties->statistics_enabled(descr_->path()) &&
-        (SortOrder::UNKNOWN != descr_->sort_order())) {
+        (is_geometry || has_sort_order)) {
       page_statistics_ = MakeStatistics<DType>(descr_, allocator_);
       chunk_statistics_ = MakeStatistics<DType>(descr_, allocator_);
     }

cpp/src/parquet/column_writer_test.cc

@@ -35,6 +35,7 @@
 #include "parquet/column_writer.h"
 #include "parquet/file_reader.h"
 #include "parquet/file_writer.h"
+#include "parquet/geometry_util_internal.h"
 #include "parquet/metadata.h"
 #include "parquet/platform.h"
 #include "parquet/properties.h"
@@ -399,6 +400,13 @@ class TestPrimitiveWriter : public PrimitiveTypedTest<TestType> {
     return metadata_accessor->key_value_metadata();
   }
 
+  std::shared_ptr<Statistics> metadata_stats() {
+    ApplicationVersion app_version(this->writer_properties_->created_by());
+    auto metadata_accessor = ColumnChunkMetaData::Make(
+        metadata_->contents(), this->descr_, default_reader_properties(), &app_version);
+    return metadata_accessor->statistics();
+  }
+
  protected:
   int64_t values_read_;
   // Keep the reader alive as for ByteArray the lifetime of the ByteArray
@@ -1774,5 +1782,191 @@ TEST_F(TestInt32Writer, WriteKeyValueMetadataEndToEnd) {
   ASSERT_EQ("bar", value);
 }
 
+// Test writing and reading geometry columns
+class TestGeometryValuesWriter : public TestPrimitiveWriter<ByteArrayType> {
+ public:
+  static const char* CRS;
+  static const char* METADATA;
+
+  void SetUpSchema(Repetition::type repetition, int num_columns) override {
+    std::vector<schema::NodePtr> fields;
+
+    for (int i = 0; i < num_columns; ++i) {
+      std::string name = TestColumnName(i);
+      std::shared_ptr<const LogicalType> logical_type =
+          GeometryLogicalType::Make(CRS, LogicalType::GeometryEdges::PLANAR,
+                                    LogicalType::GeometryEncoding::WKB, METADATA);
+      fields.push_back(schema::PrimitiveNode::Make(name, repetition, logical_type,
+                                                   ByteArrayType::type_num));
+    }
+    node_ = schema::GroupNode::Make("schema", Repetition::REQUIRED, fields);
+    schema_.Init(node_);
+  }
+
+  void GenerateData(int64_t num_values, uint32_t seed = 0) {
+    values_.resize(num_values);
+
+    buffer_.resize(num_values * WKB_POINT_SIZE);
+    uint8_t* ptr = buffer_.data();
+    for (int k = 0; k < num_values; k++) {
+      GenerateWKBPoint(ptr, k, k + 1);
+      values_[k].len = WKB_POINT_SIZE;
+      values_[k].ptr = ptr;
+      ptr += WKB_POINT_SIZE;
+    }
+
+    values_ptr_ = values_.data();
+  }
+
+  void TestWriteAndRead(ParquetVersion::type version,
+                        ParquetDataPageVersion data_page_version) {
+    this->SetUpSchema(Repetition::REQUIRED, 1);
+    this->GenerateData(SMALL_SIZE);
+    size_t num_values = this->values_.size();
+    auto writer =
+        this->BuildWriter(num_values, ColumnProperties(), version, data_page_version,
+                          /*enable_checksum*/ false);
+    writer->WriteBatch(this->values_.size(), nullptr, nullptr, this->values_.data());
+
+    writer->Close();
+    this->ReadColumn();
+    for (size_t i = 0; i < num_values; i++) {
+      const ByteArray& value = this->values_out_[i];
+      double x = 0;
+      double y = 0;
+      EXPECT_TRUE(GetWKBPointCoordinate(value, &x, &y));
+      auto expected_x = static_cast<double>(i);
+      auto expected_y = static_cast<double>(i + 1);
+      EXPECT_DOUBLE_EQ(expected_x, x);
+      EXPECT_DOUBLE_EQ(expected_y, y);
+    }
+
+    std::shared_ptr<Statistics> statistics = metadata_stats();
+    EXPECT_FALSE(statistics->HasMinMax());
+    EXPECT_TRUE(statistics->HasGeometryStatistics());
+    const GeometryStatistics* geometry_statistics = statistics->geometry_statistics();
+    std::vector<int32_t> geometry_types = geometry_statistics->GetGeometryTypes();
+    EXPECT_EQ(1, geometry_types.size());
+    EXPECT_EQ(1, geometry_types[0]);
+    EXPECT_DOUBLE_EQ(0, geometry_statistics->GetXMin());
+    EXPECT_DOUBLE_EQ(1, geometry_statistics->GetYMin());
+    EXPECT_DOUBLE_EQ(99, geometry_statistics->GetXMax());
+    EXPECT_DOUBLE_EQ(100, geometry_statistics->GetYMax());
+    EXPECT_FALSE(geometry_statistics->HasZ());
+    EXPECT_FALSE(geometry_statistics->HasM());
+
+    auto coverings = geometry_statistics->GetCoverings();
+    EXPECT_EQ(1, coverings.size());
+    EXPECT_EQ("WKB", coverings[0].first);
+    geometry::WKBGeometryBounder bounder;
+    const std::string& wkb = coverings[0].second;
+    geometry::WKBBuffer wkb_buffer(reinterpret_cast<const uint8_t*>(wkb.data()),
+                                   wkb.size());
+    bounder.ReadGeometry(&wkb_buffer);
+    bounder.Flush();
+    auto bounds = bounder.Bounds();
+    EXPECT_DOUBLE_EQ(0, bounds.min[0]);
+    EXPECT_DOUBLE_EQ(1, bounds.min[1]);
+    EXPECT_DOUBLE_EQ(99, bounds.max[0]);
+    EXPECT_DOUBLE_EQ(100, bounds.max[1]);
+  }
+
+  void TestWriteAndReadSpaced(ParquetVersion::type version,
+                              ParquetDataPageVersion data_page_version) {
+    this->SetUpSchema(Repetition::OPTIONAL, 1);
+    this->GenerateData(SMALL_SIZE);
+    size_t num_values = this->values_.size();
+
+    std::vector<int16_t> definition_levels(num_values, 1);
+    std::vector<int16_t> repetition_levels(num_values, 0);
+    std::vector<size_t> non_null_indices;
+
+    // Replace some of the generated data with NULL
+    for (size_t i = 0; i < num_values; i++) {
+      if (i % 3 == 0) {
+        definition_levels[i] = 0;
+      } else {
+        non_null_indices.push_back(i);
+      }
+    }
+
+    // Construct valid bits using definition levels
+    std::vector<uint8_t> valid_bytes(num_values);
+    std::transform(definition_levels.begin(), definition_levels.end(),
+                   valid_bytes.begin(),
+                   [&](int64_t level) { return static_cast<uint8_t>(level); });
+    std::shared_ptr<Buffer> valid_bits;
+    ASSERT_OK_AND_ASSIGN(valid_bits, ::arrow::internal::BytesToBits(valid_bytes));
+
+    auto writer =
+        this->BuildWriter(num_values, ColumnProperties(), version, data_page_version,
+                          /*enable_checksum*/ false);
+    writer->WriteBatchSpaced(this->values_.size(), definition_levels.data(),
+                             repetition_levels.data(), valid_bits->data(), 0,
+                             this->values_.data());
+
+    writer->Close();
+    this->ReadColumn();
+    size_t expected_values_read = non_null_indices.size();
+    EXPECT_EQ(expected_values_read, values_read_);
+    for (int64_t i = 0; i < values_read_; i++) {
+      const ByteArray& value = this->values_out_[i];
+      double x = 0;
+      double y = 0;
+      EXPECT_TRUE(GetWKBPointCoordinate(value, &x, &y));
+      auto expected_x = static_cast<double>(non_null_indices[i]);
+      auto expected_y = static_cast<double>(non_null_indices[i] + 1);
+      EXPECT_DOUBLE_EQ(expected_x, x);
+      EXPECT_DOUBLE_EQ(expected_y, y);
+    }
+
+    std::shared_ptr<Statistics> statistics = metadata_stats();
+    EXPECT_FALSE(statistics->HasMinMax());
+    EXPECT_TRUE(statistics->HasGeometryStatistics());
+    const GeometryStatistics* geometry_statistics = statistics->geometry_statistics();
+    std::vector<int32_t> geometry_types = geometry_statistics->GetGeometryTypes();
+    EXPECT_EQ(1, geometry_types.size());
+    EXPECT_EQ(1, geometry_types[0]);
+    EXPECT_DOUBLE_EQ(1, geometry_statistics->GetXMin());
+    EXPECT_DOUBLE_EQ(2, geometry_statistics->GetYMin());
+    EXPECT_DOUBLE_EQ(98, geometry_statistics->GetXMax());
+    EXPECT_DOUBLE_EQ(99, geometry_statistics->GetYMax());
+    EXPECT_FALSE(geometry_statistics->HasZ());
+    EXPECT_FALSE(geometry_statistics->HasM());
+  }
+};
+
+const char* TestGeometryValuesWriter::CRS =
+    R"({"id": {"authority": "OGC", "code": "CRS84"}})";
+const char* TestGeometryValuesWriter::METADATA = "test_metadata";
+
+TEST_F(TestGeometryValuesWriter, TestWriteAndReadV1) {
+  for (auto data_page_version :
+       {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) {
+    TestWriteAndRead(ParquetVersion::PARQUET_1_0, data_page_version);
+  }
+}
+
+TEST_F(TestGeometryValuesWriter, TestWriteAndReadV2) {
+  for (auto data_page_version :
+       {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) {
+    TestWriteAndRead(ParquetVersion::PARQUET_2_4, data_page_version);
+  }
+}
+
+TEST_F(TestGeometryValuesWriter, TestWriteAndReadV1Spaced) {
+  for (auto data_page_version :
+       {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) {
+    TestWriteAndReadSpaced(ParquetVersion::PARQUET_1_0, data_page_version);
+  }
+}
+
+TEST_F(TestGeometryValuesWriter, TestWriteAndReadV2Spaced) {
+  for (auto data_page_version :
+       {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) {
+    TestWriteAndReadSpaced(ParquetVersion::PARQUET_2_4, data_page_version);
+  }
+}
+
 }  // namespace test
 }  // namespace parquet