Proof-of-concept Parquet GEOMETRY logical type implementation

cpp/src/parquet/CMakeLists.txt

@@ -173,6 +173,7 @@ set(PARQUET_SRCS
     exception.cc
     file_reader.cc
     file_writer.cc
+    geometry_statistics.cc
     level_comparison.cc
     level_conversion.cc
     metadata.cc
@@ -375,6 +376,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_writer.cc

@@ -780,6 +780,9 @@ class ColumnWriterImpl {
   // Plain-encoded statistics of the whole chunk
   virtual EncodedStatistics GetChunkStatistics() = 0;
 
+  // Plain-encoded geometry statistics of the whole chunk
+  virtual EncodedGeometryStatistics GetChunkGeometryStatistics() = 0;
+
   // Merges page statistics into chunk statistics, then resets the values
   virtual void ResetPageStatistics() = 0;
 
@@ -1092,6 +1095,14 @@ int64_t ColumnWriterImpl::Close() {
     if (rows_written_ > 0 && chunk_statistics.is_set()) {
       metadata_->SetStatistics(chunk_statistics);
     }
+
+    if (descr_->logical_type() != nullptr && descr_->logical_type()->is_geometry()) {
+      EncodedGeometryStatistics geometry_stats = GetChunkGeometryStatistics();
+      if (geometry_stats.is_set()) {
+        metadata_->SetGeometryStatistics(geometry_stats);
+      }
+    }
+
     metadata_->SetKeyValueMetadata(key_value_metadata_);
     pager_->Close(has_dictionary_, fallback_);
   }
@@ -1212,10 +1223,14 @@ 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());
 
-    if (properties->statistics_enabled(descr_->path()) &&
-        (SortOrder::UNKNOWN != descr_->sort_order())) {
-      page_statistics_ = MakeStatistics<DType>(descr_, allocator_);
-      chunk_statistics_ = MakeStatistics<DType>(descr_, allocator_);
+    if (properties->statistics_enabled(descr_->path())) {
+      if (SortOrder::UNKNOWN != descr_->sort_order()) {
+        page_statistics_ = MakeStatistics<DType>(descr_, allocator_);
+        chunk_statistics_ = MakeStatistics<DType>(descr_, allocator_);
+      }
+      if (descr_->logical_type() != nullptr && descr_->logical_type()->is_geometry()) {
+        chunk_geometry_statistics_ = std::make_shared<GeometryStatistics>();
+      }
     }
     pages_change_on_record_boundaries_ =
         properties->data_page_version() == ParquetDataPageVersion::V2 ||
@@ -1363,6 +1378,12 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
     return result;
   }
 
+  EncodedGeometryStatistics GetChunkGeometryStatistics() override {
+    EncodedGeometryStatistics result;
+    if (chunk_geometry_statistics_) result = chunk_geometry_statistics_->Encode();
+    return result;
+  }
+
   void ResetPageStatistics() override {
     if (chunk_statistics_ != nullptr) {
       chunk_statistics_->Merge(*page_statistics_);
@@ -1421,6 +1442,7 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
   DictEncoder<DType>* current_dict_encoder_;
   std::shared_ptr<TypedStats> page_statistics_;
   std::shared_ptr<TypedStats> chunk_statistics_;
+  std::shared_ptr<GeometryStatistics> chunk_geometry_statistics_;
   bool pages_change_on_record_boundaries_;
 
   // If writing a sequence of ::arrow::DictionaryArray to the writer, we keep the
@@ -1607,6 +1629,11 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
     if (page_statistics_ != nullptr) {
       page_statistics_->Update(values, num_values, num_nulls);
     }
+    if constexpr (std::is_same<T, ByteArray>::value) {
+      if (chunk_geometry_statistics_ != nullptr) {
+        chunk_geometry_statistics_->Update(values, num_values, num_nulls);
+      }
+    }
   }
 
   /// \brief Write values with spaces and update page statistics accordingly.
@@ -1635,6 +1662,13 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
       page_statistics_->UpdateSpaced(values, valid_bits, valid_bits_offset,
                                      num_spaced_values, num_values, num_nulls);
     }
+    if constexpr (std::is_same<T, ByteArray>::value) {
+      if (chunk_geometry_statistics_ != nullptr) {
+        chunk_geometry_statistics_->UpdateSpaced(values, valid_bits, valid_bits_offset,
+                                                 num_spaced_values, num_values,
+                                                 num_nulls);
+      }
+    }
   }
 };
 
@@ -1712,6 +1746,12 @@ Status TypedColumnWriterImpl<DType>::WriteArrowDictionary(
     page_statistics_->IncrementNullCount(num_chunk_levels - non_null_count);
     page_statistics_->IncrementNumValues(non_null_count);
     page_statistics_->Update(*referenced_dictionary, /*update_counts=*/false);
+
+    if constexpr (std::is_same<T, ByteArray>::value) {
+      if (chunk_geometry_statistics_ != nullptr) {
+        chunk_geometry_statistics_->Update(*referenced_dictionary);
+      }
+    }
   };
 
   int64_t value_offset = 0;
@@ -2215,6 +2255,9 @@ Status TypedColumnWriterImpl<ByteArrayType>::WriteArrowDense(
       page_statistics_->IncrementNullCount(batch_size - non_null);
       page_statistics_->IncrementNumValues(non_null);
     }
+    if (chunk_geometry_statistics_ != nullptr) {
+      chunk_geometry_statistics_->Update(*data_slice);
+    }
     CommitWriteAndCheckPageLimit(batch_size, batch_num_values, batch_size - non_null,
                                  check_page);
     CheckDictionarySizeLimit();

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<GeometryStatistics> metadata_geometry_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->geometry_statistics();
+  }
+
  protected:
   int64_t values_read_;
   // Keep the reader alive as for ByteArray the lifetime of the ByteArray
@@ -1774,5 +1782,165 @@ TEST_F(TestInt32Writer, WriteKeyValueMetadataEndToEnd) {
   ASSERT_EQ("bar", value);
 }
 
+// Test writing and reading geometry columns
+class TestGeometryValuesWriter : public TestPrimitiveWriter<ByteArrayType> {
+ public:
+  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("OGC:CRS84", LogicalType::GeometryEdges::PLANAR,
+                                    LogicalType::GeometryEncoding::WKB);
+      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 * kWkbPointSize);
+    uint8_t* ptr = buffer_.data();
+    for (int k = 0; k < num_values; k++) {
+      GenerateWKBPoint(ptr, k, k + 1);
+      values_[k].len = kWkbPointSize;
+      values_[k].ptr = ptr;
+      ptr += kWkbPointSize;
+    }
+
+    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<GeometryStatistics> geometry_statistics = metadata_geometry_stats();
+    ASSERT_TRUE(geometry_statistics != nullptr);
+    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());
+  }
+
+  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<GeometryStatistics> geometry_statistics = metadata_geometry_stats();
+    ASSERT_TRUE(geometry_statistics != nullptr);
+    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());
+  }
+};
+
+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