ballista-executor 54.0.0

// 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.

//! Implementation of the Apache Arrow Flight protocol that wraps an executor.

use ballista_core::execution_plans::create_shuffle_path;
use datafusion::arrow::ipc::reader::StreamReader;
use std::convert::TryFrom;
use std::fs::File;
use std::pin::Pin;
use tokio_util::io::ReaderStream;

use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use ballista_core::error::BallistaError;
use ballista_core::execution_plans::sort_shuffle::{
    ShuffleIndex, get_index_path, is_sort_shuffle_output, stream_sort_shuffle_partition,
};
use ballista_core::serde::decode_protobuf;
use ballista_core::serde::scheduler::Action as BallistaAction;
use datafusion::arrow::ipc::CompressionType;

use arrow_flight::{
    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
    flight_service_server::FlightService,
};
use datafusion::arrow::ipc::writer::IpcWriteOptions;
use datafusion::arrow::{error::ArrowError, record_batch::RecordBatch};
use futures::{Stream, StreamExt, TryStreamExt};
use log::{debug, info};
use std::io::{BufReader, Read, Seek};
use tokio::sync::mpsc::channel;
use tokio::sync::mpsc::error::SendError;
use tokio::{sync::mpsc::Sender, task};
use tokio_stream::wrappers::ReceiverStream;
use tonic::metadata::MetadataValue;
use tonic::{Request, Response, Status, Streaming};

/// Arrow Flight service for transferring shuffle data between executors.
///
/// This service implements the Apache Arrow Flight protocol to enable efficient
/// transfer of intermediate query results (shuffle data) between executor nodes.
/// It supports both decoded streaming via `do_get` and optimized block transfer
/// via the `IO_BLOCK_TRANSPORT` action.
#[derive(Clone)]
pub struct BallistaFlightService {
    work_dir: String,
}

impl BallistaFlightService {
    /// Creates a new BallistaFlightService instance.
    pub fn new(work_dir: String) -> Self {
        Self { work_dir }
    }
}

type BoxedFlightStream<T> =
    Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;

/// shuffle file block transfer size    
const BLOCK_BUFFER_CAPACITY: usize = 8 * 1024 * 1024;

#[tonic::async_trait]
impl FlightService for BallistaFlightService {
    type DoActionStream = BoxedFlightStream<arrow_flight::Result>;
    type DoExchangeStream = BoxedFlightStream<FlightData>;
    type DoGetStream = BoxedFlightStream<FlightData>;
    type DoPutStream = BoxedFlightStream<PutResult>;
    type HandshakeStream = BoxedFlightStream<HandshakeResponse>;
    type ListActionsStream = BoxedFlightStream<ActionType>;
    type ListFlightsStream = BoxedFlightStream<FlightInfo>;

    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> Result<Response<Self::DoGetStream>, Status> {
        let ticket = request.into_inner();

        let action =
            decode_protobuf(&ticket.ticket).map_err(|e| from_ballista_err(&e))?;

        match &action {
            BallistaAction::FetchPartition {
                job_id,
                stage_id,
                partition_id,
                file_id,
                is_sort_shuffle,
                ..
            } => {
                let path = create_shuffle_path(
                    &self.work_dir,
                    job_id,
                    *stage_id,
                    *partition_id,
                    *file_id,
                    *is_sort_shuffle,
                )
                .map_err(|e| {
                    Status::internal(format!("I/O error, can't create shuffle path: {e}"))
                })?;
                debug!("FetchPartition reading partition {partition_id} from {path:?}");

                // Check if this is a sort-based shuffle output
                if is_sort_shuffle_output(&path) {
                    debug!("Detected sort-based shuffle format for {path:?}");
                    let index_path = get_index_path(path.as_path());
                    let stream =
                        stream_sort_shuffle_partition(&path, &index_path, *partition_id)
                            .map_err(|e| from_ballista_err(&e))?;

                    let schema = stream.schema();
                    // Map DataFusionError to FlightError
                    let stream =
                        stream.map_err(|e| FlightError::from(ArrowError::from(e)));

                    let write_options: IpcWriteOptions = IpcWriteOptions::default()
                        .try_with_compression(Some(CompressionType::LZ4_FRAME))
                        .map_err(|e| from_arrow_err(&e))?;
                    let flight_data_stream = FlightDataEncoderBuilder::new()
                        .with_schema(schema)
                        .with_options(write_options)
                        .build(stream)
                        .map_err(|err| Status::from_error(Box::new(err)));

                    return Ok(Response::new(
                        Box::pin(flight_data_stream) as Self::DoGetStream
                    ));
                }

                // Standard hash-based shuffle - read the entire file
                let file = File::open(&path)
                    .map_err(|e| {
                        BallistaError::General(format!(
                            "Failed to open partition file at {path:?}: {e:?}"
                        ))
                    })
                    .map_err(|e| from_ballista_err(&e))?;
                let file = BufReader::new(file);
                // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid
                let reader = unsafe {
                    StreamReader::try_new(file, None)
                        .map_err(|e| from_arrow_err(&e))?
                        .with_skip_validation(cfg!(feature = "arrow-ipc-optimizations"))
                };

                let (tx, rx) = channel(2);
                let schema = reader.schema();
                task::spawn_blocking(move || {
                    if let Err(e) = read_partition(reader, tx) {
                        log::warn!("error streaming shuffle partition: {e}");
                    }
                });

                let write_options: IpcWriteOptions = IpcWriteOptions::default()
                    .try_with_compression(Some(CompressionType::LZ4_FRAME))
                    .map_err(|e| from_arrow_err(&e))?;
                let flight_data_stream = FlightDataEncoderBuilder::new()
                    .with_schema(schema)
                    .with_options(write_options)
                    .build(ReceiverStream::new(rx))
                    .map_err(|err| Status::from_error(Box::new(err)));

                Ok(Response::new(
                    Box::pin(flight_data_stream) as Self::DoGetStream
                ))
            }
        }
    }

