Support parquet v2 format (#127)

Supports reading or writing parquet files with format v2. By default, we write files in v1 format to improve interop with common query engines. But you can specify COPY table TO '..' with (parquet_version 'v2').

Closes #104.

Author: aykut-bozkurt

.github/workflows/ci.yml

@@ -23,9 +23,9 @@ jobs:
         runs_on: [ 'ubuntu-22.04', 'ubuntu-22.04-arm' ]
         include:
           - runs_on: ubuntu-22.04
-            arch: x64
+            arch: x86_64
           - runs_on: ubuntu-22.04-arm
-            arch: arm64
+            arch: aarch64
 
     runs-on: ${{ matrix.runs_on }}
 
@@ -36,44 +36,28 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up sccache for x86_64
-        if: ${{ matrix.arch == 'x64' }}
+      - name: Set up sccache
         run: |
-          wget https://github.com/mozilla/sccache/releases/download/v$SCCACHE_VERSION/sccache-v$SCCACHE_VERSION-x86_64-unknown-linux-musl.tar.gz
-          tar -xzf sccache-v$SCCACHE_VERSION-x86_64-unknown-linux-musl.tar.gz
-          sudo mv sccache-v$SCCACHE_VERSION-x86_64-unknown-linux-musl/sccache /usr/local/bin
+          wget https://github.com/mozilla/sccache/releases/download/v$SCCACHE_VERSION/sccache-v$SCCACHE_VERSION-${{ matrix.arch }}-unknown-linux-musl.tar.gz
+          tar -xzf sccache-v$SCCACHE_VERSION-${{ matrix.arch }}-unknown-linux-musl.tar.gz
+          sudo mv sccache-v$SCCACHE_VERSION-${{ matrix.arch }}-unknown-linux-musl/sccache /usr/local/bin
           chmod +x /usr/local/bin/sccache
+          if [ "${{ matrix.arch }}" == "x86_64" ]; then
+            SCCACHE_SHA256=$SCCACHE_x86_64_SHA256
+          else
+            SCCACHE_SHA256=$SCCACHE_AARCH64_SHA256
+          fi
           echo "$SCCACHE_SHA256  /usr/local/bin/sccache" | sha256sum --check
         env:
           SCCACHE_VERSION: 0.8.1
-          SCCACHE_SHA256: "7203a4dcb3a67f3a0272366d50ede22e5faa3e2a798deaa4d1ea377b51c0ab0c"
+          SCCACHE_x86_64_SHA256: "7203a4dcb3a67f3a0272366d50ede22e5faa3e2a798deaa4d1ea377b51c0ab0c"
+          SCCACHE_AARCH64_SHA256: "36b2fd1c6c3a104ec1d526edb0533a3827c266054bf4552fb97f524beff6a612"
 
-      - name: Set up sccache for arm64
-        if: ${{ matrix.arch == 'arm64' }}
-        run: |
-          wget https://github.com/mozilla/sccache/releases/download/v$SCCACHE_VERSION/sccache-v$SCCACHE_VERSION-aarch64-unknown-linux-musl.tar.gz
-          tar -xzf sccache-v$SCCACHE_VERSION-aarch64-unknown-linux-musl.tar.gz
-          sudo mv sccache-v$SCCACHE_VERSION-aarch64-unknown-linux-musl/sccache /usr/local/bin
-          chmod +x /usr/local/bin/sccache
-          echo "$SCCACHE_SHA256  /usr/local/bin/sccache" | sha256sum --check
-        env:
-          SCCACHE_VERSION: 0.8.1
-          SCCACHE_SHA256: "36b2fd1c6c3a104ec1d526edb0533a3827c266054bf4552fb97f524beff6a612"
-
-      - name: Set up Rust for x86_64
-        if: ${{ matrix.arch == 'x64' }}
-        uses: dtolnay/rust-toolchain@stable
-        with:
-          toolchain: 1.88.0
-          target: x86_64-unknown-linux-gnu
-          components: rustfmt, clippy, llvm-tools-preview
-
-      - name: Set up Rust for arm64
-        if: ${{ matrix.arch == 'arm64' }}
+      - name: Set up Rust
         uses: dtolnay/rust-toolchain@stable
         with:
-          toolchain: 1.88.0
-          target: aarch64-unknown-linux-gnu
+          toolchain: 1.89.0
+          target: ${{ matrix.arch }}-unknown-linux-gnu
           components: rustfmt, clippy, llvm-tools-preview
 
       - name: Cache cargo registry
@@ -150,17 +134,17 @@ jobs:
           make check-lint
 
       - name: Run tests without coverage
-        if: ${{ env.PG_MAJOR != '17' || matrix.arch != 'x64' }}
+        if: ${{ env.PG_MAJOR != '17' || matrix.arch != 'x86_64' }}
         run: |
           make check
 
       - name: Run tests with coverage
-        if: ${{ env.PG_MAJOR == '17' && matrix.arch == 'x64' }}
+        if: ${{ env.PG_MAJOR == '17' && matrix.arch == 'x86_64' }}
         run: |
           make check-with-coverage
 
       - name: Upload coverage report to Codecov
-        if: ${{ env.PG_MAJOR == '17' && matrix.arch == 'x64' }}
+        if: ${{ env.PG_MAJOR == '17' && matrix.arch == 'x86_64' }}
         uses: codecov/codecov-action@v4
         with:
           fail_ci_if_error: true

README.md

@@ -310,7 +310,8 @@ Supported Google Cloud Storage uri formats are shown below:
 - `row_group_size <int64>`: the number of rows in each row group while writing Parquet files. The default row group size is `122880`,
 - `row_group_size_bytes <int64>`: the total byte size of rows in each row group while writing Parquet files. The default row group size bytes is `row_group_size * 1024`,
 - `compression <string>`: the compression format to use while writing Parquet files. The supported compression formats are `uncompressed`, `snappy`, `gzip`, `brotli`, `lz4`, `lz4raw` and `zstd`. The default compression format is `snappy`. If not specified, the compression format is determined by the file extension,
-- `compression_level <int>`: the compression level to use while writing Parquet files. The supported compression levels are only supported for `gzip`, `zstd` and `brotli` compression formats. The default compression level is `6` for `gzip (0-10)`, `1` for `zstd (1-22)` and `1` for `brotli (0-11)`.
+- `compression_level <int>`: the compression level to use while writing Parquet files. The supported compression levels are only supported for `gzip`, `zstd` and `brotli` compression formats. The default compression level is `6` for `gzip (0-10)`, `1` for `zstd (1-22)` and `1` for `brotli (0-11)`,
+- `parquet_version <string>`: writer version of the Parquet file. By default, it is set to `v1` to be more interoperable with common query engines. (some are not able to read v2 files) You can set it to `v2` to unlock some of the new encodings.
 
 `pg_parquet` supports the following options in the `COPY FROM` command:
 - `format parquet`: you need to specify this option to read or write Parquet files which does not end with `.parquet[.<compression>]` extension,

src/arrow_parquet.rs

@@ -4,6 +4,7 @@ pub(crate) mod compression;
 pub(crate) mod field_ids;
 pub(crate) mod match_by;
 pub(crate) mod parquet_reader;
+pub(crate) mod parquet_version;
 pub(crate) mod parquet_writer;
 pub(crate) mod pg_to_arrow;
 pub(crate) mod schema_parser;

src/arrow_parquet/parquet_version.rs

@@ -0,0 +1,34 @@
+use std::str::FromStr;
+
+use parquet::file::properties::WriterVersion;
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, Default, PartialEq)]
+pub(crate) enum ParquetVersion {
+    #[default]
+    V1,
+    V2,
+}
+
+impl FromStr for ParquetVersion {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "v1" => Ok(ParquetVersion::V1),
+            "v2" => Ok(ParquetVersion::V2),
+            _ => Err(format!(
+                "unrecognized parquet version: {s}. v1 or v2 is supported.",
+            )),
+        }
+    }
+}
+
+impl From<ParquetVersion> for WriterVersion {
+    fn from(value: ParquetVersion) -> Self {
+        match value {
+            ParquetVersion::V1 => WriterVersion::PARQUET_1_0,
+            ParquetVersion::V2 => WriterVersion::PARQUET_2_0,
+        }
+    }
+}

