diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -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 diff --git a/cbits/dfusion_bridge.h b/cbits/dfusion_bridge.h new file mode 100644 --- /dev/null +++ b/cbits/dfusion_bridge.h @@ -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 + +#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": } 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 */ diff --git a/dataframe-fusion.cabal b/dataframe-fusion.cabal new file mode 100644 --- /dev/null +++ b/dataframe-fusion.cabal @@ -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 diff --git a/rust/dfusion-bridge/Cargo.toml b/rust/dfusion-bridge/Cargo.toml new file mode 100644 --- /dev/null +++ b/rust/dfusion-bridge/Cargo.toml @@ -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 diff --git a/rust/dfusion-bridge/src/lib.rs b/rust/dfusion-bridge/src/lib.rs new file mode 100644 --- /dev/null +++ b/rust/dfusion-bridge/src/lib.rs @@ -0,0 +1,627 @@ +//! C ABI shim that exposes Apache DataFusion to Haskell. +//! +//! Each plan-extending function takes an opaque `*mut DfPlan` (an +//! `Arc`), 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> = RefCell::new(None); +} + +fn set_error(msg: impl Into) { + 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, +} + +pub struct DfPlan { + ctx: Arc, + 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, +} + +#[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, result: datafusion::error::Result) -> *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 = 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 { + #[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 = 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::>()?; + 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 = 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 = 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": }` 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 = 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 = 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 = keys.iter().map(|k| col(k)).collect(); + let agg_exprs: Vec = match entries + .iter() + .map(|e| Ok(decode_expr(&e.expr)?.alias(&e.name))) + .collect::, 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 = 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 = 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 { + 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 { + 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, _> = + 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 +} diff --git a/src/DataFrame/Fusion/FFI.hs b/src/DataFrame/Fusion/FFI.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Fusion/FFI.hs @@ -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 diff --git a/src/DataFrame/Fusion/Plan.hs b/src/DataFrame/Fusion/Plan.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Fusion/Plan.hs @@ -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 "" + 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)) diff --git a/src/DataFrame/Fusion/Typed.hs b/src/DataFrame/Fusion/Typed.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Fusion/Typed.hs @@ -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 diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -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)