milenkovicm commented on code in PR #2588:
URL:
https://github.com/apache/datafusion-ballista/pull/2588#discussion_r4302298271
##########
python/python/ballista/extension.py:
##########
@@ -531,80 +200,38 @@ def __repr__(self) -> str:
return
f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
-class BallistaSessionContext(SessionContext,
metaclass=RedefiningSessionContextMeta):
- """
- A session context for connecting to and querying a Ballista cluster.
-
- 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()
-
- To override DataFusion / Ballista session settings on the cluster
- (e.g. the number of target partitions used by the scheduler):
-
- >>> ctx = BallistaSessionContext(
- ... "df://localhost:50050",
- ... cluster_config={"datafusion.execution.target_partitions":
"256"},
- ... )
-
- For Jupyter notebook users:
- >>> %load_ext ballista.jupyter
- >>> %ballista connect df://localhost:50050
- >>> %sql SELECT * FROM my_table
- """
-
- 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 {}
+def BallistaSessionContext(
+ address: str,
+ config=None,
+ runtime=None,
+ cluster_config: Optional[Dict[str, str]] = None,
+):
+ 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 registered locally as a config extension below.
+ if config is None:
Review Comment:
Config is a bit of a mess at the moment, the main reason being that we can't
really tell if a user has set a value or not.
Config handling has been changed, a list of restricted options has been
copied over to python and check and default will be applied. For some of the
config values we still provide for users to override them. We do have a bit of
duplication at the moment, which we can live with for a moment.
Also new tests have been added which cover configuration.
--
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]
