Extending Ballista Scheduler And Executors

Ballista scheduler and executor provide a set of configuration options which can be used to extend their basic functionality. They allow registering new configuration extensions, object stores, logical and physical codecs …

• function registry - provides possibility to override set of build in functions.

• config producer - function which creates new SessionConfig, which can hold extended configuration options

• runtime producer - function which creates new RuntimeEnv based on provided SessionConfig.

• session builder - function which creates new SessionState for each user session

• logical codec - overrides LogicalCodec

• physical codec - overrides PhysicalCodec

Ballista executor can be configured using ExecutorProcessConfig which supports overriding function registry,runtime producer, config producer, logical codec, physical codec.

Ballista scheduler can be tuned using SchedulerConfig which supports overriding config producer, session builder, logical codec, physical codec

Example: Custom Object Store Integration

Extending basic building blocks will be demonstrated by integrating S3 object store. For this, new ObjectStoreRegistry and S3Options will be provided. ObjectStoreRegistry creates ObjectStore instances configured using S3Options, and caches them so that scans of the same bucket with the same options share one store.

For this specific task config producer, runtime producer and session builder have to be provided, and client, scheduler and executor need to be configured.

These three functions ship in ballista_core::object_store, so the snippets below are the shipped implementations rather than something you have to write from scratch. They are included from the source file at build time, so they cannot drift from it.

/// Custom [SessionConfig] constructor method
///
/// This method registers config extension [S3Options]
/// which is used to configure [ObjectStore] with ACCESS and
/// SECRET key
pub fn session_config_with_s3_support() -> SessionConfig {
    SessionConfig::new_with_ballista()
        .with_information_schema(true)
        .with_option_extension(S3Options::default())
}

/// Custom [RuntimeEnv] constructor method
///
/// It will register [CustomObjectStoreRegistry] which will
/// use configuration extension [S3Options] to configure
/// and created [ObjectStore]s
pub fn runtime_env_with_s3_support(
    session_config: &SessionConfig,
) -> Result<Arc<RuntimeEnv>> {
    let s3options = session_config
        .options()
        .extensions
        .get::<S3Options>()
        .ok_or(DataFusionError::Configuration(
            "S3 Options not set".to_string(),
        ))?;

    let runtime_env = RuntimeEnvBuilder::new()
        .with_object_store_registry(Arc::new(CustomObjectStoreRegistry::new(
            s3options.clone(),
        )))
        .build()?;

    Ok(Arc::new(runtime_env))
}

/// Custom [SessionState] with S3 support enabled
///
/// It will configure [SessionState] with provided [SessionConfig],
/// and [RuntimeEnv].
pub fn session_state_with_s3_support(
    session_config: SessionConfig,
) -> datafusion::common::Result<SessionState> {
    use crate::extension::{
        ballista_aggregate_functions, ballista_scalar_functions,
        ballista_window_functions,
    };

    let runtime_env = runtime_env_with_s3_support(&session_config)?;

    Ok(SessionStateBuilder::new()
        .with_runtime_env(runtime_env)
        .with_config(session_config)
        .with_default_features()
        .with_scalar_functions(ballista_scalar_functions())
        .with_aggregate_functions(ballista_aggregate_functions())
        .with_window_functions(ballista_window_functions())
        .build())
}

S3Options & CustomObjectStoreRegistry are implemented in ballista_core::object_store. The scheduler, executor, and client below are the runnable custom-scheduler.rs, custom-executor.rs, and custom-client.rs from examples/examples/, also included from source.

Configuring Scheduler

use ballista_core::error::BallistaError;
use ballista_core::object_store::{
    session_config_with_s3_support, session_state_with_s3_support,
};

use ballista_scheduler::cluster::BallistaCluster;
use ballista_scheduler::config::SchedulerConfig;
use ballista_scheduler::scheduler_process::start_server;
use std::net::AddrParseError;
use std::sync::Arc;

///
/// # Custom Ballista Scheduler
///
/// This example demonstrates how to crate custom ballista schedulers.
///
#[tokio::main]
async fn main() -> ballista_core::error::Result<()> {
    let _ = env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .is_test(true)
        .try_init();

    let config: SchedulerConfig = SchedulerConfig {
        // overriding default runtime producer with custom producer
        // which knows how to create S3 connections
        override_config_producer: Some(Arc::new(session_config_with_s3_support)),
        // overriding default session builder, which has custom session configuration
        // runtime environment and session state.
        override_session_builder: Some(Arc::new(session_state_with_s3_support)),
        ..Default::default()
    };

    let addr = format!("{}:{}", config.bind_host, config.bind_port);
    let addr = addr
        .parse()
        .map_err(|e: AddrParseError| BallistaError::Configuration(e.to_string()))?;

    let cluster = BallistaCluster::new_from_config(&config).await?;
    start_server(cluster, addr, Arc::new(config)).await?;

    Ok(())
}

