spiral-sdk 0.1.11

//! Host polling of native scans through async-ffi's owned waker bridge.

use std::ffi::{CString, c_char};
use std::ptr;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow_array::ffi::{FFI_ArrowArray, from_ffi};
use arrow_array::{RecordBatch, StructArray};
use arrow_schema::ffi::FFI_ArrowSchema;
use arrow_schema::{Schema, SchemaRef};
use async_ffi::ContextExt as _;
use futures::future::poll_fn;
use futures::{StreamExt as _, stream};

use super::raw::{
    ScandalClient, ScandalError as OwnedError, ScandalScan, ScandalScanPlan, scandal_plan_scan,
    scandal_plan_scan_poll, scandal_plan_scan_release, scandal_scan_poll_next,
    scandal_scan_poll_schema, scandal_scan_release,
};
use super::{Library, check, internal, json_result, json_string};
use crate::scandal::api::PlanScanRequest;
use crate::scandal::error::ScandalError;
use crate::scandal::types::ScanPlan;
use crate::scandal::{SplitRead, wire};

struct NativePlan {
    pointer: *mut ScandalScanPlan,
    _library: Arc<Library>,
}

// SAFETY: The plan is owned and polled or released by only one caller at a time.
unsafe impl Send for NativePlan {}

impl Drop for NativePlan {
    fn drop(&mut self) {
        // SAFETY: This wrapper owns the handle and no poll can outlive it.
        unsafe { scandal_plan_scan_release(self.pointer) }
    }
}

impl NativePlan {
    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<ScanPlan, ScandalError>> {
        let mut out = ptr::null_mut();
        // SAFETY: The handle and borrowed context are live for this serialized poll.
        let error = cx.with_ffi_context(|cx| unsafe {
            scandal_plan_scan_poll(self.pointer, cx, &raw mut out)
        });
        check(error)?;
        if out.is_null() {
            Poll::Pending
        } else {
            Poll::Ready(json_result(out))
        }
    }
}

pub(super) async fn plan(
    library: Arc<Library>,
    request: &PlanScanRequest,
) -> Result<ScanPlan, ScandalError> {
    let mut plan = {
        let request = json_string(request)?;
        let mut pointer = ptr::null_mut();
        // SAFETY: The client and request are live; ownership transfers to NativePlan.
        check(unsafe { scandal_plan_scan(library.client, request.as_ptr(), &raw mut pointer) })?;
        NativePlan {
            pointer,
            _library: library,
        }
    };
    poll_fn(|cx| plan.poll(cx)).await
}

type StartScan = unsafe extern "C" fn(
    *const ScandalClient,
    *const c_char,
    *mut *mut ScandalScan,
) -> *mut OwnedError;

struct NativeScan {
    pointer: *mut ScandalScan,
    schema: FFI_ArrowSchema,
    _library: Arc<Library>,
}

// SAFETY: Each scan is owned and its polls and release are serialized.
unsafe impl Send for NativeScan {}

impl Drop for NativeScan {
    fn drop(&mut self) {
        // SAFETY: No poll can outlive this sole owning wrapper.
        unsafe { scandal_scan_release(self.pointer) }
    }
}

impl NativeScan {
    fn poll_schema(&mut self, cx: &mut Context<'_>) -> Poll<Result<SchemaRef, ScandalError>> {
        let mut out = FFI_ArrowSchema::empty();
        let mut ready = false;
        // SAFETY: Context and empty output are borrowed for one serialized poll.
        let error = cx.with_ffi_context(|cx| unsafe {
            scandal_scan_poll_schema(self.pointer, cx, &raw mut ready, &raw mut out)
        });
        check(error)?;
        if ready {
            self.schema = out;
            Poll::Ready(
                Schema::try_from(&self.schema)
                    .map(Arc::new)
                    .map_err(internal),
            )
        } else {
            Poll::Pending
        }
    }

    fn poll_next(
        &mut self,
        cx: &mut Context<'_>,
        schema: &SchemaRef,
    ) -> Poll<Result<Option<RecordBatch>, ScandalError>> {
        let mut out = FFI_ArrowArray::empty();
        let mut ready = false;
        // SAFETY: The handle, context and empty output remain live through this poll.
        let error = cx.with_ffi_context(|cx| unsafe {
            scandal_scan_poll_next(self.pointer, cx, &raw mut ready, &raw mut out)
        });
        check(error)?;
        if !ready {
            return Poll::Pending;
        }
        if out.is_released() {
            return Poll::Ready(Ok(None));
        }
        // SAFETY: The native scan transfers an owned array matching its exported schema.
        let array = StructArray::from(unsafe { from_ffi(out, &self.schema) }.map_err(internal)?);
        Poll::Ready(
            RecordBatch::try_new(Arc::clone(schema), array.columns().to_vec())
                .map(Some)
                .map_err(internal),
        )
    }
}

pub(super) async fn read(
    library: Arc<Library>,
    operation: StartScan,
    request: CString,
) -> Result<SplitRead, ScandalError> {
    let mut scan = {
        let mut pointer = ptr::null_mut();
        // SAFETY: The client/request are live and the returned handle gets a sole owner.
        check(unsafe { operation(library.client, request.as_ptr(), &raw mut pointer) })?;
        NativeScan {
            pointer,
            schema: FFI_ArrowSchema::empty(),
            _library: library,
        }
    };
    let schema = poll_fn(|cx| scan.poll_schema(cx)).await?;
    let row_schema = wire::read_scan_rows(wire::execution_rows(&schema)?.as_ref())?;
    let batches = stream::try_unfold((scan, schema), |(mut scan, schema)| async move {
        let batch = poll_fn(|cx| scan.poll_next(cx, &schema)).await?;
        batch
            .map(|batch| Ok((wire::detach_execution(batch)?, (scan, schema))))
            .transpose()
    })
    .boxed();
    Ok(SplitRead {
        schema: row_schema,
        batches: wire::execution_batches(batches)
            .map(|batch| batch.and_then(wire::detach_progress))
            .boxed(),
    })
}