feat: upgrade dataframe write_parquet and write_json

The options to write_parquet changed.

write_json has a new argument that I defaulted to None. We can expose that config later.

Ref: apache/datafusion#9382

Author: Michael-J-Ward

src/dataframe.rs

@@ -20,11 +20,13 @@ use std::sync::Arc;
 use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
 use datafusion::arrow::util::pretty;
+use datafusion::config::{ParquetOptions, TableParquetOptions};
 use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::parquet::file::properties::WriterProperties;
 use datafusion::prelude::*;
+use datafusion_common::file_options::parquet_writer::ParquetWriterOptions;
 use datafusion_common::UnnestOptions;
 use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::prelude::*;
@@ -375,16 +377,15 @@ impl PyDataFrame {
             }
         };
 
-        let writer_properties = WriterProperties::builder()
-            .set_compression(compression_type)
-            .build();
+        let mut options = TableParquetOptions::default();
+        options.global.compression = Some(compression.to_string());
 
         wait_for_future(
             py,
             self.df.as_ref().clone().write_parquet(
                 path,
                 DataFrameWriteOptions::new(),
-                Option::from(writer_properties),
+                Option::from(options),
             ),
         )?;
         Ok(())
@@ -397,7 +398,7 @@ impl PyDataFrame {
             self.df
                 .as_ref()
                 .clone()
-                .write_json(path, DataFrameWriteOptions::new()),
+                .write_json(path, DataFrameWriteOptions::new(), None),
         )?;
         Ok(())
     }