algeon-cudf 0.1.2

//! Expression DSL for declarative column operations.
//!
//! Use the DSL entry points [`col`], [`lit`], and [`when`] to start
//! expressions, then extend them with arithmetic, comparisons, casts,
//! naming, aggregation expression methods, and string expression methods
//! before evaluating them against a [`DataFrame`].
//!
//! # Reduction Shape
//!
//! Expr reductions accessed through [`Expr::agg`] are broadcast aggregates.
//! They reduce their input to a single scalar internally, then broadcast that
//! scalar to the DataFrame row count so the result remains row-aligned and can
//! participate in larger expressions.
//!
//! For example, `col("x").agg().sum()` can participate in expressions like
//! `col("x") / col("x").agg().sum()`, but broadcast reductions do not inherit
//! an output name for
//! [`DataFrame::with_columns`](crate::frame::dataframe::FrameOps::with_columns).
//! Use [`.alias()`](Expr::alias) or an explicit `(name, expr)` pair to
//! materialize the broadcast result as a named column.
//!
//! For scalar reductions, use the [`Series`] reduction APIs such as
//! [`SeriesOps::sum`](crate::SeriesOps::sum) or
//! [`SeriesOps::sum_value`](crate::SeriesOps::sum_value).
//!
//! # Literal Evaluation and Memory Behavior
//!
//! [`lit`] and [`lit_null`] are lazy expression builders over a
//! [`ScalarValue`]. Creating the [`Expr`] does not allocate GPU memory and does
//! not borrow or transfer host memory.
//!
//! During [`Expr::eval`] or [`Expr::eval_in`], each literal is materialized as
//! a device-resident [`Scalar`] in the resolved output execution domain, then
//! broadcast into a newly allocated device [`Column`] with
//! the input [`DataFrame`] row count. The returned [`Series`] owns that new
//! column and retains the resolved provenance: by default the input frame's
//! stream and memory resource, or the same-device borrowed stream and memory
//! resource supplied through [`Exec`].
//!
//! Literal evaluation enqueues libcudf work and does not intentionally
//! synchronize the host before returning. Later host-visible reads of the
//! result, such as value extraction or host export, may synchronize the
//! result's execution domain.
//!
//! # Example
//!
//! ```no_run
//! use cudf::{
//!     DataFrame, Expr, Runtime,
//!     ops::expr::{col, lit},
//! };
//!
//! let rt = Runtime::new().unwrap();
//! let x = rt.series(&[1i32, 2, 3]).unwrap();
//! let y = rt.series(&[10i32, 20, 30]).unwrap();
//! let df = DataFrame::from_series([("x", x), ("y", y)]).unwrap();
//!
//! let df2 = df
//!     .on(rt.exec())
//!     .with_columns([
//!         (col("x") + col("y")).alias("sum"),
//!         (col("x") * 2i32).alias("doubled"),
//!         (col("x") / col("x").agg().sum()).alias("share_of_total"),
//!     ])
//!     .unwrap();
//! ```

use std::{cell::RefCell, sync::Arc};

use super::{
    binary::{self, BinaryOperator},
    concatenate, copying,
    datetime::{self as datetime_ops, DatetimeComponent, DatetimeRoundingFrequency},
    reduction::{self, ReduceAgg, ScanAgg, ScanType},
    replace::{self, FillPolicy},
    search,
    strings::{self as string_ops, StringSliceOptions, StripSide},
    unary::{self, UnaryOperator},
};
use crate::{
    core::{
        column::Column,
        error::{Error, ErrorCode, Result},
        scalar::{Scalar, ScalarValue},
        types::{DataType, NullPolicy, TimestampResolution, TypeId},
    },
    frame::{dataframe::DataFrame, series::Series},
    hash::{FastHashMap, fast_hash_map_with_capacity},
    internal::resolve::missing_column_error,
    runtime::{
        Exec,
        execution_domain::{EnteredExecution, ExecutionDomain},
    },
};

