Support parquet v2 format

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

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",
         ],
     );
@@ -156,6 +158,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>) {
@@ -313,6 +329,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,
     },
 };
 
@@ -66,6 +66,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).as_pg_cstr(),
@@ -76,6 +77,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) };

src/pgrx_tests/copy_options.rs

@@ -298,6 +298,22 @@ mod tests {
         test_table.assert_expected_and_result_rows();
     }
 
+    #[pg_test]
+    #[should_panic(expected = "unrecognized parquet version: vv2. v1 or v2 is supported.")]
+    fn test_invalid_parquet_version() {
+        let _file_cleanup = FileCleanup::new(LOCAL_TEST_FILE_PATH);
+
+        let mut copy_options = HashMap::new();
+        copy_options.insert(
+            "parquet_version".to_string(),
+            CopyOptionValue::StringOption("vv2".into()),
+        );
+
+        let test_table = TestTable::<i32>::new("int4".into()).with_copy_to_options(copy_options);
+        test_table.insert("INSERT INTO test_expected (a) VALUES (1), (2), (null);");
+        test_table.assert_expected_and_result_rows();
+    }
+
     #[pg_test]
     fn test_large_arrow_array_limit() {
         // disable row group size bytes limit
@@ -412,6 +428,62 @@ mod tests {
         assert_eq!(result_metadata, vec![12]);
     }
 
+    #[pg_test]
+    fn test_parquet_version() {
+        let fetch_parquet_version = || {
+            Spi::connect(|client| {
+                let parquet_file_metadata_command = format!(
+                    "select * from parquet.file_metadata('{}');",
+                    LOCAL_TEST_FILE_PATH
+                );
+
+                let mut results = Vec::new();
+                let tup_table = client
+                    .select(&parquet_file_metadata_command, None, &[])
+                    .unwrap();
+
+                for row in tup_table {
+                    let parquet_version = row["format_version"].value::<String>().unwrap().unwrap();
+                    results.push(parquet_version);
+                }
+
+                results
+            })
+        };
+
+        // try writing/reading from parquet v2 file
+        let mut copy_to_options = HashMap::new();
+        copy_to_options.insert(
+            "parquet_version".to_string(),
+            CopyOptionValue::StringOption("v2".to_string()),
+        );
+
+        let test_table = TestTable::<i32>::new("int4".into()).with_copy_to_options(copy_to_options);
+        test_table.insert("INSERT INTO test_expected (a) select i from generate_series(1,10) i;");
+        test_table.assert_expected_and_result_rows();
+
+        // v2
+        let copy_to_parquet = format!(
+            "copy (select 1 as id) to '{LOCAL_TEST_FILE_PATH}' with (parquet_version 'v2');",
+        );
+        Spi::run(&copy_to_parquet).unwrap();
+        assert_eq!(fetch_parquet_version(), vec!["2".to_string()]);
+
+        // v1
+        let copy_to_parquet = format!(
+            "copy (select 1 as id) to '{LOCAL_TEST_FILE_PATH}' with (parquet_version 'v1');",
+        );
+        Spi::run(&copy_to_parquet).unwrap();
+        assert_eq!(fetch_parquet_version(), vec!["1".to_string()]);
+
+        // V1
+        let copy_to_parquet = format!(
+            "copy (select 1 as id) to '{LOCAL_TEST_FILE_PATH}' with (parquet_version 'V1');",
+        );
+        Spi::run(&copy_to_parquet).unwrap();
+        assert_eq!(fetch_parquet_version(), vec!["1".to_string()]);
+    }
+
     #[pg_test]
     fn test_file_size_bytes() {
         let uris = [