packages feed

dataframe-fusion (empty) → 0.2.0.0

raw patch · 9 files changed

+1648/−0 lines, 9 filesdep +HUnitdep +aesondep +basebuild-type:Customsetup-changed

Dependencies added: HUnit, aeson, base, bytestring, dataframe-arrow-bridge, dataframe-core, dataframe-fusion, dataframe-lazy, dataframe-operations, directory, filepath, text

Files

+ Setup.hs view
@@ -0,0 +1,194 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}++-- Custom Setup hook for dataframe-fusion.+--+-- Two responsibilities:+--   1. Run `cargo build --release` for the dfusion-bridge staticlib before+--      Cabal's configure step looks for it (preConf), and again before the+--      build step in case the Rust source has changed.+--   2. Inject the absolute path of the staticlib into every component's+--      extraLibDirs at confHook time. We can't put a relative path in the+--      cabal file because `ghc-pkg` rejects it at registration; ${pkgroot}+--      is resolved too late for the configure-time library lookup.++import Data.Bifunctor (second)+import Distribution.Compiler (PerCompilerFlavor (..))+import Distribution.PackageDescription (+    BuildInfo (..),+    Executable (..),+    GenericPackageDescription (..),+    Library (..),+    TestSuite (..),+ )+import Distribution.Simple+import Distribution.Simple.UserHooks (UserHooks (..))+import Distribution.Types.CondTree (CondBranch (..), CondTree (..))+#if MIN_VERSION_Cabal(3,14,0)+import Distribution.Utils.Path (+    FileOrDir (Dir),+    Lib,+    Pkg,+    SymbolicPath,+    makeSymbolicPath,+ )+#endif+import Distribution.Verbosity (normal)+import System.Directory (+    canonicalizePath,+    doesFileExist,+    findExecutable,+    getHomeDirectory,+ )+import System.Exit (ExitCode (..), exitFailure)+import System.FilePath ((</>))+import System.IO (hPutStrLn, stderr)+import System.Process (proc, readCreateProcessWithExitCode)++main :: IO ()+main =+    defaultMainWithHooks+        simpleUserHooks+            { preConf = \args flags -> do+                runCargo "preConf"+                preConf simpleUserHooks args flags+            , confHook = \(gpd, hbi) flags -> do+                libDir <-+                    canonicalizePath ("rust" </> "dfusion-bridge" </> "target" </> "release")+                let gpd' = injectExtraLibDir libDir gpd+                confHook simpleUserHooks (gpd', hbi) flags+            , preBuild = \args flags -> do+                runCargo "preBuild"+                preBuild simpleUserHooks args flags+            , preRepl = \args flags -> do+                runCargo "preRepl"+                preRepl simpleUserHooks args flags+            , preTest = \args flags -> do+                runCargo "preTest"+                preTest simpleUserHooks args flags+            }++runCargo :: String -> IO ()+runCargo phase = do+    cargo <- ensureCargo phase+    let manifest = "rust/dfusion-bridge/Cargo.toml"+    hPutStrLn stderr $+        "[dataframe-fusion/Setup.hs:"+            ++ phase+            ++ "] "+            ++ cargo+            ++ " build --release --manifest-path "+            ++ manifest+    let cp = proc cargo ["build", "--release", "--manifest-path", manifest]+    (ec, out, err) <- readCreateProcessWithExitCode cp ""+    case ec of+        ExitSuccess -> return ()+        ExitFailure n -> do+            hPutStrLn stderr $+                "[dataframe-fusion/Setup.hs:"+                    ++ phase+                    ++ "] cargo failed (exit "+                    ++ show n+                    ++ ")"+            hPutStrLn stderr out+            hPutStrLn stderr err+            exitFailure++{- | Locate cargo on PATH or under ~/.cargo/bin. If neither is found,+install the stable Rust toolchain via rustup and return the path to+the freshly installed cargo.+-}+ensureCargo :: String -> IO FilePath+ensureCargo phase = do+    mCargo <- findExecutable "cargo"+    case mCargo of+        Just p -> return p+        Nothing -> do+            home <- getHomeDirectory+            let cargoBin = home </> ".cargo" </> "bin" </> "cargo"+            existing <- doesFileExist cargoBin+            if existing+                then return cargoBin+                else do+                    installRust phase+                    return cargoBin++installRust :: String -> IO ()+installRust phase = do+    hPutStrLn stderr $+        "[dataframe-fusion/Setup.hs:"+            ++ phase+            ++ "] cargo not found; installing stable Rust toolchain via rustup"+    let installer =+            proc+                "sh"+                [ "-c"+                , "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs"+                    ++ " | sh -s -- -y --default-toolchain stable --profile minimal"+                ]+    (ec, out, err) <- readCreateProcessWithExitCode installer ""+    case ec of+        ExitSuccess -> return ()+        ExitFailure n -> do+            hPutStrLn stderr $+                "[dataframe-fusion/Setup.hs:"+                    ++ phase+                    ++ "] rustup installation failed (exit "+                    ++ show n+                    ++ ")"+            hPutStrLn stderr out+            hPutStrLn stderr err+            exitFailure++injectExtraLibDir ::+    FilePath -> GenericPackageDescription -> GenericPackageDescription+injectExtraLibDir libDir gpd =+    gpd+        { condLibrary = fmap (mapTree (addToLib libDir)) (condLibrary gpd)+        , condTestSuites =+            map+                (second (mapTree (addToTest libDir)))+                (condTestSuites gpd)+        , condExecutables =+            map+                (second (mapTree (addToExe libDir)))+                (condExecutables gpd)+        }++mapTree :: (a -> a) -> CondTree v c a -> CondTree v c a+mapTree f (CondNode d c bs) = CondNode (f d) c (map mapBranch bs)+  where+    mapBranch b =+        b+            { condBranchIfTrue = mapTree f (condBranchIfTrue b)+            , condBranchIfFalse = fmap (mapTree f) (condBranchIfFalse b)+            }++addToLib :: FilePath -> Library -> Library+addToLib p lib = lib{libBuildInfo = addExtra p (libBuildInfo lib)}++addToTest :: FilePath -> TestSuite -> TestSuite+addToTest p t = t{testBuildInfo = addExtra p (testBuildInfo t)}++addToExe :: FilePath -> Executable -> Executable+addToExe p e = e{buildInfo = addExtra p (buildInfo e)}++-- TODO: mchavinda - make this linking windows compatible.+addExtra :: FilePath -> BuildInfo -> BuildInfo+addExtra p bi =+    bi+        { extraLibDirs = mkLibDir p : extraLibDirs bi+        , options = addRpath (options bi)+        }+  where+    addRpath (PerCompilerFlavor ghc ghcjs) =+        PerCompilerFlavor (("-optl-Wl,-rpath," ++ p) : ghc) ghcjs++-- | In Cabal 3.14+, extraLibDirs holds 'SymbolicPath' values, not 'FilePath'.+#if MIN_VERSION_Cabal(3,14,0)+mkLibDir :: FilePath -> SymbolicPath Pkg ('Dir Lib)+mkLibDir = makeSymbolicPath+#else+mkLibDir :: FilePath -> FilePath+mkLibDir = id+#endif
+ cbits/dfusion_bridge.h view
@@ -0,0 +1,96 @@+/*+ * dfusion_bridge.h+ *+ * C ABI for the Rust crate `dfusion-bridge`, which wraps Apache DataFusion+ * for use from Haskell. All payloads named *_json are UTF-8 encoded JSON+ * strings; expression payloads use the wire format produced by+ * DataFrame.IR.ExprJson.encodeExpr. Validity bitmaps and Arrow buffers cross+ * the boundary unchanged via the Arrow C Data Interface (see arrow_abi.h+ * in the parent dataframe package).+ */++#ifndef DFUSION_BRIDGE_H+#define DFUSION_BRIDGE_H++#include <stdint.h>++#ifdef __cplusplus+extern "C" {+#endif++typedef struct DfCtx  DfCtx;+typedef struct DfPlan DfPlan;++/* Context lifecycle. One DfCtx owns a tokio runtime and a SessionContext. */+DfCtx *df_ctx_new(void);+void   df_ctx_free(DfCtx *ctx);++/* Plan handles. df_plan_free is idempotent on NULL. */+void   df_plan_free(DfPlan *plan);++/* Thread-local last-error message. Pointer is valid until the next df_*+ * call on the same thread; copy if you need to keep it. */+const char *df_last_error(void);++/* ----- Sources ---------------------------------------------------------- */++/* Scan a CSV file. schema_json may be NULL for type inference; otherwise it+ * is a JSON object: { "fields": [["col_name", "int|double|text|bool"], ...] }.+ * Returns NULL on error (see df_last_error). */+DfPlan *df_scan_csv(DfCtx *ctx,+                    const char *path,+                    const char *schema_json);++/* ----- Operators -------------------------------------------------------- */++/* Filter rows where the boolean expression evaluates true. */+DfPlan *df_plan_filter(DfPlan *plan, const char *expr_json);++/* Keep at most `n` rows. */+DfPlan *df_plan_take(DfPlan *plan, uint64_t n);++/* Project to the named columns. names_json is a JSON array of strings. */+DfPlan *df_plan_select(DfPlan *plan, const char *names_json);++/* Add a derived column. */+DfPlan *df_plan_derive(DfPlan *plan,+                       const char *col_name,+                       const char *expr_json);++/* Sort by a list of (column, ascending) pairs.+ * orders_json shape: [{"col": "name", "asc": true}, ...] */+DfPlan *df_plan_sort_by(DfPlan *plan, const char *orders_json);++/* GroupBy + aggregate.+ *   keys_json:  JSON array of column-name strings to group by.+ *   aggs_json:  JSON array of {"name": "...", "expr": <agg-expr-json>} objects,+ *               where each agg-expr-json is a top-level "agg" node.            */+DfPlan *df_plan_groupby_aggregate(DfPlan *plan,+                                   const char *keys_json,+                                   const char *aggs_json);++/* Join two plans on key pairs.+ *   how:      "inner" | "left" | "right" | "outer"+ *   on_json:  JSON array of [left_key, right_key] string pairs.                */+DfPlan *df_plan_join(DfPlan *left,+                     DfPlan *right,+                     const char *how,+                     const char *on_json);++/* ----- Materialization -------------------------------------------------- */++/* Execute the plan, concatenate batches, and export the result via the+ * Arrow C Data Interface. *schema_out and *array_out receive the addresses+ * of newly allocated FFI_ArrowSchema and FFI_ArrowArray structs (cast to+ * uint64_t). The caller is responsible for invoking the producer's release+ * callbacks once the data has been copied (matches existing arrowToDataframe+ * in the dataframe package). Returns 0 on success, -1 on error. */+int32_t df_plan_collect(DfPlan *plan,+                        uint64_t *schema_out,+                        uint64_t *array_out);++#ifdef __cplusplus+}+#endif++#endif /* DFUSION_BRIDGE_H */
+ dataframe-fusion.cabal view
@@ -0,0 +1,106 @@+cabal-version:      3.0+name:               dataframe-fusion+version:            0.2.0.0+synopsis: Apache DataFusion backend for the dataframe typed API.++description: A typed, plan-based wrapper around Apache DataFusion. Mirrors+             the DataFrame.Typed.Lazy surface but lowers each operation to+             a DataFusion logical plan via a Rust bridge crate. Results are+             returned via the Arrow C Data Interface and reconstructed as+             ordinary TypedDataFrames in Haskell.++bug-reports: https://github.com/mchav/dataframe/issues+license:            MIT+author:             Michael Chavinda+maintainer:         mschavinda@gmail.com++copyright: (c) 2024-2026 Michael Chavinda+category: Data+build-type: Custom++extra-source-files: cbits/dfusion_bridge.h+                    rust/dfusion-bridge/Cargo.toml+                    rust/dfusion-bridge/src/lib.rs++-- The license is the same as the parent dataframe package.++custom-setup+    setup-depends: base       >= 4   && < 5,+                   Cabal      >= 3.0 && < 4,+                   directory  >= 1.3 && < 2,+                   filepath   >= 1.4 && < 2,+                   process    >= 1.6 && < 2++common warnings+    ghc-options:+        -Wincomplete-patterns+        -Wincomplete-uni-patterns+        -Wunused-imports+        -Wunused-local-binds+        -Wunused-packages++-- Mirror the meta-package's opt-out flags. dataframe-fusion depends on+-- `dataframe-arrow-bridge`, which always needs the CSV and Parquet readers,+-- so the library has to disappear in lockstep when either is disabled.+-- Pass the matching flag (e.g. `--constraint="dataframe-fusion +no-csv"`)+-- when building the meta package with `--flags="no-csv"` against the full+-- cabal.project.+flag no-csv+    default:     False+    manual:      True+    description: Skip building dataframe-fusion when the meta package's+                 CSV backend is disabled.++flag no-parquet+    default:     False+    manual:      True+    description: Skip building dataframe-fusion when the meta package's+                 Parquet backend is disabled.++library+    import: warnings+    if flag(no-csv) || flag(no-parquet)+        buildable: False+    exposed-modules:  DataFrame.Fusion.FFI+                      DataFrame.Fusion.Plan+                      DataFrame.Fusion.Typed+    build-depends:    base       >= 4    && < 5,+                      bytestring >= 0.11 && < 0.14,+                      aeson      >= 0.11 && < 3,+                      dataframe-core ^>= 2.5,+                      dataframe-lazy ^>= 2.4.1,+                      dataframe-operations ^>= 2.5,+                      dataframe-arrow-bridge ^>= 1.0,+                      text       >= 2.1  && < 3+    hs-source-dirs:   src+    include-dirs:     cbits+    includes:         dfusion_bridge.h+    install-includes: dfusion_bridge.h+    -- iconv lives in glibc on Linux; only macOS exposes it as a separate lib.+    extra-libraries:  dfusion_bridge bz2 lzma z+    if os(darwin)+        extra-libraries: iconv+        ld-options: -framework CoreFoundation -framework Security -framework SystemConfiguration+    default-language: Haskell2010++test-suite tests+    import: warnings+    type: exitcode-stdio-1.0+    main-is: Main.hs+    if flag(no-csv) || flag(no-parquet)+        buildable: False+    build-depends:    base       >= 4    && < 5,+                      dataframe-core ^>= 2.5,+                      dataframe-fusion,+                      dataframe-operations ^>= 2.5,+                      directory  >= 1.3  && < 2,+                      filepath   >= 1.4  && < 2,+                      HUnit      >= 1.6 && < 1.8,+                      text       >= 2.1  && < 3+    hs-source-dirs:   tests+    -- iconv lives in glibc on Linux; only macOS exposes it as a separate lib.+    extra-libraries:  dfusion_bridge bz2 lzma z+    if os(darwin)+        extra-libraries: iconv+        ld-options: -framework CoreFoundation -framework Security -framework SystemConfiguration+    default-language: Haskell2010
+ rust/dfusion-bridge/Cargo.toml view
@@ -0,0 +1,24 @@+[package]+name = "dfusion-bridge"+version = "0.1.0"+edition = "2021"+publish = false++[lib]+name = "dfusion_bridge"+# `cdylib` for `cabal repl` (GHCi can only load shared libraries) and for+# executables that link dynamically. `staticlib` is kept so static linking+# remains an option (e.g. for distribution as a self-contained binary).+crate-type = ["cdylib", "staticlib"]++[dependencies]+datafusion = "45"+arrow = { version = "54", features = ["ffi"] }+arrow-schema = "54"+tokio = { version = "1", features = ["rt-multi-thread", "macros"] }+serde = { version = "1", features = ["derive"] }+serde_json = "1"++[profile.release]+lto = "thin"+codegen-units = 1
+ rust/dfusion-bridge/src/lib.rs view
@@ -0,0 +1,627 @@+//! C ABI shim that exposes Apache DataFusion to Haskell.+//!+//! Each plan-extending function takes an opaque `*mut DfPlan` (an+//! `Arc<DataFusion DataFrame>`), builds a new plan node, and returns a+//! fresh `*mut DfPlan`. `df_plan_collect` materializes the plan and+//! exports the result through the Arrow C Data Interface; the Haskell side+//! consumes those pointers via `arrowToDataframe`.++use std::cell::RefCell;+use std::ffi::{c_char, CStr, CString};+use std::ptr;+use std::sync::Arc;++use arrow::array::{ArrayRef, RecordBatch, StructArray};+use arrow::compute::concat_batches;+use arrow::datatypes::{DataType, Field, Schema};+use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};+use datafusion::dataframe::DataFrame as DfDataFrame;+use datafusion::execution::context::SessionContext;+use datafusion::functions_aggregate::expr_fn::{avg, count, max as agg_max, median, min as agg_min, sum as agg_sum};+use datafusion::logical_expr::{col, lit, BinaryExpr, Expr, JoinType, Operator};+use datafusion::prelude::CsvReadOptions;+use datafusion::scalar::ScalarValue;+use serde_json::Value;+use tokio::runtime::Runtime;++thread_local! {+    static LAST_ERROR: RefCell<Option<CString>> = RefCell::new(None);+}++fn set_error(msg: impl Into<String>) {+    let s = msg.into();+    let cs = CString::new(s).unwrap_or_else(|_| CString::new("dfusion: error contains NUL").unwrap());+    LAST_ERROR.with(|cell| *cell.borrow_mut() = Some(cs));+}++fn clear_error() {+    LAST_ERROR.with(|cell| *cell.borrow_mut() = None);+}++#[no_mangle]+pub extern "C" fn df_last_error() -> *const c_char {+    LAST_ERROR.with(|cell| match &*cell.borrow() {+        Some(cs) => cs.as_ptr(),+        None => ptr::null(),+    })+}++pub struct DfCtx {+    runtime: Runtime,+    session: Arc<SessionContext>,+}++pub struct DfPlan {+    ctx: Arc<DfCtxInner>,+    df: DfDataFrame,+}++// Inner shared state so DfPlans keep the runtime/session alive even if the+// caller frees their DfCtx handle first.+struct DfCtxInner {+    runtime: Runtime,+    session: Arc<SessionContext>,+}++#[no_mangle]+pub extern "C" fn df_ctx_new() -> *mut DfCtx {+    clear_error();+    match Runtime::new() {+        Ok(rt) => {+            let session = Arc::new(SessionContext::new());+            let ctx = Box::new(DfCtx { runtime: rt, session });+            Box::into_raw(ctx)+        }+        Err(e) => {+            set_error(format!("df_ctx_new: failed to start tokio runtime: {e}"));+            ptr::null_mut()+        }+    }+}++#[no_mangle]+pub extern "C" fn df_ctx_free(ctx: *mut DfCtx) {+    if ctx.is_null() { return; }+    unsafe { drop(Box::from_raw(ctx)); }+}++#[no_mangle]+pub extern "C" fn df_plan_free(plan: *mut DfPlan) {+    if plan.is_null() { return; }+    unsafe { drop(Box::from_raw(plan)); }+}++// Wrap the result of a plan-builder operation. On error, sets last_error and+// returns null.+fn wrap_plan(ctx: Arc<DfCtxInner>, result: datafusion::error::Result<DfDataFrame>) -> *mut DfPlan {+    match result {+        Ok(df) => Box::into_raw(Box::new(DfPlan { ctx, df })),+        Err(e) => {+            set_error(format!("{e}"));+            ptr::null_mut()+        }+    }+}++unsafe fn cstr_or_err<'a>(p: *const c_char, what: &str) -> Option<&'a str> {+    if p.is_null() {+        set_error(format!("{what}: null pointer"));+        return None;+    }+    match CStr::from_ptr(p).to_str() {+        Ok(s) => Some(s),+        Err(e) => {+            set_error(format!("{what}: invalid utf-8: {e}"));+            None+        }+    }+}++#[no_mangle]+pub unsafe extern "C" fn df_scan_csv(+    ctx: *mut DfCtx,+    path: *const c_char,+    schema_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if ctx.is_null() {+        set_error("df_scan_csv: null context");+        return ptr::null_mut();+    }+    let ctx_ref = &*ctx;+    let inner = Arc::new(DfCtxInner {+        runtime: clone_runtime(),+        session: ctx_ref.session.clone(),+    });+    let path_str = match cstr_or_err(path, "df_scan_csv.path") {+        Some(s) => s.to_owned(),+        None => return ptr::null_mut(),+    };++    // Schema override (optional).+    let opts = CsvReadOptions::new();+    let opts_owned: Option<Schema> = if schema_json.is_null() {+        None+    } else {+        match cstr_or_err(schema_json, "df_scan_csv.schema_json") {+            Some(s) => match parse_schema_json(s) {+                Ok(sch) => Some(sch),+                Err(e) => {+                    set_error(format!("df_scan_csv: bad schema_json: {e}"));+                    return ptr::null_mut();+                }+            },+            None => return ptr::null_mut(),+        }+    };++    let session = ctx_ref.session.clone();+    let result = ctx_ref.runtime.block_on(async {+        let opts = match &opts_owned {+            Some(s) => opts.schema(s),+            None => opts,+        };+        session.read_csv(path_str, opts).await+    });+    wrap_plan(inner, result)+}++// We want each DfPlan to share runtime/session with siblings without copying+// the handle out from the user-owned DfCtx. Easiest path: every plan op+// captures the SessionContext + a shared Tokio runtime by Arc reference.+// Since Tokio's Runtime is not Clone, we keep a single runtime per process+// for plan-builder calls and use the session-bound runtime for execution.+fn clone_runtime() -> Runtime {+    // Plan-builder ops are essentially synchronous; create a small dedicated+    // runtime per plan handle so that spawned blocking work doesn't poison+    // the caller's context. Cheap (<1ms) compared to query execution.+    Runtime::new().expect("tokio runtime")+}++fn parse_schema_json(s: &str) -> Result<Schema, String> {+    #[derive(serde::Deserialize)]+    struct SchemaWire { fields: Vec<(String, String)> }+    let wire: SchemaWire = serde_json::from_str(s).map_err(|e| e.to_string())?;+    let fields: Vec<Field> = wire+        .fields+        .into_iter()+        .map(|(name, ty)| {+            let dt = match ty.as_str() {+                "int" | "int64" => DataType::Int64,+                "int32" => DataType::Int32,+                "double" | "float64" => DataType::Float64,+                "float" | "float32" => DataType::Float32,+                "bool" => DataType::Boolean,+                "text" | "string" | "utf8" => DataType::Utf8,+                other => return Err(format!("unsupported type tag '{other}'")),+            };+            Ok(Field::new(&name, dt, true))+        })+        .collect::<Result<_, String>>()?;+    Ok(Schema::new(fields))+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_filter(+    plan: *mut DfPlan,+    expr_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_filter: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let expr_str = match cstr_or_err(expr_json, "df_plan_filter.expr_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let val: Value = match serde_json::from_str(expr_str) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_filter: invalid json: {e}"));+            return ptr::null_mut();+        }+    };+    let expr = match decode_expr(&val) {+        Ok(e) => e,+        Err(e) => {+            set_error(format!("df_plan_filter: {e}"));+            return ptr::null_mut();+        }+    };+    let result = plan_ref.df.clone().filter(expr);+    wrap_plan(plan_ref.ctx.clone(), result)+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_take(+    plan: *mut DfPlan,+    n: u64,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_take: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let result = plan_ref.df.clone().limit(0, Some(n as usize));+    wrap_plan(plan_ref.ctx.clone(), result)+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_select(+    plan: *mut DfPlan,+    names_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_select: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let s = match cstr_or_err(names_json, "df_plan_select.names_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let names: Vec<String> = match serde_json::from_str(s) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_select: invalid json: {e}"));+            return ptr::null_mut();+        }+    };+    let exprs: Vec<Expr> = names.iter().map(|n| col(n)).collect();+    let result = plan_ref.df.clone().select(exprs);+    wrap_plan(plan_ref.ctx.clone(), result)+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_derive(+    plan: *mut DfPlan,+    col_name: *const c_char,+    expr_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_derive: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let name_str = match cstr_or_err(col_name, "df_plan_derive.col_name") {+        Some(s) => s.to_owned(),+        None => return ptr::null_mut(),+    };+    let expr_str = match cstr_or_err(expr_json, "df_plan_derive.expr_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let val: Value = match serde_json::from_str(expr_str) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_derive: invalid json: {e}"));+            return ptr::null_mut();+        }+    };+    let expr = match decode_expr(&val) {+        Ok(e) => e,+        Err(e) => {+            set_error(format!("df_plan_derive: {e}"));+            return ptr::null_mut();+        }+    };+    let result = plan_ref.df.clone().with_column(&name_str, expr);+    wrap_plan(plan_ref.ctx.clone(), result)+}++/// Group-by + aggregate. keys_json is a list of column name strings; aggs_json+/// is a list of `{"name": "alias", "expr": <agg_expr_json>}` objects, where+/// each agg_expr_json is a top-level "agg" node produced by Haskell-side+/// encodeExpr on an `Agg ...` expression.+#[no_mangle]+pub unsafe extern "C" fn df_plan_groupby_aggregate(+    plan: *mut DfPlan,+    keys_json: *const c_char,+    aggs_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_groupby_aggregate: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let keys_str = match cstr_or_err(keys_json, "df_plan_groupby_aggregate.keys_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let aggs_str = match cstr_or_err(aggs_json, "df_plan_groupby_aggregate.aggs_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let keys: Vec<String> = match serde_json::from_str(keys_str) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_groupby_aggregate: keys json: {e}"));+            return ptr::null_mut();+        }+    };+    #[derive(serde::Deserialize)]+    struct AggEntry {+        name: String,+        expr: Value,+    }+    let entries: Vec<AggEntry> = match serde_json::from_str(aggs_str) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_groupby_aggregate: aggs json: {e}"));+            return ptr::null_mut();+        }+    };+    let group_exprs: Vec<Expr> = keys.iter().map(|k| col(k)).collect();+    let agg_exprs: Vec<Expr> = match entries+        .iter()+        .map(|e| Ok(decode_expr(&e.expr)?.alias(&e.name)))+        .collect::<Result<Vec<_>, String>>()+    {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_groupby_aggregate: {e}"));+            return ptr::null_mut();+        }+    };+    let result = plan_ref.df.clone().aggregate(group_exprs, agg_exprs);+    wrap_plan(plan_ref.ctx.clone(), result)+}++/// Join two plans on a list of (left, right) key columns. `how` is one of+/// "inner", "left", "right", "outer".+#[no_mangle]+pub unsafe extern "C" fn df_plan_join(+    left: *mut DfPlan,+    right: *mut DfPlan,+    how: *const c_char,+    on_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if left.is_null() || right.is_null() {+        set_error("df_plan_join: null plan");+        return ptr::null_mut();+    }+    let left_ref = &*left;+    let right_ref = &*right;+    let how_str = match cstr_or_err(how, "df_plan_join.how") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let on_str = match cstr_or_err(on_json, "df_plan_join.on_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    let pairs: Vec<(String, String)> = match serde_json::from_str(on_str) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_join: on json: {e}"));+            return ptr::null_mut();+        }+    };+    let join_type = match how_str {+        "inner" => JoinType::Inner,+        "left"  => JoinType::Left,+        "right" => JoinType::Right,+        "outer" | "full_outer" => JoinType::Full,+        other => {+            set_error(format!("df_plan_join: unsupported how '{other}'"));+            return ptr::null_mut();+        }+    };+    let left_keys: Vec<&str> = pairs.iter().map(|(l, _)| l.as_str()).collect();+    let right_keys: Vec<&str> = pairs.iter().map(|(_, r)| r.as_str()).collect();+    // Alias both sides so DataFusion treats them as distinct relations, even+    // when both were loaded with the anonymous "?table?" qualifier.+    let left_aliased = match left_ref.df.clone().alias("l") {+        Ok(d) => d,+        Err(e) => { set_error(format!("df_plan_join: {e}")); return ptr::null_mut(); }+    };+    let right_aliased = match right_ref.df.clone().alias("r") {+        Ok(d) => d,+        Err(e) => { set_error(format!("df_plan_join: {e}")); return ptr::null_mut(); }+    };+    let result = left_aliased.join(+        right_aliased,+        join_type,+        &left_keys,+        &right_keys,+        None,+    );+    wrap_plan(left_ref.ctx.clone(), result)+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_sort_by(+    plan: *mut DfPlan,+    orders_json: *const c_char,+) -> *mut DfPlan {+    clear_error();+    if plan.is_null() {+        set_error("df_plan_sort_by: null plan");+        return ptr::null_mut();+    }+    let plan_ref = &*plan;+    let s = match cstr_or_err(orders_json, "df_plan_sort_by.orders_json") {+        Some(s) => s,+        None => return ptr::null_mut(),+    };+    #[derive(serde::Deserialize)]+    struct SortSpec { col: String, asc: bool }+    let specs: Vec<SortSpec> = match serde_json::from_str(s) {+        Ok(v) => v,+        Err(e) => {+            set_error(format!("df_plan_sort_by: invalid json: {e}"));+            return ptr::null_mut();+        }+    };+    let sort_exprs: Vec<datafusion::logical_expr::SortExpr> = specs+        .into_iter()+        .map(|sp| {+            datafusion::logical_expr::SortExpr {+                expr: col(&sp.col),+                asc: sp.asc,+                nulls_first: !sp.asc,+            }+        })+        .collect();+    let result = plan_ref.df.clone().sort(sort_exprs);+    wrap_plan(plan_ref.ctx.clone(), result)+}++fn decode_expr(v: &Value) -> Result<Expr, String> {+    let obj = v.as_object().ok_or("expected expr object")?;+    let node = obj.get("node").and_then(Value::as_str).ok_or("missing 'node'")?;+    match node {+        "col" => {+            let name = obj.get("name").and_then(Value::as_str).ok_or("col: missing name")?;+            Ok(col(name))+        }+        "lit" => {+            let out_type = obj.get("out_type").and_then(Value::as_str).ok_or("lit: missing out_type")?;+            let value = obj.get("value").ok_or("lit: missing value")?;+            decode_literal(out_type, value)+        }+        "binary" => {+            let op = obj.get("op").and_then(Value::as_str).ok_or("binary: missing op")?;+            let lhs = decode_expr(obj.get("lhs").ok_or("binary: missing lhs")?)?;+            let rhs = decode_expr(obj.get("rhs").ok_or("binary: missing rhs")?)?;+            // Wire names match DataFrame.IR.ExprJson.recognizeBinary.+            let operator = match op {+                "eq"   => Operator::Eq,+                "neq"  => Operator::NotEq,+                "lt"   => Operator::Lt,+                "leq"  => Operator::LtEq,+                "gt"   => Operator::Gt,+                "geq"  => Operator::GtEq,+                "and"  => Operator::And,+                "or"   => Operator::Or,+                "add"  => Operator::Plus,+                "sub"  => Operator::Minus,+                "mult" => Operator::Multiply,+                "divide" => Operator::Divide,+                "div"  => Operator::Divide,+                "mod"  => Operator::Modulo,+                other => return Err(format!("unsupported binary op '{other}'")),+            };+            Ok(Expr::BinaryExpr(BinaryExpr::new(Box::new(lhs), operator, Box::new(rhs))))+        }+        "if" => {+            let cond = decode_expr(obj.get("cond").ok_or("if: missing cond")?)?;+            let then_ = decode_expr(obj.get("then").ok_or("if: missing then")?)?;+            let else_ = decode_expr(obj.get("else").ok_or("if: missing else")?)?;+            // CASE WHEN cond THEN then ELSE else END+            Ok(datafusion::logical_expr::case(cond)+                .when(lit(true), then_)+                .otherwise(else_)+                .map_err(|e| e.to_string())?)+        }+        "unary" => {+            let op = obj.get("op").and_then(Value::as_str).ok_or("unary: missing op")?;+            let arg = decode_expr(obj.get("arg").ok_or("unary: missing arg")?)?;+            match op {+                "not" => Ok(!arg),+                "negate" => Ok(-arg),+                "abs" => Ok(datafusion::functions::math::abs().call(vec![arg])),+                "toDouble" => Ok(datafusion::logical_expr::cast(arg, DataType::Float64)),+                other => Err(format!("unsupported unary op '{other}'")),+            }+        }+        "agg" => {+            let name = obj.get("agg").and_then(Value::as_str).ok_or("agg: missing 'agg' name")?;+            let arg = decode_expr(obj.get("arg").ok_or("agg: missing arg")?)?;+            match name {+                "sum"     => Ok(agg_sum(arg)),+                "count"   => Ok(count(arg)),+                "mean" | "avg" => Ok(avg(arg)),+                "min"     => Ok(agg_min(arg)),+                "max"     => Ok(agg_max(arg)),+                "median"  => Ok(median(arg)),+                other     => Err(format!("unsupported aggregation '{other}'")),+            }+        }+        other => Err(format!("unknown expr node '{other}'")),+    }+}++fn decode_literal(out_type: &str, v: &Value) -> Result<Expr, String> {+    let scalar = match out_type {+        "int" | "int64" => ScalarValue::Int64(v.as_i64()),+        "int32" => ScalarValue::Int32(v.as_i64().map(|x| x as i32)),+        "double" | "float64" => ScalarValue::Float64(v.as_f64()),+        "float" | "float32" => ScalarValue::Float32(v.as_f64().map(|x| x as f32)),+        "bool" => ScalarValue::Boolean(v.as_bool()),+        "text" | "string" | "utf8" => ScalarValue::Utf8(v.as_str().map(|s| s.to_owned())),+        other => return Err(format!("lit: unsupported type tag '{other}'")),+    };+    Ok(Expr::Literal(scalar))+}++#[no_mangle]+pub unsafe extern "C" fn df_plan_collect(+    plan: *mut DfPlan,+    schema_out: *mut u64,+    array_out: *mut u64,+) -> i32 {+    clear_error();+    if plan.is_null() || schema_out.is_null() || array_out.is_null() {+        set_error("df_plan_collect: null pointer");+        return -1;+    }+    let plan_ref = &*plan;+    let batches: Result<Vec<RecordBatch>, _> =+        plan_ref.ctx.runtime.block_on(async { plan_ref.df.clone().collect().await });+    let batches = match batches {+        Ok(b) => b,+        Err(e) => {+            set_error(format!("df_plan_collect: {e}"));+            return -1;+        }+    };++    let schema = plan_ref.df.schema().as_arrow().clone();+    let combined: RecordBatch = if batches.is_empty() {+        RecordBatch::new_empty(Arc::new(schema.clone()))+    } else {+        match concat_batches(&Arc::new(schema.clone()), &batches) {+            Ok(b) => b,+            Err(e) => {+                set_error(format!("df_plan_collect: concat: {e}"));+                return -1;+            }+        }+    };++    // Convert to a top-level StructArray so the Arrow C Data Interface export+    // produces a single (schema, array) pair where children == columns.+    let struct_array: StructArray = combined.into();+    let array_ref: ArrayRef = Arc::new(struct_array);+    let array_data = array_ref.to_data();++    let ffi_array = match FFI_ArrowArray::new(&array_data) {+        a => a,+    };+    let ffi_schema = match FFI_ArrowSchema::try_from(array_ref.data_type()) {+        Ok(s) => s,+        Err(e) => {+            set_error(format!("df_plan_collect: schema export: {e}"));+            return -1;+        }+    };++    // Move both onto the heap; Haskell owns them now and is responsible for+    // calling the producer-supplied release callbacks (matches existing+    // arrowToDataframe semantics).+    let ffi_schema_box = Box::new(ffi_schema);+    let ffi_array_box = Box::new(ffi_array);+    let schema_ptr = Box::into_raw(ffi_schema_box) as u64;+    let array_ptr = Box::into_raw(ffi_array_box) as u64;+    *schema_out = schema_ptr;+    *array_out = array_ptr;+    0+}
+ src/DataFrame/Fusion/FFI.hs view
@@ -0,0 +1,80 @@+{-# LANGUAGE ForeignFunctionInterface #-}++{- | Foreign-function imports for the @dfusion-bridge@ Rust staticlib.+Pointers returned from Rust are owned by the caller and freed via+'df_ctx_free' / 'df_plan_free'. On error the @df_*@ functions return+null and set a thread-local message retrievable via 'df_last_error'.+-}+module DataFrame.Fusion.FFI (+    DfCtx,+    DfPlan,+    df_ctx_new,+    df_ctx_free,+    df_plan_free,+    df_plan_freep,+    df_last_error,+    df_scan_csv,+    df_plan_filter,+    df_plan_take,+    df_plan_select,+    df_plan_derive,+    df_plan_sort_by,+    df_plan_groupby_aggregate,+    df_plan_join,+    df_plan_collect,+) where++import Data.Word (Word64)+import Foreign.C.String (CString)+import Foreign.C.Types (CInt (..))+import Foreign.Ptr (FunPtr, Ptr)++-- Opaque handle types. Layout is irrelevant on the Haskell side; we only+-- ever pass the pointers through.+data DfCtx+data DfPlan++foreign import ccall unsafe "df_ctx_new"+    df_ctx_new :: IO (Ptr DfCtx)++foreign import ccall unsafe "df_ctx_free"+    df_ctx_free :: Ptr DfCtx -> IO ()++foreign import ccall unsafe "df_plan_free"+    df_plan_free :: Ptr DfPlan -> IO ()++-- | A FunPtr to 'df_plan_free' suitable for use as a 'ForeignPtr' finalizer.+foreign import ccall unsafe "&df_plan_free"+    df_plan_freep :: FunPtr (Ptr DfPlan -> IO ())++foreign import ccall unsafe "df_last_error"+    df_last_error :: IO CString++foreign import ccall unsafe "df_scan_csv"+    df_scan_csv :: Ptr DfCtx -> CString -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_filter"+    df_plan_filter :: Ptr DfPlan -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_take"+    df_plan_take :: Ptr DfPlan -> Word64 -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_select"+    df_plan_select :: Ptr DfPlan -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_derive"+    df_plan_derive :: Ptr DfPlan -> CString -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_sort_by"+    df_plan_sort_by :: Ptr DfPlan -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_groupby_aggregate"+    df_plan_groupby_aggregate ::+        Ptr DfPlan -> CString -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_join"+    df_plan_join ::+        Ptr DfPlan -> Ptr DfPlan -> CString -> CString -> IO (Ptr DfPlan)++foreign import ccall unsafe "df_plan_collect"+    df_plan_collect :: Ptr DfPlan -> Ptr Word64 -> Ptr Word64 -> IO CInt
+ src/DataFrame/Fusion/Plan.hs view
@@ -0,0 +1,99 @@+{-# LANGUAGE OverloadedStrings #-}++-- | Plan-handle lifecycle and error wrapping for the DataFusion bridge.+module DataFrame.Fusion.Plan (+    -- * Context+    Context,+    newContext,+    withContext,++    -- * Plan handles+    PlanHandle,+    wrapPlan,+    withPlan,++    -- * Calling Rust+    runPlanOp,+    DataFusionError (..),++    -- * Materialization+    collectArrow,+) where++import Control.Exception (Exception, throwIO)+import qualified Data.Text as T+import Foreign.C.String (peekCString)+import Foreign.ForeignPtr (ForeignPtr, newForeignPtr, withForeignPtr)+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr (Ptr, nullPtr, wordPtrToPtr)+import Foreign.Storable (peek)++import qualified DataFrame.Fusion.FFI as F+import qualified DataFrame.IO.Arrow as Arrow+import DataFrame.Internal.DataFrame (DataFrame)++-- | Process-wide DataFusion session + tokio runtime.+newtype Context = Context (Ptr F.DfCtx)++-- | A DataFusion plan handle.+newtype PlanHandle = PlanHandle (ForeignPtr F.DfPlan)++newtype DataFusionError = DataFusionError T.Text+    deriving (Show)+instance Exception DataFusionError++{- | Allocate a session. The handle stays alive until garbage-collected; in+practice that is the program lifetime.+-}+newContext :: IO Context+newContext = do+    p <- F.df_ctx_new+    if p == nullPtr+        then do+            err <- readLastError+            throwIO (DataFusionError ("newContext: " <> err))+        else return (Context p)++-- | Use the underlying context pointer for the duration of an action.+withContext :: Context -> (Ptr F.DfCtx -> IO a) -> IO a+withContext (Context p) k = k p++-- | Wrap a freshly returned plan pointer with a finalizer.+wrapPlan :: Ptr F.DfPlan -> IO PlanHandle+wrapPlan p+    | p == nullPtr = do+        err <- readLastError+        throwIO (DataFusionError ("plan op returned null: " <> err))+    | otherwise = PlanHandle <$> newForeignPtr F.df_plan_freep p++-- | Use the plan pointer for the duration of the action.+withPlan :: PlanHandle -> (Ptr F.DfPlan -> IO a) -> IO a+withPlan (PlanHandle fp) = withForeignPtr fp++-- | Run a Rust function that returns a plan pointer; throw on null.+runPlanOp :: IO (Ptr F.DfPlan) -> IO PlanHandle+runPlanOp action = action >>= wrapPlan++readLastError :: IO T.Text+readLastError = do+    cs <- F.df_last_error+    if cs == nullPtr+        then return "<no error message>"+        else T.pack <$> peekCString cs++-- | Execute a plan and import the result as an untyped DataFrame.+collectArrow :: PlanHandle -> IO DataFrame+collectArrow plan = withPlan plan $ \p ->+    alloca $ \schemaOutPtr ->+        alloca $ \arrayOutPtr -> do+            rc <- F.df_plan_collect p schemaOutPtr arrayOutPtr+            if rc /= 0+                then do+                    err <- readLastError+                    throwIO (DataFusionError ("collect: " <> err))+                else do+                    schemaAddr <- peek schemaOutPtr+                    arrayAddr <- peek arrayOutPtr+                    Arrow.arrowToDataframe+                        (wordPtrToPtr (fromIntegral schemaAddr))+                        (wordPtrToPtr (fromIntegral arrayAddr))
+ src/DataFrame/Fusion/Typed.hs view
@@ -0,0 +1,275 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{- | Typed DataFusion-backed query API. Mirrors 'DataFrame.Typed.Lazy' but runs+on Rust-side DataFusion via FFI.+-}+module DataFrame.Fusion.Typed (+    -- * Carrier+    DataFrame,++    -- * Sources+    scanCsv,++    -- * Operators+    filter,+    take,+    select,+    derive,+    sortBy,+    SortOrder (..),++    -- * Aggregation+    Grouped,+    groupBy,+    aggregate,++    -- * Joins+    innerJoin,+    leftJoin,+    rightJoin,+    fullOuterJoin,++    -- * Materialization+    run,++    -- * Re-exports+    module DataFrame.Typed.Expr,+    module DataFrame.Typed.Types,+    module DataFrame.Fusion.Plan,++    -- ** Aggregation builders (re-exported from DataFrame.Typed.Aggregate)+    AGG.as,+) where++import qualified Data.Aeson as Aeson+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as BL+import Data.Kind (Type)+import Data.Proxy (Proxy (..))+import qualified Data.Text as T+import Foreign.C.String (CString, withCString)+import Foreign.Ptr (nullPtr)+import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)+import Prelude hiding (filter, take)++import qualified DataFrame.Internal.Column as IC+import qualified DataFrame.Internal.Expression as IE+import DataFrame.Lazy (SortOrder (..))+import qualified DataFrame.Typed.Aggregate as AGG++import qualified DataFrame.Fusion.FFI as F+import DataFrame.Fusion.Plan+import DataFrame.IR.ExprJson (encodeExprToBytes)+import DataFrame.Typed.Expr+import DataFrame.Typed.Freeze (unsafeFreeze)+import DataFrame.Typed.Schema+import DataFrame.Typed.Types++{- | A query plan whose row schema is tracked at the type level. The+underlying handle is owned by Rust-side DataFusion; Haskell holds a+'ForeignPtr' that frees it on garbage collection.+-}+newtype DataFrame (cols :: [(Symbol, Type)]) = DF {unDataFrame :: PlanHandle}++{- | Scan a CSV file. The schema comes from the @cols@ phantom via 'KnownSchema';+the user does not pass it separately. DataFusion currently infers column types+from the file, so the constraint is reserved for future wire-format derivation.+-}+scanCsv ::+    forall cols.+    (KnownSchema cols) =>+    Context ->+    T.Text ->+    IO (DataFrame cols)+scanCsv ctx path =+    withContext ctx $ \cp ->+        withCString (T.unpack path) $ \cpath -> do+            ph <- runPlanOp (F.df_scan_csv cp cpath nullPtr)+            return (DF ph)++-- | Keep rows that satisfy the predicate.+filter ::+    TExpr cols Bool ->+    DataFrame cols ->+    IO (DataFrame cols)+filter (TExpr expr) (DF plan) = do+    bytes <- case encodeExprToBytes expr of+        Right bs -> return bs+        Left e -> error ("DataFrame.Fusion.Typed.filter: " <> e)+    withPlan plan $ \pp ->+        withCStringBS bytes $ \cbytes -> do+            ph <- runPlanOp (F.df_plan_filter pp cbytes)+            return (DF ph)++-- | Retain at most @n@ rows.+take ::+    Int ->+    DataFrame cols ->+    IO (DataFrame cols)+take n (DF plan) = withPlan plan $ \pp -> do+    ph <- runPlanOp (F.df_plan_take pp (fromIntegral n))+    return (DF ph)++-- | Project to the named columns. Result schema computed by 'SubsetSchema'.+select ::+    forall (names :: [Symbol]) cols.+    (AllKnownSymbol names, AssertAllPresent names cols) =>+    DataFrame cols ->+    IO (DataFrame (SubsetSchema names cols))+select (DF plan) = do+    let names = symbolVals @names+        json = BL.toStrict (Aeson.encode names)+    withPlan plan $ \pp ->+        withCStringBS json $ \cjson -> do+            ph <- runPlanOp (F.df_plan_select pp cjson)+            return (DF ph)++{- | Add a computed column, appended to the input schema (mirroring+'DataFrame.Typed.Lazy.derive'). The expression is lowered to JSON and decoded+into a DataFusion 'Expr' on the Rust side.+-}+derive ::+    forall name a cols.+    (KnownSymbol name, IC.Columnable a, AssertAbsent name cols) =>+    TExpr cols a ->+    DataFrame cols ->+    IO (DataFrame (Snoc cols '(name, a)))+derive (TExpr expr) (DF plan) = do+    bytes <- case encodeExprToBytes expr of+        Right bs -> return bs+        Left e -> error ("DataFrame.Fusion.Typed.derive: " <> e)+    let nameStr = symbolVal (Proxy @name)+    withPlan plan $ \pp ->+        withCString nameStr $ \cname ->+            withCStringBS bytes $ \cbytes -> do+                ph <- runPlanOp (F.df_plan_derive pp cname cbytes)+                return (DF ph)++-- | A grouped query: an 'DataFrame' tagged with the group-by key list.+data Grouped (keys :: [Symbol]) (cols :: [(Symbol, Type)]) = GD+    { gdKeys :: ![T.Text]+    , gdPlan :: !PlanHandle+    }++-- | Partition rows by the named keys.+groupBy ::+    forall (keys :: [Symbol]) cols.+    (AllKnownSymbol keys, AssertAllPresent keys cols) =>+    DataFrame cols ->+    Grouped keys cols+groupBy (DF plan) = GD (symbolVals @keys) plan++{- | Aggregate a grouped query. The first argument is a chain of 'AGG.as'+entries composed with @(.)@; the empty composition (@id@) yields just the+group keys.+-}+aggregate ::+    forall keys cols aggs.+    (TAgg keys cols '[] -> TAgg keys cols aggs) ->+    Grouped keys cols ->+    IO (DataFrame (Append (GroupKeyColumns keys cols) (Reverse aggs)))+aggregate build (GD keys plan) = do+    let keysJson = BL.toStrict (Aeson.encode keys)+    aggEntries <- traverse encodeAggEntry (taggToNamedExprs (build TAggNil))+    let aggsJson = BL.toStrict (Aeson.encode aggEntries)+    withPlan plan $ \pp ->+        withCStringBS keysJson $ \cKeys ->+            withCStringBS aggsJson $ \cAggs -> do+                ph <- runPlanOp (F.df_plan_groupby_aggregate pp cKeys cAggs)+                return (DF ph)+  where+    encodeAggEntry :: IE.NamedExpr -> IO Aeson.Value+    encodeAggEntry (name, IE.UExpr e) = case encodeExprToBytes e of+        Right bs -> case Aeson.decode (BL.fromStrict bs) :: Maybe Aeson.Value of+            Just v ->+                return $+                    Aeson.object+                        [ "name" Aeson..= name+                        , "expr" Aeson..= v+                        ]+            Nothing ->+                error+                    "DataFrame.Fusion.Typed.aggregate: unparseable JSON from encodeExprToBytes"+        Left err -> error ("DataFrame.Fusion.Typed.aggregate: " <> err)++{- | Inner join on a single key pair. The result schema is currently the left+schema (matching 'DataFrame.Typed.Lazy.join'); a sharper 'InnerJoinSchema'+result is on the v1.5 list.+-}+innerJoin ::+    T.Text -> T.Text -> DataFrame left -> DataFrame right -> IO (DataFrame left)+innerJoin = joinWith "inner"++leftJoin ::+    T.Text -> T.Text -> DataFrame left -> DataFrame right -> IO (DataFrame left)+leftJoin = joinWith "left"++rightJoin ::+    T.Text -> T.Text -> DataFrame left -> DataFrame right -> IO (DataFrame left)+rightJoin = joinWith "right"++fullOuterJoin ::+    T.Text -> T.Text -> DataFrame left -> DataFrame right -> IO (DataFrame left)+fullOuterJoin = joinWith "outer"++joinWith ::+    T.Text ->+    T.Text ->+    T.Text ->+    DataFrame left ->+    DataFrame right ->+    IO (DataFrame left)+joinWith how leftKey rightKey (DF leftPlan) (DF rightPlan) = do+    let onJson = BL.toStrict (Aeson.encode [[leftKey, rightKey]])+    withPlan leftPlan $ \lp ->+        withPlan rightPlan $ \rp ->+            withCString (T.unpack how) $ \cHow ->+                withCStringBS onJson $ \cOn -> do+                    ph <- runPlanOp (F.df_plan_join lp rp cHow cOn)+                    return (DF ph)++{- | Sort the result by a list of (column, direction) pairs.+The 'SortOrder' is reused from "DataFrame.Lazy.Internal.LogicalPlan" so+this signature lines up exactly with 'DataFrame.Typed.Lazy.sortBy'.+-}+sortBy ::+    [(T.Text, SortOrder)] ->+    DataFrame cols ->+    IO (DataFrame cols)+sortBy orders (DF plan) = do+    let json = BL.toStrict (Aeson.encode (map encodeOrder orders))+    withPlan plan $ \pp ->+        withCStringBS json $ \cjson -> do+            ph <- runPlanOp (F.df_plan_sort_by pp cjson)+            return (DF ph)+  where+    encodeOrder (c, o) =+        Aeson.object+            [ "col" Aeson..= c+            , "asc" Aeson..= isAsc o+            ]+    isAsc Ascending = True+    isAsc Descending = False++-- | Execute the plan and import the result as a 'TypedDataFrame'.+run ::+    forall cols.+    DataFrame cols ->+    IO (TypedDataFrame cols)+run (DF plan) = unsafeFreeze <$> collectArrow plan++{- | Pass a strict 'BS.ByteString' to a C function expecting a+NUL-terminated UTF-8 string. The bytes must not contain interior NULs;+JSON output never does.+-}+withCStringBS :: BS.ByteString -> (CString -> IO a) -> IO a+withCStringBS bs k =+    BS.useAsCString bs $ \cs ->+        if cs == nullPtr then k nullPtr else k cs
+ tests/Main.hs view
@@ -0,0 +1,147 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Main (main) where++import qualified Data.Text as T+import qualified Data.Text.IO as TIO+import System.Directory (createDirectoryIfMissing, getTemporaryDirectory)+import qualified System.Exit as Exit+import System.FilePath ((</>))+import System.IO (hPutStrLn, stderr)+import Test.HUnit (+    Test (..),+    assertEqual,+    errors,+    failures,+    runTestTT,+ )++import qualified DataFrame.Fusion.Typed as Fusion+import qualified DataFrame.Operations.Core as Core+import DataFrame.Typed.Freeze (thaw)++-- | A minimal smoke schema: id (int), score (double), name (text).+type CsvCols =+    '[ '("id", Int)+     , '("score", Double)+     , '("name", T.Text)+     ]++testCsv :: T.Text+testCsv =+    T.unlines+        [ "id,score,name"+        , "1,1.5,alice"+        , "2,2.5,bob"+        , "3,3.5,carol"+        , "4,4.5,dan"+        ]++-- | Companion CSV for the join test: department per id.+deptsCsv :: T.Text+deptsCsv =+    T.unlines+        [ "id,dept"+        , "1,eng"+        , "2,eng"+        , "3,sales"+        ]++type DeptCols =+    '[ '("id", Int)+     , '("dept", T.Text)+     ]++testFixtureDir :: IO FilePath+testFixtureDir = do+    base <- getTemporaryDirectory+    let dir = base </> "dataframe-fusion-tests"+    createDirectoryIfMissing True dir+    return dir++writeTestCsv :: IO FilePath+writeTestCsv = do+    dir <- testFixtureDir+    let path = dir </> "smoke.csv"+    TIO.writeFile path testCsv+    return path++writeDeptsCsv :: IO FilePath+writeDeptsCsv = do+    dir <- testFixtureDir+    let path = dir </> "depts.csv"+    TIO.writeFile path deptsCsv+    return path++main :: IO ()+main = do+    path <- writeTestCsv+    deptPath <- writeDeptsCsv+    ctx <- Fusion.newContext++    let scan = Fusion.scanCsv @CsvCols ctx (T.pack path)+        scanDepts = Fusion.scanCsv @DeptCols ctx (T.pack deptPath)++    let tests =+            TestList+                [ TestLabel "scanCsv + run round-trips column count" $ TestCase $ do+                    fdf <- scan+                    tdf <- Fusion.run fdf+                    let df = thaw tdf+                    assertEqual "row count" 4 (Core.nRows df)+                    assertEqual "column count" 3 (Core.nColumns df)+                , TestLabel "take limits row count" $ TestCase $ do+                    fdf <- scan+                    fdf' <- Fusion.take 2 fdf+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "row count after take 2" 2 (Core.nRows df)+                , TestLabel "select projects to fewer columns" $ TestCase $ do+                    fdf <- scan+                    fdf' <- Fusion.select @'["id", "name"] fdf+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "columns after select" 2 (Core.nColumns df)+                , TestLabel "filter drops non-matching rows" $ TestCase $ do+                    fdf <- scan+                    let pred_ = Fusion.col @"id" Fusion..>. Fusion.lit (2 :: Int)+                    fdf' <- Fusion.filter pred_ fdf+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "row count after filter id > 2" 2 (Core.nRows df)+                , TestLabel "derive adds a computed column" $ TestCase $ do+                    fdf <- scan+                    let doubled = Fusion.col @"score" Fusion..*. Fusion.lit (2.0 :: Double)+                    fdf' <- Fusion.derive @"doubled" doubled fdf+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "row count after derive" 4 (Core.nRows df)+                    assertEqual "column count after derive" 4 (Core.nColumns df)+                , TestLabel "sortBy reorders rows" $ TestCase $ do+                    fdf <- scan+                    fdf' <- Fusion.sortBy [("score", Fusion.Descending)] fdf+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "row count after sort" 4 (Core.nRows df)+                , TestLabel "groupBy + aggregate produces one row per key" $ TestCase $ do+                    fdf <- scan+                    -- group by name, sum scores (each name unique here, so 4 groups)+                    fdf' <-+                        Fusion.aggregate+                            (Fusion.as @"total" (Fusion.sum (Fusion.col @"score")))+                            (Fusion.groupBy @'["name"] fdf)+                    df <- thaw <$> Fusion.run fdf'+                    assertEqual "row count after groupBy" 4 (Core.nRows df)+                    assertEqual "column count after groupBy" 2 (Core.nColumns df)+                , TestLabel "innerJoin matches rows on id" $ TestCase $ do+                    fdf <- scan+                    fdf2 <- scanDepts+                    joined <- Fusion.innerJoin "id" "id" fdf fdf2+                    df <- thaw <$> Fusion.run joined+                    -- 3 ids in depts.csv match against 4 ids in smoke.csv -> 3 rows+                    assertEqual "row count after inner join" 3 (Core.nRows df)+                ]++    counts <- runTestTT tests+    if errors counts + failures counts == 0+        then return ()+        else do+            hPutStrLn stderr "FAILED"+            Exit.exitWith (Exit.ExitFailure 1)