algeon-cudf 0.1.2

//! Per-group row shift.
//!
//! This is the building block for SQL `LAG(value, offset, default)` and
//! `LEAD(value, offset, default)` window functions. Within each group, libcudf
//! shifts the value column by the requested offset, filling indeterminate
//! output positions with the per-column fill scalar.
//!
//! libcudf guarantees that rows within each group appear in the input row
//! order; the order of groups themselves is libcudf-defined and may differ
//! between calls. Callers that need the input row order back should add a
//! row-position column to their inputs and sort the result by that column.

use std::{ptr, sync::Arc};

use smallvec::SmallVec;

use super::{agg::NullPolicy, errors::foreign_groupby_contract_error};
use crate::{
    core::{
        column::ColumnView,
        error::{self, Error, ErrorCode, FACT_ACTUAL_COUNT, FACT_LEFT_TYPE_ID, FACT_RIGHT_TYPE_ID, Result},
        scalar::Scalar,
        table::{Table, TableView},
    },
    runtime::execution_domain::ExecutionDomain,
};

const GROUPBY_SHIFT_VALUE_HANDLE_INLINE_CAPACITY: usize = 8;

pub(crate) struct GroupByShiftOutput {
    /// Unique key groups, one row per output values row, with rows clustered
    /// by group and inter-group order libcudf-defined.
    pub keys: Table,
    /// Shifted value columns, in the same column order as the input request.
    /// Rows align with `keys` row-for-row, and within each group the row order
    /// matches the input row order.
    pub values: Table,
}

pub(crate) struct GroupByShiftRequest<'a> {
    pub values: &'a ColumnView<'a>,
    /// Row shift offset (positive shifts rows toward larger indices within
    /// the group; negative shifts toward smaller indices).
    pub offset: i32,
    /// Fill value for indeterminate output rows. Must share the column's
    /// data type. A null-valued scalar forces the output column to be
    /// nullable.
    pub fill_value: &'a Scalar,
}

pub(crate) fn groupby_shift_impl(
    keys: &TableView<'_>,
    requests: &[GroupByShiftRequest<'_>],
    null_policy: NullPolicy,
    keys_are_sorted: bool,
    output_domain: &Arc<ExecutionDomain>,
    exec_options: &cudf_sys::cudf_exec_options_t,
    operation: &'static str,
) -> Result<GroupByShiftOutput> {
    validate_request_dtypes(requests, operation)?;

    let num_value_columns_usize = requests.len();
    let num_value_columns = i32::try_from(num_value_columns_usize).map_err(|_| {
        error::limit_validation_error(
            operation,
            "number of grouped shift columns exceeds i32::MAX",
            "shift_requests",
            "request_count",
            "i32::MAX",
            num_value_columns_usize,
            "reduce the number of shift requests to fit within libcudf's i32 range",
        )
    })?;

    let value_handles: SmallVec<[cudf_sys::cudf_column_view_handle; GROUPBY_SHIFT_VALUE_HANDLE_INLINE_CAPACITY]> =
        requests.iter().map(|r| r.values.raw_handle()).collect();
    let offsets: SmallVec<[i32; GROUPBY_SHIFT_VALUE_HANDLE_INLINE_CAPACITY]> =
        requests.iter().map(|r| r.offset).collect();
    let fill_handles: SmallVec<[cudf_sys::cudf_scalar_handle; GROUPBY_SHIFT_VALUE_HANDLE_INLINE_CAPACITY]> =
        requests.iter().map(|r| r.fill_value.raw_handle()).collect();

    let mut out_keys = ptr::null_mut();
    let mut out_values = ptr::null_mut();

    let null_policy_ffi = null_policy.to_raw();
    let keys_are_sorted_ffi = i32::from(keys_are_sorted);
    let status = unsafe {
        cudf_sys::cudf_groupby_shift_ex(
            keys.raw_handle(),
            value_handles.as_ptr(),
            offsets.as_ptr(),
            fill_handles.as_ptr(),
            num_value_columns,
            null_policy_ffi,
            keys_are_sorted_ffi,
            &mut out_keys,
            &mut out_values,
            exec_options,
        )
    };

    error::check_status_in_with_details(
        status,
        operation,
        [
            ("num_value_columns", num_value_columns.to_string()),
            ("null_policy", (null_policy as i32).to_string()),
            ("keys_are_sorted", keys_are_sorted.to_string()),
            ("execution_authority", "retained_domain".to_string()),
        ],
    )?;

    let keys_table = unsafe { Table::from_raw_with_domain(out_keys, Arc::clone(output_domain)) }?;
    let values_out = unsafe { Table::from_raw_with_domain(out_values, Arc::clone(output_domain)) }?;

    let result_value_columns = values_out.num_columns().map_err(|e| e.reanchor_if_needed(operation))?;
    if result_value_columns != num_value_columns_usize {
        return Err(foreign_groupby_contract_error(
            operation,
            "FFI returned a values table with a different column count than requested",
            "result_values",
            "column_count_mismatch",
            "ensure the foreign grouped shift backend returns one output column per request",
        )
        .with_fact(FACT_ACTUAL_COUNT, result_value_columns)
        .with_detail("requested_value_columns", num_value_columns.to_string()));
    }
    let keys_row_count = keys_table.num_rows().map_err(|e| e.reanchor_if_needed(operation))?;
    let values_row_count = values_out.num_rows().map_err(|e| e.reanchor_if_needed(operation))?;
    if values_row_count != keys_row_count {
        return Err(foreign_groupby_contract_error(
            operation,
            "FFI returned keys and values tables with different row counts",
            "result_row_count",
            "row_count_mismatch",
            "ensure the foreign grouped shift backend returns keys and values with identical row counts",
        )
        .with_detail("keys_row_count", keys_row_count.to_string())
        .with_detail("values_row_count", values_row_count.to_string()));
    }

    Ok(GroupByShiftOutput { keys: keys_table, values: values_out })
}

fn validate_request_dtypes(requests: &[GroupByShiftRequest<'_>], operation: &'static str) -> Result<()> {
    for (request_index, req) in requests.iter().enumerate() {
        let value_dtype = req.values.data_type().map_err(|e| e.reanchor_if_needed(operation))?;
        let fill_dtype = req.fill_value.data_type().map_err(|e| e.reanchor_if_needed(operation))?;
        if value_dtype != fill_dtype {
            return Err(Error::new_with(
                ErrorCode::TypeMismatch,
                operation,
                format!(
                    "shift fill_value dtype ({:?}) at request {request_index} must match values dtype ({:?})",
                    fill_dtype.id(),
                    value_dtype.id()
                ),
            )
            .with_fact(FACT_LEFT_TYPE_ID, value_dtype.id())
            .with_fact(FACT_RIGHT_TYPE_ID, fill_dtype.id())
            .with_detail("resource", "shift_fill_value")
            .with_detail("reason", "dtype_mismatch")
            .with_detail("request_index", request_index.to_string())
            .with_detail("remedy", "construct fill_value with the same DataType as values"));
        }
    }
    Ok(())
}