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