apacheGH-35331: [Python] Expose Parquet sorting metadata (apache#37665)

### Rationale for this change Picking up where apache#35453 left off. Closes apache#35331 This PR builds on top of apache#37469 ### What changes are included in this PR? ### Are these changes tested? ### Are there any user-facing changes? * Closes: apache#35331 Lead-authored-by: Judah Rand <[email protected]> Co-authored-by: Will Jones <[email protected]> Signed-off-by: AlenkaF <[email protected]>
clayburn · Jan 23, 2024 · f63ad62 · f63ad62
1 parent fa7b811
commit f63ad62
Show file tree

Hide file tree

Showing 6 changed files with 394 additions and 13 deletions.
diff --git a/docs/source/python/api/formats.rst b/docs/source/python/api/formats.rst
@@ -97,6 +97,7 @@ Parquet Metadata
 
    FileMetaData
    RowGroupMetaData
+   SortingColumn
    ColumnChunkMetaData
    Statistics
    ParquetSchema

diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx
@@ -609,6 +609,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions):
             dictionary_pagesize_limit=self._properties["dictionary_pagesize_limit"],
             write_page_index=self._properties["write_page_index"],
             write_page_checksum=self._properties["write_page_checksum"],
+            sorting_columns=self._properties["sorting_columns"],
         )
 
     def _set_arrow_properties(self):
@@ -659,6 +660,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions):
             write_page_index=False,
             encryption_config=None,
             write_page_checksum=False,
+            sorting_columns=None,
         )
 
         self._set_properties()

diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd
@@ -328,11 +328,17 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil:
         optional[ParquetIndexLocation] GetColumnIndexLocation() const
         optional[ParquetIndexLocation] GetOffsetIndexLocation() const
 
+    struct CSortingColumn" parquet::SortingColumn":
+        int column_idx
+        c_bool descending
+        c_bool nulls_first
+
     cdef cppclass CRowGroupMetaData" parquet::RowGroupMetaData":
         c_bool Equals(const CRowGroupMetaData&) const
-        int num_columns()
-        int64_t num_rows()
-        int64_t total_byte_size()
+        int num_columns() const
+        int64_t num_rows() const
+        int64_t total_byte_size() const
+        vector[CSortingColumn] sorting_columns() const
         unique_ptr[CColumnChunkMetaData] ColumnChunk(int i) const
 
     cdef cppclass CFileMetaData" parquet::FileMetaData":
@@ -421,6 +427,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
             Builder* disable_dictionary()
             Builder* enable_dictionary()
             Builder* enable_dictionary(const c_string& path)
+            Builder* set_sorting_columns(vector[CSortingColumn] sorting_columns)
             Builder* disable_statistics()
             Builder* enable_statistics()
             Builder* enable_statistics(const c_string& path)
@@ -517,8 +524,8 @@ cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil:
 
     CStatus ToParquetSchema(
         const CSchema* arrow_schema,
-        const ArrowReaderProperties& properties,
-        const shared_ptr[const CKeyValueMetadata]& key_value_metadata,
+        const WriterProperties& properties,
+        const ArrowWriterProperties& arrow_properties,
         shared_ptr[SchemaDescriptor]* out)
 
 
@@ -584,7 +591,9 @@ cdef shared_ptr[WriterProperties] _create_writer_properties(
     write_batch_size=*,
     dictionary_pagesize_limit=*,
     write_page_index=*,
-    write_page_checksum=*) except *
+    write_page_checksum=*,
+    sorting_columns=*,
+) except *
 
 
 cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties(
@@ -593,7 +602,8 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties(
     allow_truncated_timestamps=*,
     writer_engine_version=*,
     use_compliant_nested_type=*,
-    store_schema=*) except *
+    store_schema=*,
+) except *
 
 cdef class ParquetSchema(_Weakrefable):
     cdef: