milenkovicm commented on code in PR #2588:
URL:
https://github.com/apache/datafusion-ballista/pull/2588#discussion_r4304019245
##########
python/python/ballista/extension.py:
##########
@@ -531,80 +203,159 @@ 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
+ Ballista does not support are overridden with Ballista defaults.
+ runtime: Optional runtime environment used by the local session.
+ cluster_config: Options applied to the local session and propagated to
+ the scheduler-side session. Keys must start with ``datafusion.`` or
+ ``ballista.``; values may be ``str``, ``int``, ``float`` or
``bool``.
+
+ Returns:
+ A ``SessionContext`` configured with the Ballista query planner.
+
+ Raises:
+ ValueError: If a ``cluster_config`` key does not start with
+ ``datafusion.`` or ``ballista.``.
+
+ Examples:
+ Connect to a scheduler and run SQL:
- This class extends DataFusion's SessionContext to work with distributed
- Ballista clusters, automatically routing query execution to the cluster
- while maintaining API compatibility with local DataFusion usage.
-
- Example:
>>> from ballista import BallistaSessionContext
>>> ctx = BallistaSessionContext("df://localhost:50050")
- >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
- >>> df.show()
+ >>> ctx.register_parquet("t", "./testdata/test.parquet")
+ >>> ctx.sql("SELECT count(*) FROM t").show()
+
+ Use the DataFrame API:
+
+ >>> from datafusion import col, lit
+ >>> df = ctx.read_parquet("./testdata/test.parquet")
+ >>> batches = df.filter(col("id") > lit(4)).limit(5).collect()
+
+ Configure settings which should reach the cluster:
+
+ >>> ctx = BallistaSessionContext(
+ ... "df://localhost:50050",
+ ... cluster_config={
+ ... "datafusion.execution.target_partitions": 16,
+ ... "datafusion.execution.batch_size": 8192,
+ ... },
+ ... )
- To override DataFusion / Ballista session settings on the cluster
- (e.g. the number of target partitions used by the scheduler):
+ Configure Ballista specific settings, such as job name shown in the
+ scheduler UI, or adaptive query planning:
>>> ctx = BallistaSessionContext(
... "df://localhost:50050",
- ... cluster_config={"datafusion.execution.target_partitions":
"256"},
+ ... cluster_config={
+ ... "ballista.job.name": "daily report",
+ ... "ballista.planner.adaptive.enabled": True,
+ ... "ballista.optimizer.broadcast_join_threshold_bytes": 10 *
1024 * 1024,
+ ... },
... )
- For Jupyter notebook users:
- >>> %load_ext ballista.jupyter
- >>> %ballista connect df://localhost:50050
- >>> %sql SELECT * FROM my_table
+ Configure the local session only:
+
+ >>> from datafusion import SessionConfig
+ >>> ctx = BallistaSessionContext(
+ ... "df://localhost:50050",
+ ... config=SessionConfig().with_information_schema(True),
+ ... )
"""
+ cluster_config = _normalize_cluster_config(cluster_config)
+ config = _produce_valid_config(config, cluster_config)
- def __init__(
- self,
- address: str,
- config=None,
- runtime=None,
- cluster_config: Optional[Dict[str, str]] = None,
- ):
- self.cluster_config = (
- {str(k): str(v) for k, v in cluster_config.items()}
- if cluster_config is not None
- else None
- )
- # Apply overrides to the local SessionContext too: some settings
- # (e.g. datafusion.execution.listing_table_factory_infer_partitions)
- # are consulted during local table registration / planning, before
- # the plan is shipped to the scheduler. Ballista-namespaced keys
- # are not understood by the local SessionConfig and are forwarded
- # to the scheduler only.
- if self.cluster_config:
- if config is None:
- config = SessionConfig()
- for key, value in self.cluster_config.items():
- if key.startswith("ballista."):
- continue
- config = config.set(key, value)
- super().__init__(config, runtime)
- self.address = address
- self.session_id_internal = super().session_id()
-
- @property
- def session_id(self):
- return self.session_id_internal
-
- def get_tables(self) -> Optional[dict[str, List[str]]]:
- """Get tables and their respective schemas (in terms of database
schema)."""
- try:
- catalog = self.catalog()
- schema_names = list(catalog.schema_names())
- if schema_names:
- tables_info = {}
- for schema_name in schema_names:
- tables_info[schema_name] = list(
- catalog.schema(name=schema_name).table_names()
- )
- return tables_info
- except (AttributeError, NotImplementedError) as e:
- warnings.warn(f"Could not retrieve tables from catalog: {e}")
- pass
- return {}
+ source_ctx = SessionContext(config, runtime)
+ configured_ctx = with_ballista_query_planner(
+ source_ctx,
+ address,
+ cluster_config,
+ )
+ return configured_ctx
+
+
+# DataFusion options overridden by Ballista, mirrors
+# `SessionConfigExt::ballista_restricted_configuration` in ballista-core
+#
+# value is `True` if option is restricted and cant be overridden,
+# if value is `False` if users can override it using dict parameter
+# in dictionary.
+#
+# TODO: At the moment values are duplicated,
+_BALLISTA_RESTRICTED_DATAFUSION_OPTIONS: Dict[str, bool] = {
+ "datafusion.optimizer.enable_round_robin_repartition": True,
+ "datafusion.execution.parquet.schema_force_view_types": False,
+ "datafusion.sql_parser.map_string_types_to_utf8view": False,
+ "datafusion.optimizer.hash_join_single_partition_threshold": False,
+ "datafusion.optimizer.hash_join_single_partition_threshold_rows": False,
+ "datafusion.optimizer.prefer_hash_join": True,
+ "datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery": True,
+ "datafusion.optimizer.enable_dynamic_filter_pushdown": True,
+}
+
+
+def _config_value_to_str(value: ConfigValue) -> str:
+ # `str(True)` is "True", config parsers expect lowercase
+ if isinstance(value, bool):
+ return "true" if value else "false"
+ 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():
Review Comment:
good catch, added target_partitions to restricted options.
I believe config needs to be revisited but will do it in a follow up PR
--
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]
