sail python SparkConnectServer的处理简单说明

sail 基于rust 开发，对于python包的支持基于了比较火的框架pyo3，目前包含了python （主要是spark connect server 启动的）以及python udf 的，以下简单说明下SparkConnectServer 的

服务启动的处理

基于pyo3 暴露了SparkConnectServer 类，当然还提供了一个cli 方便启动的

• connect server

use pyo3::prelude::*;

pub(crate) mod server;

pub(super) fn register_module(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    let module = PyModule::new(parent.py(), "spark")?;
    module.add_class::<server::SparkConnectServer>()?;
    parent.add_submodule(&module)?;
    Ok(())
}

python 参考使用

from pysail.spark import SparkConnectServer
from pyspark.sql import SparkSession

server = SparkConnectServer()
server.start()
_, port = server.listening_address

spark = SparkSession.builder.remote(f"sc://localhost:{port}").getOrCreate()

spark.sql("SELECT 1 + 1").show()

• cli 模式

主要是方便通过python 包快速启动sail conenct server 服务，主要目的是测试

#[pyfunction]
pub(crate) fn main(py: Python<'_>, args: Vec<String>) -> PyResult<()> {
    py.allow_threads(move || {
        sail_cli::runner::main(args).map_err(|e| PyErr::new::<PyRuntimeError, _>(format!("{e}")))
    })
}

pyproject.toml 文件

[project.scripts]
sail = "pysail.cli:main"

使用

RUST_LOG=debug sail spark server

说明

SparkConnectServer 处理实际上就是包装sail rust server 的能力为python包，机制上基于了pyo3，代码并不多

参考资料

crates/sail-common/src/config/application.rs

https://github.com/lakehq/sail/tree/main/crates/sail-python

https://github.com/lakehq/sail/tree/main/crates/sail-python-udf

https://github.com/PyO3/pyo3