src/arrow_parquet/parquet_writer.rs

@@ -94,6 +94,7 @@ impl ParquetWriterContext {
             .set_statistics_enabled(EnabledStatistics::Page)
             .set_compression(compression.into())
             .set_max_row_group_size(options.row_group_size as usize)
+            .set_writer_version(options.parquet_version.into())
             .set_created_by("pg_parquet".to_string());
 
         let geometry_columns_metadata_value = geoparquet_metadata_json_from_tupledesc(tupledesc);

src/parquet_copy_hook/copy_to_split_dest_receiver.rs

@@ -9,6 +9,7 @@ use pgrx::prelude::*;
 use crate::arrow_parquet::{
     compression::{PgParquetCompression, INVALID_COMPRESSION_LEVEL},
     field_ids::FieldIds,
+    parquet_version::ParquetVersion,
     parquet_writer::{DEFAULT_ROW_GROUP_SIZE, DEFAULT_ROW_GROUP_SIZE_BYTES},
 };
 
@@ -39,6 +40,7 @@ pub(crate) struct CopyToParquetOptions {
     pub(crate) row_group_size_bytes: i64,
     pub(crate) compression: PgParquetCompression,
     pub(crate) compression_level: i32,
+    pub(crate) parquet_version: ParquetVersion,
 }
 
 impl CopyToParquetSplitDestReceiver {
@@ -212,6 +214,7 @@ pub extern "C-unwind" fn create_copy_to_parquet_split_dest_receiver(
     row_group_size_bytes: *const i64,
     compression: *const PgParquetCompression,
     compression_level: *const i32,
+    parquet_version: *const ParquetVersion,
 ) -> *mut DestReceiver {
     let file_size_bytes = if file_size_bytes.is_null() {
         INVALID_FILE_SIZE_BYTES
@@ -251,13 +254,20 @@ pub extern "C-unwind" fn create_copy_to_parquet_split_dest_receiver(
         unsafe { *compression_level }
     };
 
+    let parquet_version = if parquet_version.is_null() {
+        ParquetVersion::default()
+    } else {
+        unsafe { *parquet_version }
+    };
+
     let options = CopyToParquetOptions {
         file_size_bytes,
         field_ids,
         row_group_size,
         row_group_size_bytes,
         compression,
         compression_level,
+        parquet_version,
     };
 
     let mut split_dest =

src/parquet_copy_hook/copy_utils.rs

@@ -20,6 +20,7 @@ use crate::{
         compression::{all_supported_compressions, PgParquetCompression},
         field_ids,
         match_by::MatchBy,
+        parquet_version::ParquetVersion,
         parquet_writer::{DEFAULT_ROW_GROUP_SIZE, DEFAULT_ROW_GROUP_SIZE_BYTES},
         uri_utils::ParsedUriInfo,
     },
@@ -44,6 +45,7 @@ pub(crate) fn validate_copy_to_options(p_stmt: &PgBox<PlannedStmt>, uri_info: &P
             "row_group_size_bytes",
             "compression",
             "compression_level",
+            "parquet_version",
             "freeze",
         ],
     );
@@ -153,6 +155,20 @@ pub(crate) fn validate_copy_to_options(p_stmt: &PgBox<PlannedStmt>, uri_info: &P
 
         compression.ensure_compression_level(compression_level);
     }
+
+    let parquet_version_option = copy_stmt_get_option(p_stmt, "parquet_version");
+
+    if !parquet_version_option.is_null() {
+        let parquet_version = unsafe { defGetString(parquet_version_option.as_ptr()) };
+
+        let parquet_version = unsafe {
+            CStr::from_ptr(parquet_version)
+                .to_str()
+                .expect("parquet_version option is not a valid CString")
+        };
+
+        ParquetVersion::from_str(parquet_version).unwrap_or_else(|e| panic!("{}", e));
+    }
 }
 
 pub(crate) fn validate_copy_from_options(p_stmt: &PgBox<PlannedStmt>) {
@@ -307,6 +323,23 @@ pub(crate) fn copy_to_stmt_compression_level(
     }
 }
 
+pub(crate) fn copy_to_stmt_parquet_version(p_stmt: &PgBox<PlannedStmt>) -> ParquetVersion {
+    let parquet_version_option = copy_stmt_get_option(p_stmt, "parquet_version");
+
+    if parquet_version_option.is_null() {
+        ParquetVersion::default()
+    } else {
+        let parquet_version = unsafe { defGetString(parquet_version_option.as_ptr()) };
+        let parquet_version = unsafe {
+            CStr::from_ptr(parquet_version)
+                .to_str()
+                .expect("parquet_version option is not a valid CString")
+        };
+
+        ParquetVersion::from_str(parquet_version).unwrap_or_else(|e| panic!("{}", e))
+    }
+}
+
 pub(crate) fn copy_from_stmt_create_option_list(p_stmt: &PgBox<PlannedStmt>) -> PgList<DefElem> {
     let mut new_copy_options = PgList::<DefElem>::new();
 

src/parquet_copy_hook/hook.rs

@@ -24,7 +24,7 @@ use super::{
     copy_to_split_dest_receiver::free_copy_to_parquet_split_dest_receiver,
     copy_utils::{
         copy_to_stmt_compression, copy_to_stmt_field_ids, copy_to_stmt_file_size_bytes,
-        validate_copy_from_options, validate_copy_to_options,
+        copy_to_stmt_parquet_version, validate_copy_from_options, validate_copy_to_options,
     },
 };
 
@@ -64,6 +64,7 @@ fn process_copy_to_parquet(
     let row_group_size_bytes = copy_to_stmt_row_group_size_bytes(p_stmt);
     let compression = copy_to_stmt_compression(p_stmt, &uri_info);
     let compression_level = copy_to_stmt_compression_level(p_stmt, &uri_info);
+    let parquet_version = copy_to_stmt_parquet_version(p_stmt);
 
     let parquet_split_dest = create_copy_to_parquet_split_dest_receiver(
         uri_as_string(&uri_info.uri).as_pg_cstr(),
@@ -74,6 +75,7 @@ fn process_copy_to_parquet(
         &row_group_size_bytes,
         &compression,
         &compression_level.unwrap_or(INVALID_COMPRESSION_LEVEL),
+        &parquet_version,
     );
 
     let parquet_split_dest = unsafe { PgBox::from_pg(parquet_split_dest) };