andygrove commented on code in PR #2637:
URL:
https://github.com/apache/datafusion-ballista/pull/2637#discussion_r4406552366
##########
python/python/ballista/extension.py:
##########
@@ -204,10 +204,11 @@ def __repr__(self) -> str:
def BallistaSessionContext(
Review Comment:
`cluster_config` has been the documented way to send settings to the cluster
since 53.0.0, so removing it will break existing scripts on upgrade. Could we
keep accepting it for one more release? Something like merging it into `config`
and raising a `DeprecationWarning` should be just a few lines, and then we can
drop it in 57.
##########
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:
Small idea: since a passed `SessionConfig` gets its Ballista-managed options
reset, a `warnings.warn` here when `config` is a `SessionConfig` could point
people to `BallistaSessionConfig` or a dict. That would make the
`with_target_partitions` surprise much easier to discover than the docs alone.
##########
python/python/tests/test_session_config.py:
##########
@@ -0,0 +1,309 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+from ballista._internal_ballista import (
+ BallistaConfigExtension,
+ ballista_datafusion_config_defaults,
+)
+from ballista.extension import (
+ _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS,
+ BallistaSessionConfig,
+)
+from datafusion import SessionConfig, SessionContext
+
+
+def _settings(config: SessionConfig) -> dict:
+ ctx = SessionContext(config)
+ ctx.sql("SET datafusion.catalog.information_schema = true").collect()
+ rows = ctx.sql("SELECT name, value FROM
information_schema.df_settings").to_pylist()
+ return {row["name"]: row["value"] for row in rows}
+
+
+def test_restricted_options_set_without_config():
+ defaults = ballista_datafusion_config_defaults()
+ settings = _settings(BallistaSessionConfig().build())
+
+ for key in _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS:
+ assert settings[key] == defaults[key], key
+
+
+def test_restricted_options_override_session_config():
+ config = SessionConfig().set("datafusion.optimizer.prefer_hash_join",
"true")
+ settings = _settings(BallistaSessionConfig(config).build())
+
+ assert settings["datafusion.optimizer.prefer_hash_join"] == "false"
+
+
+def test_restricted_options_override_set():
+ settings = _settings(
+ BallistaSessionConfig()
+ .set("datafusion.optimizer.prefer_hash_join", True)
+ .set("datafusion.optimizer.enable_round_robin_repartition", True)
+ .build()
+ )
+
+ assert settings["datafusion.optimizer.prefer_hash_join"] == "false"
+ assert settings["datafusion.optimizer.enable_round_robin_repartition"] ==
"false"
+
+
+def test_restricted_options_override_dict():
+ key = "datafusion.optimizer.prefer_hash_join"
+ assert _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS[key]
+
+ settings = _settings(BallistaSessionConfig({key: True}).build())
+
+ assert settings[key] == "false"
+
+
+def test_round_robin_repartition_disabled_without_config():
+ settings = _settings(BallistaSessionConfig().build())
+
+ assert settings["datafusion.optimizer.enable_round_robin_repartition"] ==
"false"
+
+
+def test_non_restricted_options_from_session_config():
+ config = SessionConfig().with_target_partitions(3).with_batch_size(1234)
+ settings = _settings(BallistaSessionConfig(config).build())
+
+ assert settings["datafusion.execution.target_partitions"] == "16"
+ assert settings["datafusion.execution.batch_size"] == "1234"
+
+
+def test_target_partitions_default():
+ settings = _settings(BallistaSessionConfig().build())
+
+ assert settings["datafusion.execution.target_partitions"] == "16"
+
+
+def test_target_partitions_overridden_by_dict():
+ config = {"datafusion.execution.target_partitions": 128}
+ settings = _settings(BallistaSessionConfig(config).build())
+
+ assert settings["datafusion.execution.target_partitions"] == "128"
+
+
+def test_target_partitions_overridden_by_set_on_session_config():
+ config = SessionConfig().with_batch_size(1234)
+ settings = _settings(
+ BallistaSessionConfig(config)
+ .set("datafusion.execution.target_partitions", 128)
+ .build()
+ )
+
+ assert settings["datafusion.execution.target_partitions"] == "128"
+ assert settings["datafusion.execution.batch_size"] == "1234"
+
+
+def test_with_target_partitions():
+ config = BallistaSessionConfig()
+ assert config.with_target_partitions(3) is config
+
+ settings = _settings(config.build())
+
+ assert settings["datafusion.execution.target_partitions"] == "3"
+
+
+def test_dict_keeps_non_restricted_options():
+ key = "datafusion.sql_parser.map_string_types_to_utf8view"
+ assert not _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS[key]
+ assert ballista_datafusion_config_defaults()[key] == "false"
+
+ settings = _settings(BallistaSessionConfig({key: True}).build())
+
+ assert settings[key] == "true"
+
+
+def test_session_config_overrides_non_restricted_options():
+ key = "datafusion.sql_parser.map_string_types_to_utf8view"
+ config = SessionConfig().set(key, "true")
+
+ settings = _settings(BallistaSessionConfig(config).build())
+
+ assert settings[key] == "false"
+
+
+def test_ballista_defaults_set_without_config():
+ expected =
_settings(SessionConfig().with_extension(BallistaConfigExtension()))
+ settings = _settings(BallistaSessionConfig().build())
+
+ ballista_defaults = {k: v for k, v in expected.items() if
k.startswith("ballista.")}
+ assert ballista_defaults
+ for key, value in ballista_defaults.items():
+ assert settings[key] == value, key
+
+
+def test_dict_accepts_ballista_and_datafusion_options():
+ settings = _settings(
+ BallistaSessionConfig(
+ {
+ "datafusion.execution.batch_size": 1234,
+ "ballista.job.name": "test_job",
+ "ballista.shuffle.force_remote_read": True,
+ }
+ ).build()
+ )
+
+ assert settings["datafusion.execution.batch_size"] == "1234"
+ assert settings["ballista.job.name"] == "test_job"
+ assert settings["ballista.shuffle.force_remote_read"] == "true"
+
+
+def test_set_overrides_ballista_defaults():
+ key = "ballista.shuffle.max_concurrent_read_requests"
+ assert _settings(BallistaSessionConfig().build())[key] != "7"
+
+ settings = _settings(BallistaSessionConfig().set(key, 7).build())
+
+ assert settings[key] == "7"
+
+
[email protected](
+ "method, value, key, expected",
+ [
+ ("with_job_name", "test_job", "ballista.job.name", "test_job"),
+ (
+ "with_shuffle_reader_maximum_concurrent_requests",
+ 7,
+ "ballista.shuffle.max_concurrent_read_requests",
+ "7",
+ ),
+ (
+ "with_shuffle_reader_force_remote_read",
+ True,
+ "ballista.shuffle.force_remote_read",
+ "true",
+ ),
+ (
+ "with_adaptive_query_planner",
+ True,
+ "ballista.planner.adaptive.enabled",
+ "true",
+ ),
+ (
+ "with_broadcast_join_threshold_bytes",
+ 1024,
+ "ballista.optimizer.broadcast_join_threshold_bytes",
+ "1024",
+ ),
+ (
+ "with_coalesce_small_partition_factor",
+ 0.5,
+ "ballista.planner.coalesce.small_partition_factor",
+ "0.5",
+ ),
+ (
+ "with_shuffle_compression_codec",
+ "zstd",
+ "ballista.shuffle.compression.codec",
+ "zstd",
+ ),
+ (
+ "with_grpc_client_http2_keepalive_interval_seconds",
+ 42,
+ "ballista.client.grpc_http2_keepalive_interval_seconds",
+ "42",
+ ),
+ (
+ "with_share_file_statistics_cache",
+ False,
+ "ballista.scheduler.share_file_statistics_cache",
+ "false",
+ ),
+ ],
+)
+def test_typed_setters(method, value, key, expected):
Review Comment:
It might be worth adding a test that calls every `with_*` setter. Right now
9 of the 43 are covered. Since `set` validates `ballista.*` keys, calling each
one (any value works, since values aren't checked) would catch a setter whose
key drifts from the Rust config.
--
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]
