andygrove commented on code in PR #2588: URL: https://github.com/apache/datafusion-ballista/pull/2588#discussion_r4303460680
########## python/src/ffi_planner_workaround.rs: ########## @@ -0,0 +1,287 @@ +// 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. + +//! Workaround for missing file-format support in `datafusion-ffi` 55.2.0. +//! +//! datafusion-python replaces the codecs on an installed [`FFI_QueryPlanner`] +//! with its own, so Ballista decodes every logical plan through a +//! `ForeignLogicalExtensionCodec`. That type's `try_decode_file_format` is a +//! hard-coded `not_impl_err!`, which breaks any plan carrying a +//! `FileFormatFactory` (`LogicalPlan::Copy`, i.e. `DataFrame.write_*`). +//! +//! The file formats in those plans are written by the session's terminal +//! `DefaultLogicalExtensionCodec` without a chain envelope, so they can be +//! decoded locally. We replace the planner's `create_physical_plan` entry +//! with one that builds a temporary planner handle whose logical codec is a +//! local [`FileFormatFallbackCodec`] around the session's codec, and forwards +//! the call to `datafusion-ffi`'s own entry on that handle. The `clone` entry +//! is replaced too, so the override survives the clone datafusion-python +//! makes when it swaps codecs. +//! +//! This relies on the private field layouts of `FFI_QueryPlanner` and +//! `FFI_SessionRef`, mirrored below. It is only valid for the pinned +//! `datafusion-ffi = "=55.2.0"`. +//! +//! TODO: remove once `datafusion-ffi` delegates file formats to +//! `DefaultLogicalExtensionCodec` in `ForeignLogicalExtensionCodec`. + +use async_ffi::FfiFuture; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::TableProvider; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::file_format::FileFormatFactory; +use datafusion::execution::context::{QueryPlanner, SessionContext}; +use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::logical_expr::{ + AggregateUDF, Extension, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF, +}; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_ffi::util::FFI_Result; +use datafusion_proto::logical_plan::{ + DefaultLogicalExtensionCodec, LogicalExtensionCodec, +}; +use stabby::vec::Vec as SVec; +use std::ffi::c_void; +use std::mem::MaybeUninit; +use std::sync::{Arc, OnceLock}; + +/// Same size and alignment as `datafusion_ffi::session::FFI_SessionRef` +/// (55.2.0), which is `pub(crate)`. It is only passed through by value, so +/// ownership moves to the callee and nothing here drops it. The struct is far +/// larger than two registers, so every supported C ABI passes it in memory +/// regardless of field types. +#[repr(C)] +struct FfiSessionRefBytes { Review Comment: Not a blocker, just something to think about. This module reinterprets `FFI_QueryPlanner` and `FFI_SessionRef` through hand-mirrored private layouts, so it is only sound for exactly `datafusion-ffi =55.2.0`. `FFI_QueryPlanner` has a compile-time size and alignment assert, but `FfiSessionRefBytes` can't have one since the real type is `pub(crate)`, and the pointer checks in `layout()` are `debug_assert`s so release builds skip them. A layout mismatch would show up as UB rather than an error. Could we link an upstream issue in the module doc for the `TODO: remove once ...` and add a short note next to the `=55.2.0` pin in `Cargo.toml` that it must not be bumped without re-checking these layouts? If upstream is likely to land the fix soon, documenting `DataFrame.write_*` as unsupported for now would be another option, but I can see the case for keeping the workaround if write support matters for this release. ########## 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: I think the default `target_partitions` changed with this loop. On `main`, `BallistaSessionContext(addr)` with no config ended up with Ballista's 16 (`new_with_ballista()` sets it). Here `datafusion.execution.target_partitions` is not in `_BALLISTA_RESTRICTED_DATAFUSION_OPTIONS`, so it falls through and the plain `SessionConfig()` value (the client's core count) is kept. `DistributedQueryExec::execute` then sends every local entry to the scheduler, so the cluster plans with the client's core count. A 10-core laptop and a 64-core box would get different plans for the same query. Could we seed the defaults for keys the user didn't set, for example by treating `target_partitions` as a non-restricted option with the Ballista default applied unless it appears in the user's dict? A test that checks `target_partitions` for `config=None` would pin it down. The note in the upgrade guide ("`config` applies to the local session only") and the matching `config` docstring and README lines also need a tweak, since everything in the local config is forwarded to the scheduler on this path. -- 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]