struct ExprEvalContext<'df> {
    df: &'df DataFrame,
    output_domain: Arc<ExecutionDomain>,
    exec_options: EnteredExecution,
    operation: &'static str,
    normalized_columns: RefCell<FastHashMap<String, Series>>,
}

impl<'df> ExprEvalContext<'df> {
    fn new(df: &'df DataFrame, overrides: Exec<'_>, operation: &'static str) -> Result<Self> {
        let output_domain = overrides.resolve_output_domain(df.execution_domain())?;
        let exec_options = output_domain.enter_execution()?;
        Ok(Self {
            df,
            output_domain,
            exec_options,
            operation,
            normalized_columns: RefCell::new(fast_hash_map_with_capacity(df.num_columns())),
        })
    }

    fn with_output_domain(
        df: &'df DataFrame,
        output_domain: Arc<ExecutionDomain>,
        operation: &'static str,
    ) -> Result<Self> {
        let exec_options = output_domain.enter_execution()?;
        Ok(Self {
            df,
            output_domain,
            exec_options,
            operation,
            normalized_columns: RefCell::new(fast_hash_map_with_capacity(df.num_columns())),
        })
    }

    fn required_exec_options(&self) -> &cudf_sys::cudf_exec_options_t {
        self.exec_options.options()
    }

    fn output_domain(&self) -> Arc<ExecutionDomain> {
        Arc::clone(&self.output_domain)
    }

    fn column(&self, name: &str) -> Result<Series> {
        if let Some(series) = self.normalized_columns.borrow().get(name) {
            return Ok(series.clone());
        }

        let column =
            self.df.columns_map().get(name).cloned().ok_or_else(|| missing_column_error(self.operation, name))?;
        let normalized = if ExecutionDomain::same_provenance(column.execution_domain(), &self.output_domain) {
            column
        } else {
            // copy_into_domain installs both fence directions internally.
            Arc::new(
                column
                    .copy_into_domain(Arc::clone(&self.output_domain), self.operation)
                    .map_err(|error| error.with_column(name.to_string()))?,
            )
        };
        let normalized = Series::from_column_arc(normalized).map_err(|error| error.with_column(name.to_string()))?;
        self.normalized_columns.borrow_mut().insert(name.to_string(), normalized.clone());
        Ok(normalized)
    }

    fn literal(&self, value: &ScalarValue) -> Result<Series> {
        let scalar = Scalar::from_scalar_value_in_domain(value, &self.output_domain)?;
        let col = Column::from_scalar(&scalar, self.df.num_rows())?;
        Series::from_column(col)
    }

    fn wrap_column(&self, column: Column) -> Result<Series> {
        Series::from_column(column)
    }
}

pub(crate) fn eval_many_in_context<'a, I>(
    df: &DataFrame,
    overrides: Exec<'_>,
    operation: &'static str,
    exprs: I,
) -> Result<(Arc<ExecutionDomain>, Vec<Series>)>
where
    I: IntoIterator<Item = &'a Expr>,
{
    let ctx = ExprEvalContext::new(df, overrides, operation)?;
    let output_domain = ctx.output_domain();
    let results = exprs.into_iter().map(|expr| expr.eval_with_context(&ctx)).collect::<Result<_>>()?;
    Ok((output_domain, results))
}

pub(crate) fn eval_in_output_domain(
    df: &DataFrame,
    output_domain: Arc<ExecutionDomain>,
    operation: &'static str,
    expr: &Expr,
) -> Result<Series> {
    let ctx = ExprEvalContext::with_output_domain(df, output_domain, operation)?;
    expr.eval_with_context(&ctx)
}

mod agg;
mod core;
mod datetime;
mod lit;
mod strings;
mod when;

pub use core::{Expr, IntoExpr, col, concat};

pub use agg::ExprAggNamespace;
pub use datetime::ExprDatetimeNamespace;
pub use lit::{lit, lit_null};
pub use strings::{ExprStringNamespace, StringOperator};
pub use when::{ThenBuilder, WhenBuilder, when};