diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Revision history for dataframe +## 1.2.0.0 +### Breaking changes +* Remove `Eq` and `Ord` instances for `Expr` — they were buggy under cross-type comparison. Use `eqExpr` / `compareExpr` from `DataFrame.Internal.Expression` directly. +* Lazy executor: `ExecutorConfig` removed; `execute` now takes only a `PhysicalPlan`. `CsvSource` carries the `CsvReader` per scan, so the executor no longer has to know about the CSV implementation. + +### New features +* New `dataframe-fusion` package: typed Apache DataFusion-backed query API (`DataFrame.Fusion.Typed`). +* Can now use any CSV reader for the lazy pipeline. +* New `readCsvWithSchema :: Schema -> FilePath -> IO DataFrame` for schema-driven CSV reads, plus a `CsvReader` type alias exported from `DataFrame.IO.CSV`. +* Numeric comparison operators (`.==`, `./=`, `.<`, `.>`, `.<=`, `.>=`) now widen their operands to a common numeric type, so e.g. `Expr Double .== Expr Int` typechecks. Mirrors the existing arithmetic widening. + +### Internal +* Bump `cabal-version` to 3.0; add `dataframe-fastcsv` to the lint targets; nix flake and `cabal.project` improvements. +* `dataframe-persistent` license corrected to MIT in cabal (#201). + ## 1.1.2.1 * Add `over` and `median` functions to typed API * Fix bug opening web plots in MacOS diff --git a/cbits/dataframe_arrow.h b/cbits/dataframe_arrow.h deleted file mode 100644 --- a/cbits/dataframe_arrow.h +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/** - * Execute a JSON-encoded query plan and return Arrow C Data Interface pointers. - * - * plan_json – null-terminated UTF-8 JSON (see DataFrame.IR for schema) - * schema_out – receives ArrowSchema* cast to uint64_t - * array_out – receives ArrowArray* cast to uint64_t - * - * Returns 0 on success, -1 on error (message written to stderr). - * - * Plan ops: ReadCsv, ReadTsv, FromArrow, Select, GroupBy, Sort, Limit - * - * Example (GroupBy): - * {"op":"GroupBy","keys":["Sex"], - * "aggregations":[{"name":"n","agg":"count","col":"Sex"}], - * "input":{"op":"ReadCsv","path":"data/titanic.csv"}} - */ -int dfExecutePlan(const char* plan_json, - uint64_t* schema_out, - uint64_t* array_out); - -#ifdef __cplusplus -} -#endif diff --git a/cbits/rts_init.c b/cbits/rts_init.c deleted file mode 100644 --- a/cbits/rts_init.c +++ /dev/null @@ -1,18 +0,0 @@ -/* Auto-initialise / finalise the GHC RTS when the shared library is - loaded / unloaded. Required because cabal's native-shared libraries - do not run hs_init automatically on macOS. */ - -#include -#include "HsFFI.h" - -__attribute__((constructor)) -static void df_lib_init(void) -{ - hs_init(0, NULL); -} - -__attribute__((destructor)) -static void df_lib_fini(void) -{ - hs_exit(); -} diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.0 name: dataframe -version: 1.1.2.1 +version: 1.2.0.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -17,8 +17,6 @@ tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 extra-doc-files: CHANGELOG.md README.md extra-source-files: cbits/arrow_abi.h - cbits/dataframe_arrow.h - cbits/rts_init.c tests/data/typing/texts.txt tests/data/typing/texts_with_empties.txt tests/data/typing/texts_with_empties_and_nullish.txt @@ -113,9 +111,11 @@ DataFrame.Typed.Lazy, DataFrame.Typed build-depends: base >= 4 && <5, + async >= 2.2 && < 3, deepseq >= 1 && < 2, aeson >= 0.11.0.0 && < 3, array >= 0.5.4.0 && < 0.6, + temporary >= 1.3 && < 2, attoparsec >= 0.12 && < 0.15, bytestring >= 0.11 && < 0.13, bytestring-lexing >= 0.5 && < 0.6, @@ -147,32 +147,29 @@ hs-source-dirs: src default-language: Haskell2010 -foreign-library dataframe-arrow - type: native-shared +library arrow-bridge + import: warnings + visibility: public hs-source-dirs: ffi - other-modules: DataFrame.FFI - DataFrame.IO.Arrow + exposed-modules: DataFrame.IO.Arrow DataFrame.IR + DataFrame.IR.ExprJson build-depends: base >= 4 && < 5, - dataframe ^>= 1.1, + dataframe, text >= 2.0 && < 3, aeson >= 0.11 && < 3, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, vector ^>= 0.13 - c-sources: cbits/rts_init.c include-dirs: cbits - includes: arrow_abi.h dataframe_arrow.h - install-includes: arrow_abi.h dataframe_arrow.h - ghc-options: -threaded default-language: Haskell2010 executable dataframe-benchmark-example import: warnings main-is: Benchmark.hs build-depends: base >= 4 && < 5, - dataframe ^>= 1.1, + dataframe >= 1 && < 2, random >= 1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13, @@ -184,7 +181,7 @@ import: warnings main-is: Synthesis.hs build-depends: base >= 4 && < 5, - dataframe ^>= 1.1, + dataframe >= 1 && < 2, random >= 1 && < 2, text >= 2.0 && < 3 hs-source-dirs: app @@ -211,7 +208,7 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, - dataframe ^>= 1.1, + dataframe >= 1 && < 2, directory >= 1.3.0.0 && < 2, random >= 1 && < 2, text >= 2.0 && < 3, @@ -228,7 +225,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, process >= 1.6 && < 2, - dataframe ^>= 1.1, + dataframe >= 1 && < 2, random >= 1 && < 2, default-language: Haskell2010 ghc-options: @@ -274,7 +271,8 @@ Monad build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, - dataframe ^>= 1.1, + dataframe >= 1 && < 2, + directory >= 1.3 && < 2, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, random-shuffle >= 0.0.4 && < 1, diff --git a/ffi/DataFrame/FFI.hs b/ffi/DataFrame/FFI.hs deleted file mode 100644 --- a/ffi/DataFrame/FFI.hs +++ /dev/null @@ -1,49 +0,0 @@ -{-# LANGUAGE ForeignFunctionInterface #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - --- | C-exported entry points for the Python bindings. -module DataFrame.FFI where - -import Control.Exception (SomeException, try) -import Data.Word (Word64) -import Foreign (Ptr, castPtr, poke, ptrToWordPtr) -import Foreign.C.String (CString) -import Foreign.C.Types (CInt (..)) -import System.IO (hPrint, hPutStrLn, stderr) - -import qualified Data.Aeson as Aeson -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL - -import DataFrame.IO.Arrow (dataframeToArrow) -import DataFrame.IR (executePlan) - -foreign export ccall "dfExecutePlan" - dfExecutePlan :: CString -> Ptr Word64 -> Ptr Word64 -> IO CInt - -{- | Execute a JSON-encoded query plan and return Arrow C Data Interface - pointers to Python via output parameters. - - @planCS@ – JSON plan string (see DataFrame.IR for schema) - @schemaOut@ – receives the ArrowSchema* address as a Word64 - @arrayOut@ – receives the ArrowArray* address as a Word64 - - Returns 0 on success, -1 on error (error message written to stderr). --} -dfExecutePlan :: CString -> Ptr Word64 -> Ptr Word64 -> IO CInt -dfExecutePlan planCS schemaOut arrayOut = do - planBytes <- BS.packCString planCS - result <- try @SomeException $ do - node <- - either - fail - return - (Aeson.eitherDecode (BL.fromStrict planBytes)) - df <- executePlan node - (sPtr, aPtr) <- dataframeToArrow df - poke schemaOut (fromIntegral (ptrToWordPtr (castPtr sPtr))) - poke arrayOut (fromIntegral (ptrToWordPtr (castPtr aPtr))) - case result of - Left ex -> hPrint stderr ex >> return (CInt (-1)) - Right _ -> return (CInt 0) diff --git a/ffi/DataFrame/IO/Arrow.hs b/ffi/DataFrame/IO/Arrow.hs --- a/ffi/DataFrame/IO/Arrow.hs +++ b/ffi/DataFrame/IO/Arrow.hs @@ -11,7 +11,6 @@ dataframeToArrow, columnToArrow, arrowToDataframe, - -- exported only to force linking of release-callback symbols releaseSchemaImpl, releaseArrayImpl, ) where @@ -24,13 +23,9 @@ import qualified Data.Vector.Unboxed as VU import qualified DataFrame.Internal.Column as DI -import Control.Monad (foldM_, forM, join, void, when, zipWithM_) -import Data.Bits (popCount, setBit, testBit) -import Data.Int (Int32, Int64) -import Data.Maybe (fromMaybe, isNothing) +import Control.Monad (foldM_, forM, join, when, zipWithM_) import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) -import Data.Word (Word8) -import Foreign hiding (void) +import Foreign import Foreign.C.String (CString, newCString, peekCString) import Type.Reflection (typeRep) @@ -225,7 +220,7 @@ sPtr <- makeLeafSchema "g" colName aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) return (sPtr, aPtr) -columnToArrow colName (BoxedColumn _ (vec :: V.Vector a)) +columnToArrow colName (BoxedColumn Nothing (vec :: V.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = do let n = V.length vec bss = map TE.encodeUtf8 (V.toList vec) @@ -253,7 +248,7 @@ [nullPtr, castPtr offPtr, castPtr charsPtr] (free offPtr >> free charsPtr) return (sPtr, aPtr) -columnToArrow colName (BoxedColumn _ (vec :: V.Vector a)) +columnToArrow colName (BoxedColumn Nothing (vec :: V.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do let n = V.length vec dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double) @@ -261,7 +256,7 @@ sPtr <- makeLeafSchema "g" colName aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) return (sPtr, aPtr) -columnToArrow colName (BoxedColumn _ (vec :: V.Vector a)) +columnToArrow colName (BoxedColumn Nothing (vec :: V.Vector a)) | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do let n = V.length vec dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64) @@ -402,11 +397,6 @@ topArray `at` _arrayPrivateData $ castStablePtrToPtr cleanupA return (topSchema, topArray) - --- | Test whether bit i is set in a validity bitmap. -bitmapIsSet :: Ptr Word8 -> Int -> IO Bool -bitmapIsSet bitmapPtr i = - testBit <$> peekElemOff bitmapPtr (i `div` 8) <*> pure (i `mod` 8) {- | Import an Arrow RecordBatch from raw C Data Interface pointers. Copies all data into GC-managed Haskell vectors, then calls the diff --git a/ffi/DataFrame/IR.hs b/ffi/DataFrame/IR.hs --- a/ffi/DataFrame/IR.hs +++ b/ffi/DataFrame/IR.hs @@ -1,7 +1,9 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} @@ -15,6 +17,7 @@ ) where import Data.Aeson (FromJSON (..), withObject, (.:)) +import qualified Data.Aeson as Aeson import Data.Aeson.Types (Parser) import qualified Data.ByteString as BS import Data.Int (Int16, Int32, Int64, Int8) @@ -26,7 +29,7 @@ ) import qualified Data.Vector as V import qualified Data.Vector.Unboxed as VU -import Data.Word (Word, Word16, Word32, Word64, Word8) +import Data.Word (Word16, Word32, Word64, Word8) import Foreign (wordPtrToPtr) import Type.Reflection (SomeTypeRep (..), eqTypeRep, typeRep) @@ -34,17 +37,28 @@ import qualified DataFrame.Functions as Functions import DataFrame.IO.Arrow (arrowToDataframe) import DataFrame.IO.CSV ( + CsvReader, defaultReadOptions, readSeparated, readTsv, + writeCsv, ) +import DataFrame.IO.JSON (readJSON) +import qualified DataFrame.IO.Parquet as Parquet +import DataFrame.IR.ExprJson (SomeExpr (..), decodeExprAny, decodeExprAt) import DataFrame.Internal.Column (Column (..), Columnable) import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn) import DataFrame.Internal.Expression (Expr (..), NamedExpr) -import DataFrame.Operations.Aggregation (aggregate, groupBy) +import DataFrame.Internal.Schema (Schema, makeSchema, schemaType) +import qualified DataFrame.Lazy as Lazy +import DataFrame.Operations.Aggregation (aggregate, distinct, groupBy) +import DataFrame.Operations.Core (insertVector, renameMany) +import DataFrame.Operations.Join (JoinType (..), join) import DataFrame.Operations.Permutation (SortOrder (..), sortBy) -import DataFrame.Operations.Subset (select) +import qualified DataFrame.Operations.Statistics as Stats +import DataFrame.Operations.Subset (exclude, filterWhere, range, select) import qualified DataFrame.Operations.Subset as Subset +import DataFrame.Operations.Transformations (derive) import DataFrame.Operators ((.=)) -- --------------------------------------------------------------------------- @@ -67,6 +81,33 @@ | GroupBy [T.Text] [AggSpec] PlanNode | Sort [T.Text] Bool PlanNode | Limit Int PlanNode + | -- | predicate JSON, child plan + Filter Aeson.Value PlanNode + | -- | column name, expr JSON, child plan + Derive T.Text Aeson.Value PlanNode + | Exclude [T.Text] PlanNode + | Rename [(T.Text, T.Text)] PlanNode + | Distinct PlanNode + | TakeLast Int PlanNode + | Drop Int PlanNode + | DropLast Int PlanNode + | Range Int Int PlanNode + | -- | joinType ("inner"|"left"|"right"|"outer"), shared key columns, left, right + Join T.Text [T.Text] PlanNode PlanNode + | Describe PlanNode + | -- | first column, second column, child plan + Correlation T.Text T.Text PlanNode + | Frequencies T.Text PlanNode + | ReadParquet FilePath + | ReadJson FilePath + | -- | path, separator (single character), child plan; runs as a terminal op + WriteCsv FilePath PlanNode + | {- | path, schema (column name → type-tag map). Reads via the lazy + engine with predicate / projection pushdown; subsequent ops + currently still run eagerly on the materialized result. + -} + ScanCsv FilePath [(T.Text, T.Text)] + | ScanParquet FilePath [(T.Text, T.Text)] deriving (Show) -- --------------------------------------------------------------------------- @@ -91,29 +132,122 @@ "GroupBy" -> GroupBy <$> o .: "keys" <*> o .: "aggregations" <*> o .: "input" "Sort" -> Sort <$> o .: "cols" <*> o .: "ascending" <*> o .: "input" "Limit" -> Limit <$> o .: "n" <*> o .: "input" + "Filter" -> Filter <$> o .: "predicate" <*> o .: "input" + "Derive" -> Derive <$> o .: "name" <*> o .: "expr" <*> o .: "input" + "Exclude" -> Exclude <$> o .: "cols" <*> o .: "input" + "Rename" -> Rename <$> o .: "pairs" <*> o .: "input" + "Distinct" -> Distinct <$> o .: "input" + "TakeLast" -> TakeLast <$> o .: "n" <*> o .: "input" + "Drop" -> Drop <$> o .: "n" <*> o .: "input" + "DropLast" -> DropLast <$> o .: "n" <*> o .: "input" + "Range" -> Range <$> o .: "start" <*> o .: "end" <*> o .: "input" + "Join" -> + Join + <$> o .: "how" + <*> o .: "on" + <*> o .: "left" + <*> o .: "right" + "Describe" -> Describe <$> o .: "input" + "Correlation" -> + Correlation + <$> o .: "first" + <*> o .: "second" + <*> o .: "input" + "Frequencies" -> Frequencies <$> o .: "col" <*> o .: "input" + "ReadParquet" -> ReadParquet <$> o .: "path" + "ReadJson" -> ReadJson <$> o .: "path" + "WriteCsv" -> WriteCsv <$> o .: "path" <*> o .: "input" + "ScanCsv" -> ScanCsv <$> o .: "path" <*> o .: "schema" + "ScanParquet" -> ScanParquet <$> o .: "path" <*> o .: "schema" _ -> fail $ "DataFrame.IR: unknown op: " ++ T.unpack op -executePlan :: PlanNode -> IO DataFrame -executePlan (ReadCsv path) = +executePlan :: CsvReader -> PlanNode -> IO DataFrame +executePlan _reader (ReadCsv path) = readSeparated defaultReadOptions path -executePlan (ReadTsv path) = +executePlan _reader (ReadTsv path) = readTsv path -executePlan (FromArrow schemaAddr arrayAddr) = +executePlan _reader (FromArrow schemaAddr arrayAddr) = arrowToDataframe (wordPtrToPtr (fromIntegral schemaAddr)) (wordPtrToPtr (fromIntegral arrayAddr)) -executePlan (Select cols node) = - select cols <$> executePlan node -executePlan (GroupBy keys aggs node) = do - df <- executePlan node +executePlan reader (Select cols node) = + select cols <$> executePlan reader node +executePlan reader (GroupBy keys aggs node) = do + df <- executePlan reader node nes <- mapM (buildNamedExpr df) aggs return $ aggregate nes (groupBy keys df) -executePlan (Sort cols ascending node) = do - df <- executePlan node +executePlan reader (Sort cols ascending node) = do + df <- executePlan reader node let orders = map (\c -> mkSortOrder ascending c (unsafeGetColumn c df)) cols return $ sortBy orders df -executePlan (Limit k node) = - Subset.take k <$> executePlan node +executePlan reader (Limit k node) = + Subset.take k <$> executePlan reader node +executePlan reader (Filter predJson node) = do + df <- executePlan reader node + case decodeExprAt @Bool predJson of + Right pred_ -> return $ filterWhere pred_ df + Left err -> ioError $ userError $ "DataFrame.IR.Filter: " <> err +executePlan reader (Derive name exprJson node) = do + df <- executePlan reader node + case decodeExprAny exprJson of + Right (SomeExpr _trep expr) -> return $ derive name expr df + Left err -> ioError $ userError $ "DataFrame.IR.Derive: " <> err +executePlan reader (Exclude cols node) = + exclude cols <$> executePlan reader node +executePlan reader (Rename pairs node) = + renameMany pairs <$> executePlan reader node +executePlan reader (Distinct node) = + distinct <$> executePlan reader node +executePlan reader (TakeLast n node) = + Subset.takeLast n <$> executePlan reader node +executePlan reader (Drop n node) = + Subset.drop n <$> executePlan reader node +executePlan reader (DropLast n node) = + Subset.dropLast n <$> executePlan reader node +executePlan reader (Range start end node) = + range (start, end) <$> executePlan reader node +executePlan reader (Join how on leftPlan rightPlan) = do + left <- executePlan reader leftPlan + right <- executePlan reader rightPlan + jt <- case how of + "inner" -> return INNER + "left" -> return LEFT + "right" -> return RIGHT + "outer" -> return FULL_OUTER + "full_outer" -> return FULL_OUTER + other -> + ioError . userError $ + "DataFrame.IR.Join: unknown join type " <> T.unpack other + return $ join jt on left right +executePlan reader (Describe node) = Stats.summarize <$> executePlan reader node +executePlan reader (Correlation a b node) = do + df <- executePlan reader node + let r = Stats.correlation a b df + valueCol = case r of + Just d -> V.singleton d + Nothing -> V.singleton (0 / 0 :: Double) -- NaN when correlation is undefined + -- Return a single-row frame: { first, second, correlation } + return $ + insertVector "first" (V.singleton a) $ + insertVector "second" (V.singleton b) $ + insertVector "correlation" valueCol mempty +executePlan reader (Frequencies colName node) = do + df <- executePlan reader node + runFrequencies colName df +executePlan _reader (ReadParquet path) = Parquet.readParquet path +executePlan _reader (ReadJson path) = readJSON path +executePlan reader (WriteCsv path node) = do + df <- executePlan reader node + writeCsv path df + -- Return the same frame so callers that pipe through .collect() get the + -- written data back too. Callers that don't care can discard. + return df +executePlan reader (ScanCsv path schemaPairs) = do + schema <- buildSchema schemaPairs + Lazy.runDataFrame (Lazy.scanCsvWith reader schema (T.pack path)) +executePlan _reader (ScanParquet path schemaPairs) = do + schema <- buildSchema schemaPairs + Lazy.runDataFrame (Lazy.scanParquet schema (T.pack path)) {- | Build a SortOrder from a column's runtime type. Uses type dispatch to recover Ord for known column types. @@ -155,11 +289,81 @@ "count" -> countExpr name colName (unsafeGetColumn colName df) "sum" -> sumExpr name colName (unsafeGetColumn colName df) "mean" -> meanExpr name colName (unsafeGetColumn colName df) + "min" -> minMaxExpr Functions.minimum name colName (unsafeGetColumn colName df) + "max" -> minMaxExpr Functions.maximum name colName (unsafeGetColumn colName df) + "median" -> doubleStatExpr Functions.median name colName (unsafeGetColumn colName df) + "variance" -> doubleStatExpr Functions.variance name colName (unsafeGetColumn colName df) + "std" -> doubleStatExpr stdDevExpr name colName (unsafeGetColumn colName df) other -> ioError $ userError $ "DataFrame.IR: unknown aggregation '" ++ T.unpack other ++ "'" +-- | Variance → standard deviation; sqrt of the underlying variance Expr. +stdDevExpr :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double +stdDevExpr e = sqrt (Functions.variance e) + +-- | Build a 'Schema' from a list of (col, type-tag) pairs sent over the wire. +buildSchema :: [(T.Text, T.Text)] -> IO Schema +buildSchema pairs = do + sch <- mapM resolve pairs + return (makeSchema sch) + where + resolve (name, tag) = case tag of + "int" -> return (name, schemaType @Int) + "int8" -> return (name, schemaType @Int8) + "int16" -> return (name, schemaType @Int16) + "int32" -> return (name, schemaType @Int32) + "int64" -> return (name, schemaType @Int64) + "double" -> return (name, schemaType @Double) + "float" -> return (name, schemaType @Float) + "bool" -> return (name, schemaType @Bool) + "text" -> return (name, schemaType @T.Text) + "string" -> return (name, schemaType @String) + other -> + ioError . userError $ + "DataFrame.IR.buildSchema: unsupported schema type tag '" + ++ T.unpack other + ++ "' for column '" + ++ T.unpack name + ++ "'" + +-- | Dispatch 'frequencies' on the column's runtime element type. +runFrequencies :: T.Text -> DataFrame -> IO DataFrame +runFrequencies colName df = dispatchType (columnTypeRep (unsafeGetColumn colName df)) + where + columnTypeRep :: Column -> SomeTypeRep + columnTypeRep (UnboxedColumn _ (_ :: VU.Vector a)) = SomeTypeRep (typeRep @a) + columnTypeRep (BoxedColumn _ (_ :: V.Vector a)) = SomeTypeRep (typeRep @a) + + fr :: forall a. (Columnable a, Ord a) => IO DataFrame + fr = return $ Stats.frequencies (Col @a colName) df + + dispatchType :: SomeTypeRep -> IO DataFrame + dispatchType (SomeTypeRep tr) + | Just HRefl <- eqTypeRep tr (typeRep @Int) = fr @Int + | Just HRefl <- eqTypeRep tr (typeRep @Int8) = fr @Int8 + | Just HRefl <- eqTypeRep tr (typeRep @Int16) = fr @Int16 + | Just HRefl <- eqTypeRep tr (typeRep @Int32) = fr @Int32 + | Just HRefl <- eqTypeRep tr (typeRep @Int64) = fr @Int64 + | Just HRefl <- eqTypeRep tr (typeRep @Word) = fr @Word + | Just HRefl <- eqTypeRep tr (typeRep @Word8) = fr @Word8 + | Just HRefl <- eqTypeRep tr (typeRep @Word16) = fr @Word16 + | Just HRefl <- eqTypeRep tr (typeRep @Word32) = fr @Word32 + | Just HRefl <- eqTypeRep tr (typeRep @Word64) = fr @Word64 + | Just HRefl <- eqTypeRep tr (typeRep @Integer) = fr @Integer + | Just HRefl <- eqTypeRep tr (typeRep @Double) = fr @Double + | Just HRefl <- eqTypeRep tr (typeRep @Float) = fr @Float + | Just HRefl <- eqTypeRep tr (typeRep @Bool) = fr @Bool + | Just HRefl <- eqTypeRep tr (typeRep @Char) = fr @Char + | Just HRefl <- eqTypeRep tr (typeRep @T.Text) = fr @T.Text + | Just HRefl <- eqTypeRep tr (typeRep @String) = fr @String + | otherwise = + ioError . userError $ + "DataFrame.IR.Frequencies: unsupported column type for '" + ++ T.unpack colName + ++ "'" + countExpr :: T.Text -> T.Text -> Column -> IO NamedExpr countExpr name colName (UnboxedColumn Nothing (_ :: VU.Vector a)) = return $ name .= count (Col @a colName) countExpr name colName (UnboxedColumn (Just _) (_ :: VU.Vector a)) = return $ name .= count (Col @(Maybe a) colName) @@ -187,7 +391,7 @@ return $ name .= sumMaybe (Col @(Maybe Int) colName) | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = return $ name .= sumMaybe (Col @(Maybe Double) colName) -sumExpr name colName _ = +sumExpr _ colName _ = ioError $ userError $ "DataFrame.IR: sum: unsupported column type for '" ++ T.unpack colName ++ "'" @@ -211,7 +415,59 @@ return $ name .= meanMaybe (Col @(Maybe Double) colName) | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = return $ name .= meanMaybe (Col @(Maybe Int) colName) -meanExpr name colName _ = +meanExpr _ colName _ = ioError $ userError $ "DataFrame.IR: mean: unsupported column type for '" ++ T.unpack colName ++ "'" + +-- | min / max — preserve column type, require Ord. +minMaxExpr :: + (forall a. (Columnable a, Ord a) => Expr a -> Expr a) -> + T.Text -> + T.Text -> + Column -> + IO NamedExpr +minMaxExpr op name colName (UnboxedColumn Nothing (_ :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= op (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= op (Col @Double colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Float) = + return $ name .= op (Col @Float colName) +minMaxExpr op name colName (BoxedColumn Nothing (_ :: V.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = + return $ name .= op (Col @T.Text colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= op (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= op (Col @Double colName) +minMaxExpr _ _ colName _ = + ioError . userError $ + "DataFrame.IR: min/max: unsupported column type for '" + ++ T.unpack colName + ++ "'" + +-- | median / variance / std — return Double, require Real + Unbox. +doubleStatExpr :: + (forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double) -> + T.Text -> + T.Text -> + Column -> + IO NamedExpr +doubleStatExpr op name colName (UnboxedColumn Nothing (_ :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= op (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= op (Col @Double colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Float) = + return $ name .= op (Col @Float colName) +doubleStatExpr op name colName (BoxedColumn Nothing (_ :: V.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= op (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= op (Col @Double colName) +doubleStatExpr _ _ colName _ = + ioError . userError $ + "DataFrame.IR: median/variance/std: unsupported column type for '" + ++ T.unpack colName + ++ "'" diff --git a/ffi/DataFrame/IR/ExprJson.hs b/ffi/DataFrame/IR/ExprJson.hs new file mode 100644 --- /dev/null +++ b/ffi/DataFrame/IR/ExprJson.hs @@ -0,0 +1,701 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | JSON wire format for 'Expr' values. + + The Python bindings build expression trees as JSON; this module decodes them + back into 'Expr' GADTs (typed) and encodes them in the reverse direction + (used for shipping fitted decision trees back to Python). + + Encoded shape: + + > { "node": "col" , "out_type": "double", "name": "x" } + > { "node": "lit" , "out_type": "int" , "value": 42 } + > { "node": "unary" , "out_type": "double", "op": "toDouble" + > , "arg_type": "int" , "arg": } + > { "node": "binary", "out_type": "bool" , "op": "leq" + > , "arg_type": "double", "lhs": , "rhs": } + > { "node": "if" , "out_type": "text" , "cond": + > , "then": , "else": } + + Operator names follow 'binaryName' / 'unaryName' from + "DataFrame.Internal.Expression". Nullable variants ("nulladd", "nullor", …) + are accepted on decode and treated as their non-null equivalents — Python + data crosses the FFI boundary as non-null Arrow buffers, so this is lossless + for the supported workflow. +-} +module DataFrame.IR.ExprJson ( + SomeExpr (..), + encodeExpr, + encodeExprToBytes, + decodeExprAny, + decodeExprAt, + parseSomeExpr, + typeTagOf, + withTypeTag, + withOrdTypeTag, + withNumTypeTag, + withFracTypeTag, + withRealTypeTag, +) where + +import Data.Aeson (object, (.:), (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.Types as Aeson +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BL +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import Data.Word (Word16, Word32, Word64, Word8) +import Type.Reflection (TypeRep, Typeable, typeRep) + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression ( + AggStrategy (..), + BinaryOp (binaryName), + Expr (..), + UnaryOp (unaryName), + ) +import qualified DataFrame.Internal.Expression as F +import DataFrame.Operators ( + ifThenElse, + (.&&.), + (./=.), + (.<.), + (.<=.), + (.==.), + (.>.), + (.>=.), + (.||.), + ) + +{- | Existential wrapper around a typed expression decoded from JSON. +TODO: mchavinda - Maybe consolidate with UExpr from the main package. +-} +data SomeExpr where + SomeExpr :: (Columnable a) => TypeRep a -> Expr a -> SomeExpr + +-- | Map a Haskell type to its wire-format tag string. +typeTagOf :: forall a. (Typeable a) => Maybe T.Text +typeTagOf + | Just _ <- testEquality (typeRep @a) (typeRep @Int) = Just "int" + | Just _ <- testEquality (typeRep @a) (typeRep @Int8) = Just "int8" + | Just _ <- testEquality (typeRep @a) (typeRep @Int16) = Just "int16" + | Just _ <- testEquality (typeRep @a) (typeRep @Int32) = Just "int32" + | Just _ <- testEquality (typeRep @a) (typeRep @Int64) = Just "int64" + | Just _ <- testEquality (typeRep @a) (typeRep @Word) = Just "word" + | Just _ <- testEquality (typeRep @a) (typeRep @Word8) = Just "word8" + | Just _ <- testEquality (typeRep @a) (typeRep @Word16) = Just "word16" + | Just _ <- testEquality (typeRep @a) (typeRep @Word32) = Just "word32" + | Just _ <- testEquality (typeRep @a) (typeRep @Word64) = Just "word64" + | Just _ <- testEquality (typeRep @a) (typeRep @Integer) = Just "integer" + | Just _ <- testEquality (typeRep @a) (typeRep @Double) = Just "double" + | Just _ <- testEquality (typeRep @a) (typeRep @Float) = Just "float" + | Just _ <- testEquality (typeRep @a) (typeRep @Bool) = Just "bool" + | Just _ <- testEquality (typeRep @a) (typeRep @Char) = Just "char" + | Just _ <- testEquality (typeRep @a) (typeRep @T.Text) = Just "text" + | Just _ <- testEquality (typeRep @a) (typeRep @String) = Just "string" + | otherwise = Nothing + +-- | Dispatch on a type tag with a 'Columnable' continuation. +withTypeTag :: + T.Text -> + (forall a. (Columnable a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withTypeTag t k = case t of + "int" -> k (Proxy @Int) + "int8" -> k (Proxy @Int8) + "int16" -> k (Proxy @Int16) + "int32" -> k (Proxy @Int32) + "int64" -> k (Proxy @Int64) + "word" -> k (Proxy @Word) + "word8" -> k (Proxy @Word8) + "word16" -> k (Proxy @Word16) + "word32" -> k (Proxy @Word32) + "word64" -> k (Proxy @Word64) + "integer" -> k (Proxy @Integer) + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + "bool" -> k (Proxy @Bool) + "char" -> k (Proxy @Char) + "text" -> k (Proxy @T.Text) + "string" -> k (Proxy @String) + _ -> fail $ "DataFrame.IR.ExprJson: unknown type tag: " <> T.unpack t + +-- | Subset that supports 'Ord' (for comparison ops). +withOrdTypeTag :: + T.Text -> + (forall a. (Columnable a, Ord a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withOrdTypeTag t k = case t of + "int" -> k (Proxy @Int) + "int8" -> k (Proxy @Int8) + "int16" -> k (Proxy @Int16) + "int32" -> k (Proxy @Int32) + "int64" -> k (Proxy @Int64) + "word" -> k (Proxy @Word) + "word8" -> k (Proxy @Word8) + "word16" -> k (Proxy @Word16) + "word32" -> k (Proxy @Word32) + "word64" -> k (Proxy @Word64) + "integer" -> k (Proxy @Integer) + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + "bool" -> k (Proxy @Bool) + "char" -> k (Proxy @Char) + "text" -> k (Proxy @T.Text) + "string" -> k (Proxy @String) + _ -> + fail $ "DataFrame.IR.ExprJson: type does not support ordering: " <> T.unpack t + +-- | Subset that supports 'Num' (arithmetic). +withNumTypeTag :: + T.Text -> + (forall a. (Columnable a, Num a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withNumTypeTag t k = case t of + "int" -> k (Proxy @Int) + "int8" -> k (Proxy @Int8) + "int16" -> k (Proxy @Int16) + "int32" -> k (Proxy @Int32) + "int64" -> k (Proxy @Int64) + "word" -> k (Proxy @Word) + "word8" -> k (Proxy @Word8) + "word16" -> k (Proxy @Word16) + "word32" -> k (Proxy @Word32) + "word64" -> k (Proxy @Word64) + "integer" -> k (Proxy @Integer) + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + _ -> + fail $ "DataFrame.IR.ExprJson: type does not support arithmetic: " <> T.unpack t + +-- | Subset that supports 'Fractional' (division). +withFracTypeTag :: + T.Text -> + (forall a. (Columnable a, Fractional a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withFracTypeTag t k = case t of + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + _ -> + fail $ + "DataFrame.IR.ExprJson: type does not support fractional division: " + <> T.unpack t + +-- | Subset that supports 'Real' (used by toDouble source types). +withRealTypeTag :: + T.Text -> + (forall a. (Columnable a, Real a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withRealTypeTag t k = case t of + "int" -> k (Proxy @Int) + "int8" -> k (Proxy @Int8) + "int16" -> k (Proxy @Int16) + "int32" -> k (Proxy @Int32) + "int64" -> k (Proxy @Int64) + "word" -> k (Proxy @Word) + "word8" -> k (Proxy @Word8) + "word16" -> k (Proxy @Word16) + "word32" -> k (Proxy @Word32) + "word64" -> k (Proxy @Word64) + "integer" -> k (Proxy @Integer) + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + _ -> + fail $ + "DataFrame.IR.ExprJson: type does not support Real (toDouble source): " + <> T.unpack t + +-- | Subset that supports 'Floating'. +withFloatingTypeTag :: + T.Text -> + (forall a. (Columnable a, Floating a) => Proxy a -> Aeson.Parser r) -> + Aeson.Parser r +withFloatingTypeTag t k = case t of + "double" -> k (Proxy @Double) + "float" -> k (Proxy @Float) + _ -> + fail $ "DataFrame.IR.ExprJson: type does not support Floating: " <> T.unpack t + +{- | Encode an 'Expr' to a JSON value. Returns 'Left' on unsupported +constructors (Agg, Over, CastWith, CastExprWith) or unsupported operator +names (binaryUdf, unaryUdf, …). +-} +encodeExpr :: forall a. (Columnable a) => Expr a -> Either String Aeson.Value +encodeExpr expr = case expr of + Col name -> do + outTag <- requireTypeTag @a + Right $ + object + [ "node" .= ("col" :: T.Text) + , "out_type" .= outTag + , "name" .= name + ] + Lit v -> do + outTag <- requireTypeTag @a + litVal <- encodeLit @a v + Right $ + object + [ "node" .= ("lit" :: T.Text) + , "out_type" .= outTag + , "value" .= litVal + ] + Unary op (arg :: Expr b) -> do + outTag <- requireTypeTag @a + argTag <- requireTypeTag @b + opTag <- recognizeUnary (unaryName op) + argEnc <- encodeExpr arg + Right $ + object + [ "node" .= ("unary" :: T.Text) + , "out_type" .= outTag + , "op" .= opTag + , "arg_type" .= argTag + , "arg" .= argEnc + ] + Binary op (lhs :: Expr c) (rhs :: Expr b) -> do + outTag <- requireTypeTag @a + argTag <- requireTypeTag @c + opTag <- recognizeBinary (binaryName op) + lEnc <- encodeExpr lhs + rEnc <- encodeExpr rhs + Right $ + object + [ "node" .= ("binary" :: T.Text) + , "out_type" .= outTag + , "op" .= opTag + , "arg_type" .= argTag + , "lhs" .= lEnc + , "rhs" .= rEnc + ] + If cond th el -> do + outTag <- requireTypeTag @a + cEnc <- encodeExpr cond + tEnc <- encodeExpr th + eEnc <- encodeExpr el + Right $ + object + [ "node" .= ("if" :: T.Text) + , "out_type" .= outTag + , "cond" .= cEnc + , "then" .= tEnc + , "else" .= eEnc + ] + CastWith{} -> + Left + "DataFrame.IR.ExprJson.encodeExpr: CastWith is not supported in the wire format" + CastExprWith{} -> + Left + "DataFrame.IR.ExprJson.encodeExpr: CastExprWith is not supported in the wire format" + Agg (strat :: F.AggStrategy a b) (inner :: F.Expr b) -> do + outTag <- requireTypeTag @a + argTag <- requireTypeTag @b + innerEnc <- encodeExpr inner + Right $ + object + [ "node" .= ("agg" :: T.Text) + , "out_type" .= outTag + , "agg" .= aggStrategyName strat + , "arg_type" .= argTag + , "arg" .= innerEnc + ] + Over names inner -> do + outTag <- requireTypeTag @a + innerEnc <- encodeExpr inner + Right $ + object + [ "node" .= ("over" :: T.Text) + , "out_type" .= outTag + , "partition_by" .= names + , "arg" .= innerEnc + ] + where + requireTypeTag :: forall x. (Typeable x) => Either String T.Text + requireTypeTag = case typeTagOf @x of + Just t -> Right t + Nothing -> + Left $ + "DataFrame.IR.ExprJson.encodeExpr: unsupported type: " <> show (typeRep @x) + +-- | Wire-format name for an aggregation strategy. +aggStrategyName :: AggStrategy a b -> T.Text +aggStrategyName (CollectAgg n _) = n +aggStrategyName (FoldAgg n _ _) = n +aggStrategyName (MergeAgg n _ _ _ _) = n + +-- | Encode an 'Expr' to a strict 'BS.ByteString' of JSON. +encodeExprToBytes :: + forall a. (Columnable a) => Expr a -> Either String BS.ByteString +encodeExprToBytes e = BL.toStrict . Aeson.encode <$> encodeExpr e + +encodeLit :: forall a. (Columnable a) => a -> Either String Aeson.Value +encodeLit v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + Right $ Aeson.toJSON (v :: Int) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) = + Right $ Aeson.toJSON (v :: Int8) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) = + Right $ Aeson.toJSON (v :: Int16) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) = + Right $ Aeson.toJSON (v :: Int32) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) = + Right $ Aeson.toJSON (v :: Int64) + | Just Refl <- testEquality (typeRep @a) (typeRep @Word) = + Right $ Aeson.toJSON (v :: Word) + | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) = + Right $ Aeson.toJSON (v :: Word8) + | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) = + Right $ Aeson.toJSON (v :: Word16) + | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) = + Right $ Aeson.toJSON (v :: Word32) + | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) = + Right $ Aeson.toJSON (v :: Word64) + | Just Refl <- testEquality (typeRep @a) (typeRep @Integer) = + Right $ Aeson.toJSON (v :: Integer) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + Right $ Aeson.toJSON (v :: Double) + | Just Refl <- testEquality (typeRep @a) (typeRep @Float) = + Right $ Aeson.toJSON (v :: Float) + | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = + Right $ Aeson.toJSON (v :: Bool) + | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = + Right $ Aeson.toJSON (v :: T.Text) + | Just Refl <- testEquality (typeRep @a) (typeRep @Char) = + Right $ Aeson.toJSON [v :: Char] + | Just Refl <- testEquality (typeRep @a) (typeRep @String) = + Right $ Aeson.toJSON (v :: String) + | otherwise = + Left $ + "DataFrame.IR.ExprJson.encodeLit: unsupported type: " <> show (typeRep @a) + +-- | Map a 'binaryName' string to a wire-format op name. Errors on opaque UDF names. +recognizeBinary :: T.Text -> Either String T.Text +recognizeBinary n = case n of + "add" -> Right "add" + "sub" -> Right "sub" + "mult" -> Right "mult" + "divide" -> Right "divide" + "eq" -> Right "eq" + "neq" -> Right "neq" + "lt" -> Right "lt" + "leq" -> Right "leq" + "gt" -> Right "gt" + "geq" -> Right "geq" + "and" -> Right "and" + "or" -> Right "or" + -- Nullable-aware aliases (lossless on non-null inputs) + "nulladd" -> Right "add" + "nullsub" -> Right "sub" + "nullmul" -> Right "mult" + "nulldiv" -> Right "divide" + "nulland" -> Right "and" + "nullor" -> Right "or" + "exponentiate" -> Right "exponentiate" + "logBase" -> Right "logBase" + "div" -> Right "div" + "mod" -> Right "mod" + "pow" -> Right "pow" + other -> + Left $ + "DataFrame.IR.ExprJson: unsupported binary op (cannot serialize): " + <> T.unpack other + +-- | Map a 'unaryName' string to a wire-format op name. Errors on opaque UDF names. +recognizeUnary :: T.Text -> Either String T.Text +recognizeUnary n = case n of + "not" -> Right "not" + "negate" -> Right "negate" + "abs" -> Right "abs" + "signum" -> Right "signum" + "exp" -> Right "exp" + "sqrt" -> Right "sqrt" + "log" -> Right "log" + "sin" -> Right "sin" + "cos" -> Right "cos" + "tan" -> Right "tan" + "asin" -> Right "asin" + "acos" -> Right "acos" + "atan" -> Right "atan" + "sinh" -> Right "sinh" + "cosh" -> Right "cosh" + "asinh" -> Right "asinh" + "acosh" -> Right "acosh" + "atanh" -> Right "atanh" + "toDouble" -> Right "toDouble" + other -> + Left $ + "DataFrame.IR.ExprJson: unsupported unary op (cannot serialize): " + <> T.unpack other + +{- | Decode a JSON value into a 'SomeExpr'. The output type comes from the +@out_type@ field on the root node. +-} +decodeExprAny :: Aeson.Value -> Either String SomeExpr +decodeExprAny = Aeson.parseEither parseSomeExpr + +-- | Decode a JSON value, asserting the result type matches @a@. +decodeExprAt :: + forall a. (Columnable a) => Aeson.Value -> Either String (Expr a) +decodeExprAt v = do + SomeExpr trep expr <- decodeExprAny v + case testEquality trep (typeRep @a) of + Just Refl -> Right expr + Nothing -> + Left $ + "DataFrame.IR.ExprJson.decodeExprAt: expected " + <> show (typeRep @a) + <> " but got " + <> show trep + +-- | The Aeson.Parser entry point — useful when composing with bigger parsers. +parseSomeExpr :: Aeson.Value -> Aeson.Parser SomeExpr +parseSomeExpr = Aeson.withObject "Expr" $ \o -> do + node <- o .: "node" :: Aeson.Parser T.Text + outType <- o .: "out_type" :: Aeson.Parser T.Text + case node of + "col" -> do + name <- o .: "name" :: Aeson.Parser T.Text + withTypeTag outType $ \(_ :: Proxy a) -> + return $ SomeExpr (typeRep @a) (Col @a name) + "lit" -> do + rawVal <- o .: "value" + withTypeTag outType $ \(_ :: Proxy a) -> do + litVal <- decodeLit @a rawVal + return $ SomeExpr (typeRep @a) (Lit @a litVal) + "if" -> do + rawCond <- o .: "cond" + rawThen <- o .: "then" + rawElse <- o .: "else" + cond <- parseExprAt @Bool rawCond + withTypeTag outType $ \(_ :: Proxy a) -> do + thenE <- parseExprAt @a rawThen + elseE <- parseExprAt @a rawElse + return $ SomeExpr (typeRep @a) (ifThenElse cond thenE elseE) + "unary" -> do + op <- o .: "op" :: Aeson.Parser T.Text + argType <- o .: "arg_type" :: Aeson.Parser T.Text + rawArg <- o .: "arg" + parseUnary op outType argType rawArg + "binary" -> do + op <- o .: "op" :: Aeson.Parser T.Text + argType <- o .: "arg_type" :: Aeson.Parser T.Text + rawLhs <- o .: "lhs" + rawRhs <- o .: "rhs" + parseBinary op outType argType rawLhs rawRhs + other -> fail $ "DataFrame.IR.ExprJson: unknown node kind: " <> T.unpack other + +parseExprAt :: forall a. (Columnable a) => Aeson.Value -> Aeson.Parser (Expr a) +parseExprAt v = do + SomeExpr trep expr <- parseSomeExpr v + case testEquality trep (typeRep @a) of + Just Refl -> return expr + Nothing -> + fail $ + "DataFrame.IR.ExprJson: expected " + <> show (typeRep @a) + <> " but got " + <> show trep + +decodeLit :: forall a. (Columnable a) => Aeson.Value -> Aeson.Parser a +decodeLit v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = parseAs @Int v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int8) = parseAs @Int8 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int16) = parseAs @Int16 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int32) = parseAs @Int32 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Int64) = parseAs @Int64 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Word) = parseAs @Word v + | Just Refl <- testEquality (typeRep @a) (typeRep @Word8) = parseAs @Word8 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Word16) = parseAs @Word16 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Word32) = parseAs @Word32 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Word64) = parseAs @Word64 v + | Just Refl <- testEquality (typeRep @a) (typeRep @Integer) = parseAs @Integer v + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = parseAs @Double v + | Just Refl <- testEquality (typeRep @a) (typeRep @Float) = parseAs @Float v + | Just Refl <- testEquality (typeRep @a) (typeRep @Bool) = parseAs @Bool v + | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = parseAs @T.Text v + | Just Refl <- testEquality (typeRep @a) (typeRep @Char) = parseChar v + | Just Refl <- testEquality (typeRep @a) (typeRep @String) = parseAs @String v + | otherwise = + fail $ + "DataFrame.IR.ExprJson.decodeLit: unsupported type: " <> show (typeRep @a) + +parseAs :: forall b. (Aeson.FromJSON b) => Aeson.Value -> Aeson.Parser b +parseAs = Aeson.parseJSON + +parseChar :: Aeson.Value -> Aeson.Parser Char +parseChar v = do + s <- Aeson.parseJSON v :: Aeson.Parser T.Text + case T.unpack s of + [c] -> return c + _ -> + fail $ + "DataFrame.IR.ExprJson: expected single-character string, got: " <> show s + +requireTag :: T.Text -> T.Text -> Aeson.Parser () +requireTag actual expected + | actual == expected = return () + | otherwise = + fail $ + "DataFrame.IR.ExprJson: type mismatch — expected " + <> T.unpack expected + <> " but got " + <> T.unpack actual + +requireSame :: T.Text -> T.Text -> Aeson.Parser () +requireSame a b + | a == b = return () + | otherwise = + fail $ + "DataFrame.IR.ExprJson: type mismatch — " + <> T.unpack a + <> " vs " + <> T.unpack b + +parseUnary :: + T.Text -> T.Text -> T.Text -> Aeson.Value -> Aeson.Parser SomeExpr +parseUnary op outType argType rawArg = case op of + "not" -> do + requireTag outType "bool" + requireTag argType "bool" + arg <- parseExprAt @Bool rawArg + return $ SomeExpr (typeRep @Bool) (F.not arg) + "negate" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + arg <- parseExprAt @a rawArg + return $ SomeExpr (typeRep @a) (negate arg) + "abs" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + arg <- parseExprAt @a rawArg + return $ SomeExpr (typeRep @a) (abs arg) + "signum" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + arg <- parseExprAt @a rawArg + return $ SomeExpr (typeRep @a) (signum arg) + "toDouble" -> do + requireTag outType "double" + withRealTypeTag argType $ \(_ :: Proxy a) -> do + arg <- parseExprAt @a rawArg + return $ SomeExpr (typeRep @Double) (F.toDouble arg) + "exp" -> floatingUnary outType argType rawArg exp + "sqrt" -> floatingUnary outType argType rawArg sqrt + "log" -> floatingUnary outType argType rawArg log + "sin" -> floatingUnary outType argType rawArg sin + "cos" -> floatingUnary outType argType rawArg cos + "tan" -> floatingUnary outType argType rawArg tan + "asin" -> floatingUnary outType argType rawArg asin + "acos" -> floatingUnary outType argType rawArg acos + "atan" -> floatingUnary outType argType rawArg atan + "sinh" -> floatingUnary outType argType rawArg sinh + "cosh" -> floatingUnary outType argType rawArg cosh + "asinh" -> floatingUnary outType argType rawArg asinh + "acosh" -> floatingUnary outType argType rawArg acosh + "atanh" -> floatingUnary outType argType rawArg atanh + other -> fail $ "DataFrame.IR.ExprJson: unsupported unary op: " <> T.unpack other + +floatingUnary :: + T.Text -> + T.Text -> + Aeson.Value -> + (forall x. (Columnable x, Floating x) => Expr x -> Expr x) -> + Aeson.Parser SomeExpr +floatingUnary outType argType rawArg f = withFloatingTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + arg <- parseExprAt @a rawArg + return $ SomeExpr (typeRep @a) (f arg) + +parseBinary :: + T.Text -> + T.Text -> + T.Text -> + Aeson.Value -> + Aeson.Value -> + Aeson.Parser SomeExpr +parseBinary op outType argType rawLhs rawRhs = case op of + "eq" -> do + requireTag outType "bool" + withTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l .==. r) + "neq" -> do + requireTag outType "bool" + withTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l ./=. r) + "lt" -> do + requireTag outType "bool" + withOrdTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l .<. r) + "leq" -> do + requireTag outType "bool" + withOrdTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l .<=. r) + "gt" -> do + requireTag outType "bool" + withOrdTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l .>. r) + "geq" -> do + requireTag outType "bool" + withOrdTypeTag argType $ \(_ :: Proxy a) -> do + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @Bool) (l .>=. r) + "and" -> do + requireTag outType "bool" + l <- parseExprAt @Bool rawLhs + r <- parseExprAt @Bool rawRhs + return $ SomeExpr (typeRep @Bool) (l .&&. r) + "or" -> do + requireTag outType "bool" + l <- parseExprAt @Bool rawLhs + r <- parseExprAt @Bool rawRhs + return $ SomeExpr (typeRep @Bool) (l .||. r) + "add" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @a) (l + r) + "sub" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @a) (l - r) + "mult" -> withNumTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @a) (l * r) + "divide" -> withFracTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @a) (l / r) + "exponentiate" -> withFloatingTypeTag outType $ \(_ :: Proxy a) -> do + requireSame outType argType + l <- parseExprAt @a rawLhs + r <- parseExprAt @a rawRhs + return $ SomeExpr (typeRep @a) (l ** r) + "nulladd" -> parseBinary "add" outType argType rawLhs rawRhs + "nullsub" -> parseBinary "sub" outType argType rawLhs rawRhs + "nullmul" -> parseBinary "mult" outType argType rawLhs rawRhs + "nulldiv" -> parseBinary "divide" outType argType rawLhs rawRhs + "nulland" -> parseBinary "and" outType argType rawLhs rawRhs + "nullor" -> parseBinary "or" outType argType rawLhs rawRhs + other -> fail $ "DataFrame.IR.ExprJson: unsupported binary op: " <> T.unpack other diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -12,7 +12,7 @@ import qualified DataFrame.Functions as F import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) -import DataFrame.Internal.Expression (Expr (..), eSize, getColumns) +import DataFrame.Internal.Expression (Expr (..), eSize, eqExpr, getColumns) import DataFrame.Internal.Interpreter (interpret) import DataFrame.Internal.Statistics (percentileOrd') import DataFrame.Internal.Types @@ -23,9 +23,9 @@ import Control.Monad (guard) import Data.Function (on) #if MIN_VERSION_base(4,20,0) -import Data.List (maximumBy, minimumBy, nub, sort, sortBy) +import Data.List (maximumBy, minimumBy, nub, nubBy, sort, sortBy) #else -import Data.List (foldl', maximumBy, minimumBy, nub, sort, sortBy) +import Data.List (foldl', maximumBy, minimumBy, nub, nubBy, sort, sortBy) #endif import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map.Strict as M @@ -149,7 +149,7 @@ data Tree a = Leaf !a | Branch !(Expr Bool) !(Tree a) !(Tree a) - deriving (Eq, Show) + deriving (Show) treeDepth :: Tree a -> Int treeDepth (Leaf _) = 0 @@ -171,7 +171,7 @@ fitDecisionTree cfg (Col target) df = let conds = - nub $ + nubBy eqExpr $ numericConditions cfg (exclude [target] df) ++ generateConditionsOld cfg (exclude [target] df) @@ -522,15 +522,15 @@ in Branch cond left' right' -pruneExpr :: forall a. (Columnable a, Eq a) => Expr a -> Expr a +pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a pruneExpr (If cond trueBranch falseBranch) = let t = pruneExpr trueBranch f = pruneExpr falseBranch - in if t == f + in if eqExpr t f then t else case (t, f) of - (If condInner tInner _, _) | cond == condInner -> If cond tInner f - (_, If condInner _ fInner) | cond == condInner -> If cond t fInner + (If condInner tInner _, _) | eqExpr cond condInner -> If cond tInner f + (_, If condInner _ fInner) | eqExpr cond condInner -> If cond t fInner _ -> If cond t f pruneExpr (Unary op e) = Unary op (pruneExpr e) pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) @@ -624,8 +624,8 @@ numExprCols (NMaybeDouble e) = getColumns e numExprEq :: NumExpr -> NumExpr -> Bool -numExprEq (NDouble e1) (NDouble e2) = e1 == e2 -numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = e1 == e2 +numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2 +numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2 numExprEq _ _ = False combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] @@ -759,7 +759,7 @@ combinedExprs = do e1 <- prevExprs e2 <- baseExprs - guard (e1 /= e2) + guard (Prelude.not (eqExpr e1 e2)) [F.and e1 e2, F.or e1 e2] generateConditionsOld :: TreeConfig -> DataFrame -> [Expr Bool] @@ -770,7 +770,7 @@ genConds colName = case unsafeGetColumn colName df of (BoxedColumn Nothing (column :: V.Vector a)) -> case withOrdFrom @a ords (map (Lit . (`percentileOrd'` column)) [1, 25, 75, 99]) of - Just ps -> map (F.lift2 (==) (Col @a colName)) ps + Just ps -> map (\p -> Col @a colName .==. p) ps Nothing -> [] (BoxedColumn (Just _) (column :: V.Vector a)) -> case sFloating @a of STrue -> [] -- handled by numericCols / numericExprs @@ -780,7 +780,7 @@ case withOrdFrom @a ords (map (Lit . Just . (`percentileOrd'` column)) [1, 25, 75, 99]) of - Just ps -> map (F.lift2 (==) (Col @(Maybe a) colName)) ps + Just ps -> map (\p -> Col @(Maybe a) colName .==. p) ps Nothing -> [] (UnboxedColumn _ (_ :: VU.Vector a)) -> [] @@ -803,7 +803,7 @@ ) -> case testEquality (typeRep @a) (typeRep @b) of Nothing -> [] - Just Refl -> [F.lift2 (==) (Col @a l) (Col @a r)] + Just Refl -> [Col @a l .==. Col @a r] (UnboxedColumn _ (_ :: VU.Vector a), UnboxedColumn _ (_ :: VU.Vector b)) -> [] ( BoxedColumn (Just _) (_ :: V.Vector a) , BoxedColumn (Just _) (_ :: V.Vector b) @@ -811,11 +811,11 @@ Nothing -> [] Just Refl -> case testEquality (typeRep @a) (typeRep @T.Text) of Nothing -> - case withOrdFrom @a ords [F.lift2 (<=) (Col @(Maybe a) l) (Col r)] of + case withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r] of Just leExprs -> - leExprs ++ [F.lift2 (==) (Col @(Maybe a) l) (Col r)] - Nothing -> [F.lift2 (==) (Col @(Maybe a) l) (Col r)] - Just Refl -> [F.lift2 (==) (Col @(Maybe a) l) (Col r)] + leExprs ++ [Col @(Maybe a) l .==. Col @(Maybe a) r] + Nothing -> [Col @(Maybe a) l .==. Col @(Maybe a) r] + Just Refl -> [Col @(Maybe a) l .==. Col @(Maybe a) r] _ -> [] in concatMap genConds (columnNames df) ++ columnConds @@ -885,7 +885,7 @@ TreeConfig -> T.Text -> [Expr Bool] -> DataFrame -> Maybe (Expr Bool) findBestSplit = findBestGreedySplit @a -pruneTree :: forall a. (Columnable a, Eq a) => Expr a -> Expr a +pruneTree :: forall a. (Columnable a) => Expr a -> Expr a pruneTree = pruneExpr -- | A tree where each leaf stores a class-probability distribution. @@ -958,7 +958,7 @@ fitProbTree cfg (Col target) df = let conds = - nub $ + nubBy eqExpr $ numericConditions cfg (exclude [target] df) ++ generateConditionsOld cfg (exclude [target] df) initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs --- a/src/DataFrame/IO/CSV.hs +++ b/src/DataFrame/IO/CSV.hs @@ -291,6 +291,28 @@ readCsv :: FilePath -> IO DataFrame readCsv = readSeparated defaultReadOptions +type CsvReader = Schema -> FilePath -> IO DataFrame + +{- | Schema-driven attoparsec CSV reader. Coerces each column to the +type declared in 'Schema'; columns absent from the schema fall back to +the default inference path. Defined in terms of 'readSeparated' with +the 'TypeSpec' filled in. + +@ +import qualified DataFrame as D +df <- D.readCsvWithSchema schema "input.csv" +@ +-} +readCsvWithSchema :: CsvReader +readCsvWithSchema schema = + readSeparated + defaultReadOptions + { typeSpec = + SpecifyTypes + (M.toList (elements schema)) + (typeSpec defaultReadOptions) + } + {- | Read CSV file from path and load it into a dataframe. ==== __Example__ diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -80,7 +80,7 @@ , safeColumns :: Bool -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema. } - deriving (Eq, Show) + deriving (Show) {- | Default Parquet read options. diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -287,33 +287,30 @@ exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e -instance (Ord a, Columnable a) => Ord (Expr a) where - compare l r = compareExpr (normalize l) (normalize r) - -instance (Eq a, Columnable a) => Eq (Expr a) where - (==) l r = eqNormalized (normalize l) (normalize r) - where - exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool - exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of - Just Refl -> e1 == e2 - Nothing -> False - eqNormalized :: Expr a -> Expr a -> Bool - eqNormalized (Col n1) (Col n2) = n1 == n2 - eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2 - eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2 - eqNormalized (Lit v1) (Lit v2) = v1 == v2 - eqNormalized (If c1 t1 e1) (If c2 t2 e2) = - c1 == c2 && t1 `exprEq` t2 && e1 `exprEq` e2 - eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2 - eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b - eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) = - n1 == n2 && e1 `exprEq` e2 - eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2 - eqNormalized _ _ = False +eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool +eqExpr l r = eqNormalized (normalize l) (normalize r) + where + exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool + exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of + Just Refl -> eqExpr e1 e2 + Nothing -> False + eqNormalized :: Expr a -> Expr a -> Bool + eqNormalized (Col n1) (Col n2) = n1 == n2 + eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2 + eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2 + eqNormalized (Lit v1) (Lit v2) = v1 == v2 + eqNormalized (If c1 t1 e1) (If c2 t2 e2) = + eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2 + eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2 + eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b + eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) = + n1 == n2 && e1 `exprEq` e2 + eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) = + n1 == n2 && e1 `exprEq` e2 + eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) = + n1 == n2 && e1 `exprEq` e2 + eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2 + eqNormalized _ _ = False replaceExpr :: forall a b c. @@ -321,7 +318,7 @@ Expr a -> Expr b -> Expr c -> Expr c replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of Just Refl -> case testEquality (typeRep @a) (typeRep @c) of - Just Refl -> if old == expr then new else replace' + Just Refl -> if eqExpr old expr then new else replace' Nothing -> expr Nothing -> replace' where diff --git a/src/DataFrame/Internal/Nullable.hs b/src/DataFrame/Internal/Nullable.hs --- a/src/DataFrame/Internal/Nullable.hs +++ b/src/DataFrame/Internal/Nullable.hs @@ -63,6 +63,7 @@ -- * Numeric widening NumericWidenOp (..), widenArithOp, + widenCmpOp, WidenResult, -- * Division widening (integral × integral → Double) @@ -405,6 +406,16 @@ b -> Promote a b widenArithOp f x y = f (widen1 @a @b x) (widen2 @a @b y) + +-- | Apply a comparison function after widening both operands to their common type. +widenCmpOp :: + forall a b. + (NumericWidenOp a b) => + (Promote a b -> Promote a b -> Bool) -> + a -> + b -> + Bool +widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y) -- | Result type of a widening binary operator, accounting for nullable wrappers. type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b)) diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs --- a/src/DataFrame/Lazy/Internal/DataFrame.hs +++ b/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -8,15 +8,12 @@ module DataFrame.Lazy.Internal.DataFrame where import qualified Data.Text as T +import DataFrame.IO.CSV (CsvReader, readCsvWithSchema) import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import DataFrame.Internal.Schema (Schema) -import DataFrame.Lazy.Internal.Executor ( - ExecutorConfig (..), - defaultExecutorConfig, - execute, - ) +import DataFrame.Lazy.Internal.Executor (execute) import DataFrame.Lazy.Internal.LogicalPlan ( DataSource (..), LogicalPlan (..), @@ -46,11 +43,11 @@ {- | Execute the lazy query: optimise the logical plan, then stream-execute the resulting physical plan, returning a fully-materialised 'D.DataFrame'. +The CSV reader (default: attoparsec) is set per scan via 'scanCsv' / +'scanCsvWith'. -} runDataFrame :: LazyDataFrame -> IO D.DataFrame -runDataFrame ldf = do - let physPlan = Opt.optimize (batchSize ldf) (plan ldf) - execute physPlan defaultExecutorConfig{defaultBatchSize = batchSize ldf} +runDataFrame ldf = execute (Opt.optimize (batchSize ldf) (plan ldf)) -- --------------------------------------------------------------------------- -- Builders that construct the logical plan tree @@ -60,19 +57,32 @@ fromDataFrame :: D.DataFrame -> LazyDataFrame fromDataFrame df = LazyDataFrame{plan = SourceDF df, batchSize = 1_000_000} --- | Scan a CSV file with the default comma separator. +{- | Scan a CSV file with the default comma separator and the in-tree +attoparsec reader. For the SIMD reader use 'scanCsvWith'. +-} scanCsv :: Schema -> T.Text -> LazyDataFrame -scanCsv schema path = +scanCsv = scanCsvWith readCsvWithSchema + +{- | Like 'scanCsv' but with an explicit CSV reader (e.g. the SIMD reader +@fastReadCsvWithSchema@ from @dataframe-fastcsv@). +-} +scanCsvWith :: CsvReader -> Schema -> T.Text -> LazyDataFrame +scanCsvWith reader schema path = LazyDataFrame - { plan = Scan (CsvSource (T.unpack path) ',') schema + { plan = Scan (CsvSource (T.unpack path) ',' reader) schema , batchSize = 1_000_000 } --- | Scan a character-separated file. +-- | Scan a character-separated file with the default attoparsec reader. scanSeparated :: Char -> Schema -> T.Text -> LazyDataFrame -scanSeparated sep schema path = +scanSeparated = scanSeparatedWith readCsvWithSchema + +-- | Like 'scanSeparated' but with an explicit CSV reader. +scanSeparatedWith :: + CsvReader -> Char -> Schema -> T.Text -> LazyDataFrame +scanSeparatedWith reader sep schema path = LazyDataFrame - { plan = Scan (CsvSource (T.unpack path) sep) schema + { plan = Scan (CsvSource (T.unpack path) sep reader) schema , batchSize = 1_000_000 } diff --git a/src/DataFrame/Lazy/Internal/Executor.hs b/src/DataFrame/Lazy/Internal/Executor.hs --- a/src/DataFrame/Lazy/Internal/Executor.hs +++ b/src/DataFrame/Lazy/Internal/Executor.hs @@ -1,7 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE NumericUnderscores #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} @@ -16,32 +15,34 @@ expressions support it. -} module DataFrame.Lazy.Internal.Executor ( - ExecutorConfig (..), - defaultExecutorConfig, + CsvReader, execute, foldBatches, ) where -import Control.Concurrent (forkIO) +import Control.Concurrent (forkIO, getNumCapabilities) +import Control.Concurrent.Async (mapConcurrently) import Control.Concurrent.STM (atomically) import Control.Concurrent.STM.TBQueue (newTBQueueIO, readTBQueue, writeTBQueue) import Control.DeepSeq (force) import Control.Exception (evaluate) -import Control.Monad (filterM, when) +import Control.Monad (filterM, forM, forM_, when) import qualified Data.ByteString as BS import Data.IORef import qualified Data.Map as M +import qualified Data.Maybe import qualified Data.Set as S import qualified Data.Text as T import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import qualified Data.Vector.Unboxed as VU +import Data.Word (Word8) +import DataFrame.IO.CSV (CsvReader) import qualified DataFrame.IO.Parquet as Parquet import qualified DataFrame.Internal.Column as C import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import DataFrame.Internal.Schema (elements) import qualified DataFrame.Lazy.IO.Binary as Bin -import qualified DataFrame.Lazy.IO.CSV as LCSV import DataFrame.Lazy.Internal.LogicalPlan (DataSource (..), SortOrder (..)) import DataFrame.Lazy.Internal.PhysicalPlan import qualified DataFrame.Operations.Aggregation as Agg @@ -51,32 +52,13 @@ import qualified DataFrame.Operations.Permutation as Perm import qualified DataFrame.Operations.Subset as Sub import qualified DataFrame.Operations.Transformations as Trans -import System.Directory (doesDirectoryExist) +import System.Directory (doesDirectoryExist, removeFile) import System.FilePath (()) import System.FilePath.Glob (glob) -import System.IO (hClose) +import System.IO.Temp (emptySystemTempFile) import Type.Reflection (typeRep) -- --------------------------------------------------------------------------- --- Configuration --- --------------------------------------------------------------------------- - -data ExecutorConfig = ExecutorConfig - { memoryBudgetBytes :: !Int - -- ^ Per-node spill threshold (currently informational; not enforced yet). - , spillDirectory :: FilePath - , defaultBatchSize :: !Int - } - -defaultExecutorConfig :: ExecutorConfig -defaultExecutorConfig = - ExecutorConfig - { memoryBudgetBytes = 512 * 1_048_576 -- 512 MiB - , spillDirectory = "/tmp" - , defaultBatchSize = 1_000_000 - } - --- --------------------------------------------------------------------------- -- Stream abstraction -- --------------------------------------------------------------------------- @@ -102,16 +84,16 @@ {- | Execute a physical plan, returning the complete result as a single 'DataFrame'. -} -execute :: PhysicalPlan -> ExecutorConfig -> IO D.DataFrame -execute plan cfg = buildStream plan cfg >>= collectStream +execute :: PhysicalPlan -> IO D.DataFrame +execute plan = buildStream plan >>= collectStream {- | Fold a function over every batch produced by a physical plan. The fold is strict in the accumulator; each batch is discarded after folding. -} foldBatches :: - (b -> D.DataFrame -> IO b) -> b -> PhysicalPlan -> ExecutorConfig -> IO b -foldBatches f seed plan cfg = do - stream <- buildStream plan cfg + (b -> D.DataFrame -> IO b) -> b -> PhysicalPlan -> IO b +foldBatches f seed plan = do + stream <- buildStream plan let loop !acc = do mb <- pullBatch stream case mb of @@ -125,14 +107,14 @@ -- Per-operator stream builders -- --------------------------------------------------------------------------- -buildStream :: PhysicalPlan -> ExecutorConfig -> IO Stream +buildStream :: PhysicalPlan -> IO Stream -- Scan ----------------------------------------------------------------------- -buildStream (PhysicalScan (CsvSource path sep) cfg) _ = - executeCsvScan path sep cfg -buildStream (PhysicalScan (ParquetSource path) cfg) _ = +buildStream (PhysicalScan (CsvSource path sep reader) cfg) = + executeCsvScan path sep reader cfg +buildStream (PhysicalScan (ParquetSource path) cfg) = executeParquetScan path cfg -buildStream (PhysicalSpill child path) execCfg = do - df <- execute child execCfg +buildStream (PhysicalSpill child path) = do + df <- execute child Bin.spillToDisk path df df' <- Bin.readSpilled path ref <- newIORef (Just df') @@ -143,32 +125,32 @@ return mb ) -- Filter --------------------------------------------------------------------- -buildStream (PhysicalFilter p child) execCfg = do - childStream <- buildStream child execCfg +buildStream (PhysicalFilter p child) = do + childStream <- buildStream child return . Stream $ ( do mb <- pullBatch childStream return $ fmap (Sub.filterWhere p) mb ) -- Project -------------------------------------------------------------------- -buildStream (PhysicalProject cols child) execCfg = do - childStream <- buildStream child execCfg +buildStream (PhysicalProject cols child) = do + childStream <- buildStream child return . Stream $ ( do mb <- pullBatch childStream return $ fmap (Sub.select cols) mb ) -- Derive --------------------------------------------------------------------- -buildStream (PhysicalDerive name uexpr child) execCfg = do - childStream <- buildStream child execCfg +buildStream (PhysicalDerive name uexpr child) = do + childStream <- buildStream child return . Stream $ ( do mb <- pullBatch childStream return $ fmap (Trans.deriveMany [(name, uexpr)]) mb ) -- Limit ---------------------------------------------------------------------- -buildStream (PhysicalLimit n child) execCfg = do - childStream <- buildStream child execCfg +buildStream (PhysicalLimit n child) = do + childStream <- buildStream child countRef <- newIORef (0 :: Int) return . Stream $ ( do @@ -185,8 +167,8 @@ return $ Just (Sub.take toTake df) ) -- Sort (blocking) ------------------------------------------------------------ -buildStream (PhysicalSort cols child) execCfg = do - df <- execute child execCfg +buildStream (PhysicalSort cols child) = do + df <- execute child let sortOrds = fmap toPermSortOrder cols let sorted = Perm.sortBy sortOrds df ref <- newIORef (Just sorted) @@ -197,34 +179,31 @@ return mb ) -- HashAggregate -------------------------------------------------------------- -buildStream (PhysicalHashAggregate keys aggs child) execCfg = do - childStream <- buildStream child execCfg +buildStream (PhysicalHashAggregate keys aggs child) = do + childStream <- buildStream child if all (isStreamableAgg . snd) aggs then do - -- Streaming partial aggregation: O(|groups|) memory + -- Parallel streaming partial aggregation: + -- * N workers, each pulls batches from the child stream and + -- maintains its own local accumulator. + -- * Once the stream is drained, the N partials are merged + -- sequentially using the same merge expression. + -- * O(|groups| × N) memory in flight, then O(|groups|). let (partialAggs, mergeAggs, finalizer) = buildAggPlan aggs - accRef <- newIORef (Nothing :: Maybe D.DataFrame) - let loop = do - mb <- pullBatch childStream - case mb of - Nothing -> return () - Just batch -> do - -- Force to NF so the batch DataFrame can be GC'd immediately. - -- evaluate . force breaks the thunk chain that would otherwise - -- keep every batch (~60 MB each) alive until the end = OOM. - !partial <- - evaluate . force $ Agg.aggregate partialAggs (Agg.groupBy keys batch) - mAcc <- readIORef accRef - !newAcc <- case mAcc of - Nothing -> return partial - Just acc -> - evaluate . force $ - Agg.aggregate mergeAggs $ - Agg.groupBy keys (acc <> partial) - writeIORef accRef (Just newAcc) - loop - loop - mFinal <- fmap (fmap finalizer) (readIORef accRef) + nCaps <- getNumCapabilities + let workers = max 1 nCaps + partials <- + mapConcurrently + (\_ -> workerLoop childStream keys partialAggs mergeAggs) + [1 .. workers] + mFinal <- + let nonEmpty = Data.Maybe.catMaybes partials + in case nonEmpty of + [] -> return Nothing + [single] -> return (Just (finalizer single)) + (a : rest) -> do + !merged <- mergePartials keys mergeAggs a rest + return (Just (finalizer merged)) ref <- newIORef mFinal return . Stream $ do mb <- readIORef ref @@ -240,9 +219,8 @@ writeIORef ref Nothing return mb -- SourceDF (split pre-loaded DataFrame into batches) ------------------------- -buildStream (PhysicalSourceDF df) execCfg = do - let bs = defaultBatchSize execCfg - total = Core.nRows df +buildStream (PhysicalSourceDF bs df) = do + let total = Core.nRows df posRef <- newIORef (0 :: Int) return . Stream $ do i <- readIORef posRef @@ -254,14 +232,14 @@ writeIORef posRef (i + n) return (Just batch) -- HashJoin — streaming probe (INNER/LEFT) or blocking fallback ---------------- -buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) execCfg = +buildStream (PhysicalHashJoin jt leftKey rightKey leftPlan rightPlan) = case jt of Join.INNER -> streamingHashJoin assembleInnerBatch Join.LEFT -> streamingHashJoin assembleLeftBatch _ -> do -- Blocking fallback for RIGHT / FULL_OUTER - leftDf <- execute leftPlan execCfg - rightDf <- execute rightPlan execCfg + leftDf <- execute leftPlan + rightDf <- execute rightPlan let result = performJoin jt leftKey rightKey leftDf rightDf ref <- newIORef (Just result) return . Stream $ do @@ -271,7 +249,7 @@ where streamingHashJoin assembleFn = do -- Materialise build (right) side once and build the compact index. - rightDf <- execute rightPlan execCfg + rightDf <- execute rightPlan let rightDf' = if leftKey == rightKey then rightDf @@ -281,7 +259,7 @@ rightHashes = Join.buildHashColumn [joinKey] rightDf' ci = Join.buildCompactIndex rightHashes -- Stream probe (left) side batch by batch. - leftStream <- buildStream leftPlan execCfg + leftStream <- buildStream leftPlan return . Stream $ do mBatch <- pullBatch leftStream case mBatch of @@ -307,9 +285,9 @@ assembleInnerBatch = Join.assembleInner -- SortMergeJoin (blocking on both sides) ------------------------------------- -buildStream (PhysicalSortMergeJoin jt leftKey rightKey leftPlan rightPlan) execCfg = do - leftDf <- execute leftPlan execCfg - rightDf <- execute rightPlan execCfg +buildStream (PhysicalSortMergeJoin jt leftKey rightKey leftPlan rightPlan) = do + leftDf <- execute leftPlan + rightDf <- execute rightPlan let result = performJoin jt leftKey rightKey leftDf rightDf ref <- newIORef (Just result) return . Stream $ @@ -326,6 +304,51 @@ {- | True when an aggregate expression can be computed incrementally (i.e., partial results can be merged without materialising all rows). -} + +{- | One worker's loop: pull batches off the shared child stream until +exhausted, building up a per-worker accumulator. +-} +workerLoop :: + Stream -> + [T.Text] -> + [E.NamedExpr] -> + [E.NamedExpr] -> + IO (Maybe D.DataFrame) +workerLoop childStream keys partialAggs mergeAggs = loop Nothing + where + loop !acc = do + mb <- pullBatch childStream + case mb of + Nothing -> return acc + Just batch -> do + !partial <- + evaluate . force $ + Agg.aggregate partialAggs (Agg.groupBy keys batch) + !next <- case acc of + Nothing -> return (Just partial) + Just a -> do + !merged <- + evaluate . force $ + Agg.aggregate mergeAggs (Agg.groupBy keys (a <> partial)) + return (Just merged) + loop next + +-- | Merge a head accumulator with the rest of the workers' partials. +mergePartials :: + [T.Text] -> + [E.NamedExpr] -> + D.DataFrame -> + [D.DataFrame] -> + IO D.DataFrame +mergePartials keys mergeAggs = go + where + go !acc [] = return acc + go !acc (p : ps) = do + !merged <- + evaluate . force $ + Agg.aggregate mergeAggs (Agg.groupBy keys (acc <> p)) + go merged ps + isStreamableAgg :: E.UExpr -> Bool isStreamableAgg (E.UExpr (E.Agg (E.CollectAgg _ _) _)) = False isStreamableAgg (E.UExpr (E.Agg (E.FoldAgg _ Nothing (_ :: a -> b -> a)) _)) = @@ -523,26 +546,37 @@ -- CSV scan implementation -- --------------------------------------------------------------------------- -{- | CSV scan with pipeline parallelism: a dedicated reader thread fills a -bounded queue while the caller's thread applies pushdown predicates and -delivers batches to the rest of the pipeline. The queue depth of 8 keeps -at most eight raw batches in flight, bounding memory while hiding I/O latency. +{- | CSV scan, SIMD-parallel. + +The file is read once into memory, split at newline boundaries into N +ByteString slices (N = RTS capabilities), and each slice is parsed in +parallel with the SIMD reader from "DataFrame.IO.CSV.Fast" via the +in-memory entry point — no temp-file roundtrip. The resulting per-chunk +DataFrames are sliced into batches and a dedicated thread feeds them +into a bounded queue. Pushdown predicates are applied per batch by the +consumer. -} -executeCsvScan :: FilePath -> Char -> ScanConfig -> IO Stream -executeCsvScan path sep cfg = do - (handle, colSpec) <- LCSV.openCsvStream sep (scanSchema cfg) path - -- Queue carries raw batches; Nothing is the end-of-stream sentinel. - -- Depth 2: each batch holds ~60 MB (1M Text + Double columns); 8 would be ~480 MB. - queue <- newTBQueueIO 2 +executeCsvScan :: FilePath -> Char -> CsvReader -> ScanConfig -> IO Stream +executeCsvScan path _sep reader cfg = do + nCaps <- getNumCapabilities + chunkPaths <- splitCsvAtNewlines (max 1 nCaps) path + + -- Each chunk parses in parallel via the reader carried on the + -- 'CsvSource' plan node. Parsing and queue-feeding stay disjoint to + -- avoid 14 producers all hammering a shared TBQueue (STM contention + -- dominates throughput). + let schema = scanSchema cfg + batchSz = scanBatchSize cfg + chunkDfs <- mapConcurrently (reader schema) chunkPaths + mapM_ removeFile chunkPaths + + -- Bounded queue with a single writer, N concurrent readers. + queue <- newTBQueueIO (fromIntegral (max 4 (2 * nCaps))) _ <- forkIO $ do - let loop lo = do - result <- LCSV.readBatch sep colSpec (scanBatchSize cfg) lo handle - case result of - Nothing -> - hClose handle >> atomically (writeTBQueue queue Nothing) - Just (df, lo') -> - atomically (writeTBQueue queue (Just df)) >> loop lo' - loop BS.empty + forM_ chunkDfs $ \df -> + forM_ (sliceIntoBatches batchSz df) $ \b -> + atomically (writeTBQueue queue (Just b)) + atomically (writeTBQueue queue Nothing) return . Stream $ ( do mb <- atomically (readTBQueue queue) @@ -555,6 +589,46 @@ Just p -> Sub.filterWhere p df in return (Just df') ) + +-- | Slice a 'DataFrame' into row-bounded batches of at most @n@ rows. +sliceIntoBatches :: Int -> D.DataFrame -> [D.DataFrame] +sliceIntoBatches n df = + let total = Core.nRows df + starts = [0, n .. total - 1] + in [Sub.range (s, min (s + n) total) df | s <- starts] + +{- | Split a CSV file at newline boundaries into @n@ temp files, each +carrying the original header followed by an aligned-at-newlines slice +of the body. Returns the temp file paths; the caller is responsible +for removing them after use. The path-based 'fastReadCsvWithSchema' +mmap's each file, so we get OS-paged reads instead of a single +monolithic 'BS.readFile' of the whole input. +-} +splitCsvAtNewlines :: Int -> FilePath -> IO [FilePath] +splitCsvAtNewlines n path = do + bs <- BS.readFile path + let (header, rest) = BS.break (== nl) bs + body = BS.drop 1 rest + bodyLen = BS.length body + rawOffsets = [(bodyLen * i) `div` n | i <- [0 .. n]] + snapped = 0 : map (snap body) (init (drop 1 rawOffsets)) ++ [bodyLen] + ranges = zip snapped (drop 1 snapped) + slices = + [ BS.take (hi - lo) (BS.drop lo body) + | (lo, hi) <- ranges + , hi > lo + ] + forM slices $ \chunk -> do + p <- emptySystemTempFile "lazy_csv_chunk_.csv" + BS.writeFile p (header <> BS.singleton nl <> chunk) + return p + where + nl :: Word8 + nl = 0x0A + snap body off = + case BS.elemIndex nl (BS.drop off body) of + Just i -> off + i + 1 + Nothing -> BS.length body -- --------------------------------------------------------------------------- -- Join helper diff --git a/src/DataFrame/Lazy/Internal/LogicalPlan.hs b/src/DataFrame/Lazy/Internal/LogicalPlan.hs --- a/src/DataFrame/Lazy/Internal/LogicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/LogicalPlan.hs @@ -3,6 +3,7 @@ module DataFrame.Lazy.Internal.LogicalPlan where import qualified Data.Text as T +import DataFrame.IO.CSV (CsvReader) import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.Expression as E import DataFrame.Internal.Schema (Schema) @@ -10,10 +11,14 @@ -- | Data source for a scan node. data DataSource - = -- | path, separator - CsvSource FilePath Char + = -- | path, separator, CSV reader (e.g. attoparsec or SIMD) + CsvSource FilePath Char CsvReader | ParquetSource FilePath - deriving (Show) + +instance Show DataSource where + show (CsvSource path sep _) = + "CsvSource " ++ show path ++ " " ++ show sep ++ " " + show (ParquetSource path) = "ParquetSource " ++ show path -- | Sort direction used in Sort nodes and the public API. data SortOrder = Ascending | Descending diff --git a/src/DataFrame/Lazy/Internal/Optimizer.hs b/src/DataFrame/Lazy/Internal/Optimizer.hs --- a/src/DataFrame/Lazy/Internal/Optimizer.hs +++ b/src/DataFrame/Lazy/Internal/Optimizer.hs @@ -171,13 +171,13 @@ -} toPhysical :: Int -> LogicalPlan -> PhysicalPlan -- Special case: Filter directly on a Scan → push into ScanConfig. -toPhysical batchSz (Filter p (Scan (CsvSource path sep) schema)) = +toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) = PhysicalScan - (CsvSource path sep) + (CsvSource path sep reader) (ScanConfig batchSz sep schema (Just p)) -toPhysical batchSz (Scan (CsvSource path sep) schema) = +toPhysical batchSz (Scan (CsvSource path sep reader) schema) = PhysicalScan - (CsvSource path sep) + (CsvSource path sep reader) (ScanConfig batchSz sep schema Nothing) toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) = PhysicalScan @@ -206,4 +206,4 @@ PhysicalSort cols (toPhysical batchSz child) toPhysical batchSz (Limit n child) = PhysicalLimit n (toPhysical batchSz child) -toPhysical _ (SourceDF df) = PhysicalSourceDF df +toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df diff --git a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs --- a/src/DataFrame/Lazy/Internal/PhysicalPlan.hs +++ b/src/DataFrame/Lazy/Internal/PhysicalPlan.hs @@ -31,6 +31,6 @@ | PhysicalLimit Int PhysicalPlan | -- | Materialize child to a binary file on disk (used for build sides). PhysicalSpill PhysicalPlan FilePath - | -- | Emit an already-loaded DataFrame as a stream of batches. - PhysicalSourceDF D.DataFrame + | -- | Emit an already-loaded DataFrame as a stream of batches of size @n@. + PhysicalSourceDF Int D.DataFrame deriving (Show) diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs --- a/src/DataFrame/Operators.hs +++ b/src/DataFrame/Operators.hs @@ -33,6 +33,7 @@ WidenResultDiv, divArithOp, widenArithOp, + widenCmpOp, ) import DataFrame.Internal.Types (Promote, PromoteDiv) @@ -198,53 +199,76 @@ -- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing) --- | Nullable-aware equality. Returns @Maybe Bool@ when either operand is nullable. +{- | Nullable-aware equality. Widens numeric operands to their common type, +so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either +operand is nullable. +-} (.==) :: - (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(.==) = lift2Decorated (nullCmpOp (==)) "eq" (Just ".==") True 4 +(.==) = lift2Decorated (applyNull2 (widenCmpOp (==))) "eq" (Just ".==") True 4 --- | Nullable-aware inequality. +-- | Nullable-aware inequality. Widens numeric operands to their common type. (./=) :: - (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(./=) = lift2Decorated (nullCmpOp (/=)) "neq" (Just "./=") True 4 +(./=) = lift2Decorated (applyNull2 (widenCmpOp (/=))) "neq" (Just "./=") True 4 --- | Nullable-aware less-than. +-- | Nullable-aware less-than. Widens numeric operands to their common type. (.<) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(.<) = lift2Decorated (nullCmpOp (<)) "lt" (Just ".<") False 4 +(.<) = lift2Decorated (applyNull2 (widenCmpOp (<))) "lt" (Just ".<") False 4 --- | Nullable-aware greater-than. +-- | Nullable-aware greater-than. Widens numeric operands to their common type. (.>) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(.>) = lift2Decorated (nullCmpOp (>)) "gt" (Just ".>") False 4 +(.>) = lift2Decorated (applyNull2 (widenCmpOp (>))) "gt" (Just ".>") False 4 --- | Nullable-aware less-than-or-equal. +{- | Nullable-aware less-than-or-equal. Widens numeric operands to their +common type, so @Expr Double .<= Expr Int@ typechecks. +-} (.<=) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(.<=) = lift2Decorated (nullCmpOp (<=)) "leq" (Just ".<=") False 4 +(.<=) = lift2Decorated (applyNull2 (widenCmpOp (<=))) "leq" (Just ".<=") False 4 --- | Nullable-aware greater-than-or-equal. +-- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type. (.>=) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => Expr a -> Expr b -> Expr (NullCmpResult a b) -(.>=) = lift2Decorated (nullCmpOp (>=)) "geq" (Just ".>=") False 4 +(.>=) = lift2Decorated (applyNull2 (widenCmpOp (>=))) "geq" (Just ".>=") False 4 (.&&.) :: Expr Bool -> Expr Bool -> Expr Bool (.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3 diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -20,6 +20,7 @@ import DataFrame.Internal.Expression ( Expr (..), eSize, + eqExpr, ) import DataFrame.Internal.Interpreter (interpret) import DataFrame.Internal.Statistics @@ -48,7 +49,7 @@ [ p .<= q | p <- filter (not . isLiteral) ps , q <- ps - , p /= q + , Prelude.not (eqExpr p q) ] ++ [ F.not p | p <- conds @@ -56,8 +57,8 @@ expandedConds = conds ++ newConds - ++ [p .&& q | p <- newConds, q <- conds, p /= q] - ++ [p .|| q | p <- newConds, q <- conds, p /= q] + ++ [p .&& q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)] + ++ [p .|| q | p <- newConds, q <- conds, Prelude.not (eqExpr p q)] in pickTopNBool df labels (deduplicate df expandedConds) @@ -114,7 +115,7 @@ , (j, q) <- zip [(0 :: Int) ..] existingPrograms , Prelude.not (isLiteral p && isLiteral q) , Prelude.not (isConditional p || isConditional q) - , p /= q + , Prelude.not (eqExpr p q) , i > j ] ++ [ F.max p q @@ -122,7 +123,7 @@ , (j, q) <- zip [(0 :: Int) ..] existingPrograms , Prelude.not (isLiteral p && isLiteral q) , Prelude.not (isConditional p || isConditional q) - , p /= q + , Prelude.not (eqExpr p q) , i > j ] ++ [ F.ifThenElse cond r s @@ -130,7 +131,7 @@ , r <- existingPrograms , s <- existingPrograms , Prelude.not (isConditional r || isConditional s) - , r /= s + , Prelude.not (eqExpr r s) ] else [] ) @@ -146,7 +147,7 @@ , q <- existingPrograms , Prelude.not (isLiteral p && isLiteral q) , Prelude.not (isConditional p || isConditional q) - , p /= q + , Prelude.not (eqExpr p q) ] isLiteral :: Expr a -> Bool @@ -163,7 +164,7 @@ DataFrame -> [Expr a] -> [(Expr a, TypedColumn a)] -deduplicate df = go [] . L.nub . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2)) +deduplicate df = go [] . L.nubBy eqExpr . L.sortBy (\e1 e2 -> compare (eSize e1) (eSize e2)) where go _ [] = [] go seen (x : xs) diff --git a/src/DataFrame/Typed/Expr.hs b/src/DataFrame/Typed/Expr.hs --- a/src/DataFrame/Typed/Expr.hs +++ b/src/DataFrame/Typed/Expr.hs @@ -156,6 +156,7 @@ WidenResultDiv, divArithOp, widenArithOp, + widenCmpOp, ) import DataFrame.Internal.Types (Promote, PromoteDiv) @@ -471,59 +472,86 @@ -- Nullable-aware comparison operators (three-valued logic) ------------------------------------------------------------------------------- --- | Nullable-aware equality. Returns @Maybe Bool@ when either operand is nullable. +{- | Nullable-aware equality. Widens numeric operands to their common type, +so @TExpr cols Double .== TExpr cols Int@ typechecks. Returns @Maybe Bool@ +when either operand is nullable. +-} (.==) :: - (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.==) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (==)) "eq" (Just "==") True 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (==))) "eq" (Just "==") True 4) a b) --- | Nullable-aware inequality. +-- | Nullable-aware inequality. Widens numeric operands to their common type. (./=) :: - (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Eq (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (./=) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (/=)) "neq" (Just "/=") True 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (/=))) "neq" (Just "/=") True 4) a b) --- | Nullable-aware less-than. +-- | Nullable-aware less-than. Widens numeric operands to their common type. (.<) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.<) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (<)) "lt" (Just "<") False 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<))) "lt" (Just "<") False 4) a b) --- | Nullable-aware less-than-or-equal. +-- | Nullable-aware less-than-or-equal. Widens numeric operands to their common type. (.<=) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.<=) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (<=)) "leq" (Just "<=") False 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (<=))) "leq" (Just "<=") False 4) a b) --- | Nullable-aware greater-than-or-equal. +-- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type. (.>=) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.>=) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (>=)) "geq" (Just ">=") False 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>=))) "geq" (Just ">=") False 4) a b) --- | Nullable-aware greater-than. +-- | Nullable-aware greater-than. Widens numeric operands to their common type. (.>) :: - (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b Bool (NullCmpResult a b) + , Ord (Promote (BaseType a) (BaseType b)) + ) => TExpr cols a -> TExpr cols b -> TExpr cols (NullCmpResult a b) (.>) (TExpr a) (TExpr b) = - TExpr (Binary (MkBinaryOp (nullCmpOp (>)) "gt" (Just ">") False 4) a b) + TExpr + (Binary (MkBinaryOp (applyNull2 (widenCmpOp (>))) "gt" (Just ">") False 4) a b) not :: TExpr cols Bool -> TExpr cols Bool not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) diff --git a/tests/DecisionTree.hs b/tests/DecisionTree.hs --- a/tests/DecisionTree.hs +++ b/tests/DecisionTree.hs @@ -9,7 +9,7 @@ import DataFrame.DecisionTree import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI -import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Expression (Expr (..), eqExpr) import DataFrame.Internal.Interpreter (interpret) import DataFrame.Operators @@ -634,8 +634,8 @@ probExprsLeaf = TestCase $ do let pt = Leaf (M.fromList [("A", 0.75), ("B", 0.25)]) :: ProbTree T.Text pe = probExprs pt - assertEqual "A expr is Lit 0.75" (Lit 0.75) (pe M.! "A") - assertEqual "B expr is Lit 0.25" (Lit 0.25) (pe M.! "B") + assertBool "A expr is Lit 0.75" (eqExpr (Lit 0.75) (pe M.! "A")) + assertBool "B expr is Lit 0.25" (eqExpr (Lit 0.25) (pe M.! "B")) -- probExprs: class absent from one leaf gets Lit 0.0 on that side probExprsMissingClass :: Test @@ -647,14 +647,12 @@ (Leaf (M.fromList [("B", 1.0)])) :: ProbTree T.Text pe = probExprs pt - assertEqual + assertBool "A expr: If cond (Lit 1.0) (Lit 0.0)" - (F.ifThenElse splitCond (Lit 1.0) (Lit 0.0)) - (pe M.! "A") - assertEqual + (eqExpr (F.ifThenElse splitCond (Lit 1.0) (Lit 0.0)) (pe M.! "A")) + assertBool "B expr: If cond (Lit 0.0) (Lit 1.0)" - (F.ifThenElse splitCond (Lit 0.0) (Lit 1.0)) - (pe M.! "B") + (eqExpr (F.ifThenElse splitCond (Lit 0.0) (Lit 1.0)) (pe M.! "B")) -- probExprs: keys equal all classes that appear across any leaf probExprsAllClasses :: Test diff --git a/tests/data/unstable_csv/empty_file.csv b/tests/data/unstable_csv/empty_file.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/empty_file.csv diff --git a/tests/data/unstable_csv/utf8_bom.csv b/tests/data/unstable_csv/utf8_bom.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/utf8_bom.csv @@ -0,0 +1,3 @@ +name,value +Alice,1 +Bob,2