andygrove commented on code in PR #2588:
URL:
https://github.com/apache/datafusion-ballista/pull/2588#discussion_r4371934779
##########
python/python/ballista/extension.py:
##########
@@ -531,80 +203,160 @@ def __repr__(self) -> str:
return
f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
-class BallistaSessionContext(SessionContext,
metaclass=RedefiningSessionContextMeta):
+def BallistaSessionContext(
+ address: str,
+ config: Optional[Union[SessionConfig, Dict[str, str]]] = None,
+ runtime: Optional[RuntimeEnvBuilder] = None,
+ cluster_config: Optional[Dict[str, ConfigValue]] = None,
+) -> SessionContext:
"""
- A session context for connecting to and querying a Ballista cluster.
+ Create a DataFusion ``SessionContext`` that executes queries on a Ballista
cluster.
+
+ The returned object is a regular ``datafusion.SessionContext``, so the
whole
+ DataFusion Python API (``sql``, ``read_parquet``, ``register_*``, DataFrame
+ operations, ...) is available. Planning happens locally; physical plans are
+ sent to the Ballista scheduler at ``address`` for distributed execution.
+
+ Args:
+ address: Ballista scheduler URL, for example
``"df://localhost:50050"``.
+ config: Local session configuration, either a ``SessionConfig`` or a
+ dictionary of ``datafusion.*`` options. It is used for local
planning
+ only and is not propagated to the cluster. DataFusion options which
Review Comment:
Nit: I think this isn't quite true anymore. `DistributedQueryExec::execute`
sends every local `datafusion.*` entry to the scheduler, so whatever is in
`config` reaches the cluster too. The "Configure the local session only"
example below has the same issue. Thanks for already dropping the matching line
from the upgrade guide.
##########
python/Cargo.toml:
##########
@@ -30,27 +30,46 @@ include = ["/src", "/ballista", "/LICENSE.txt",
"pyproject.toml", "Cargo.toml",
publish = false
[dependencies]
-# ballista = { path = "../ballista/client", version = "55.0.0" }
-# ballista-core = { path = "../ballista/core", version = "55.0.0" }
-# ballista-executor = { path = "../ballista/executor", version = "55.0.0",
default-features = false }
-# ballista-scheduler = { path = "../ballista/scheduler", version = "55.0.0",
default-features = false }
-ballista = { version = "=55.0.0" }
-ballista-core = { version = "=55.0.0" }
-ballista-executor = { version = "=55.0.0", default-features = false }
-ballista-scheduler = { version = "=55.0.0", default-features = false }
+async-ffi = { version = "0.5" }
+async-trait = { version = "0.1.92" }
+ballista = { path = "../ballista/client", version = "55.0.0" }
+ballista-core = { path = "../ballista/core", version = "55.0.0" }
+ballista-executor = { path = "../ballista/executor", version = "55.0.0",
default-features = false }
+ballista-scheduler = { path = "../ballista/scheduler", version = "55.0.0",
default-features = false }
+
+# ballista = { version = "=55.0.0" }
+# ballista-core = { version = "=55.0.0" }
+# ballista-executor = { version = "=55.0.0", default-features = false }
+# ballista-scheduler = { version = "=55.0.0", default-features = false }
+
+# Version se to 55.2.0 as Workaround for missing file-format
Review Comment:
Nit: small typo, "Version se to" should be "Version set to" (and there's a
double space before "Workaround").
##########
python/python/tests/test_config.py:
##########
@@ -0,0 +1,194 @@
+# 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,
+ _produce_valid_config,
+)
+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(_produce_valid_config())
+
+ for key in _BALLISTA_RESTRICTED_DATAFUSION_OPTIONS:
+ assert settings[key] == defaults[key], key
+
+
+def test_restricted_options_override_user_config():
+ config = SessionConfig().set("datafusion.optimizer.prefer_hash_join",
"true")
+ settings = _settings(_produce_valid_config(config))
+
+ assert settings["datafusion.optimizer.prefer_hash_join"] == "false"
+
+
+def test_round_robin_repartition_disabled_without_config():
+ settings = _settings(_produce_valid_config())
+
+ assert settings["datafusion.optimizer.enable_round_robin_repartition"] ==
"false"
+
+
+def test_round_robin_repartition_disabled_with_user_config():
+ config = SessionConfig().set(
+ "datafusion.optimizer.enable_round_robin_repartition", "true"
+ )
+ settings = _settings(_produce_valid_config(config))
+
+ assert settings["datafusion.optimizer.enable_round_robin_repartition"] ==
"false"
+
+
+def test_non_restricted_options_kept_from_user_config():
+ config = SessionConfig().with_target_partitions(3).with_batch_size(1234)
+ settings = _settings(_produce_valid_config(config))
+
+ assert settings["datafusion.execution.target_partitions"] == "16"
Review Comment:
Nit for the config follow-up: `SessionConfig().with_target_partitions(3)`
quietly ends up as 16 because only dict values count as user overrides. It's
documented in the upgrade guide, but I could see people tripping over it since
`with_target_partitions` is the natural thing to reach for.
--
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]