To keep the scheduler’s own command-line interface, parse ballista_scheduler::config::Config with clap and convert it, then apply the overrides:

let opt = Config::parse();
let mut config: SchedulerConfig = opt.try_into()?;
config.override_config_producer = Some(Arc::new(session_config_with_s3_support));
config.override_session_builder = Some(Arc::new(session_state_with_s3_support));

Configuring Executor

use ballista_core::object_store::{
    runtime_env_with_s3_support, session_config_with_s3_support,
};

use ballista_executor::executor_process::{
    ExecutorProcessConfig, start_executor_process,
};
use std::sync::Arc;
///
/// # Custom Ballista Executor
///
/// This example demonstrates how to crate custom ballista executors.
///
#[tokio::main]
async fn main() -> ballista_core::error::Result<()> {
    let _ = env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .is_test(true)
        .try_init();

    let config: ExecutorProcessConfig = ExecutorProcessConfig {
        // overriding default config producer with custom producer
        // which has required S3 configuration options
        override_config_producer: Some(Arc::new(session_config_with_s3_support)),
        // overriding default runtime producer with custom producer
        // which knows how to create S3 connections
        override_runtime_producer: Some(Arc::new(runtime_env_with_s3_support)),
        ..Default::default()
    };

    start_executor_process(Arc::new(config)).await
}

As with the scheduler, to keep the executor’s own command line, parse ballista_executor::config::Config with clap, convert it to ExecutorProcessConfig with try_into(), and then apply the overrides.

Configuring Client

use ballista::extension::SessionContextExt;
use ballista_core::object_store::state_with_s3_support;
use datafusion::error::Result;
use datafusion::{assert_batches_eq, prelude::SessionContext};

/// bucket name to be used for this example
const S3_BUCKET: &str = "ballista";
/// S3 access key
const S3_ACCESS_KEY_ID: &str = "BALLISTA";
/// S3 secret key
const S3_SECRET_KEY: &str = "BALLISTASECRET";
///
/// # Extending Ballista
///
/// This example demonstrates how to extend ballista scheduler and executor registering new object store registry.
/// It uses a local [RustFS](https://github.com/rustfs/rustfs) server to act as S3 object store.
///
/// Ballista will be extended providing custom session configuration, runtime environment and session state.
///
/// RustFS can be started, along with the bucket this example writes to:
///
/// ```bash
/// docker run --rm -d --name rustfs -p 9000:9000 -e RUSTFS_ACCESS_KEY=BALLISTA -e RUSTFS_SECRET_KEY=BALLISTASECRET rustfs/rustfs:1.0.0
/// docker exec rustfs mkdir /data/ballista
/// ```
/// After RustFS, we need to start `custom-scheduler`
///
/// ```bash
/// cargo run --example custom-scheduler
/// ```
///
/// and `custom-executor`
///
/// ```bash
/// cargo run --example custom-executor
/// ```
///
/// ```bash
/// cargo run --example custom-client
/// ```
#[tokio::main]
async fn main() -> Result<()> {
    let test_data = ballista_examples::test_util::examples_test_data();

    // new sessions state with required custom session configuration and runtime environment
    let state = state_with_s3_support()?;

    let ctx: SessionContext =
        SessionContext::remote_with_state("df://localhost:50050", state).await?;

    // session config has relevant S3 options registered and exposed.
    // S3 configuration options can be changed using `SET` statement
    ctx.sql("SET s3.allow_http = true").await?.show().await?;

    ctx.sql(&format!("SET s3.access_key_id = '{S3_ACCESS_KEY_ID}'"))
        .await?
        .show()
        .await?;
    ctx.sql(&format!("SET s3.secret_access_key = '{S3_SECRET_KEY}'"))
        .await?
        .show()
        .await?;
    ctx.sql("SET s3.endpoint = 'http://localhost:9000'")
        .await?
        .show()
        .await?;

    ctx.register_parquet(
        "test",
        &format!("{test_data}/alltypes_plain.parquet"),
        Default::default(),
    )
    .await?;

    let write_dir_path = &format!("s3://{S3_BUCKET}/write_test.parquet");

    ctx.sql("select * from test")
        .await?
        .write_parquet(write_dir_path, Default::default(), Default::default())
        .await?;

    ctx.register_parquet("written_table", write_dir_path, Default::default())
        .await?;

    let result = ctx
        .sql("select id, string_col, timestamp_col from written_table where id > 4")
        .await?
        .collect()
        .await?;

    let expected = [
        "+----+------------+---------------------+",
        "| id | string_col | timestamp_col       |",
        "+----+------------+---------------------+",
        "| 5  | 31         | 2009-03-01T00:01:00 |",
        "| 6  | 30         | 2009-04-01T00:00:00 |",
        "| 7  | 31         | 2009-04-01T00:01:00 |",
        "+----+------------+---------------------+",
    ];

    assert_batches_eq!(expected, &result);
    Ok(())
}