martin-g commented on code in PR #2039:
URL:
https://github.com/apache/datafusion-ballista/pull/2039#discussion_r3585870554
##########
ballista/core/src/config.rs:
##########
@@ -585,6 +597,21 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED)
}
+ /// Returns compression codec that will be used during write stage of
shuffle
+ pub fn shuffle_compression_codec(
+ &self,
+ ) -> std::result::Result<Option<CompressionType>, String> {
+ match self
+ .get_string_setting(BALLISTA_SHUFFLE_COMPRESSION_CODEC)
+ .as_str()
+ {
+ "lz4" => Ok(Some(CompressionType::LZ4_FRAME)),
+ "zstd" => Ok(Some(CompressionType::ZSTD)),
+ "none" => Ok(None),
+ other => Err(format!("Invalid compression codec: {}", other)),
Review Comment:
```suggestion
other => Err(format!("Invalid compression codec: {other}")),
```
nit: In my other projects I use more strict lint rules and this is reported
as a warning. It would be nice to be prepared if this lint rule is enabled here
too one day.
##########
ballista/core/src/utils.rs:
##########
@@ -177,6 +177,23 @@ pub fn default_config_producer() -> SessionConfig {
SessionConfig::new_with_ballista()
}
+/// Creates [IpcWriteOptions] using the compression codec configured in the
BallistaConfig
+pub fn create_new_write_options(
+ config: Arc<BallistaConfig>,
+) -> std::result::Result<IpcWriteOptions, DataFusionError> {
+ let compression_codec = config
+ .shuffle_compression_codec()
+ .map_err(DataFusionError::Configuration)?;
+
+ let options = IpcWriteOptions::default()
+ .try_with_compression(compression_codec)
+ .map_err(|err| {
+ DataFusionError::Internal(format!("Failed to set compression
codec: {}", err))
Review Comment:
```suggestion
DataFusionError::Internal(format!("Failed to set compression
codec: {err}"))
```
##########
ballista/core/src/config.rs:
##########
@@ -585,6 +597,21 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED)
}
+ /// Returns compression codec that will be used during write stage of
shuffle
+ pub fn shuffle_compression_codec(
+ &self,
+ ) -> std::result::Result<Option<CompressionType>, String> {
+ match self
+ .get_string_setting(BALLISTA_SHUFFLE_COMPRESSION_CODEC)
+ .as_str()
+ {
+ "lz4" => Ok(Some(CompressionType::LZ4_FRAME)),
+ "zstd" => Ok(Some(CompressionType::ZSTD)),
+ "none" => Ok(None),
+ other => Err(format!("Invalid compression codec: {}", other)),
Review Comment:
Also, I think it would be nice to list the supported values in the error
message.
##########
ballista/core/src/config.rs:
##########
@@ -334,6 +338,14 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
DataType::Utf8,
Some("".to_string()),
),
+ ConfigEntry::new(
+ BALLISTA_SHUFFLE_COMPRESSION_CODEC.to_string(),
+ "Compression codec specification \
+ used in the shuffle process. Possible values: \
+ none, lz4, zstd. Defaults to lz4 to preserve current
behaviour".to_string(),
+ DataType::Utf8,
+ Some("lz4".to_string()),
+ )
Review Comment:
```suggestion
),
```
##########
ballista/core/src/config.rs:
##########
@@ -585,6 +597,21 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_PROPAGATE_EMPTY_ENABLED)
}
+ /// Returns compression codec that will be used during write stage of
shuffle
+ pub fn shuffle_compression_codec(
+ &self,
+ ) -> std::result::Result<Option<CompressionType>, String> {
+ match self
+ .get_string_setting(BALLISTA_SHUFFLE_COMPRESSION_CODEC)
+ .as_str()
Review Comment:
nit: For better UX - trim + to_lowercase the value ?!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
