andygrove commented on code in PR #2637:
URL:
https://github.com/apache/datafusion-ballista/pull/2637#discussion_r4437306070
##########
python/python/ballista/extension.py:
##########
@@ -318,44 +319,349 @@ def _config_value_to_str(value: ConfigValue) -> str:
return str(value)
-def _normalize_cluster_config(
- cluster_config: Optional[Dict[str, ConfigValue]],
-) -> Optional[Dict[str, str]]:
- if cluster_config is None:
- return None
- return {str(k): _config_value_to_str(v) for k, v in cluster_config.items()}
-
-
-def _produce_valid_config(
- config: Optional[Union[SessionConfig, Dict[str, str]]] = None,
- cluster_config: Optional[Dict[str, ConfigValue]] = None,
-) -> SessionConfig:
- cluster_config = _normalize_cluster_config(cluster_config)
-
- ballista_settings = {
- k: v for k, v in (cluster_config or {}).items() if
k.startswith("ballista.")
- }
-
- # options explicitly set by the user, only known when config is a dict
- user_options = config if isinstance(config, dict) else {}
-
- if config is None:
- config = SessionConfig()
- elif isinstance(config, dict):
- config = SessionConfig(config)
-
- for key, value in ballista_datafusion_config_defaults().items():
- if (
- key.startswith("datafusion.")
- and key in _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS
- ):
- restricted = _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS[key]
- # keep user value for options which are not restricted
- if not restricted and key in user_options:
- continue
+class BallistaSessionConfig:
+ """
+ Builder for a Ballista-ready DataFusion ``SessionConfig``.
+
+ Mirrors Rust ``BallistaConfig``: typed ``with_*`` setters for Ballista
+ options, plus ``set`` for any ``datafusion.*`` or ``ballista.*`` option.
+ ``build`` produces a ``SessionConfig`` with Ballista defaults applied and
+ ``BallistaConfigExtension`` registered.
+
+ DataFusion options which Ballista overrides are set to Ballista defaults:
+ restricted options are always overridden, non-restricted options are kept
+ only when set with ``set`` or a dict.
+ Options inside a passed ``SessionConfig`` can not be inspected, so all of
+ them are overridden with Ballista defaults.
+
+ For example ``datafusion.execution.target_partitions`` defaults to ``16``,
+ change it with ``with_target_partitions``, ``set`` or a dict, not with
+ ``SessionConfig.with_target_partitions``.
+
+ Args:
+ config: Initial configuration, either a ``SessionConfig`` or a
dictionary
+ of ``datafusion.*`` and ``ballista.*`` options.
+
+ Examples:
+ >>> from ballista import BallistaSessionConfig
+ >>> config = (
+ ... BallistaSessionConfig({"datafusion.execution.batch_size":
8192})
+ ... .with_target_partitions(32)
+ ... .with_job_name("daily report")
+ ... .with_adaptive_query_planner(True)
+ ... .set("ballista.shuffle.compression.codec", "zstd")
+ ... .build()
+ ... )
+ """
+
+ def __init__(
+ self,
+ config: Optional[Union[SessionConfig, Dict[str, ConfigValue]]] = None,
Review Comment:
Sorry, that was a bit terse! What I meant: if someone passes
`config=SessionConfig().with_target_partitions(64)`, they silently get 16 on
the cluster because we can't see inside the `SessionConfig` and reset it to
Ballista defaults. That's documented now, but people coming from DataFusion
will likely hit it before they read the docs. So the idea was a `warnings.warn`
(a plain `UserWarning`) in `BallistaSessionConfig.__init__` when it's given a
`SessionConfig`, saying something like "options in a SessionConfig are reset to
Ballista defaults, use a dict or BallistaSessionConfig.set to override them".
Totally fine to skip it if you think it's too noisy though, it's not a blocker.
##########
python/python/ballista/extension.py:
##########
@@ -318,44 +327,374 @@ def _config_value_to_str(value: ConfigValue) -> str:
return str(value)
-def _normalize_cluster_config(
- cluster_config: Optional[Dict[str, ConfigValue]],
-) -> Optional[Dict[str, str]]:
- if cluster_config is None:
- return None
- return {str(k): _config_value_to_str(v) for k, v in cluster_config.items()}
+class BallistaSessionConfig:
+ """
+ Builder for a Ballista-ready DataFusion ``SessionConfig``.
+ Mirrors Rust ``BallistaConfig``: typed ``with_*`` setters for Ballista
+ options, plus ``set`` for any ``datafusion.*`` or ``ballista.*`` option.
+ ``build`` produces a ``SessionConfig`` with Ballista defaults applied and
+ ``BallistaConfigExtension`` registered.
-def _produce_valid_config(
- config: Optional[Union[SessionConfig, Dict[str, str]]] = None,
- cluster_config: Optional[Dict[str, ConfigValue]] = None,
-) -> SessionConfig:
- cluster_config = _normalize_cluster_config(cluster_config)
-
- ballista_settings = {
- k: v for k, v in (cluster_config or {}).items() if
k.startswith("ballista.")
- }
-
- # options explicitly set by the user, only known when config is a dict
- user_options = config if isinstance(config, dict) else {}
-
- if config is None:
- config = SessionConfig()
- elif isinstance(config, dict):
- config = SessionConfig(config)
-
- for key, value in ballista_datafusion_config_defaults().items():
- if (
- key.startswith("datafusion.")
- and key in _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS
- ):
- restricted = _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS[key]
- # keep user value for options which are not restricted
- if not restricted and key in user_options:
- continue
+ DataFusion options which Ballista overrides are set to Ballista defaults:
+ restricted options are always overridden, non-restricted options are kept
+ only when set with ``set`` or a dict.
+ Options inside a passed ``SessionConfig`` can not be inspected, so all of
+ them are overridden with Ballista defaults.
+
+ For example ``datafusion.execution.target_partitions`` defaults to ``16``,
+ change it with ``with_target_partitions``, ``set`` or a dict, not with
+ ``SessionConfig.with_target_partitions``.
+
+ Args:
+ config: Initial configuration, either a ``SessionConfig`` or a
dictionary
+ of ``datafusion.*`` and ``ballista.*`` options.
+
+ Examples:
+ >>> from ballista import BallistaSessionConfig
+ >>> config = (
+ ... BallistaSessionConfig({"datafusion.execution.batch_size":
8192})
+ ... .with_target_partitions(32)
+ ... .with_job_name("daily report")
+ ... .with_adaptive_query_planner(True)
+ ... .set("ballista.shuffle.compression.codec", "zstd")
+ ... .build()
+ ... )
+ """
+
+ def __init__(
+ self,
+ config: Optional[Union[SessionConfig, Dict[str, ConfigValue]]] = None,
+ ):
+ self._base: Optional[SessionConfig] = None
+ self._datafusion_options: Dict[str, str] = {}
+ self._ballista_settings: Dict[str, str] = {}
+
+ if isinstance(config, SessionConfig):
+ self._base = config
+ elif config is not None:
+ for key, value in config.items():
+ self.set(key, value)
+
+ def set(self, key: str, value: ConfigValue) -> "BallistaSessionConfig":
+ """
+ Set a ``datafusion.*`` or ``ballista.*`` option.
+
+ Raises:
+ ValueError: If the key has an unsupported prefix, or is an unknown
+ ``ballista.*`` key.
+ """
+ value = _config_value_to_str(value)
+ if key.startswith("ballista."):
+ # validates the key, values are not validated by ballista
+ BallistaConfigExtension({key: value})
+ self._ballista_settings[key] = value
+ elif key.startswith("datafusion."):
+ self._datafusion_options[key] = value
+ else:
+ raise ValueError(
+ f"unsupported configuration key: {key}, "
+ "expected `datafusion.` or `ballista.` prefix"
+ )
+ return self
+
+ def build(self) -> SessionConfig:
+ """Produce a ``SessionConfig`` with ``BallistaConfigExtension``
registered."""
+ config = self._base if self._base is not None else SessionConfig()
Review Comment:
Small one: `SessionConfig.set` and `with_extension` modify the config in
place in datafusion-python, so `build()` ends up writing Ballista defaults into
the `SessionConfig` the user passed in. If they reuse it for a plain
`SessionContext` afterwards they'd get Ballista's target_partitions and the
extension. Might be worth starting from a copy of the base. Not a blocker.
--
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]
