diff --git a/Cargo.toml b/Cargo.toml index e51f4ddea..a28521594 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,5 +61,5 @@ name = "datafusion_python" crate-type = ["cdylib", "rlib"] [profile.release] -lto = true -codegen-units = 1 +lto = "thin" +#codegen-units = 1 diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 61cb09438..2478a01ce 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -26,6 +26,7 @@ TYPE_CHECKING, Any, Iterable, + List, Literal, Optional, Union, @@ -56,6 +57,7 @@ from enum import Enum +from datafusion._internal import InsertOp # excerpt from deltalake # https://github.com/apache/datafusion-python/pull/981#discussion_r1905619163 @@ -875,6 +877,7 @@ def except_all(self, other: DataFrame) -> DataFrame: """ return DataFrame(self.df.except_all(other.df)) + @overload def write_csv(self, path: str | pathlib.Path, with_header: bool = False) -> None: """Execute the :py:class:`DataFrame` and write the results to a CSV file. @@ -883,6 +886,19 @@ def write_csv(self, path: str | pathlib.Path, with_header: bool = False) -> None with_header: If true, output the CSV header row. """ self.df.write_csv(str(path), with_header) + + @overload + def write_csv(self, path: str | pathlib.Path, with_header: bool = False, insert_operation: InsertOp = InsertOp.Append, single_file_output: bool = False, partition_by: Optional[List[str]] = None,) -> None: + """Execute the :py:class:`DataFrame` and write the results to a CSV file. + + Args: + path: Path of the CSV file to write. + with_header: If true, output the CSV header row. + insert_operation: The operation to perform on the CSV file(Append, Overwrite, Replace). + single_file_output: If true, write the CSV file as a single file. + partition_by: The columns to partition the CSV file by. + """ + self.df.write_csv(str(path), with_header, insert_operation, single_file_output, partition_by or []) @overload def write_parquet( @@ -911,8 +927,11 @@ def write_parquet( def write_parquet( self, path: str | pathlib.Path, - compression: Union[str, Compression, ParquetWriterOptions] = Compression.ZSTD, + compression: Union[str, Compression] = Compression.ZSTD, compression_level: int | None = None, + insert_operation: InsertOp = InsertOp.Append, + single_file_output: bool = False, + partition_by: Optional[List[str]] = None, ) -> None: """Execute the :py:class:`DataFrame` and write the results to a Parquet file. @@ -931,12 +950,16 @@ def write_parquet( compression_level: Compression level to use. For ZSTD, the recommended range is 1 to 22, with the default being 4. Higher levels provide better compression but slower speed. + insert_operation: The operation to perform on the Parquet file(Append, Overwrite, Replace). + single_file_output: If true, write the Parquet file as a single file. + partition_by: The columns to partition the Parquet file by. """ + if isinstance(compression, ParquetWriterOptions): if compression_level is not None: msg = "compression_level should be None when using ParquetWriterOptions" raise ValueError(msg) - self.write_parquet_with_options(path, compression) + self.write_parquet_with_options(path, compression, insert_operation, single_file_output, partition_by or []) return if isinstance(compression, str): @@ -948,10 +971,14 @@ def write_parquet( ): compression_level = compression.get_default_level() - self.df.write_parquet(str(path), compression.value, compression_level) + self.df.write_parquet(str(path), compression.value, compression_level, insert_operation, single_file_output, partition_by or []) def write_parquet_with_options( - self, path: str | pathlib.Path, options: ParquetWriterOptions + self, path: str | pathlib.Path, + options: ParquetWriterOptions, + insert_operation: InsertOp = InsertOp.Append, + single_file_output: bool = False, + partition_by: Optional[List[str]] = None, ) -> None: """Execute the :py:class:`DataFrame` and write the results to a Parquet file. @@ -1000,15 +1027,21 @@ def write_parquet_with_options( str(path), options_internal, column_specific_options_internal, + insert_operation, + single_file_output, + partition_by, ) - def write_json(self, path: str | pathlib.Path) -> None: + def write_json(self, path: str | pathlib.Path, insert_operation: InsertOp = InsertOp.Append, single_file_output: bool = False, partition_by: Optional[List[str]] = None) -> None: """Execute the :py:class:`DataFrame` and write the results to a JSON file. Args: path: Path of the JSON file to write. + insert_operation: The operation to perform on the JSON file(Append, Overwrite, Replace). + single_file_output: If true, write the JSON file as a single file. + partition_by: The columns to partition the JSON file by. """ - self.df.write_json(str(path)) + self.df.write_json(str(path), insert_operation, single_file_output, partition_by or []) def to_arrow_table(self) -> pa.Table: """Execute the :py:class:`DataFrame` and convert it into an Arrow Table. diff --git a/python/datafusion/object_store.py b/python/datafusion/object_store.py index 6298526f5..b6cba8c12 100644 --- a/python/datafusion/object_store.py +++ b/python/datafusion/object_store.py @@ -24,4 +24,7 @@ MicrosoftAzure = object_store.MicrosoftAzure Http = object_store.Http -__all__ = ["AmazonS3", "GoogleCloud", "Http", "LocalFileSystem", "MicrosoftAzure"] +RetryConfig = object_store.RetryConfig +ClientOptions = object_store.ClientOptions + +__all__ = ["AmazonS3", "GoogleCloud", "Http", "LocalFileSystem", "MicrosoftAzure", "RetryConfig", "ClientOptions"] diff --git a/src/dataframe.rs b/src/dataframe.rs index 1437f5f82..19da41a99 100644 --- a/src/dataframe.rs +++ b/src/dataframe.rs @@ -34,6 +34,7 @@ use datafusion::dataframe::{DataFrame, DataFrameWriteOptions}; use datafusion::datasource::TableProvider; use datafusion::error::DataFusionError; use datafusion::execution::SendableRecordBatchStream; +use datafusion::logical_expr::dml::InsertOp; use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use datafusion::prelude::*; use datafusion_ffi::table_provider::FFI_TableProvider; @@ -58,6 +59,27 @@ use crate::{ expr::{sort_expr::PySortExpr, PyExpr}, }; +#[derive(Clone, Copy, PartialEq)] +#[pyclass(name = "InsertOp", module = "datafusion", eq, eq_int)] +pub enum PyInsertOp { + #[pyo3(name = "Append")] + Append, + #[pyo3(name = "Overwrite")] + Overwrite, + #[pyo3(name = "Replace")] + Replace, +} + +impl From for InsertOp { + fn from(op: PyInsertOp) -> Self { + match op { + PyInsertOp::Append => InsertOp::Append, + PyInsertOp::Overwrite => InsertOp::Overwrite, + PyInsertOp::Replace => InsertOp::Replace, + } + } +} + // https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116 // - we have not decided on the table_provider approach yet // this is an interim implementation @@ -743,19 +765,22 @@ impl PyDataFrame { } /// Write a `DataFrame` to a CSV file. - fn write_csv(&self, path: &str, with_header: bool, py: Python) -> PyDataFusionResult<()> { + fn write_csv(&self, path: &str, with_header: bool, insert_operation: PyInsertOp, single_file_output: bool, partition_by: Vec, py: Python) -> PyDataFusionResult<()> { let csv_options = CsvOptions { has_header: Some(with_header), ..Default::default() }; - wait_for_future( + let _ = wait_for_future( py, self.df.as_ref().clone().write_csv( path, - DataFrameWriteOptions::new(), + DataFrameWriteOptions::new() + .with_insert_operation(insert_operation.into()) + .with_single_file_output(single_file_output) + .with_partition_by(partition_by), Some(csv_options), ), - )??; + )?; Ok(()) } @@ -763,13 +788,19 @@ impl PyDataFrame { #[pyo3(signature = ( path, compression="zstd", - compression_level=None + compression_level=None, + insert_operation=PyInsertOp::Append, + single_file_output=false, + partition_by=vec![], ))] fn write_parquet( &self, path: &str, compression: &str, compression_level: Option, + insert_operation: PyInsertOp, + single_file_output: bool, + partition_by: Vec, py: Python, ) -> PyDataFusionResult<()> { fn verify_compression_level(cl: Option) -> Result { @@ -813,7 +844,10 @@ impl PyDataFrame { py, self.df.as_ref().clone().write_parquet( path, - DataFrameWriteOptions::new(), + DataFrameWriteOptions::new() + .with_insert_operation(insert_operation.into()) + .with_single_file_output(single_file_output) + .with_partition_by(partition_by), Option::from(options), ), )??; @@ -821,11 +855,22 @@ impl PyDataFrame { } /// Write a `DataFrame` to a Parquet file, using advanced options. + #[pyo3(signature = ( + path, + options, + column_specific_options, + insert_operation=PyInsertOp::Append, + single_file_output=false, + partition_by=vec![], + ))] fn write_parquet_with_options( &self, path: &str, options: PyParquetWriterOptions, column_specific_options: HashMap, + insert_operation: PyInsertOp, + single_file_output: bool, + partition_by: Vec, py: Python, ) -> PyDataFusionResult<()> { let table_options = TableParquetOptions { @@ -841,7 +886,10 @@ impl PyDataFrame { py, self.df.as_ref().clone().write_parquet( path, - DataFrameWriteOptions::new(), + DataFrameWriteOptions::new() + .with_insert_operation(insert_operation.into()) + .with_single_file_output(single_file_output) + .with_partition_by(partition_by), Option::from(table_options), ), )??; @@ -849,14 +897,18 @@ impl PyDataFrame { } /// Executes a query and writes the results to a partitioned JSON file. - fn write_json(&self, path: &str, py: Python) -> PyDataFusionResult<()> { - wait_for_future( + fn write_json(&self, path: &str, insert_operation: PyInsertOp, single_file_output: bool, partition_by: Vec, py: Python) -> PyDataFusionResult<()> { + let _ = wait_for_future( py, self.df .as_ref() .clone() - .write_json(path, DataFrameWriteOptions::new(), None), - )??; + .write_json(path, DataFrameWriteOptions::new() + .with_insert_operation(insert_operation.into()) + .with_single_file_output(single_file_output) + .with_partition_by(partition_by), + None), + )?; Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 29d3f41da..799161876 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -97,6 +97,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; let catalog = PyModule::new(py, "catalog")?; catalog::init_module(&catalog)?; diff --git a/src/store.rs b/src/store.rs index 1e5fab472..ffcd2774a 100644 --- a/src/store.rs +++ b/src/store.rs @@ -16,6 +16,9 @@ // under the License. use std::sync::Arc; +use std::time::Duration; + +use object_store::{ClientOptions, RetryConfig}; use pyo3::prelude::*; @@ -27,6 +30,99 @@ use object_store::local::LocalFileSystem; use pyo3::exceptions::PyValueError; use url::Url; +#[pyclass(name = "RetryConfig", module = "datafusion.store", subclass)] +#[derive(FromPyObject)] +pub struct PyRetryConfig { + pub max_retries: usize, + pub retry_timeout: u64, +} + +#[pymethods] +impl PyRetryConfig { + #[new] + #[pyo3(signature = (max_retries=None, retry_timeout=None))] + fn py_new(max_retries: Option, retry_timeout: Option) -> Self { + let default_retry = RetryConfig::default(); + Self { + max_retries: max_retries.unwrap_or(default_retry.max_retries), + retry_timeout: retry_timeout.unwrap_or(default_retry.retry_timeout.as_secs()), + } + } + + #[getter] + fn get_max_retries(&self) -> usize { + self.max_retries + } + + #[getter] + fn get_retry_timeout(&self) -> u64 { + self.retry_timeout + } +} + +impl From for RetryConfig { + fn from(config: PyRetryConfig) -> Self { + RetryConfig{max_retries:config.max_retries, retry_timeout:Duration::from_secs(config.retry_timeout), ..Default::default()} + } +} + +impl Default for PyRetryConfig { + fn default() -> Self { + let default_retry = RetryConfig::default(); + Self { + max_retries: default_retry.max_retries, + retry_timeout: default_retry.retry_timeout.as_secs(), + } + } +} + +#[pyclass(name = "ClientOptions", module = "datafusion.store", subclass)] +#[derive(FromPyObject)] +pub struct PyClientOptions { + pub connect_timeout: Duration, + pub timeout: Duration, +} + +impl Default for PyClientOptions { + fn default() -> Self { + Self { + connect_timeout: Duration::from_secs(30), + timeout: Duration::from_secs(5), + } + } +} + +#[pymethods] +impl PyClientOptions { + #[new] + #[pyo3(signature = (connect_timeout=None, timeout=None))] + fn py_new(connect_timeout: Option, timeout: Option) -> Self { + let default_client_options = Self::default(); + Self { + connect_timeout: connect_timeout.map_or(default_client_options.connect_timeout, Duration::from_secs), + timeout: timeout.map_or(default_client_options.timeout, Duration::from_secs), + } + } + + #[getter] + fn get_connect_timeout(&self) -> Duration { + self.connect_timeout + } + + #[getter] + fn get_timeout(&self) -> Duration { + self.timeout + } +} + +impl From for ClientOptions { + fn from(options: PyClientOptions) -> Self { + ClientOptions::new() + .with_connect_timeout(options.connect_timeout) + .with_timeout(options.timeout) + } +} + #[derive(FromPyObject)] pub enum StorageContexts { AmazonS3(PyAmazonS3Context), @@ -174,7 +270,7 @@ pub struct PyAmazonS3Context { #[pymethods] impl PyAmazonS3Context { #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (bucket_name, region=None, access_key_id=None, secret_access_key=None, endpoint=None, allow_http=false, imdsv1_fallback=false))] + #[pyo3(signature = (bucket_name, region=None, access_key_id=None, secret_access_key=None, endpoint=None, retry_config=None, client_options=None, allow_http=false, imdsv1_fallback=false))] #[new] fn new( bucket_name: String, @@ -182,7 +278,8 @@ impl PyAmazonS3Context { access_key_id: Option, secret_access_key: Option, endpoint: Option, - //retry_config: RetryConfig, + retry_config: Option, + client_options: Option, allow_http: bool, imdsv1_fallback: bool, ) -> Self { @@ -211,7 +308,8 @@ impl PyAmazonS3Context { let store = builder .with_bucket_name(bucket_name.clone()) - //.with_retry_config(retry_config) #TODO: add later + .with_retry(retry_config.map_or(RetryConfig::default(), RetryConfig::from)) + .with_client_options(client_options.map_or(ClientOptions::default(), ClientOptions::from)) .with_allow_http(allow_http) .build() .expect("failed to build AmazonS3"); @@ -255,5 +353,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) }