    async fn get_schema(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<SchemaResult>, Status> {
        Err(Status::unimplemented("get_schema"))
    }

    async fn get_flight_info(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        Err(Status::unimplemented("get_flight_info"))
    }

    async fn handshake(
        &self,
        _request: Request<Streaming<HandshakeRequest>>,
    ) -> Result<Response<Self::HandshakeStream>, Status> {
        let token = uuid::Uuid::new_v4();
        info!("do_handshake token={}", token);

        let result = HandshakeResponse {
            protocol_version: 0,
            payload: token.as_bytes().to_vec().into(),
        };
        let result = Ok(result);
        let output = futures::stream::iter(vec![result]);
        let str = format!("Bearer {token}");
        let mut resp: Response<
            Pin<Box<dyn Stream<Item = Result<_, Status>> + Send + 'static>>,
        > = Response::new(Box::pin(output));
        let md = MetadataValue::try_from(str)
            .map_err(|_| Status::invalid_argument("authorization not parsable"))?;
        resp.metadata_mut().insert("authorization", md);
        Ok(resp)
    }

    async fn list_flights(
        &self,
        _request: Request<Criteria>,
    ) -> Result<Response<Self::ListFlightsStream>, Status> {
        Err(Status::unimplemented("list_flights"))
    }

    async fn do_put(
        &self,
        request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoPutStream>, Status> {
        let mut request = request.into_inner();

        while let Some(data) = request.next().await {
            let _data = data?;
        }

        Err(Status::unimplemented("do_put"))
    }

    async fn do_action(
        &self,
        request: Request<Action>,
    ) -> Result<Response<Self::DoActionStream>, Status> {
        let action = request.into_inner();

        match action.r#type.as_str() {
            // Block transfer will transfer arrow ipc file block by block
            // without decoding or decompressing, this will provide less resource utilization
            // as file are not decoded nor decompressed/compressed. Usually this would transfer less data across
            // as files are better compressed due to its size.
            //
            // For further discussion regarding performance implications, refer to:
            // https://github.com/apache/datafusion-ballista/issues/1315
            "IO_BLOCK_TRANSPORT" => {
                let action =
                    decode_protobuf(&action.body).map_err(|e| from_ballista_err(&e))?;

                match &action {
                    BallistaAction::FetchPartition {
                        job_id,
                        stage_id,
                        partition_id,
                        file_id,
                        is_sort_shuffle,
                        ..
                    } => {
                        let path = create_shuffle_path(
                            &self.work_dir,
                            job_id,
                            *stage_id,
                            *partition_id,
                            *file_id,
                            *is_sort_shuffle,
                        )
                        .map_err(|e| {
                            Status::internal(format!(
                                "I/O error, can't create shuffle path: {e}"
                            ))
                        })?;

                        debug!("FetchPartition reading {path:?}");

                        let stream = if is_sort_shuffle_output(&path) {
                            // Sort-shuffle: stream the leading schema-header
                            // bytes followed by the requested partition's
                            // byte range. The receiver (BlockDataStream) walks
                            // the resulting concatenated IPC streams.
                            stream_sort_shuffle_block(&path, *partition_id).await?
                        } else {
                            // Hash-shuffle: file contains exactly one partition.
                            stream_whole_file(&path).await?
                        };

                        Ok(Response::new(stream))
                    }
                }
            }
            action_type => Err(Status::unimplemented(format!(
                "do_action does not implement: {}",
                action_type
            ))),
        }
    }

    async fn list_actions(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::ListActionsStream>, Status> {
        let actions = vec![Ok(ActionType {
            r#type: "IO_BLOCK_TRANSFER".to_owned(),
            description: "optimized shuffle data transfer".to_owned(),
        })];

        Ok(Response::new(
            Box::pin(futures::stream::iter(actions)) as Self::ListActionsStream
        ))
    }

    async fn do_exchange(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoExchangeStream>, Status> {
        Err(Status::unimplemented("do_exchange"))
    }

    async fn poll_flight_info(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<PollInfo>, Status> {
        Err(Status::unimplemented("poll_flight_info"))
    }
}

async fn stream_whole_file(
    path: &std::path::Path,
) -> Result<<BallistaFlightService as FlightService>::DoActionStream, Status> {
    let file = tokio::fs::File::open(path)
        .await
        .map_err(|e| Status::internal(format!("Failed to open file: {e}")))?;
    debug!(
        "streaming file: {:?} with size: {}",
        path,
        file.metadata().await?.len()
    );
    let file_stream = ReaderStream::with_capacity(file, BLOCK_BUFFER_CAPACITY);
    Ok(Box::pin(file_stream.map(|result| {
        result
            .map(|bytes| arrow_flight::Result { body: bytes })
            .map_err(|e| Status::internal(format!("I/O error: {e}")))
    })))
}

async fn stream_sort_shuffle_block(
    data_path: &std::path::Path,
    partition_id: usize,
) -> Result<<BallistaFlightService as FlightService>::DoActionStream, Status> {
    use tokio::io::{AsyncReadExt, AsyncSeekExt};

    let index_path = get_index_path(data_path);
    let index =
        ShuffleIndex::read_from_file(&index_path).map_err(|e| from_ballista_err(&e))?;

    if partition_id >= index.partition_count() {
        return Err(Status::out_of_range(format!(
            "partition_id {partition_id} not found in index (max: {})",
            index.partition_count()
        )));
    }

    // The leading [0, header_end) bytes hold the schema-header IPC stream.
    // We always prepend it to the partition's byte range so the receiver
    // recovers the schema even when the partition is empty.
    let header_end = index.header_end_offset() as u64;
    let (start, end) = index.get_partition_range(partition_id);
    let (start, end) = (start as u64, end as u64);

    // One open + dup gives us two independent file cursors over the same
    // inode: one reads the header from offset 0, the other seeks to the
    // partition. `chain` and `take` consume their readers by value, so two
    // cursors are unavoidable here.
    let header_file = tokio::fs::File::open(data_path)
        .await
        .map_err(|e| Status::internal(format!("Failed to open file: {e}")))?;
    let mut partition_file = header_file
        .try_clone()
        .await
        .map_err(|e| Status::internal(format!("dup file handle: {e}")))?;
    partition_file
        .seek(std::io::SeekFrom::Start(start))
        .await
        .map_err(|e| Status::internal(format!("seek partition: {e}")))?;

    let combined = header_file
        .take(header_end)
        .chain(partition_file.take(end - start));
    let file_stream = ReaderStream::with_capacity(combined, BLOCK_BUFFER_CAPACITY);
    Ok(Box::pin(file_stream.map(|result| {
        result
            .map(|bytes| arrow_flight::Result { body: bytes })
            .map_err(|e| Status::internal(format!("I/O error: {e}")))
    })))
}

fn read_partition<T>(
    reader: StreamReader<std::io::BufReader<T>>,
    tx: Sender<Result<RecordBatch, FlightError>>,
) -> Result<(), FlightError>
where
    T: Read + Seek,
{
    if tx.is_closed() {
        return Err(FlightError::Tonic(Box::new(Status::internal(
            "Can't send a batch, channel is closed",
        ))));
    }

    for batch in reader {
        tx.blocking_send(batch.map_err(|err| err.into()))
            .map_err(|err| {
                if let SendError(Err(err)) = err {
                    err
                } else {
                    FlightError::Tonic(Box::new(Status::internal(format!(
                        "Can't send a batch, something went wrong: {err:?}"
                    ))))
                }
            })?
    }
    Ok(())
}

fn from_arrow_err(e: &ArrowError) -> Status {
    Status::internal(format!("ArrowError: {e:?}"))
}

fn from_ballista_err(e: &ballista_core::error::BallistaError) -> Status {
    Status::internal(format!("Ballista Error: {e:?}"))
}