diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Revision history for dataframe +## 1.0.0.0 +* Fix mappend to respect schema of empty columns. +* Add cast operators that force column schema +* Add null aware operators so some operations are easier. +* Add arrow shim with python example. +* Add numeric promotion for numeric operations. +* Add column provenance tracking +* Add stratified sampling +* Read files from hugging face + ## 0.7.0.0 * This release adds A LOT of AI code to the repo (which we'll now pause in favour of refactoring, testing, and completeness for 1.0) * The lazy reader now has a custom binary format that it spills to. This almost halved the time it takes to run the 1 billion row challenge. The lazy evaluation now also supports Parquet. diff --git a/app/LazyBenchmark.hs b/app/LazyBenchmark.hs --- a/app/LazyBenchmark.hs +++ b/app/LazyBenchmark.hs @@ -224,7 +224,7 @@ -- Demonstrates that the executor reads only the first batch. runQuery "Q1 — preview first 20 rows (no filter)" $ L.runDataFrame $ - L.limit 20 $ + L.take 20 $ L.scanCsv schema pathT -- Q2: Filter + limit. @@ -232,8 +232,8 @@ -- ~512 matches in the first batch and stops — reads only one batch. runQuery "Q2 — filter (x > 0.999), limit 20" $ L.runDataFrame $ - L.limit 20 $ - L.filter (col @Double "x" .> lit 0.999) $ + L.take 20 $ + L.filter (col @Double "x" .> lit (0.999 :: Double)) $ L.scanCsv schema pathT -- Q3: Filter + derive + select + limit. @@ -241,10 +241,10 @@ -- Predicate pushdown moves the filter into the scan batch loop. runQuery "Q3 — filter (x > 0.999), derive z = x*y, select [id,z], limit 20" $ L.runDataFrame $ - L.limit 20 $ + L.take 20 $ L.select ["id", "z"] $ L.derive "z" (col @Double "x" * col @Double "y") $ - L.filter (col @Double "x" .> lit 0.999) $ + L.filter (col @Double "x" .> lit (0.999 :: Double)) $ L.scanCsv schema pathT -- Q4: Filter fusion demo. @@ -253,9 +253,9 @@ -- We limit to keep result size manageable. runQuery "Q4 — filter fusion: (x > 0.5) . (y > 0.5), limit 20" $ L.runDataFrame $ - L.limit 20 $ - L.filter (col @Double "y" .> lit 0.5) $ - L.filter (col @Double "x" .> lit 0.5) $ + L.take 20 $ + L.filter (col @Double "y" .> lit (0.5 :: Double)) $ + L.filter (col @Double "x" .> lit (0.5 :: Double)) $ L.scanCsv schema pathT -- Q5: Full scan, heavy filter, count results. @@ -269,7 +269,7 @@ ) $ L.runDataFrame $ L.select ["id", "x"] - $ L.filter (col @Double "x" .> lit 0.999) + $ L.filter (col @Double "x" .> lit (0.999 :: Double)) $ L.scanCsv schema pathT putStrLn "\nDone." diff --git a/app/Synthesis.hs b/app/Synthesis.hs --- a/app/Synthesis.hs +++ b/app/Synthesis.hs @@ -53,13 +53,13 @@ fitDecisionTree ( defaultTreeConfig { maxTreeDepth = 5 - , minSamplesSplit = 10 + , minSamplesSplit = 5 , minLeafSize = 3 , taoIterations = 100 , synthConfig = defaultSynthConfig - { complexityPenalty = 0.00 - , maxExprDepth = 2 + { complexityPenalty = 0.1 + , maxExprDepth = 3 , disallowedCombinations = [ (F.name age, F.name fare) , ("passenger_class", "number_of_siblings_and_spouses") @@ -99,12 +99,32 @@ computeAccuracy df = let tp = - fromIntegral $ D.nRows (D.filterWhere (survived .== 1 .&& prediction .== 1) df) + fromIntegral $ + D.nRows + ( D.filterWhere + (survived .== F.lit (1 :: Int) .&& prediction .== F.lit (1 :: Int)) + df + ) tn = - fromIntegral $ D.nRows (D.filterWhere (survived .== 0 .&& prediction .== 0) df) + fromIntegral $ + D.nRows + ( D.filterWhere + (survived .== F.lit (0 :: Int) .&& prediction .== F.lit (0 :: Int)) + df + ) fp = - fromIntegral $ D.nRows (D.filterWhere (survived .== 0 .&& prediction .== 1) df) + fromIntegral $ + D.nRows + ( D.filterWhere + (survived .== F.lit (0 :: Int) .&& prediction .== F.lit (1 :: Int)) + df + ) fn = - fromIntegral $ D.nRows (D.filterWhere (survived .== 1 .&& prediction .== 0) df) + fromIntegral $ + D.nRows + ( D.filterWhere + (survived .== F.lit (1 :: Int) .&& prediction .== F.lit (0 :: Int)) + df + ) in (tp + tn) / (tp + tn + fp + fn) diff --git a/cbits/arrow_abi.h b/cbits/arrow_abi.h new file mode 100644 --- /dev/null +++ b/cbits/arrow_abi.h @@ -0,0 +1,44 @@ +/* Arrow C Data Interface structs (verbatim from specification). + See https://arrow.apache.org/docs/format/CDataInterface.html */ + +#pragma once + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +struct ArrowSchema { + /* Array type description */ + const char *format; + const char *name; + const char *metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema **children; + struct ArrowSchema *dictionary; + + void (*release)(struct ArrowSchema *); + void *private_data; +}; + +struct ArrowArray { + /* Array data description */ + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void **buffers; + struct ArrowArray **children; + struct ArrowArray *dictionary; + + void (*release)(struct ArrowArray *); + /* Opaque producer-specific data */ + void *private_data; +}; + +#ifdef __cplusplus +} +#endif diff --git a/cbits/dataframe_arrow.h b/cbits/dataframe_arrow.h new file mode 100644 --- /dev/null +++ b/cbits/dataframe_arrow.h @@ -0,0 +1,30 @@ +#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 new file mode 100644 --- /dev/null +++ b/cbits/rts_init.c @@ -0,0 +1,18 @@ +/* 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/data/sharded/part-0.parquet b/data/sharded/part-0.parquet new file mode 100644 Binary files /dev/null and b/data/sharded/part-0.parquet differ diff --git a/data/sharded/part-1.parquet b/data/sharded/part-1.parquet new file mode 100644 Binary files /dev/null and b/data/sharded/part-1.parquet differ diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe -version: 0.7.0.0 +version: 1.0.0.0 synopsis: A fast, safe, and intuitive DataFrame library. @@ -17,7 +17,11 @@ 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/process_csv.h + cbits/arrow_abi.h + cbits/dataframe_arrow.h + cbits/rts_init.c data/titanic/*.csv + data/sharded/*.parquet tests/data/*.csv tests/data/*.tsv tests/data/*.parquet @@ -53,8 +57,10 @@ DataFrame.Internal.Types, DataFrame.Internal.Expression, DataFrame.Internal.Interpreter, + DataFrame.Internal.Nullable, DataFrame.Internal.Parsing, DataFrame.Internal.Column, + DataFrame.Internal.Binary, DataFrame.Internal.Statistics, DataFrame.Display.Terminal.PrettyPrint, DataFrame.Display.Terminal.Colours, @@ -86,6 +92,7 @@ DataFrame.IO.Parquet.Compression, DataFrame.IO.Parquet.Encoding, DataFrame.IO.Parquet.Page, + DataFrame.IO.Parquet.Seeking, DataFrame.IO.Parquet.Time, DataFrame.IO.Parquet.Types, DataFrame.Lazy.IO.CSV, @@ -138,6 +145,9 @@ stm >= 2.5 && < 3, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, + http-conduit >= 2.3 && < 3, + streamly-core, + streamly-bytestring, hs-source-dirs: src c-sources: cbits/process_csv.c @@ -146,11 +156,32 @@ install-includes: process_csv.h default-language: Haskell2010 +foreign-library dataframe-arrow + type: native-shared + hs-source-dirs: ffi + other-modules: DataFrame.FFI + DataFrame.IO.Arrow + DataFrame.IR + build-depends: + base >= 4 && < 5, + dataframe ^>= 1, + 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 >= 0.5 && < 1, + dataframe ^>= 1, random >= 1 && < 2, time >= 1.12 && < 2, vector ^>= 0.13, @@ -162,7 +193,7 @@ import: warnings main-is: Synthesis.hs build-depends: base >= 4 && < 5, - dataframe >= 0.5 && < 1, + dataframe ^>= 1, random >= 1 && < 2, text >= 2.0 && < 3 hs-source-dirs: app @@ -189,7 +220,7 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, - dataframe >= 0.5 && < 1, + dataframe ^>= 1, directory >= 1.3.0.0 && < 2, random >= 1 && < 2, text >= 2.0 && < 3, @@ -206,7 +237,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, process >= 1.6 && < 2, - dataframe >= 0.5 && < 1, + dataframe ^>= 1, random >= 1 && < 2, default-language: Haskell2010 ghc-options: @@ -219,6 +250,7 @@ type: exitcode-stdio-1.0 main-is: Main.hs other-modules: Assertions, + DecisionTree, Functions, GenDataFrame, Internal.Parsing, @@ -232,6 +264,8 @@ Operations.InsertColumn, Operations.Join, Operations.Merge, + Operations.Nullable, + Operations.Provenance, Operations.ReadCsv, Operations.Shuffle, Operations.Sort, @@ -245,7 +279,8 @@ Monad build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.13, - dataframe >= 0.5 && < 1, + dataframe ^>= 1, + bytestring >= 0.11 && < 0.13, directory >= 1.3.0.0 && < 2, HUnit ^>= 1.6, QuickCheck >= 2 && < 3, @@ -258,4 +293,3 @@ containers >= 0.6.7 && < 0.9 hs-source-dirs: tests default-language: Haskell2010 - diff --git a/ffi/DataFrame/FFI.hs b/ffi/DataFrame/FFI.hs new file mode 100644 --- /dev/null +++ b/ffi/DataFrame/FFI.hs @@ -0,0 +1,49 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/ffi/DataFrame/IO/Arrow.hs @@ -0,0 +1,606 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE ForeignFunctionInterface #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Convert a 'DataFrame' to Arrow C Data Interface structs for zero-copy + transfer to Python (or any other Arrow consumer). +-} +module DataFrame.IO.Arrow ( + dataframeToArrow, + columnToArrow, + arrowToDataframe, + -- exported only to force linking of release-callback symbols + releaseSchemaImpl, + releaseArrayImpl, +) where + +import qualified Data.ByteString as BS +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Vector as V +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 (setBit, testBit) +import Data.Int (Int32, Int64) +import Data.Maybe (fromMaybe, isNothing) +import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) +import Data.Word (Word8) +import Foreign hiding (void) +import Foreign.C.String (CString, newCString, peekCString) +import Type.Reflection (typeRep) + +import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Operations.Core (fromNamedColumns) + +-- --------------------------------------------------------------------------- +-- Opaque phantom types for the Arrow structs +-- --------------------------------------------------------------------------- + +data ArrowSchema +data ArrowArray + +arrowSchemaSize :: Int +arrowSchemaSize = 72 -- 9 × 8 bytes + +arrowArraySize :: Int +arrowArraySize = 80 -- 10 × 8 bytes + +-- ArrowSchema field byte offsets +_schemaFormat + , _schemaName + , _schemaMetadata + , _schemaFlags + , _schemaNChildren + , _schemaChildren + , _schemaDictionary + , _schemaRelease + , _schemaPrivateData :: + Int +_schemaFormat = 0 +_schemaName = 8 +_schemaMetadata = 16 +_schemaFlags = 24 +_schemaNChildren = 32 +_schemaChildren = 40 +_schemaDictionary = 48 +_schemaRelease = 56 +_schemaPrivateData = 64 + +-- ArrowArray field byte offsets +_arrayLength + , _arrayNullCount + , _arrayOffset + , _arrayNBuffers + , _arrayNChildren + , _arrayBuffers + , _arrayChildren + , _arrayDictionary + , _arrayRelease + , _arrayPrivateData :: + Int +_arrayLength = 0 +_arrayNullCount = 8 +_arrayOffset = 16 +_arrayNBuffers = 24 +_arrayNChildren = 32 +_arrayBuffers = 40 +_arrayChildren = 48 +_arrayDictionary = 56 +_arrayRelease = 64 +_arrayPrivateData = 72 + +-- --------------------------------------------------------------------------- +-- Helpers +-- --------------------------------------------------------------------------- + +-- Write a Storable value at a byte offset from a base pointer. +at :: (Storable a) => Ptr b -> Int -> a -> IO () +at p off = poke (castPtr (p `plusPtr` off)) + +-- Read a Storable value at a byte offset from a base pointer. +readAt :: (Storable a) => Ptr b -> Int -> IO a +readAt p off = peek (castPtr (p `plusPtr` off)) + +-- --------------------------------------------------------------------------- +-- Release callbacks (self-import trick for compile-time-constant FunPtr) +-- --------------------------------------------------------------------------- + +foreign export ccall "df_release_schema" + releaseSchemaImpl :: Ptr ArrowSchema -> IO () + +foreign import ccall "&df_release_schema" + pReleaseSchema :: FunPtr (Ptr ArrowSchema -> IO ()) + +foreign export ccall "df_release_array" + releaseArrayImpl :: Ptr ArrowArray -> IO () + +foreign import ccall "&df_release_array" + pReleaseArray :: FunPtr (Ptr ArrowArray -> IO ()) + +-- Dynamic wrappers to call producer's release callbacks after copying. +foreign import ccall "dynamic" + callRelSchema :: FunPtr (Ptr ArrowSchema -> IO ()) -> Ptr ArrowSchema -> IO () + +foreign import ccall "dynamic" + callRelArray :: FunPtr (Ptr ArrowArray -> IO ()) -> Ptr ArrowArray -> IO () + +releaseSchemaImpl :: Ptr ArrowSchema -> IO () +releaseSchemaImpl p = do + rawPriv <- peek (castPtr (p `plusPtr` _schemaPrivateData) :: Ptr (Ptr ())) + let sp = castPtrToStablePtr rawPriv :: StablePtr (IO ()) + join (deRefStablePtr sp) + freeStablePtr sp + -- Arrow spec: release callback must set release to NULL to signal completion. + -- p here is Arrow C++'s internal copy of the struct (not our mallocBytes + -- allocation); our original allocation is freed inside the cleanup closure. + p `at` _schemaRelease $ (nullFunPtr :: FunPtr (Ptr ArrowSchema -> IO ())) + +releaseArrayImpl :: Ptr ArrowArray -> IO () +releaseArrayImpl p = do + rawPriv <- peek (castPtr (p `plusPtr` _arrayPrivateData) :: Ptr (Ptr ())) + let sp = castPtrToStablePtr rawPriv :: StablePtr (IO ()) + join (deRefStablePtr sp) + freeStablePtr sp + -- Same reasoning as releaseSchemaImpl. + p `at` _arrayRelease $ (nullFunPtr :: FunPtr (Ptr ArrowArray -> IO ())) + +makeLeafSchema :: String -> T.Text -> IO (Ptr ArrowSchema) +makeLeafSchema fmt colName = do + p <- mallocBytes arrowSchemaSize + fmtStr <- newCString fmt + nameStr <- newCString (T.unpack colName) + p `at` _schemaFormat $ fmtStr + p `at` _schemaName $ nameStr + p `at` _schemaMetadata $ (nullPtr :: Ptr ()) + p `at` _schemaFlags $ (0 :: Int64) + p `at` _schemaNChildren $ (0 :: Int64) + p `at` _schemaChildren $ (nullPtr :: Ptr ()) + p `at` _schemaDictionary $ (nullPtr :: Ptr ()) + p `at` _schemaRelease $ pReleaseSchema + -- Capture p so our original mallocBytes allocation is freed when release runs. + cleanup <- newStablePtr (free fmtStr >> free nameStr >> free p) + p `at` _schemaPrivateData $ castStablePtrToPtr cleanup + return p + +makeLeafArray :: Int -> Int64 -> [Ptr ()] -> IO () -> IO (Ptr ArrowArray) +makeLeafArray nRows nullCnt bufPtrs extraCleanup = do + p <- mallocBytes arrowArraySize + let nb = length bufPtrs + bufArr <- mallocArray nb :: IO (Ptr (Ptr ())) + zipWithM_ (pokeElemOff bufArr) [0 ..] bufPtrs + p `at` _arrayLength $ (fromIntegral nRows :: Int64) + p `at` _arrayNullCount $ nullCnt + p `at` _arrayOffset $ (0 :: Int64) + p `at` _arrayNBuffers $ (fromIntegral nb :: Int64) + p `at` _arrayNChildren $ (0 :: Int64) + p `at` _arrayBuffers $ bufArr + p `at` _arrayChildren $ (nullPtr :: Ptr ()) + p `at` _arrayDictionary $ (nullPtr :: Ptr ()) + p `at` _arrayRelease $ pReleaseArray + -- Capture p so our original mallocBytes allocation is freed when release runs. + cleanup <- newStablePtr (free bufArr >> extraCleanup >> free p) + p `at` _arrayPrivateData $ castStablePtrToPtr cleanup + return p + +buildBitmap :: V.Vector (Maybe a) -> IO (Ptr Word8, Int) +buildBitmap vec = do + let n = V.length vec + bitmapBytes = max 1 ((n + 7) `div` 8) + nullCount = V.length (V.filter isNothing vec) + bitmapPtr <- mallocBytes bitmapBytes :: IO (Ptr Word8) + mapM_ (\i -> pokeElemOff bitmapPtr i (0 :: Word8)) [0 .. bitmapBytes - 1] + V.iforM_ vec $ \i mv -> + case mv of + Nothing -> return () + Just _ -> do + let byteIdx = i `div` 8 + bitIdx = i `mod` 8 + b <- peekElemOff bitmapPtr byteIdx + pokeElemOff bitmapPtr byteIdx (setBit b bitIdx) + return (bitmapPtr, nullCount) + +columnToArrow :: T.Text -> Column -> IO (Ptr ArrowSchema, Ptr ArrowArray) +columnToArrow colName (UnboxedColumn (vec :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do + let n = VU.length vec + dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64) + VU.imapM_ (\i v -> pokeElemOff dataPtr i (fromIntegral v)) vec + sPtr <- makeLeafSchema "l" colName + aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (UnboxedColumn (vec :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do + let n = VU.length vec + dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double) + VU.imapM_ (pokeElemOff dataPtr) vec + sPtr <- makeLeafSchema "g" colName + aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (BoxedColumn (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) + cumOff = scanl (+) 0 (map BS.length bss) + total = last cumOff + offPtr <- mallocArray (n + 1) :: IO (Ptr Int32) + zipWithM_ + (\i o -> pokeElemOff offPtr i (fromIntegral o :: Int32)) + [0 ..] + cumOff + charsPtr <- mallocBytes (max 1 total) :: IO (Ptr Word8) + foldM_ + ( \pos bs -> do + BS.useAsCStringLen bs $ \(src, len) -> + copyBytes (charsPtr `plusPtr` pos) (castPtr src) len + return (pos + BS.length bs) + ) + 0 + bss + sPtr <- makeLeafSchema "u" colName + aPtr <- + makeLeafArray + n + 0 + [nullPtr, castPtr offPtr, castPtr charsPtr] + (free offPtr >> free charsPtr) + return (sPtr, aPtr) +columnToArrow colName (BoxedColumn (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) + V.imapM_ (pokeElemOff dataPtr) vec + sPtr <- makeLeafSchema "g" colName + aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (BoxedColumn (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) + V.imapM_ (\i v -> pokeElemOff dataPtr i (fromIntegral v)) vec + sPtr <- makeLeafSchema "l" colName + aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (OptionalColumn (vec :: V.Vector (Maybe a))) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do + let n = V.length vec + (bitmapPtr, nullCount) <- buildBitmap vec + dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64) + V.imapM_ + ( \i mv -> + pokeElemOff dataPtr i (fromIntegral (fromMaybe 0 mv) :: Int64) + ) + vec + sPtr <- makeLeafSchema "l" colName + aPtr <- + makeLeafArray + n + (fromIntegral nullCount) + [castPtr bitmapPtr, castPtr dataPtr] + (free bitmapPtr >> free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (OptionalColumn (vec :: V.Vector (Maybe a))) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do + let n = V.length vec + (bitmapPtr, nullCount) <- buildBitmap vec + dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double) + V.imapM_ + ( \i mv -> + pokeElemOff dataPtr i (maybe 0.0 realToFrac mv :: Double) + ) + vec + sPtr <- makeLeafSchema "g" colName + aPtr <- + makeLeafArray + n + (fromIntegral nullCount) + [castPtr bitmapPtr, castPtr dataPtr] + (free bitmapPtr >> free dataPtr) + return (sPtr, aPtr) +columnToArrow colName (OptionalColumn (vec :: V.Vector (Maybe a))) + | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = do + let n = V.length vec + texts = V.toList vec + bss = map (maybe BS.empty TE.encodeUtf8) texts + cumOff = scanl (+) 0 (map BS.length bss) + total = last cumOff + (bitmapPtr, nullCount) <- buildBitmap vec + offPtr <- mallocArray (n + 1) :: IO (Ptr Int32) + zipWithM_ + (\i o -> pokeElemOff offPtr i (fromIntegral o :: Int32)) + [0 ..] + cumOff + charsPtr <- mallocBytes (max 1 total) :: IO (Ptr Word8) + foldM_ + ( \pos bs -> do + BS.useAsCStringLen bs $ \(src, len) -> + copyBytes (charsPtr `plusPtr` pos) (castPtr src) len + return (pos + BS.length bs) + ) + 0 + bss + sPtr <- makeLeafSchema "u" colName + aPtr <- + makeLeafArray + n + (fromIntegral nullCount) + [castPtr bitmapPtr, castPtr offPtr, castPtr charsPtr] + (free bitmapPtr >> free offPtr >> free charsPtr) + return (sPtr, aPtr) +columnToArrow colName _ = + error $ + "DataFrame.IO.Arrow.columnToArrow: unsupported column type for '" + ++ T.unpack colName + ++ "'" + +dataframeToArrow :: DataFrame -> IO (Ptr ArrowSchema, Ptr ArrowArray) +dataframeToArrow df = do + let idxToName = M.fromList [(v, k) | (k, v) <- M.toList (columnIndices df)] + ncols = M.size (columnIndices df) + colsInOrder = + [ (idxToName M.! i, columns df V.! i) + | i <- [0 .. ncols - 1] + ] + + childPairs <- forM colsInOrder (uncurry columnToArrow) + let childSPtrs = map fst childPairs + childAPtrs = map snd childPairs + + let nRows = case colsInOrder of + [] -> 0 + _ -> DI.columnLength (snd (head colsInOrder)) + topSchema <- mallocBytes arrowSchemaSize + fmtStr <- newCString "+s" + nameStr <- newCString "" + childSArr <- mallocArray ncols :: IO (Ptr (Ptr ArrowSchema)) + zipWithM_ (pokeElemOff childSArr) [0 ..] childSPtrs + topSchema `at` _schemaFormat $ fmtStr + topSchema `at` _schemaName $ nameStr + topSchema `at` _schemaMetadata $ (nullPtr :: Ptr ()) + topSchema `at` _schemaFlags $ (0 :: Int64) + topSchema `at` _schemaNChildren $ (fromIntegral ncols :: Int64) + topSchema `at` _schemaChildren $ childSArr + topSchema `at` _schemaDictionary $ (nullPtr :: Ptr ()) + topSchema `at` _schemaRelease $ pReleaseSchema + -- Do NOT loop over children here: Arrow C++ zeroes children[i]->release + -- during import, so reading it would yield a null function pointer. + -- Children are released independently by Arrow C++; their own cleanup + -- closures free their buffers and struct memory. + cleanupS <- newStablePtr $ do + free childSArr + free fmtStr + free nameStr + free topSchema -- free our original mallocBytes allocation + topSchema `at` _schemaPrivateData $ castStablePtrToPtr cleanupS + + -- ── Top-level struct array ────────────────────────────────────────────── + topArray <- mallocBytes arrowArraySize + childAArr <- mallocArray ncols :: IO (Ptr (Ptr ArrowArray)) + zipWithM_ (pokeElemOff childAArr) [0 ..] childAPtrs + topBufArr <- mallocArray 1 :: IO (Ptr (Ptr ())) + pokeElemOff topBufArr 0 nullPtr + topArray `at` _arrayLength $ (fromIntegral nRows :: Int64) + topArray `at` _arrayNullCount $ (0 :: Int64) + topArray `at` _arrayOffset $ (0 :: Int64) + topArray `at` _arrayNBuffers $ (1 :: Int64) + topArray `at` _arrayNChildren $ (fromIntegral ncols :: Int64) + topArray `at` _arrayBuffers $ topBufArr + topArray `at` _arrayChildren $ childAArr + topArray `at` _arrayDictionary $ (nullPtr :: Ptr ()) + topArray `at` _arrayRelease $ pReleaseArray + -- Same reasoning as cleanupS: Arrow C++ manages children independently. + cleanupA <- newStablePtr $ do + free childAArr + free topBufArr + free topArray -- free our original mallocBytes allocation + 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 + producer's release callbacks. +-} +arrowToDataframe :: Ptr () -> Ptr () -> IO DataFrame +arrowToDataframe rawSchema rawArray = do + let schemaPtr = castPtr rawSchema :: Ptr ArrowSchema + arrayPtr = castPtr rawArray :: Ptr ArrowArray + nCols <- readAt schemaPtr _schemaNChildren :: IO Int64 + childSArr <- readAt schemaPtr _schemaChildren :: IO (Ptr (Ptr ArrowSchema)) + childAArr <- readAt arrayPtr _arrayChildren :: IO (Ptr (Ptr ArrowArray)) + cols <- forM [0 .. fromIntegral nCols - 1] $ \i -> do + cs <- peekElemOff childSArr i + ca <- peekElemOff childAArr i + readArrowColumn cs ca + -- Call producer's release callbacks after all data has been copied. + relA <- readAt arrayPtr _arrayRelease :: IO (FunPtr (Ptr ArrowArray -> IO ())) + when (relA /= nullFunPtr) $ callRelArray relA arrayPtr + relS <- + readAt schemaPtr _schemaRelease :: IO (FunPtr (Ptr ArrowSchema -> IO ())) + when (relS /= nullFunPtr) $ callRelSchema relS schemaPtr + return $ fromNamedColumns cols + +readArrowColumn :: Ptr ArrowSchema -> Ptr ArrowArray -> IO (T.Text, Column) +readArrowColumn schemaPtr arrayPtr = do + fmtStr <- (readAt schemaPtr _schemaFormat :: IO CString) >>= peekCString + nameStr <- (readAt schemaPtr _schemaName :: IO CString) >>= peekCString + let name = T.pack nameStr + len <- readAt arrayPtr _arrayLength :: IO Int64 + nullCnt <- readAt arrayPtr _arrayNullCount :: IO Int64 + bufArr <- readAt arrayPtr _arrayBuffers :: IO (Ptr (Ptr ())) + let n = fromIntegral len + col <- case fmtStr of + "l" -> readInt64Col n nullCnt bufArr + "i" -> readInt32Col n nullCnt bufArr + "g" -> readFloat64Col n nullCnt bufArr + "f" -> readFloat32Col n nullCnt bufArr + "U" -> readLargeUtf8Col n nullCnt bufArr + "u" -> readUtf8Col n nullCnt bufArr + _ -> + error $ + "DataFrame.IO.Arrow.readArrowColumn: unsupported format '" + ++ fmtStr + ++ "' for column '" + ++ nameStr + ++ "'" + return (name, col) + +readInt64Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readInt64Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + dataVoid <- peekElemOff bufArr 1 + let dataPtr = castPtr dataVoid :: Ptr Int64 + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then do + v <- peekElemOff dataPtr i + return (Just (fromIntegral v :: Int)) + else return Nothing + return $ OptionalColumn (vec :: V.Vector (Maybe Int)) + else do + vec <- VU.generateM n $ \i -> do + v <- peekElemOff dataPtr i + return (fromIntegral v :: Int) + return $ UnboxedColumn (vec :: VU.Vector Int) + +readInt32Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readInt32Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + dataVoid <- peekElemOff bufArr 1 + let dataPtr = castPtr dataVoid :: Ptr Int32 + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then do + v <- peekElemOff dataPtr i + return (Just (fromIntegral v :: Int)) + else return Nothing + return $ OptionalColumn (vec :: V.Vector (Maybe Int)) + else do + vec <- VU.generateM n $ \i -> do + v <- peekElemOff dataPtr i + return (fromIntegral v :: Int) + return $ UnboxedColumn (vec :: VU.Vector Int) + +readFloat64Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readFloat64Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + dataVoid <- peekElemOff bufArr 1 + let dataPtr = castPtr dataVoid :: Ptr Double + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then Just <$> (peekElemOff dataPtr i :: IO Double) + else return Nothing + return $ OptionalColumn (vec :: V.Vector (Maybe Double)) + else do + vec <- VU.generateM n (peekElemOff dataPtr) + return $ UnboxedColumn (vec :: VU.Vector Double) + +readFloat32Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readFloat32Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + dataVoid <- peekElemOff bufArr 1 + let dataPtr = castPtr dataVoid :: Ptr Float + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then do + v <- peekElemOff dataPtr i + return (Just (realToFrac v :: Double)) + else return Nothing + return $ OptionalColumn (vec :: V.Vector (Maybe Double)) + else do + vec <- VU.generateM n $ \i -> do + v <- peekElemOff dataPtr i + return (realToFrac v :: Double) + return $ UnboxedColumn (vec :: VU.Vector Double) + +-- | Read a large_string (format "U") column with int64 offsets. +readLargeUtf8Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readLargeUtf8Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + offsetVoid <- peekElemOff bufArr 1 + charVoid <- peekElemOff bufArr 2 + let offsetPtr = castPtr offsetVoid :: Ptr Int64 + charPtr = castPtr charVoid :: Ptr Word8 + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then do + start <- fromIntegral <$> peekElemOff offsetPtr i + end <- fromIntegral <$> peekElemOff offsetPtr (i + 1) + bs <- + BS.packCStringLen + (castPtr (charPtr `plusPtr` start), end - start) + return $ Just (TE.decodeUtf8 bs) + else return Nothing + return $ OptionalColumn vec + else do + vec <- V.generateM n $ \i -> do + start <- fromIntegral <$> peekElemOff offsetPtr i + end <- fromIntegral <$> peekElemOff offsetPtr (i + 1) + bs <- + BS.packCStringLen + (castPtr (charPtr `plusPtr` start), end - start) + return $ TE.decodeUtf8 bs + return $ BoxedColumn vec + +-- | Read a utf8 (format "u") column with int32 offsets. +readUtf8Col :: Int -> Int64 -> Ptr (Ptr ()) -> IO Column +readUtf8Col n nullCnt bufArr = do + bitmapVoid <- peekElemOff bufArr 0 + offsetVoid <- peekElemOff bufArr 1 + charVoid <- peekElemOff bufArr 2 + let offsetPtr = castPtr offsetVoid :: Ptr Int32 + charPtr = castPtr charVoid :: Ptr Word8 + if nullCnt > 0 + then do + let bitmapPtr = castPtr bitmapVoid :: Ptr Word8 + vec <- V.generateM n $ \i -> do + valid <- bitmapIsSet bitmapPtr i + if valid + then do + start <- fromIntegral <$> peekElemOff offsetPtr i + end <- fromIntegral <$> peekElemOff offsetPtr (i + 1) + bs <- + BS.packCStringLen + (castPtr (charPtr `plusPtr` start), end - start) + return $ Just (TE.decodeUtf8 bs) + else return Nothing + return $ OptionalColumn vec + else do + vec <- V.generateM n $ \i -> do + start <- fromIntegral <$> peekElemOff offsetPtr i + end <- fromIntegral <$> peekElemOff offsetPtr (i + 1) + bs <- + BS.packCStringLen + (castPtr (charPtr `plusPtr` start), end - start) + return $ TE.decodeUtf8 bs + return $ BoxedColumn vec diff --git a/ffi/DataFrame/IR.hs b/ffi/DataFrame/IR.hs new file mode 100644 --- /dev/null +++ b/ffi/DataFrame/IR.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Intermediate Representation for DataFrame query plans. + JSON-decodable plan tree + interpreter. +-} +module DataFrame.IR ( + PlanNode (..), + AggSpec (..), + executePlan, +) where + +import Data.Aeson (FromJSON (..), withObject, (.:)) +import Data.Aeson.Types (Parser) +import qualified Data.Text as T +import Data.Type.Equality ( + TestEquality (testEquality), + type (:~:) (Refl), + ) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Data.Word (Word64) +import Foreign (wordPtrToPtr) +import Type.Reflection (typeRep) + +import DataFrame.Functions (count, mean, meanMaybe, sumMaybe) +import qualified DataFrame.Functions as Functions +import DataFrame.IO.Arrow (arrowToDataframe) +import DataFrame.IO.CSV ( + defaultReadOptions, + readSeparated, + readTsv, + ) +import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..), NamedExpr) +import DataFrame.Operations.Aggregation (aggregate, groupBy) +import DataFrame.Operations.Permutation (SortOrder (..), sortBy) +import DataFrame.Operations.Subset (select) +import qualified DataFrame.Operations.Subset as Subset +import DataFrame.Operators ((.=)) + +-- --------------------------------------------------------------------------- +-- IR types +-- --------------------------------------------------------------------------- + +data AggSpec = AggSpec + { aggName :: T.Text + , aggFn :: T.Text + , aggCol :: T.Text + } + deriving (Show) + +data PlanNode + = ReadCsv FilePath + | ReadTsv FilePath + | -- | schema_addr array_addr + FromArrow Word64 Word64 + | Select [T.Text] PlanNode + | GroupBy [T.Text] [AggSpec] PlanNode + | Sort [T.Text] Bool PlanNode + | Limit Int PlanNode + deriving (Show) + +-- --------------------------------------------------------------------------- +-- JSON decoding +-- --------------------------------------------------------------------------- + +instance FromJSON AggSpec where + parseJSON = withObject "AggSpec" $ \o -> + AggSpec + <$> o .: "name" + <*> o .: "agg" + <*> o .: "col" + +instance FromJSON PlanNode where + parseJSON = withObject "PlanNode" $ \o -> do + op <- o .: "op" :: Parser T.Text + case op of + "ReadCsv" -> ReadCsv <$> o .: "path" + "ReadTsv" -> ReadTsv <$> o .: "path" + "FromArrow" -> FromArrow <$> o .: "schema" <*> o .: "array" + "Select" -> Select <$> o .: "cols" <*> o .: "input" + "GroupBy" -> GroupBy <$> o .: "keys" <*> o .: "aggregations" <*> o .: "input" + "Sort" -> Sort <$> o .: "cols" <*> o .: "ascending" <*> o .: "input" + "Limit" -> Limit <$> o .: "n" <*> o .: "input" + _ -> fail $ "DataFrame.IR: unknown op: " ++ T.unpack op + +executePlan :: PlanNode -> IO DataFrame +executePlan (ReadCsv path) = + readSeparated defaultReadOptions path +executePlan (ReadTsv path) = + readTsv path +executePlan (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 + nes <- mapM (buildNamedExpr df) aggs + return $ aggregate nes (groupBy keys df) +executePlan (Sort cols ascending node) = do + df <- executePlan 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 + +-- | Build a SortOrder from a column's runtime type. +mkSortOrder :: Bool -> T.Text -> Column -> SortOrder +mkSortOrder True n (UnboxedColumn (_ :: VU.Vector a)) = Asc (Col @a n) +mkSortOrder False n (UnboxedColumn (_ :: VU.Vector a)) = Desc (Col @a n) +mkSortOrder True n (BoxedColumn (_ :: V.Vector a)) = Asc (Col @a n) +mkSortOrder False n (BoxedColumn (_ :: V.Vector a)) = Desc (Col @a n) +mkSortOrder True n (OptionalColumn (_ :: V.Vector (Maybe a))) = Asc (Col @(Maybe a) n) +mkSortOrder False n (OptionalColumn (_ :: V.Vector (Maybe a))) = Desc (Col @(Maybe a) n) + +-- | Dispatch aggregation by fn name and runtime column type. +buildNamedExpr :: DataFrame -> AggSpec -> IO NamedExpr +buildNamedExpr df (AggSpec name fn colName) = + case fn of + "count" -> countExpr name colName (unsafeGetColumn colName df) + "sum" -> sumExpr name colName (unsafeGetColumn colName df) + "mean" -> meanExpr name colName (unsafeGetColumn colName df) + other -> + ioError $ + userError $ + "DataFrame.IR: unknown aggregation '" ++ T.unpack other ++ "'" + +countExpr :: T.Text -> T.Text -> Column -> IO NamedExpr +countExpr name colName (UnboxedColumn (_ :: VU.Vector a)) = return $ name .= count (Col @a colName) +countExpr name colName (BoxedColumn (_ :: V.Vector a)) = return $ name .= count (Col @a colName) +countExpr name colName (OptionalColumn (_ :: V.Vector (Maybe a))) = return $ name .= count (Col @(Maybe a) colName) + +sumExpr :: T.Text -> T.Text -> Column -> IO NamedExpr +sumExpr name colName (UnboxedColumn (_ :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= Functions.sum (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= Functions.sum (Col @Double colName) +sumExpr name colName (BoxedColumn (_ :: V.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= Functions.sum (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= Functions.sum (Col @Double colName) +sumExpr name colName (OptionalColumn (_ :: V.Vector (Maybe a))) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= sumMaybe (Col @(Maybe Int) colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= sumMaybe (Col @(Maybe Double) colName) +sumExpr name colName _ = + ioError $ + userError $ + "DataFrame.IR: sum: unsupported column type for '" ++ T.unpack colName ++ "'" + +meanExpr :: T.Text -> T.Text -> Column -> IO NamedExpr +meanExpr name colName (UnboxedColumn (_ :: VU.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= mean (Col @Int colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= mean (Col @Double colName) +meanExpr name colName (BoxedColumn (_ :: V.Vector a)) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= mean (Col @Double colName) +meanExpr name colName (OptionalColumn (_ :: V.Vector (Maybe a))) + | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = + return $ name .= meanMaybe (Col @(Maybe Double) colName) + | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = + return $ name .= meanMaybe (Col @(Maybe Int) colName) +meanExpr name colName _ = + ioError $ + userError $ + "DataFrame.IR: mean: unsupported column type for '" ++ T.unpack colName ++ "'" diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -219,6 +219,9 @@ module UnstableCSV, module Parquet, + -- * Type conversion + module Typing, + -- * Operations module Subset, module Transformations, @@ -298,6 +301,7 @@ toRowVector, ) import DataFrame.Internal.Schema as Schema ( + makeSchema, schemaType, ) import DataFrame.Operations.Aggregation as Aggregation ( @@ -359,8 +363,11 @@ sample, select, selectBy, + stratifiedSample, + stratifiedSplit, take, takeLast, ) import DataFrame.Operations.Transformations as Transformations +import DataFrame.Operations.Typing as Typing import DataFrame.Operators as Operators diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -1,5 +1,6 @@ {-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -12,7 +13,7 @@ import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) import DataFrame.Internal.Expression (Expr (..), eSize, getColumns) import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Internal.Statistics (percentile', percentileOrd') +import DataFrame.Internal.Statistics (percentileOrd') import DataFrame.Internal.Types import DataFrame.Operations.Core (columnNames, nRows) import DataFrame.Operations.Subset (exclude, filterWhere) @@ -21,9 +22,12 @@ import Control.Monad (guard) import Data.Containers.ListUtils (nubOrd) import Data.Function (on) +#if MIN_VERSION_base(4,20,0) +import Data.List (maximumBy, minimumBy, sort, sortBy) +#else import Data.List (foldl', maximumBy, minimumBy, sort, sortBy) +#endif import qualified Data.Map.Strict as M -import Data.Maybe import qualified Data.Text as T import Data.Type.Equality import qualified Data.Vector as V @@ -548,38 +552,108 @@ (boolExpansion (synthConfig cfg)) ) +-- | Unifies non-nullable and nullable Double expressions for feature generation. +data NumExpr + = NDouble !(Expr Double) + | NMaybeDouble !(Expr (Maybe Double)) + +numExprCols :: NumExpr -> [T.Text] +numExprCols (NDouble e) = getColumns e +numExprCols (NMaybeDouble e) = getColumns e + +numExprEq :: NumExpr -> NumExpr -> Bool +numExprEq (NDouble e1) (NDouble e2) = e1 == e2 +numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = e1 == e2 +numExprEq _ _ = False + +combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] +combineNumExprs (NDouble e1) (NDouble e2) = + [ NDouble (e1 .+ e2) + , NDouble (e1 .- e2) + , NDouble (e1 .* e2) + , NDouble + (F.ifThenElse (e2 ./= F.lit (0 :: Double)) (e1 ./ e2) (F.lit (0 :: Double))) + ] +combineNumExprs (NDouble e1) (NMaybeDouble e2) = + [ NMaybeDouble (e1 .+ e2) + , NMaybeDouble (e1 .- e2) + , NMaybeDouble (e1 .* e2) + , NMaybeDouble + ( F.ifThenElse + (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) + (e1 ./ e2) + (F.lit (Nothing :: Maybe Double)) + ) + ] +combineNumExprs (NMaybeDouble e1) (NDouble e2) = + [ NMaybeDouble (e1 .+ e2) + , NMaybeDouble (e1 .- e2) + , NMaybeDouble (e1 .* e2) + , NMaybeDouble + ( F.ifThenElse + (e2 ./= F.lit (0 :: Double)) + (e1 ./ e2) + (F.lit (Nothing :: Maybe Double)) + ) + ] +combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) = + [ NMaybeDouble (e1 .+ e2) + , NMaybeDouble (e1 .- e2) + , NMaybeDouble (e1 .* e2) + , NMaybeDouble + ( F.ifThenElse + (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) + (e1 ./ e2) + (F.lit (Nothing :: Maybe Double)) + ) + ] + numericConditions :: TreeConfig -> DataFrame -> [Expr Bool] numericConditions = generateNumericConds generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool] generateNumericConds cfg df = do expr <- numericExprsWithTerms (synthConfig cfg) df - let thresholds = map (\p -> percentile p expr df) (percentiles cfg) + let thresholds = numericThresholds expr threshold <- thresholds - [ expr .<= F.lit threshold - , expr .>= F.lit threshold - , expr .< F.lit threshold - , expr .> F.lit threshold + numericCondsFromExpr expr threshold + where + numericThresholds (NDouble e) = map (\p -> percentile p e df) (percentiles cfg) + numericThresholds (NMaybeDouble e) = map (\p -> percentile p (F.fromMaybe 0 e) df) (percentiles cfg) + + numericCondsFromExpr (NDouble e) t = + [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] + numericCondsFromExpr (NMaybeDouble e) t = + [ F.fromMaybe False (e .<= F.lit t) + , F.fromMaybe False (e .>= F.lit t) + , F.fromMaybe False (e .< F.lit t) + , F.fromMaybe False (e .> F.lit t) ] -numericExprsWithTerms :: SynthConfig -> DataFrame -> [Expr Double] +numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr] numericExprsWithTerms cfg df = concatMap (numericExprs cfg df [] 0) [0 .. maxExprDepth cfg] -numericCols :: DataFrame -> [Expr Double] +numericCols :: DataFrame -> [NumExpr] numericCols df = concatMap extract (columnNames df) where extract col = case unsafeGetColumn col df of UnboxedColumn (_ :: VU.Vector b) -> case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [Col col] + Just Refl -> [NDouble (Col col)] Nothing -> case sIntegral @b of - STrue -> [F.toDouble (Col @b col)] + STrue -> [NDouble (F.toDouble (Col @b col))] SFalse -> [] + OptionalColumn (_ :: V.Vector (Maybe b)) -> + case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [NMaybeDouble (Col @(Maybe b) col)] + Nothing -> case sIntegral @b of + STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) col))] + SFalse -> [] _ -> [] numericExprs :: - SynthConfig -> DataFrame -> [Expr Double] -> Int -> Int -> [Expr Double] + SynthConfig -> DataFrame -> [NumExpr] -> Int -> Int -> [NumExpr] numericExprs cfg df prevExprs depth maxDepth | depth == 0 = baseExprs ++ numericExprs cfg df baseExprs (depth + 1) maxDepth | depth >= maxDepth = [] @@ -592,16 +666,16 @@ | otherwise = do e1 <- prevExprs e2 <- baseExprs - let cols = getColumns e1 <> getColumns e2 + let cols = numExprCols e1 <> numExprCols e2 guard - ( e1 /= e2 + ( not (numExprEq e1 e2) && not ( any (\(l, r) -> l `elem` cols && r `elem` cols) (disallowedCombinations cfg) ) ) - [e1 + e2, e1 - e2, e1 * e2, F.ifThenElse (e2 ./= 0) (e1 / e2) 0] + combineNumExprs e1 e2 boolExprs :: DataFrame -> [Expr Bool] -> [Expr Bool] -> Int -> Int -> [Expr Bool] @@ -625,42 +699,14 @@ genConds colName = case unsafeGetColumn colName df of (BoxedColumn (col :: V.Vector a)) -> let ps = map (Lit . (`percentileOrd'` col)) [1, 25, 75, 99] - in map (Col @a colName .==) ps + in map (F.lift2 (==) (Col @a colName)) ps (OptionalColumn (col :: V.Vector (Maybe a))) -> case sFloating @a of - STrue -> - let doubleCol = - VU.convert - (V.map fromJust (V.filter isJust (V.map (fmap (realToFrac @a @Double)) col))) - in zipWith - ($) - [ (Col @(Maybe a) colName .==) - , (Col @(Maybe a) colName .<=) - , (Col @(Maybe a) colName .>=) - ] - ( Lit Nothing - : map - (Lit . Just . realToFrac . (`percentile'` doubleCol)) - (percentiles cfg) - ) + STrue -> [] -- handled by numericCols / numericExprs SFalse -> case sIntegral @a of - STrue -> - let doubleCol = - VU.convert - (V.map fromJust (V.filter isJust (V.map (fmap (fromIntegral @a @Double)) col))) - in zipWith - ($) - [ (Col @(Maybe a) colName .==) - , (Col @(Maybe a) colName .<=) - , (Col @(Maybe a) colName .>=) - ] - ( Lit Nothing - : map - (Lit . Just . round . (`percentile'` doubleCol)) - (percentiles cfg) - ) + STrue -> [] -- handled by numericCols / numericExprs SFalse -> map - ((Col @(Maybe a) colName .==) . Lit . (`percentileOrd'` col)) + (F.lift2 (==) (Col @(Maybe a) colName) . Lit . (`percentileOrd'` col)) [1, 25, 75, 99] (UnboxedColumn (_ :: VU.Vector a)) -> [] @@ -681,15 +727,18 @@ (BoxedColumn (col1 :: V.Vector a), BoxedColumn (_ :: V.Vector b)) -> case testEquality (typeRep @a) (typeRep @b) of Nothing -> [] - Just Refl -> [Col @a l .== Col @a r] + Just Refl -> [F.lift2 (==) (Col @a l) (Col @a r)] (UnboxedColumn (_ :: VU.Vector a), UnboxedColumn (_ :: VU.Vector b)) -> [] ( OptionalColumn (_ :: V.Vector (Maybe a)) , OptionalColumn (_ :: V.Vector (Maybe b)) ) -> case testEquality (typeRep @a) (typeRep @b) of Nothing -> [] Just Refl -> case testEquality (typeRep @a) (typeRep @T.Text) of - Nothing -> [Col @(Maybe a) l .<= Col r, Col @(Maybe a) l .== Col r] - Just Refl -> [Col @(Maybe a) l .== Col r] + Nothing -> + [ F.lift2 (<=) (Col @(Maybe a) l) (Col r) + , F.lift2 (==) (Col @(Maybe a) l) (Col r) + ] + Just Refl -> [F.lift2 (==) (Col @(Maybe a) l) (Col r)] _ -> [] in concatMap genConds (columnNames df) ++ columnConds @@ -759,3 +808,117 @@ pruneTree :: forall a. (Columnable a, Eq a) => Expr a -> Expr a pruneTree = pruneExpr + +-- | A tree where each leaf stores a class-probability distribution. +type ProbTree a = Tree (M.Map a Double) + +-- | Compute normalised class probabilities from a subset of training rows. +probsFromIndices :: + forall a. + (Columnable a) => + T.Text -> + DataFrame -> + V.Vector Int -> + M.Map a Double +probsFromIndices target df indices = + case interpret @a df (Col target) of + Left _ -> M.empty + Right (TColumn col) -> + case toVector @a col of + Left _ -> M.empty + Right vals -> + let counts = + V.foldl' + (\acc i -> M.insertWith (+) (vals V.! i) 1 acc) + M.empty + indices + total = fromIntegral (V.length indices) :: Double + in M.map (\c -> fromIntegral c / total) counts + +{- | Annotate a fitted 'Tree a' with class distributions by routing the + training data through it. The split conditions are preserved; only the + leaf values change from a majority label to a probability map. +-} +buildProbTree :: + forall a. + (Columnable a) => + Tree a -> + T.Text -> + DataFrame -> + V.Vector Int -> + ProbTree a +buildProbTree (Leaf _) target df indices = + Leaf (probsFromIndices @a target df indices) +buildProbTree (Branch cond left right) target df indices = + let (indicesL, indicesR) = partitionIndices cond df indices + in Branch + cond + (buildProbTree @a left target df indicesL) + (buildProbTree @a right target df indicesR) + +{- | Fit a TAO decision tree and return one @Expr Double@ per class. + + Each @(c, e)@ pair in the result map means: evaluate @e@ on a 'DataFrame' + row to get the predicted probability of class @c@. You can insert these + as new columns with 'derive' or evaluate them with 'interpret'. + + Example: + @ + let pes = fitProbTree \@T.Text cfg (Col \"species\") trainDf + -- pes M.! \"setosa\" :: Expr Double + df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes + @ +-} +fitProbTree :: + forall a. + (Columnable a) => + TreeConfig -> + Expr a -> -- target column, e.g. @Col \"label\"@ + DataFrame -> + M.Map a (Expr Double) +fitProbTree cfg (Col target) df = + let + conds = + nubOrd $ + numericConditions cfg (exclude [target] df) + ++ generateConditionsOld cfg (exclude [target] df) + initialTree = buildGreedyTree @a cfg (maxTreeDepth cfg) target conds df + indices = V.enumFromN 0 (nRows df) + optimizedTree = taoOptimize @a cfg target conds df indices initialTree + pruned = pruneDead optimizedTree + in + probExprs (buildProbTree @a pruned target df indices) +fitProbTree _ expr _ = + error $ "Cannot create prob tree for compound expression: " ++ show expr + +{- | Convert a 'ProbTree' into one 'Expr Double' per class. + + Each @(c, e)@ pair means: evaluate @e@ on a 'DataFrame' row to get the + predicted probability of class @c@. You can insert these as new columns + with 'derive' or evaluate them with 'interpret'. + + Example: + @ + let pt = fitProbTree \@T.Text cfg (Col \"species\") trainDf + pes = probExprs pt + -- pes M.! \"setosa\" :: Expr Double + df' = M.foldlWithKey' (\\d cls e -> D.derive (cls <> \"_prob\") e d) testDf pes + @ +-} +probExprs :: + forall a. + (Columnable a) => + ProbTree a -> + M.Map a (Expr Double) +probExprs tree = + let classes = nubOrd (allClasses tree) + in M.fromList [(c, classExpr c tree) | c <- classes] + where + allClasses :: ProbTree a -> [a] + allClasses (Leaf m) = M.keys m + allClasses (Branch _ l r) = allClasses l ++ allClasses r + + classExpr :: a -> ProbTree a -> Expr Double + classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m) + classExpr c (Branch cond l r) = + F.ifThenElse cond (classExpr c l) (classExpr c r) diff --git a/src/DataFrame/Functions.hs b/src/DataFrame/Functions.hs --- a/src/DataFrame/Functions.hs +++ b/src/DataFrame/Functions.hs @@ -1,5 +1,4 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} @@ -11,6 +10,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where @@ -29,12 +29,14 @@ import Control.Monad import Control.Monad.IO.Class import qualified Data.Char as Char +import Data.Either import Data.Function import Data.Functor import Data.Int import qualified Data.List as L import qualified Data.Map as M import qualified Data.Maybe as Maybe +import qualified Data.Set as S import qualified Data.Text as T import Data.Time import qualified Data.Vector as V @@ -43,6 +45,14 @@ import qualified DataFrame.IO.CSV as CSV import qualified DataFrame.IO.Parquet as Parquet import DataFrame.IO.Parquet.Thrift +import DataFrame.IO.Parquet.Types (columnNullCount) +import DataFrame.Internal.Nullable ( + BaseType, + NullLift1Op (applyNull1), + NullLift1Result, + NullLift2Op (applyNull2), + NullLift2Result, + ) import DataFrame.Operators import Debug.Trace (trace) import Language.Haskell.TH @@ -72,32 +82,107 @@ } ) -liftDecorated :: - (Columnable a, Columnable b) => - (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b -liftDecorated f name rep = Unary (MkUnaryOp{unaryFn = f, unaryName = name, unarySymbol = rep}) +{- | Lift a unary function over a nullable or non-nullable column expression. +When the input is @Maybe a@, 'Nothing' short-circuits (like 'fmap'). +When the input is plain @a@, the function is applied directly. -lift2Decorated :: - (Columnable c, Columnable b, Columnable a) => - (c -> b -> a) -> - T.Text -> - Maybe T.Text -> - Bool -> - Int -> - Expr c -> +The return type is inferred via 'NullLift1Result': no annotation needed. +-} +nullLift :: + (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => + (BaseType a -> r) -> + Expr a -> + Expr (NullLift1Result a r) +nullLift f = + Unary + (MkUnaryOp{unaryFn = applyNull1 f, unaryName = "nullLift", unarySymbol = Nothing}) + +{- | Lift a binary function over nullable or non-nullable column expressions. +Any 'Nothing' operand short-circuits to 'Nothing' in the result. + +The return type is inferred via 'NullLift2Result': no annotation needed. +-} +nullLift2 :: + (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => + (BaseType a -> BaseType b -> r) -> + Expr a -> Expr b -> - Expr a -lift2Decorated f name rep comm prec = + Expr (NullLift2Result a b r) +nullLift2 f = Binary ( MkBinaryOp - { binaryFn = f - , binaryName = name - , binarySymbol = rep - , binaryCommutative = comm - , binaryPrecedence = prec + { binaryFn = applyNull2 f + , binaryName = "nullLift2" + , binarySymbol = Nothing + , binaryCommutative = False + , binaryPrecedence = 0 } ) +{- | Lenient numeric \/ text coercion returning @Maybe a@. Looks up column +@name@ and coerces its values to @a@. Values that cannot be converted +(parse failures, type mismatches) become 'Nothing'; successfully converted +values are wrapped in 'Just'. Existing 'Nothing' in optional source columns +stays as 'Nothing'. +-} +cast :: forall a. (Columnable a) => T.Text -> Expr (Maybe a) +cast name = CastWith name "cast" (either (const Nothing) Just) + +{- | Lenient coercion that substitutes a default for unconvertible values. +Looks up column @name@, coerces its values to @a@, and uses @def@ wherever +conversion fails or the source value is 'Nothing'. +-} +castWithDefault :: forall a. (Columnable a) => a -> T.Text -> Expr a +castWithDefault def name = + CastWith name ("castWithDefault:" <> T.pack (show def)) (fromRight def) + +{- | Lenient coercion returning @Either T.Text a@. Successfully converted +values are 'Right'; values that cannot be parsed are kept as 'Left' with +their original string representation, so the caller can inspect or handle +them downstream. Existing 'Nothing' in optional source columns becomes +@Left \"null\"@. +-} +castEither :: forall a. (Columnable a) => T.Text -> Expr (Either T.Text a) +castEither name = CastWith name "castEither" (either (Left . T.pack) Right) + +{- | Lenient coercion for assertedly non-nullable columns. +Substitutes @error@ for @Nothing@, so it will crash at evaluation time if +any @Nothing@ is actually encountered. For non-nullable and +fully-populated nullable columns no cost is paid. +-} +unsafeCast :: forall a. (Columnable a) => T.Text -> Expr a +unsafeCast name = + CastWith + name + "unsafeCast" + (fromRight (error "unsafeCast: unexpected Nothing in column")) + +castExpr :: + forall b src. (Columnable b, Columnable src) => Expr src -> Expr (Maybe b) +castExpr = CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) + +castExprWithDefault :: + forall b src. (Columnable b, Columnable src) => b -> Expr src -> Expr b +castExprWithDefault def = + CastExprWith @b @b @src + ("castExprWithDefault:" <> T.pack (show def)) + (fromRight def) + +castExprEither :: + forall b src. + (Columnable b, Columnable src) => Expr src -> Expr (Either T.Text b) +castExprEither = + CastExprWith @b @(Either T.Text b) @src + "castExprEither" + (either (Left . T.pack) Right) + +unsafeCastExpr :: + forall b src. (Columnable b, Columnable src) => Expr src -> Expr b +unsafeCastExpr = + CastExprWith @b @b @src + "unsafeCastExpr" + (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) + toDouble :: (Columnable a, Real a) => Expr a -> Expr Double toDouble = Unary @@ -115,20 +200,22 @@ mod :: (Integral a, Columnable a) => Expr a -> Expr a -> Expr a mod = lift2Decorated Prelude.mod "mod" Nothing False 7 -eq :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -eq = (.==) +eq :: (Columnable a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +eq = lift2Decorated (==) "eq" (Just "==") True 4 -lt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -lt = (.<) +lt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +lt = lift2Decorated (<) "lt" (Just "<") False 4 -gt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -gt = (.>) +gt :: (Columnable a, Ord a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +gt = lift2Decorated (>) "gt" (Just ">") False 4 -leq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -leq = (.<=) +leq :: + (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +leq = lift2Decorated (<=) "leq" (Just "<=") False 4 -geq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -geq = (.>=) +geq :: + (Columnable a, Ord a, Eq a, a ~ BaseType a) => Expr a -> Expr a -> Expr Bool +geq = lift2Decorated (>=) "geq" (Just ">=") False 4 and :: Expr Bool -> Expr Bool -> Expr Bool and = (.&&) @@ -222,7 +309,7 @@ zScore c = (c - mean c) / stddev c pow :: (Columnable a, Num a) => Expr a -> Int -> Expr a -pow = (.^^) +pow expr i = lift2Decorated (^) "max" (Just "^") True 8 expr (Lit i) relu :: (Columnable a, Num a, Ord a) => Expr a -> Expr a relu = liftDecorated (Prelude.max 0) "relu" Nothing @@ -256,13 +343,13 @@ whenPresent :: forall a b. (Columnable a, Columnable b) => (a -> b) -> Expr (Maybe a) -> Expr (Maybe b) -whenPresent f = lift (fmap f) +whenPresent f = liftDecorated (fmap f) "whenPresent" Nothing whenBothPresent :: forall a b c. (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Expr (Maybe a) -> Expr (Maybe b) -> Expr (Maybe c) -whenBothPresent f = lift2 (\l r -> f <$> l <*> r) +whenBothPresent f = lift2Decorated (\l r -> f <$> l <*> r) "whenBothPresent" Nothing False 0 recode :: forall a b. @@ -340,7 +427,7 @@ forall a b m. (Columnable a, Columnable (m a), Monad m, Columnable b, Columnable (m b)) => (a -> m b) -> Expr (m a) -> Expr (m b) -bind f = lift (>>= f) +bind f = liftDecorated (>>= f) "bind" Nothing -- See Section 2.4 of the Haskell Report https://www.haskell.org/definition/haskell2010.pdf isReservedId :: T.Text -> Bool @@ -452,33 +539,43 @@ declareColumnsFromParquetFile :: String -> DecsQ declareColumnsFromParquetFile path = do isDir <- liftIO $ doesDirectoryExist path - let pat = if isDir then path "*.parquet" else path - matches <- liftIO $ glob pat - files <- liftIO $ filterM (fmap Prelude.not . doesDirectoryExist) matches - df <- - liftIO $ - foldM - ( \acc p -> do - (metadata, _) <- liftIO (Parquet.readMetadataFromPath path) - let d = schemaToEmptyDataFrame (schema metadata) - pure $ acc <> d - ) + metas <- liftIO $ mapM (fmap fst . Parquet.readMetadataFromPath) files + let nullableCols :: S.Set T.Text + nullableCols = + S.fromList + [ T.pack (last colPath) + | meta <- metas + , rg <- rowGroups meta + , cc <- rowGroupColumns rg + , let cm = columnMetaData cc + colPath = columnPathInSchema cm + , Prelude.not (null colPath) + , columnNullCount (columnStatistics cm) > 0 + ] + let df = + foldl + (\acc meta -> acc <> schemaToEmptyDataFrame nullableCols (schema meta)) DataFrame.Internal.DataFrame.empty - files + metas declareColumns df -schemaToEmptyDataFrame :: [SchemaElement] -> DataFrame -schemaToEmptyDataFrame elems = +schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame +schemaToEmptyDataFrame nullableCols elems = let leafElems = filter (\e -> numChildren e == 0) elems - in fromNamedColumns (map schemaElemToColumn leafElems) + in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems) -schemaElemToColumn :: SchemaElement -> (T.Text, Column) -schemaElemToColumn elem = +schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column) +schemaElemToColumn nullableCols elem = let name = elementName elem - in (name, emptyColumnForType (elementType elem)) + isNull = name `S.member` nullableCols + col = + if isNull + then emptyNullableColumnForType (elementType elem) + else emptyColumnForType (elementType elem) + in (name, col) emptyColumnForType :: TType -> Column emptyColumnForType = \case @@ -491,6 +588,19 @@ FLOAT -> fromList @Float [] DOUBLE -> fromList @Double [] STRING -> fromList @T.Text [] + other -> error $ "Unsupported parquet type for column: " <> show other + +emptyNullableColumnForType :: TType -> Column +emptyNullableColumnForType = \case + BOOL -> fromList @(Maybe Bool) [] + BYTE -> fromList @(Maybe Word8) [] + I16 -> fromList @(Maybe Int16) [] + I32 -> fromList @(Maybe Int32) [] + I64 -> fromList @(Maybe Int64) [] + I96 -> fromList @(Maybe Int64) [] + FLOAT -> fromList @(Maybe Float) [] + DOUBLE -> fromList @(Maybe Double) [] + STRING -> fromList @(Maybe T.Text) [] other -> error $ "Unsupported parquet type for column: " <> show other declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ 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 @@ -169,7 +169,7 @@ data TypeSpec = InferFromSample Int - | SpecifyTypes [(T.Text, SchemaType)] + | SpecifyTypes [(T.Text, SchemaType)] TypeSpec | NoInference -- | CSV read parameters. @@ -202,14 +202,16 @@ shouldInferFromSample :: TypeSpec -> Bool shouldInferFromSample (InferFromSample _) = True +shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback shouldInferFromSample _ = False schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType -schemaTypeMap (SpecifyTypes xs) = M.fromList xs +schemaTypeMap (SpecifyTypes xs _) = M.fromList xs schemaTypeMap _ = M.empty typeInferenceSampleSize :: TypeSpec -> Int typeInferenceSampleSize (InferFromSample n) = n +typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback typeInferenceSampleSize _ = 0 defaultReadOptions :: ReadOptions @@ -221,7 +223,8 @@ , dateFormat = "%Y-%m-%d" , columnSeparator = ',' , numColumns = Nothing - , missingIndicators = [] + , missingIndicators = + ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"] } {- | Read CSV file from path and load it into a dataframe. @@ -310,12 +313,13 @@ frozenCols <- V.mapM (finalizeBuilderColumn opts) builderColsV let numRows = maybe 0 columnLength (frozenCols V.!? 0) - return $ - DataFrame - frozenCols - (M.fromList (zip columnNames [0 ..])) - (numRows, V.length frozenCols) - M.empty -- TODO give typed column references + let df = + DataFrame + frozenCols + (M.fromList (zip columnNames [0 ..])) + (numRows, V.length frozenCols) + M.empty -- TODO give typed column references + pure $ parseWithTypes (safeRead opts) (schemaTypeMap (typeSpec opts)) df initializeColumns :: [T.Text] -> [BL.ByteString] -> ReadOptions -> IO [BuilderColumn] @@ -325,7 +329,7 @@ -- Return Nothing for columns that should be inferred from BS shouldInfer = case typeSpec opts of InferFromSample _ -> True - SpecifyTypes _ -> True + SpecifyTypes _ fallback -> shouldInferFromSample fallback NoInference -> False lookupType name = M.lookup name typeMap initColumn :: T.Text -> Maybe SchemaType -> IO BuilderColumn @@ -370,12 +374,12 @@ Nothing -> appendPagedUnboxedVector gv 0.0 >> appendPagedUnboxedVector valid 0 BuilderText gv valid -> do let !val = T.strip (TE.decodeUtf8Lenient bs') - if isNullish val || val `elem` missing + if val `elem` missing then appendPagedVector gv T.empty >> appendPagedUnboxedVector valid 0 else appendPagedVector gv val >> appendPagedUnboxedVector valid 1 BuilderBS gv valid -> do let !bs'' = C.strip bs' - if isNullishBS bs'' || TE.decodeUtf8Lenient bs'' `elem` missing + if TE.decodeUtf8Lenient bs'' `elem` missing then appendPagedVector gv BS.empty >> appendPagedUnboxedVector valid 0 else appendPagedVector gv bs'' >> appendPagedUnboxedVector valid 1 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 @@ -6,39 +6,53 @@ module DataFrame.IO.Parquet where -import Control.Exception (throw) +import Control.Exception (throw, try) import Control.Monad -import Data.Bits import qualified Data.ByteString as BSO import Data.Either import Data.IORef import Data.Int import qualified Data.List as L -import qualified Data.Map as M +import qualified Data.Map.Strict as M import qualified Data.Set as S import qualified Data.Text as T import Data.Text.Encoding import Data.Time import Data.Time.Clock.POSIX (posixSecondsToUTCTime) -import Data.Word import DataFrame.Errors (DataFrameException (ColumnNotFoundException)) +import DataFrame.Internal.Binary (littleEndianWord32) import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr, getColumns) import qualified DataFrame.Operations.Core as DI import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS -import System.FilePath.Glob (glob) +import System.FilePath.Glob (compile, glob, match) +import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:)) import DataFrame.IO.Parquet.Dictionary import DataFrame.IO.Parquet.Levels import DataFrame.IO.Parquet.Page import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types -import System.Directory (doesDirectoryExist) +import Network.HTTP.Simple ( + getResponseBody, + getResponseStatusCode, + httpBS, + parseRequest, + setRequestHeader, + ) +import System.Directory ( + doesDirectoryExist, + getHomeDirectory, + getTemporaryDirectory, + ) +import System.Environment (lookupEnv) import qualified Data.Vector.Unboxed as VU +import DataFrame.IO.Parquet.Seeking import System.FilePath (()) +import System.IO (IOMode (ReadMode)) -- Options ----------------------------------------------------------------- @@ -131,8 +145,19 @@ p : go (sChildren n) ps False readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame -readParquetWithOpts opts path = do - (fileMetadata, contents) <- readMetadataFromPath path +readParquetWithOpts opts path + | isHFUri path = do + paths <- fetchHFParquetFiles path + let optsNoRange = opts{rowRange = Nothing} + dfs <- mapM (_readParquetWithOpts Nothing optsNoRange) paths + pure (applyRowRange opts (mconcat dfs)) + | otherwise = _readParquetWithOpts Nothing opts path + +-- | Internal function to pass testing parameters +_readParquetWithOpts :: + ForceNonSeekable -> ParquetReadOptions -> FilePath -> IO DataFrame +_readParquetWithOpts extraConfig opts path = withFileBufferedOrSeekable extraConfig path ReadMode $ \file -> do + fileMetadata <- readMetadataFromHandle file let columnPaths = getColumnPaths (drop 1 $ schema fileMetadata) let columnNames = map fst columnPaths let leafNames = map (last . T.splitOn ".") columnNames @@ -205,7 +230,11 @@ else colDataPageOffset let colLength = columnTotalCompressedSize metadata - let columnBytes = BSO.take (fromIntegral colLength) (BSO.drop (fromIntegral colStart) contents) + columnBytes <- + seekAndReadBytes + (Just (AbsoluteSeek, fromIntegral colStart)) + (fromIntegral colLength) + file pages <- readAllPages (columnCodec metadata) columnBytes @@ -237,13 +266,13 @@ Nothing -> do mc <- DI.newMutableColumn totalRows column DI.copyIntoMutableColumn mc 0 column - modifyIORef colMutMap (M.insert colFullName mc) - modifyIORef colOffMap (M.insert colFullName (DI.columnLength column)) + modifyIORef' colMutMap (M.insert colFullName mc) + modifyIORef' colOffMap (M.insert colFullName (DI.columnLength column)) Just mc -> do off <- (M.! colFullName) <$> readIORef colOffMap DI.copyIntoMutableColumn mc off column - modifyIORef colOffMap (M.adjust (+ DI.columnLength column) colFullName) - modifyIORef lTypeMap (M.insert colFullName lType) + modifyIORef' colOffMap (M.adjust (+ DI.columnLength column) colFullName) + modifyIORef' lTypeMap (M.insert colFullName lType) finalMutMap <- readIORef colMutMap finalColMap <- @@ -282,23 +311,29 @@ @ -} readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame -readParquetFilesWithOpts opts path = do - isDir <- doesDirectoryExist path +readParquetFilesWithOpts opts path + | isHFUri path = do + files <- fetchHFParquetFiles path + let optsWithoutRowRange = opts{rowRange = Nothing} + dfs <- mapM (_readParquetWithOpts Nothing optsWithoutRowRange) files + pure (applyRowRange opts (mconcat dfs)) + | otherwise = do + isDir <- doesDirectoryExist path - let pat = if isDir then path "*.parquet" else path + let pat = if isDir then path "*.parquet" else path - matches <- glob pat + matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches + files <- filterM (fmap not . doesDirectoryExist) matches - case files of - [] -> - error $ - "readParquetFiles: no parquet files found for " ++ path - _ -> do - let optsWithoutRowRange = opts{rowRange = Nothing} - dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files - pure (applyRowRange opts (mconcat dfs)) + case files of + [] -> + error $ + "readParquetFiles: no parquet files found for " ++ path + _ -> do + let optsWithoutRowRange = opts{rowRange = Nothing} + dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files + pure (applyRowRange opts (mconcat dfs)) -- Options application ----------------------------------------------------- @@ -322,28 +357,35 @@ -- File and metadata parsing ----------------------------------------------- +-- | read the file in memory at once, parse magicString and return the entire file ByteString readMetadataFromPath :: FilePath -> IO (FileMetadata, BSO.ByteString) readMetadataFromPath path = do contents <- BSO.readFile path - let (size, magicString) = contents `seq` readMetadataSizeFromFooter contents + let (size, magicString) = readMetadataSizeFromFooter contents when (magicString /= "PAR1") $ error "Invalid Parquet file" meta <- readMetadata contents size pure (meta, contents) -readMetadataSizeFromFooter :: BSO.ByteString -> (Int, BSO.ByteString) -readMetadataSizeFromFooter contents = +-- | read from the end of the file, parse magicString and return the entire file ByteString +readMetadataFromHandle :: FileBufferedOrSeekable -> IO FileMetadata +readMetadataFromHandle sh = do + footerBs <- readLastBytes (fromIntegral footerSize) sh + let (size, magicString) = readMetadataSizeFromFooterSlice footerBs + when (magicString /= "PAR1") $ error "Invalid Parquet file" + readMetadataByHandleMetaSize sh size + +-- | Takes the last 8 bit of the file to parse metadata size and magic string +readMetadataSizeFromFooterSlice :: BSO.ByteString -> (Int, BSO.ByteString) +readMetadataSizeFromFooterSlice contents = let - footerOffSet = BSO.length contents - 8 - sizeBytes = - map - (fromIntegral @Word8 @Int32 . BSO.index contents) - [footerOffSet .. footerOffSet + 3] - size = fromIntegral $ L.foldl' (.|.) 0 $ zipWith shift sizeBytes [0, 8, 16, 24] - magicStringBytes = map (BSO.index contents) [footerOffSet + 4 .. footerOffSet + 7] - magicString = BSO.pack magicStringBytes + size = fromIntegral (littleEndianWord32 contents) + magicString = BSO.take 4 (BSO.drop 4 contents) in (size, magicString) +readMetadataSizeFromFooter :: BSO.ByteString -> (Int, BSO.ByteString) +readMetadataSizeFromFooter = readMetadataSizeFromFooterSlice . BSO.takeEnd 8 + -- Schema navigation ------------------------------------------------------- getColumnPaths :: [SchemaElement] -> [(T.Text, Int)] @@ -582,3 +624,155 @@ applyScale :: Int32 -> Int32 -> Double applyScale scale rawValue = fromIntegral rawValue / (10 ^ scale) + +-- HuggingFace support ----------------------------------------------------- + +data HFRef = HFRef + { hfOwner :: T.Text + , hfDataset :: T.Text + , hfGlob :: T.Text + } + +data HFParquetFile = HFParquetFile + { hfpUrl :: T.Text + , hfpConfig :: T.Text + , hfpSplit :: T.Text + , hfpFilename :: T.Text + } + deriving (Show) + +instance FromJSON HFParquetFile where + parseJSON = withObject "HFParquetFile" $ \o -> + HFParquetFile + <$> o .: "url" + <*> o .: "config" + <*> o .: "split" + <*> o .: "filename" + +newtype HFParquetResponse = HFParquetResponse {hfParquetFiles :: [HFParquetFile]} + +instance FromJSON HFParquetResponse where + parseJSON = withObject "HFParquetResponse" $ \o -> + HFParquetResponse <$> o .: "parquet_files" + +isHFUri :: FilePath -> Bool +isHFUri = L.isPrefixOf "hf://" + +parseHFUri :: FilePath -> Either String HFRef +parseHFUri path = + let stripped = drop (length ("hf://datasets/" :: String)) path + in case T.splitOn "/" (T.pack stripped) of + (owner : dataset : rest) + | not (null rest) -> + Right $ HFRef owner dataset (T.intercalate "/" rest) + _ -> + Left $ "Invalid hf:// URI (expected hf://datasets/owner/dataset/glob): " ++ path + +getHFToken :: IO (Maybe BSO.ByteString) +getHFToken = do + envToken <- lookupEnv "HF_TOKEN" + case envToken of + Just t -> pure (Just (encodeUtf8 (T.pack t))) + Nothing -> do + home <- getHomeDirectory + let tokenPath = home ".cache" "huggingface" "token" + result <- try (BSO.readFile tokenPath) :: IO (Either IOError BSO.ByteString) + case result of + Right bs -> pure (Just (BSO.takeWhile (/= 10) bs)) + Left _ -> pure Nothing + +{- | Extract the repo-relative path from a HuggingFace download URL. +URL format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/{ref}/{path} +Returns the {path} portion (e.g. "data/train-00000-of-00001.parquet"). +-} +hfUrlRepoPath :: HFParquetFile -> String +hfUrlRepoPath f = + case T.breakOn "/resolve/" (hfpUrl f) of + (_, rest) + | not (T.null rest) -> + -- Drop "/resolve/", then drop the ref component (up to and including "/") + T.unpack $ T.drop 1 $ T.dropWhile (/= '/') $ T.drop (T.length "/resolve/") rest + _ -> + T.unpack (hfpConfig f) T.unpack (hfpSplit f) T.unpack (hfpFilename f) + +matchesGlob :: T.Text -> HFParquetFile -> Bool +matchesGlob g f = match (compile (T.unpack g)) (hfUrlRepoPath f) + +resolveHFUrls :: Maybe BSO.ByteString -> HFRef -> IO [HFParquetFile] +resolveHFUrls mToken ref = do + let dataset = hfOwner ref <> "/" <> hfDataset ref + let apiUrl = "https://datasets-server.huggingface.co/parquet?dataset=" ++ T.unpack dataset + req0 <- parseRequest apiUrl + let req = case mToken of + Nothing -> req0 + Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 + resp <- httpBS req + let status = getResponseStatusCode resp + when (status /= 200) $ + ioError $ + userError $ + "HuggingFace API returned status " + ++ show status + ++ " for dataset " + ++ T.unpack dataset + case eitherDecodeStrict (getResponseBody resp) of + Left err -> ioError $ userError $ "Failed to parse HF API response: " ++ err + Right hfResp -> pure $ filter (matchesGlob (hfGlob ref)) (hfParquetFiles hfResp) + +downloadHFFiles :: Maybe BSO.ByteString -> [HFParquetFile] -> IO [FilePath] +downloadHFFiles mToken files = do + tmpDir <- getTemporaryDirectory + forM files $ \f -> do + -- Derive a collision-resistant temp name from the URL path components + let fname = case (hfpConfig f, hfpSplit f) of + (c, s) | T.null c && T.null s -> T.unpack (hfpFilename f) + (c, s) -> T.unpack c <> "_" <> T.unpack s <> "_" <> T.unpack (hfpFilename f) + let destPath = tmpDir fname + req0 <- parseRequest (T.unpack (hfpUrl f)) + let req = case mToken of + Nothing -> req0 + Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 + resp <- httpBS req + let status = getResponseStatusCode resp + when (status /= 200) $ + ioError $ + userError $ + "Failed to download " ++ T.unpack (hfpUrl f) ++ " (HTTP " ++ show status ++ ")" + BSO.writeFile destPath (getResponseBody resp) + pure destPath + +-- | True when the path contains glob wildcard characters. +hasGlob :: T.Text -> Bool +hasGlob = T.any (\c -> c == '*' || c == '?' || c == '[') + +{- | Build the direct HF repo download URL for a path with no wildcards. +Format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/main/{path} +-} +directHFUrl :: HFRef -> T.Text +directHFUrl ref = + "https://huggingface.co/datasets/" + <> hfOwner ref + <> "/" + <> hfDataset ref + <> "/resolve/main/" + <> hfGlob ref + +fetchHFParquetFiles :: FilePath -> IO [FilePath] +fetchHFParquetFiles uri = do + ref <- case parseHFUri uri of + Left err -> ioError (userError err) + Right r -> pure r + mToken <- getHFToken + if hasGlob (hfGlob ref) + then do + hfFiles <- resolveHFUrls mToken ref + when (null hfFiles) $ + ioError $ + userError $ + "No parquet files found for " ++ uri + downloadHFFiles mToken hfFiles + else do + -- Direct repo file download — no datasets-server needed + let url = directHFUrl ref + let filename = last $ T.splitOn "/" (hfGlob ref) + downloadHFFiles mToken [HFParquetFile url "" "" filename] diff --git a/src/DataFrame/IO/Parquet/Binary.hs b/src/DataFrame/IO/Parquet/Binary.hs --- a/src/DataFrame/IO/Parquet/Binary.hs +++ b/src/DataFrame/IO/Parquet/Binary.hs @@ -15,44 +15,6 @@ import qualified Foreign.Ptr as Foreign import qualified Foreign.Storable as Foreign -littleEndianWord32 :: BS.ByteString -> Word32 -littleEndianWord32 bytes - | BS.length bytes >= 4 = - foldr - (.|.) - 0 - ( zipWith - (\b i -> fromIntegral b `shiftL` i) - (BS.unpack $ BS.take 4 bytes) - [0, 8, 16, 24] - ) - | otherwise = - littleEndianWord32 (BS.take 4 $ bytes `BS.append` BS.pack [0, 0, 0, 0]) - -littleEndianWord64 :: BS.ByteString -> Word64 -littleEndianWord64 bytes = - foldr - (.|.) - 0 - ( zipWith - (\b i -> fromIntegral b `shiftL` i) - (BS.unpack $ BS.take 8 bytes) - [0, 8 ..] - ) - -littleEndianInt32 :: BS.ByteString -> Int32 -littleEndianInt32 = fromIntegral . littleEndianWord32 - -word64ToLittleEndian :: Word64 -> BS.ByteString -word64ToLittleEndian w = - BS.map - (\i -> fromIntegral (w `shiftR` fromIntegral i)) - (BS.pack [0, 8, 16, 24, 32, 40, 48, 56]) - -word32ToLittleEndian :: Word32 -> BS.ByteString -word32ToLittleEndian w = - BS.map (\i -> fromIntegral (w `shiftR` fromIntegral i)) (BS.pack [0, 8, 16, 24]) - readUVarInt :: BS.ByteString -> (Word64, BS.ByteString) readUVarInt xs = loop xs 0 0 0 where @@ -100,7 +62,7 @@ readAndAdvance bufferPos buffer = do pos <- readIORef bufferPos let b = BS.index buffer pos - modifyIORef bufferPos (+ 1) + modifyIORef' bufferPos (+ 1) return b readVarIntFromBuffer :: (Integral a) => BS.ByteString -> IORef Int -> IO a diff --git a/src/DataFrame/IO/Parquet/Dictionary.hs b/src/DataFrame/IO/Parquet/Dictionary.hs --- a/src/DataFrame/IO/Parquet/Dictionary.hs +++ b/src/DataFrame/IO/Parquet/Dictionary.hs @@ -17,11 +17,15 @@ import qualified Data.Vector as V import qualified Data.Vector.Mutable as VM import qualified Data.Vector.Unboxed as VU -import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Encoding import DataFrame.IO.Parquet.Levels import DataFrame.IO.Parquet.Time import DataFrame.IO.Parquet.Types +import DataFrame.Internal.Binary ( + littleEndianInt32, + littleEndianWord32, + littleEndianWord64, + ) import qualified DataFrame.Internal.Column as DI import GHC.Float diff --git a/src/DataFrame/IO/Parquet/Encoding.hs b/src/DataFrame/IO/Parquet/Encoding.hs --- a/src/DataFrame/IO/Parquet/Encoding.hs +++ b/src/DataFrame/IO/Parquet/Encoding.hs @@ -1,14 +1,18 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} module DataFrame.IO.Parquet.Encoding where import Data.Bits import qualified Data.ByteString as BS import qualified Data.ByteString.Unsafe as BSU +#if !MIN_VERSION_base(4,20,0) import Data.List (foldl') +#endif import qualified Data.Vector.Unboxed as VU import Data.Word -import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Binary (readUVarInt) +import DataFrame.Internal.Binary (littleEndianWord32) ceilLog2 :: Int -> Int ceilLog2 x diff --git a/src/DataFrame/IO/Parquet/Levels.hs b/src/DataFrame/IO/Parquet/Levels.hs --- a/src/DataFrame/IO/Parquet/Levels.hs +++ b/src/DataFrame/IO/Parquet/Levels.hs @@ -5,10 +5,10 @@ import Data.List import qualified Data.Text as T -import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Encoding import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types +import DataFrame.Internal.Binary (littleEndianWord32) readLevelsV1 :: Int -> Int -> Int -> BS.ByteString -> ([Int], [Int], BS.ByteString) diff --git a/src/DataFrame/IO/Parquet/Page.hs b/src/DataFrame/IO/Parquet/Page.hs --- a/src/DataFrame/IO/Parquet/Page.hs +++ b/src/DataFrame/IO/Parquet/Page.hs @@ -14,6 +14,11 @@ import DataFrame.IO.Parquet.Binary import DataFrame.IO.Parquet.Thrift import DataFrame.IO.Parquet.Types +import DataFrame.Internal.Binary ( + littleEndianInt32, + littleEndianWord32, + littleEndianWord64, + ) import GHC.Float import qualified Snappy diff --git a/src/DataFrame/IO/Parquet/Seeking.hs b/src/DataFrame/IO/Parquet/Seeking.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet/Seeking.hs @@ -0,0 +1,145 @@ +{- | This module contains low-level utilities around file seeking + +potentially also contains all Streamly related low-level utilities. + +later this module can be renamed / moved to an internal module. +-} +module DataFrame.IO.Parquet.Seeking ( + SeekableHandle (getSeekableHandle), + SeekMode (..), + FileBufferedOrSeekable (..), + ForceNonSeekable, + advanceBytes, + mkFileBufferedOrSeekable, + mkSeekableHandle, + readLastBytes, + seekAndReadBytes, + seekAndStreamBytes, + withFileBufferedOrSeekable, +) where + +import Control.Monad +import Control.Monad.IO.Class +import qualified Data.ByteString as BS +import Data.IORef +import Data.Int +import Data.Word +import Streamly.Data.Stream (Stream) +import qualified Streamly.Data.Stream as S +import qualified Streamly.External.ByteString as SBS +import qualified Streamly.FileSystem.Handle as SHandle +import System.IO + +{- | This handle carries a proof that it must be seekable. +Note: Handle and SeekableHandle are not thread safe, should not be +shared across threads, beaware when running parallel/concurrent code. + +Not seekable: + - stdin / stdout + - pipes / FIFOs + +But regular files are always seekable. Parquet fundamentally wants random +access, a non-seekable source will not support effecient access without +buffering the entire file. +-} +newtype SeekableHandle = SeekableHandle {getSeekableHandle :: Handle} + +{- | If we truely want to support non-seekable files, we need to also consider the case +to buffer the entire file in memory. + +Not thread safe, contains mutable reference (as Handle already is). + +If we need concurrent / parallel parsing or something, we need to read into ByteString +first, not sharing the same handle. +-} +data FileBufferedOrSeekable + = FileBuffered !(IORef Int64) !BS.ByteString + | FileSeekable !SeekableHandle + +-- | Smart constructor for SeekableHandle +mkSeekableHandle :: Handle -> IO (Maybe SeekableHandle) +mkSeekableHandle h = do + seekable <- hIsSeekable h + pure $ if seekable then Just (SeekableHandle h) else Nothing + +-- | For testing only +type ForceNonSeekable = Maybe Bool + +{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case +if possible. +-} +mkFileBufferedOrSeekable :: + ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable +mkFileBufferedOrSeekable forceNonSeek h = do + seekable <- hIsSeekable h + if not seekable || forceNonSeek == Just True + then FileBuffered <$> newIORef 0 <*> BS.hGetContents h + else pure $ FileSeekable $ SeekableHandle h + +{- | With / bracket pattern for FileBufferedOrSeekable + +Warning: do not return the FileBufferedOrSeekable outside the scope of the action as +it will be closed. +-} +withFileBufferedOrSeekable :: + ForceNonSeekable -> + FilePath -> + IOMode -> + (FileBufferedOrSeekable -> IO a) -> + IO a +withFileBufferedOrSeekable forceNonSeek path ioMode action = withFile path ioMode $ \h -> do + fbos <- mkFileBufferedOrSeekable forceNonSeek h + action fbos + +-- | Read from the end, useful for reading metadata without loading entire file +readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString +readLastBytes n (FileSeekable sh) = do + let h = getSeekableHandle sh + hSeek h SeekFromEnd (negate n) + S.fold SBS.write (SHandle.read h) +readLastBytes n (FileBuffered i bs) = do + writeIORef i (fromIntegral $ BS.length bs) + when (n > fromIntegral (BS.length bs)) $ error "lastBytes: n > length bs" + pure $ BS.drop (BS.length bs - fromIntegral n) bs + +-- | Note: this does not guarantee n bytes (if it ends early) +advanceBytes :: Int -> FileBufferedOrSeekable -> IO BS.ByteString +advanceBytes = seekAndReadBytes Nothing + +-- | Note: this does not guarantee n bytes (if it ends early) +seekAndReadBytes :: + Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> IO BS.ByteString +seekAndReadBytes mSeek len f = seekAndStreamBytes mSeek len f >>= S.fold SBS.write + +{- | Warning: the stream produced from this function accesses to the mutable handler. +if multiple streams are pulled from the same handler at the same time, chaos happen. +Make sure there is only one stream running at one time for each SeekableHandle, +and streams are not read again when they are not used anymore. +-} +seekAndStreamBytes :: + (MonadIO m) => + Maybe (SeekMode, Integer) -> Int -> FileBufferedOrSeekable -> m (Stream m Word8) +seekAndStreamBytes mSeek len f = do + liftIO $ + case mSeek of + Nothing -> pure () + Just (seekMode, seekTo) -> fSeek f seekMode seekTo + pure $ S.take len $ fRead f + +fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO () +fSeek (FileSeekable (SeekableHandle h)) seekMode seekTo = hSeek h seekMode seekTo +fSeek (FileBuffered i bs) AbsoluteSeek seekTo = writeIORef i (fromIntegral seekTo) +fSeek (FileBuffered i bs) RelativeSeek seekTo = modifyIORef' i (+ fromIntegral seekTo) +fSeek (FileBuffered i bs) SeekFromEnd seekTo = writeIORef i (fromIntegral $ BS.length bs + fromIntegral seekTo) + +fRead :: (MonadIO m) => FileBufferedOrSeekable -> Stream m Word8 +fRead (FileSeekable (SeekableHandle h)) = SHandle.read h +fRead (FileBuffered i bs) = S.concatEffect $ do + pos <- liftIO $ readIORef i + pure $ + S.mapM + ( \x -> do + liftIO (modifyIORef' i (+ 1)) + pure x + ) + (S.unfold SBS.reader (BS.drop (fromIntegral pos) bs)) diff --git a/src/DataFrame/IO/Parquet/Thrift.hs b/src/DataFrame/IO/Parquet/Thrift.hs --- a/src/DataFrame/IO/Parquet/Thrift.hs +++ b/src/DataFrame/IO/Parquet/Thrift.hs @@ -21,6 +21,7 @@ import qualified Data.Vector.Unboxed as VU import Data.Word import DataFrame.IO.Parquet.Binary +import DataFrame.IO.Parquet.Seeking import DataFrame.IO.Parquet.Types import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn) @@ -329,6 +330,17 @@ let elemType = toTType sizeAndType replicateM_ sizeOnly (skipFieldData elemType buf pos) +{- | This avoids reading entire bytestring at once: it uses the seekable handle + seeks it to the end of the file to read the metadata +-} +readMetadataByHandleMetaSize :: FileBufferedOrSeekable -> Int -> IO FileMetadata +readMetadataByHandleMetaSize sh metaSize = do + let lastFieldId = 0 + bs <- readLastBytes (fromIntegral $ metaSize + footerSize) sh + bufferPos <- newIORef 0 + readFileMetaData defaultMetadata bs bufferPos lastFieldId + +-- | metadata starts from (L - 8 - meta_size) to L - 8 - 1. readMetadata :: BS.ByteString -> Int -> IO FileMetadata readMetadata contents size = do let metadataStartPos = BS.length contents - footerSize - size diff --git a/src/DataFrame/IO/Parquet/Time.hs b/src/DataFrame/IO/Parquet/Time.hs --- a/src/DataFrame/IO/Parquet/Time.hs +++ b/src/DataFrame/IO/Parquet/Time.hs @@ -6,7 +6,12 @@ import Data.Time import Data.Word -import DataFrame.IO.Parquet.Binary +import DataFrame.Internal.Binary ( + littleEndianWord32, + littleEndianWord64, + word32ToLittleEndian, + word64ToLittleEndian, + ) int96ToUTCTime :: BS.ByteString -> UTCTime int96ToUTCTime bytes diff --git a/src/DataFrame/IO/Unstable/CSV.hs b/src/DataFrame/IO/Unstable/CSV.hs --- a/src/DataFrame/IO/Unstable/CSV.hs +++ b/src/DataFrame/IO/Unstable/CSV.hs @@ -53,7 +53,7 @@ typeInferenceSampleSize, ) import DataFrame.Internal.DataFrame (DataFrame (..)) -import DataFrame.Operations.Typing (parseFromExamples) +import DataFrame.Operations.Typing (ParseOptions (..), parseFromExamples) readSeparatedDefaultFast :: Word8 -> FilePath -> IO DataFrame readSeparatedDefaultFast separator = @@ -124,12 +124,14 @@ if shouldInferFromSample (typeSpec opts) then typeInferenceSampleSize (typeSpec opts) else 0 - in parseFromExamples - (missingIndicators opts) - n - (safeRead opts) - (dateFormat opts) - col + parseOpts = + ParseOptions + { missingValues = missingIndicators opts + , sampleSize = n + , parseSafe = safeRead opts + , parseDateFormat = dateFormat opts + } + in parseFromExamples parseOpts col generateColumn col = parseTypes $ Vector.fromListN diff --git a/src/DataFrame/Internal/Binary.hs b/src/DataFrame/Internal/Binary.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Binary.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE BangPatterns #-} + +module DataFrame.Internal.Binary where + +import Data.Bits (Bits (unsafeShiftL, (.|.))) +import Data.ByteString (toStrict) +import qualified Data.ByteString as BS +import Data.ByteString.Builder (toLazyByteString, word32LE, word64LE) +import qualified Data.ByteString.Unsafe as BS +import Data.Int (Int32) +import Data.Word (Word32, Word64, Word8) + +littleEndianWord32 :: BS.ByteString -> Word32 +littleEndianWord32 bytes + | len >= 4 = + assembleWord32 + (BS.unsafeIndex bytes 0) + (BS.unsafeIndex bytes 1) + (BS.unsafeIndex bytes 2) + (BS.unsafeIndex bytes 3) + | otherwise = + assembleWord32 + (byteAtOrZero len bytes 0) + (byteAtOrZero len bytes 1) + (byteAtOrZero len bytes 2) + (byteAtOrZero len bytes 3) + where + len = BS.length bytes +{-# INLINE littleEndianWord32 #-} + +littleEndianWord64 :: BS.ByteString -> Word64 +littleEndianWord64 bytes + | len >= 8 = + assembleWord64 + (BS.index bytes 0) + (BS.index bytes 1) + (BS.index bytes 2) + (BS.index bytes 3) + (BS.index bytes 4) + (BS.index bytes 5) + (BS.index bytes 6) + (BS.index bytes 7) + | otherwise = + assembleWord64 + (byteAtOrZero len bytes 0) + (byteAtOrZero len bytes 1) + (byteAtOrZero len bytes 2) + (byteAtOrZero len bytes 3) + (byteAtOrZero len bytes 4) + (byteAtOrZero len bytes 5) + (byteAtOrZero len bytes 6) + (byteAtOrZero len bytes 7) + where + len = BS.length bytes +{-# INLINE littleEndianWord64 #-} + +littleEndianInt32 :: BS.ByteString -> Int32 +littleEndianInt32 = fromIntegral . littleEndianWord32 +{-# INLINE littleEndianInt32 #-} + +word64ToLittleEndian :: Word64 -> BS.ByteString +word64ToLittleEndian = toStrict . toLazyByteString . word64LE +{-# INLINE word64ToLittleEndian #-} + +word32ToLittleEndian :: Word32 -> BS.ByteString +word32ToLittleEndian = toStrict . toLazyByteString . word32LE +{-# INLINE word32ToLittleEndian #-} + +byteAtOrZero :: Int -> BS.ByteString -> Int -> Word8 +byteAtOrZero len bytes i + | i >= 0 && i < len = BS.unsafeIndex bytes i + | otherwise = 0 +{-# INLINE byteAtOrZero #-} + +assembleWord32 :: Word8 -> Word8 -> Word8 -> Word8 -> Word32 +assembleWord32 !b0 !b1 !b2 !b3 = + fromIntegral b0 + .|. (fromIntegral b1 `unsafeShiftL` 8) + .|. (fromIntegral b2 `unsafeShiftL` 16) + .|. (fromIntegral b3 `unsafeShiftL` 24) +{-# INLINE assembleWord32 #-} + +assembleWord64 :: + Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word64 +assembleWord64 !b0 !b1 !b2 !b3 !b4 !b5 !b6 !b7 = + fromIntegral b0 + .|. (fromIntegral b1 `unsafeShiftL` 8) + .|. (fromIntegral b2 `unsafeShiftL` 16) + .|. (fromIntegral b3 `unsafeShiftL` 24) + .|. (fromIntegral b4 `unsafeShiftL` 32) + .|. (fromIntegral b5 `unsafeShiftL` 40) + .|. (fromIntegral b6 `unsafeShiftL` 48) + .|. (fromIntegral b7 `unsafeShiftL` 56) +{-# INLINE assembleWord64 #-} diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -18,7 +18,6 @@ module DataFrame.Internal.Column where -import qualified Data.Set as S import qualified Data.Text as T import qualified Data.Vector as VB import qualified Data.Vector.Algorithms.Merge as VA @@ -50,6 +49,10 @@ UnboxedColumn :: (Columnable a, VU.Unbox a) => VU.Vector a -> Column OptionalColumn :: (Columnable a) => VB.Vector (Maybe a) -> Column +{- | A mutable companion struct to dataframe columns. + +Used mostly as an intermediate structure for I/O. +-} data MutableColumn where MBoxedColumn :: (Columnable a) => VBM.IOVector a -> MutableColumn MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn @@ -57,6 +60,9 @@ {- | A TypedColumn is a wrapper around our type-erased column. It is used to type check expressions on columns. + +Note: there is no guarantee that the Phanton type is the +same as the underlying vector type. -} data TypedColumn a where TColumn :: (Columnable a) => Column -> TypedColumn a @@ -99,15 +105,13 @@ -- | Checks if a column is of a given type values. hasElemType :: forall a. (Columnable a) => Column -> Bool -hasElemType (BoxedColumn (column :: VB.Vector b)) = fromMaybe False $ do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure True -hasElemType (UnboxedColumn (column :: VU.Vector b)) = fromMaybe False $ do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure True -hasElemType (OptionalColumn (column :: VB.Vector b)) = fromMaybe False $ do - Refl <- testEquality (typeRep @a) (typeRep @b) - pure True +hasElemType = \case + BoxedColumn (column :: VB.Vector b) -> check (typeRep @b) + UnboxedColumn (column :: VU.Vector b) -> check (typeRep @b) + OptionalColumn (column :: VB.Vector b) -> check (typeRep @b) + where + check :: forall (b :: Type). TypeRep b -> Bool + check = isJust . testEquality (typeRep @a) -- | An internal/debugging function to get the column type of a column. columnVersionString :: Column -> String @@ -254,6 +258,7 @@ [a] -> Column fromList = toColumnRep @(KindOf a) . VB.fromList +-- An internal helper for type errors throwTypeMismatch :: forall (a :: Type) (b :: Type). (Typeable a, Typeable b) => Either DataFrameException Column @@ -263,7 +268,7 @@ MkTypeErrorContext { userType = Right (typeRep @b) , expectedType = Right (typeRep @a) - , callingFunctionName = Just "mapColumn" + , callingFunctionName = Nothing , errorColumnName = Nothing } @@ -274,7 +279,7 @@ (b -> c) -> Column -> Either DataFrameException Column mapColumn f = \case BoxedColumn (col :: VB.Vector a) -> run col - OptionalColumn (col :: VB.Vector a) -> run col + OptionalColumn (col :: VB.Vector (Maybe a)) -> runOpt col UnboxedColumn (col :: VU.Vector a) -> runUnboxed col where run :: forall a. (Typeable a) => VB.Vector a -> Either DataFrameException Column @@ -282,6 +287,17 @@ Just Refl -> Right (fromVector @c (VB.map f col)) Nothing -> throwTypeMismatch @a @b + -- For OptionalColumn: first try exact match (b = Maybe a, user maps over Maybe + -- values directly), then try lenient fmap fallback (b = a, auto-fmap over inner). + runOpt :: + forall a. + (Columnable a) => VB.Vector (Maybe a) -> Either DataFrameException Column + runOpt col = case testEquality (typeRep @(Maybe a)) (typeRep @b) of + Just Refl -> Right (fromVector @c (VB.map f col)) + Nothing -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Right (OptionalColumn (VB.map (fmap f) col)) + Nothing -> throwTypeMismatch @(Maybe a) @b + runUnboxed :: forall a. (Typeable a, VU.Unbox a) => VU.Vector a -> Either DataFrameException Column @@ -299,7 +315,7 @@ (Int -> b -> c) -> Column -> Either DataFrameException Column imapColumn f = \case BoxedColumn (col :: VB.Vector a) -> run col - OptionalColumn (col :: VB.Vector a) -> run col + OptionalColumn (col :: VB.Vector (Maybe a)) -> runOpt col UnboxedColumn (col :: VU.Vector a) -> runUnboxed col where run :: forall a. (Typeable a) => VB.Vector a -> Either DataFrameException Column @@ -307,6 +323,15 @@ Just Refl -> Right (fromVector @c (VB.imap f col)) Nothing -> throwTypeMismatch @a @b + runOpt :: + forall a. + (Columnable a) => VB.Vector (Maybe a) -> Either DataFrameException Column + runOpt col = case testEquality (typeRep @(Maybe a)) (typeRep @b) of + Just Refl -> Right (fromVector @c (VB.imap f col)) + Nothing -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Right (OptionalColumn (VB.imap (fmap . f) col)) + Nothing -> throwTypeMismatch @(Maybe a) @b + runUnboxed :: forall a. (Typeable a, VU.Unbox a) => VU.Vector a -> Either DataFrameException Column @@ -318,16 +343,16 @@ -- | O(1) Gets the number of elements in the column. columnLength :: Column -> Int -columnLength (BoxedColumn xs) = VG.length xs -columnLength (UnboxedColumn xs) = VG.length xs -columnLength (OptionalColumn xs) = VG.length xs +columnLength (BoxedColumn xs) = VB.length xs +columnLength (UnboxedColumn xs) = VU.length xs +columnLength (OptionalColumn xs) = VB.length xs {-# INLINE columnLength #-} -- | O(n) Gets the number of elements in the column. numElements :: Column -> Int -numElements (BoxedColumn xs) = VG.length xs -numElements (UnboxedColumn xs) = VG.length xs -numElements (OptionalColumn xs) = VG.foldl' (\acc x -> acc + fromEnum (isJust x)) 0 xs +numElements (BoxedColumn xs) = VB.length xs +numElements (UnboxedColumn xs) = VU.length xs +numElements (OptionalColumn xs) = VB.foldl' (\acc x -> acc + fromEnum (isJust x)) 0 xs {-# INLINE numElements #-} -- | O(n) Takes the first n values of a column. @@ -349,13 +374,6 @@ sliceColumn start n (OptionalColumn xs) = OptionalColumn $ VG.slice start n xs {-# INLINE sliceColumn #-} --- | O(n) Selects the elements at a given set of indices. May change the order. -atIndices :: S.Set Int -> Column -> Column -atIndices indexes (BoxedColumn column) = BoxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (OptionalColumn column) = OptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column -atIndices indexes (UnboxedColumn column) = UnboxedColumn $ VU.ifilter (\i _ -> i `S.member` indexes) column -{-# INLINE atIndices #-} - -- | O(n) Selects the elements at a given set of indices. Does not change the order. atIndicesStable :: VU.Vector Int -> Column -> Column atIndicesStable indexes (BoxedColumn column) = @@ -396,16 +414,6 @@ in if idx < 0 then Nothing else v `VB.unsafeIndex` idx {-# INLINE gatherWithSentinel #-} -atIndicesWithNulls :: VB.Vector (Maybe Int) -> Column -> Column -atIndicesWithNulls indices column = - case column of - BoxedColumn col -> - OptionalColumn $ VB.map (fmap (col VB.!)) indices - UnboxedColumn col -> - OptionalColumn $ VB.map (fmap (col VU.!)) indices - OptionalColumn col -> - OptionalColumn $ VB.map (\ix -> ix >>= (col VB.!)) indices - -- | Internal helper to get indices in a boxed vector. getIndices :: VU.Vector Int -> VB.Vector a -> VB.Vector a getIndices indices xs = VB.generate (VU.length indices) (\i -> xs VB.! (indices VU.! i)) @@ -448,19 +456,19 @@ -- | An internal function that returns a vector of how indexes change after a column is sorted. sortedIndexes :: Bool -> Column -> VU.Vector Int sortedIndexes asc = \case - BoxedColumn column -> sortWorker column - UnboxedColumn column -> sortWorker column - OptionalColumn column -> sortWorker column + BoxedColumn column -> sortWorker VG.convert column + UnboxedColumn column -> sortWorker id column + OptionalColumn column -> sortWorker VG.convert column where sortWorker :: (VG.Vector v a, Ord a, VG.Vector v (Int, a), VG.Vector v Int) => - v a -> VU.Vector Int - sortWorker column = runST $ do + (v Int -> VU.Vector Int) -> v a -> VU.Vector Int + sortWorker finalize column = runST $ do withIndexes <- VG.thaw $ VG.indexed column let cmp = if asc then compare else flip compare VA.sortBy (\(_, b) (_, b') -> cmp b b') withIndexes sorted <- VG.unsafeFreeze withIndexes - return $ VG.convert $ VG.map fst sorted + return $ finalize $ VG.map fst sorted {-# INLINE sortedIndexes #-} -- | Fold (right) column with index. @@ -468,262 +476,83 @@ forall a b. (Columnable a, Columnable b) => (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b -ifoldrColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.ifoldr f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "ifoldrColumn" - , errorColumnName = Nothing - } - ) -ifoldrColumn f acc c@(OptionalColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.ifoldr f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "ifoldrColumn" - , errorColumnName = Nothing - } - ) -ifoldrColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.ifoldr f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "ifoldrColumn" - , errorColumnName = Nothing - } - ) - -foldlColumn :: - forall a b. - (Columnable a, Columnable b) => - (b -> a -> b) -> b -> Column -> Either DataFrameException b -foldlColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumn" - , errorColumnName = Nothing - } - ) -foldlColumn f acc c@(OptionalColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumn" - , errorColumnName = Nothing - } - ) -foldlColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumn" - , errorColumnName = Nothing - } - ) - -foldlColumnWith :: - forall a b. - (Columnable a) => - (b -> a -> b) -> b -> Column -> Either DataFrameException b -foldlColumnWith f acc (BoxedColumn (column :: VB.Vector d)) = - case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column +ifoldrColumn f acc = \case + BoxedColumn column -> foldrWorker column + UnboxedColumn column -> foldrWorker column + OptionalColumn column -> foldrWorker column + where + foldrWorker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException b + foldrWorker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> pure $ VG.ifoldr f acc vec Nothing -> Left $ TypeMismatchException ( MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumnWith" + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "ifoldrColumn" , errorColumnName = Nothing } ) -foldlColumnWith f acc (OptionalColumn (column :: VB.Vector d)) = - case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column + +foldlColumn :: + forall a b. + (Columnable a, Columnable b) => + (b -> a -> b) -> b -> Column -> Either DataFrameException b +foldlColumn f acc = \case + BoxedColumn column -> foldlWorker column + UnboxedColumn column -> foldlWorker column + OptionalColumn column -> foldlWorker column + where + foldlWorker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException b + foldlWorker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> pure $ VG.foldl' f acc vec Nothing -> Left $ TypeMismatchException ( MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumnWith" + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "ifoldrColumn" , errorColumnName = Nothing } ) -foldlColumnWith f acc (UnboxedColumn (column :: VU.Vector d)) = - case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl' f acc column + +foldl1Column :: + forall a. + (Columnable a) => + (a -> a -> a) -> Column -> Either DataFrameException a +foldl1Column f = \case + BoxedColumn column -> foldl1Worker column + UnboxedColumn column -> foldl1Worker column + OptionalColumn column -> foldl1Worker column + where + foldl1Worker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException a + foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> pure $ VG.foldl1' f vec Nothing -> Left $ TypeMismatchException ( MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldlColumnWith" + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "foldl1Column" , errorColumnName = Nothing } ) -foldl1Column :: - forall a. - (Columnable a) => - (a -> a -> a) -> Column -> Either DataFrameException a -foldl1Column f c@(BoxedColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl1' f column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldl1Column" - , errorColumnName = Nothing - } - ) -foldl1Column f c@(OptionalColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl1' f column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldl1Column" - , errorColumnName = Nothing - } - ) -foldl1Column f c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> pure $ VG.foldl1' f column - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldl1Column" - , errorColumnName = Nothing - } - ) - -{- | O(n) Fold a column over groups without materialising a sorted copy. -Instead of backpermuting the column (expensive: O(n) allocation + random reads), -this iterates @valueIndices@ sequentially and accesses the column at each index. -Avoids one 10M-element allocation per column per aggregation expression. --} -foldDirectGroups :: - forall b acc. - (Columnable b) => - (acc -> b -> acc) -> - acc -> - Column -> - VU.Vector Int -> -- valueIndices (sorted row order) - VU.Vector Int -> -- offsets (group boundaries) - Either DataFrameException (VB.Vector acc) -foldDirectGroups f seed col valueIndices offsets - | VU.length offsets <= 1 = Right VB.empty - | otherwise = - let !nGroups = VU.length offsets - 1 - in case col of - UnboxedColumn (vec :: VU.Vector d) -> - case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - in go s e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VU.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldDirectGroups" - , errorColumnName = Nothing - } - BoxedColumn (vec :: VB.Vector d) -> - case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - in go s e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldDirectGroups" - , errorColumnName = Nothing - } - OptionalColumn (vec :: VB.Vector (Maybe d)) -> - case testEquality (typeRep @b) (typeRep @(Maybe d)) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - in go s e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @(Maybe d)) - , callingFunctionName = Just "foldDirectGroups" - , errorColumnName = Nothing - } -{-# INLINEABLE foldDirectGroups #-} - {- | O(n) Seedless fold over groups using the first element of each group as seed. Like 'foldDirectGroups' but for the case where no initial accumulator is available. -} @@ -734,87 +563,43 @@ Column -> VU.Vector Int -> VU.Vector Int -> - Either DataFrameException (VB.Vector a) + Either DataFrameException Column foldl1DirectGroups f col valueIndices offsets - | VU.length offsets <= 1 = Right VB.empty - | otherwise = - let !nGroups = VU.length offsets - 1 - in case col of - UnboxedColumn (vec :: VU.Vector d) -> - case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - !seed = VU.unsafeIndex vec (VU.unsafeIndex valueIndices s) - in go (s + 1) e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VU.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldl1DirectGroups" - , errorColumnName = Nothing - } - BoxedColumn (vec :: VB.Vector d) -> - case testEquality (typeRep @a) (typeRep @d) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - !seed = VB.unsafeIndex vec (VU.unsafeIndex valueIndices s) - in go (s + 1) e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldl1DirectGroups" - , errorColumnName = Nothing - } - OptionalColumn (vec :: VB.Vector (Maybe d)) -> - case testEquality (typeRep @a) (typeRep @(Maybe d)) of - Just Refl -> - Right $ - VB.generate nGroups foldGroup - where - foldGroup k = - let !s = VU.unsafeIndex offsets k - !e = VU.unsafeIndex offsets (k + 1) - !seed = VB.unsafeIndex vec (VU.unsafeIndex valueIndices s) - in go (s + 1) e seed - go !i !e !acc - | i >= e = acc - | otherwise = - go (i + 1) e $! - f acc (VB.unsafeIndex vec (VU.unsafeIndex valueIndices i)) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @(Maybe d)) - , callingFunctionName = Just "foldl1DirectGroups" - , errorColumnName = Nothing - } + | VU.length offsets <= 1 = pure $ fromVector @a VB.empty + | otherwise = case col of + UnboxedColumn (vec :: VU.Vector d) -> UnboxedColumn <$> foldl1Worker vec + BoxedColumn (vec :: VB.Vector d) -> BoxedColumn <$> foldl1Worker vec + OptionalColumn (vec :: VB.Vector (Maybe d)) -> OptionalColumn <$> foldl1Worker vec + where + foldl1Worker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException (v c) + foldl1Worker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> + Right $ + VG.generate (VU.length offsets - 1) foldGroup + where + foldGroup k = + let !s = VU.unsafeIndex offsets k + !e = VU.unsafeIndex offsets (k + 1) + !seed = VG.unsafeIndex vec (VU.unsafeIndex valueIndices s) + in go (s + 1) e seed + go !i !e !acc + | i >= e = acc + | otherwise = + go (i + 1) e $! + f acc (VG.unsafeIndex vec (VU.unsafeIndex valueIndices i)) + Nothing -> + Left $ + TypeMismatchException + MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "foldl1DirectGroups" + , errorColumnName = Nothing + } {-# INLINEABLE foldl1DirectGroups #-} {- | O(n) fold over groups by scanning the column LINEARLY. @@ -836,70 +621,36 @@ foldLinearGroups f seed col rowToGroup nGroups | nGroups == 0 = Right (fromVector @acc VB.empty) | otherwise = case col of - UnboxedColumn (vec :: VU.Vector d) -> - case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - Right $ - unsafePerformIO $ - runWith - ( \readAt writeAt -> - VU.iforM_ vec $ \row x -> do - let !k = VU.unsafeIndex rowToGroup row - cur <- readAt k - writeAt k $! f cur x - ) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldLinearGroups" - , errorColumnName = Nothing - } - BoxedColumn (vec :: VB.Vector d) -> - case testEquality (typeRep @b) (typeRep @d) of - Just Refl -> - Right $ - unsafePerformIO $ - runWith - ( \readAt writeAt -> - VB.iforM_ vec $ \row x -> do - let !k = VU.unsafeIndex rowToGroup row - cur <- readAt k - writeAt k $! f cur x - ) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @d) - , callingFunctionName = Just "foldLinearGroups" - , errorColumnName = Nothing - } - OptionalColumn (vec :: VB.Vector (Maybe d)) -> - case testEquality (typeRep @b) (typeRep @(Maybe d)) of - Just Refl -> - Right $ - unsafePerformIO $ - runWith - ( \readAt writeAt -> - VB.iforM_ vec $ \row x -> do - let !k = VU.unsafeIndex rowToGroup row - cur <- readAt k - writeAt k $! f cur x - ) - Nothing -> - Left $ - TypeMismatchException - MkTypeErrorContext - { userType = Right (typeRep @b) - , expectedType = Right (typeRep @(Maybe d)) - , callingFunctionName = Just "foldLinearGroups" - , errorColumnName = Nothing - } + UnboxedColumn (vec :: VU.Vector d) -> foldLinearWorker vec + BoxedColumn (vec :: VB.Vector d) -> foldLinearWorker vec + OptionalColumn (vec :: VB.Vector (Maybe d)) -> foldLinearWorker vec where + foldLinearWorker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException Column + foldLinearWorker vec = case testEquality (typeRep @b) (typeRep @c) of + Just Refl -> + Right $ + unsafePerformIO $ + runWith + ( \readAt writeAt -> + VG.iforM_ vec $ \row x -> do + let !k = VG.unsafeIndex rowToGroup row + cur <- readAt k + writeAt k $! f cur x + ) + Nothing -> + Left $ + TypeMismatchException + MkTypeErrorContext + { userType = Right (typeRep @b) + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "foldLinearGroups" + , errorColumnName = Nothing + } + -- \| Allocate accumulators, run the traversal, return a frozen Column. -- When @acc@ is unboxable, uses an unboxed mutable vector (no pointer -- indirection per read/write) and returns UnboxedColumn directly — @@ -918,51 +669,31 @@ {-# INLINEABLE foldLinearGroups #-} headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a -headColumn (BoxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> - if VG.null col - then Left (EmptyDataSetException "headColumn") - else pure (VG.head col) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "headColumn" - , errorColumnName = Nothing - } - ) -headColumn (UnboxedColumn (col :: VU.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> - if VG.null col - then Left (EmptyDataSetException "headColumn") - else pure (VG.head col) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "headColumn" - , errorColumnName = Nothing - } - ) -headColumn (OptionalColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> - if VG.null col - then Left (EmptyDataSetException "headColumn") - else pure (VG.head col) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "headColumn" - , errorColumnName = Nothing - } - ) +headColumn = \case + BoxedColumn col -> headWorker col + UnboxedColumn col -> headWorker col + OptionalColumn col -> headWorker col + where + headWorker :: + forall c v. + (Typeable c, VG.Vector v c) => + v c -> + Either DataFrameException a + headWorker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> + if VG.null vec + then Left (EmptyDataSetException "headColumn") + else pure (VG.head vec) + Nothing -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Right (typeRep @c) + , callingFunctionName = Just "headColumn" + , errorColumnName = Nothing + } + ) -- | An internal, column version of zip. zipColumns :: Column -> Column -> Column @@ -1033,7 +764,7 @@ zipWithColumns f (UnboxedColumn (column :: VU.Vector d)) (UnboxedColumn (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of Just Refl -> case testEquality (typeRep @b) (typeRep @e) of Just Refl -> pure $ case sUnbox @c of - STrue -> fromUnboxedVector (VU.zipWith f column other) + STrue -> UnboxedColumn (VU.zipWith f column other) SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other) Nothing -> Left $ @@ -1055,6 +786,8 @@ , errorColumnName = Nothing } ) +-- TODO: mchavinda - reuse pattern from interpret where we augment the +-- error at the end. zipWithColumns f left right = case toVector @a left of Left (TypeMismatchException context) -> Left $ @@ -1172,52 +905,31 @@ Returns Nothing if the columns are of different types. -} concatColumns :: Column -> Column -> Either DataFrameException Column -concatColumns (OptionalColumn left) (OptionalColumn right) = case testEquality (typeOf left) (typeOf right) of - Just Refl -> pure (OptionalColumn $ left <> right) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeOf right) - , expectedType = Right (typeOf left) - , callingFunctionName = Just "concatColumns" - , errorColumnName = Nothing - } - ) -concatColumns (BoxedColumn left) (BoxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Just Refl -> pure (BoxedColumn $ left <> right) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeOf right) - , expectedType = Right (typeOf left) - , callingFunctionName = Just "concatColumns" - , errorColumnName = Nothing - } - ) -concatColumns (UnboxedColumn left) (UnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of - Just Refl -> pure (UnboxedColumn $ left <> right) - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeOf right) - , expectedType = Right (typeOf left) - , callingFunctionName = Just "concatColumns" - , errorColumnName = Nothing - } - ) -concatColumns left right = - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeOf right) - , expectedType = Right (typeOf left) - , callingFunctionName = Just "concatColumns" - , errorColumnName = Nothing - } - ) +concatColumns left right = case (left, right) of + (OptionalColumn l, OptionalColumn r) -> case testEquality (typeOf l) (typeOf r) of + Just Refl -> pure (OptionalColumn $ l <> r) + Nothing -> Left (mismatchErr (typeOf r) (typeOf l)) + (BoxedColumn l, BoxedColumn r) -> case testEquality (typeOf l) (typeOf r) of + Just Refl -> pure (BoxedColumn $ l <> r) + Nothing -> Left (mismatchErr (typeOf r) (typeOf l)) + (UnboxedColumn l, UnboxedColumn r) -> case testEquality (typeOf l) (typeOf r) of + Just Refl -> pure (UnboxedColumn $ l <> r) + Nothing -> Left (mismatchErr (typeOf r) (typeOf l)) + _ -> Left (mismatchErr (typeOf right) (typeOf left)) + where + mismatchErr :: + forall (x :: Type) (y :: Type). TypeRep x -> TypeRep y -> DataFrameException + mismatchErr ta tb = + withTypeable ta $ + withTypeable tb $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right ta + , expectedType = Right tb + , callingFunctionName = Just "concatColumns" + , errorColumnName = Nothing + } + ) {- | Concatenates two columns. @@ -1375,41 +1087,24 @@ toVector :: forall a v. (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a) -toVector column@(OptionalColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVector" - , errorColumnName = Nothing - } - ) -toVector (BoxedColumn (col :: VB.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col - Nothing -> - Left $ - TypeMismatchException - ( MkTypeErrorContext - { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) - , callingFunctionName = Just "toVector" - , errorColumnName = Nothing - } - ) -toVector (UnboxedColumn (col :: VU.Vector b)) = - case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Right $ VG.convert col +toVector = \case + OptionalColumn col -> toVectorWorker col + BoxedColumn col -> toVectorWorker col + UnboxedColumn col -> toVectorWorker col + where + toVectorWorker :: + forall c w. + (Typeable c, VG.Vector w c) => + w c -> + Either DataFrameException (v a) + toVectorWorker vec = case testEquality (typeRep @a) (typeRep @c) of + Just Refl -> Right $ VG.convert vec Nothing -> Left $ TypeMismatchException ( MkTypeErrorContext { userType = Right (typeRep @a) - , expectedType = Right (typeRep @b) + , expectedType = Right (typeRep @c) , callingFunctionName = Just "toVector" , errorColumnName = Nothing } 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 @@ -56,6 +56,18 @@ data Expr a where Col :: (Columnable a) => T.Text -> Expr a + CastWith :: + (Columnable a, Columnable b) => + T.Text -> + T.Text -> + (Either String a -> b) -> + Expr b + CastExprWith :: + (Columnable a, Columnable b, Columnable src) => + T.Text -> + (Either String a -> b) -> + Expr src -> + Expr b Lit :: (Columnable a) => a -> Expr a Unary :: (Columnable a, Columnable b) => UnaryOp b a -> Expr b -> Expr a @@ -228,6 +240,8 @@ instance (Show a) => Show (Expr a) where show :: forall a. (Show a) => Expr a -> String show (Col name) = "(col @" ++ show (typeRep @a) ++ " " ++ show name ++ ")" + show (CastWith name tag _) = "(castWith " ++ show tag ++ " " ++ show name ++ ")" + show (CastExprWith tag _ inner) = "(castExprWith " ++ show tag ++ " " ++ show inner ++ ")" show (Lit value) = "(lit (" ++ show value ++ "))" show (If cond l r) = "(ifThenElse " ++ show cond ++ " " ++ show l ++ " " ++ show r ++ ")" show (Unary op value) = "(" ++ T.unpack (unaryName op) ++ " " ++ show value ++ ")" @@ -239,6 +253,8 @@ normalize :: (Eq a, Ord a, Show a, Typeable a) => Expr a -> Expr a normalize expr = case expr of Col name -> Col name + CastWith n t f -> CastWith n t f + CastExprWith t f e -> CastExprWith t f (normalize e) Lit val -> Lit val If cond th el -> If (normalize cond) (normalize th) (normalize el) Unary op e -> Unary op (normalize e) @@ -261,6 +277,8 @@ where exprKey :: Expr a -> String exprKey (Col name) = "0:" ++ T.unpack name + exprKey (CastWith name tag _) = "0CW:" ++ T.unpack name ++ ":" ++ T.unpack tag + exprKey (CastExprWith tag _ _) = "0CE:" ++ T.unpack tag exprKey (Lit val) = "1:" ++ show val exprKey (If c t e) = "2:" ++ exprKey c ++ exprKey t ++ exprKey e exprKey (Unary op e) = "3:" ++ T.unpack (unaryName op) ++ exprKey e @@ -278,6 +296,8 @@ 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 @@ -295,6 +315,8 @@ compare :: Expr a -> Expr a -> Ordering compare e1 e2 = case (e1, e2) of (Col n1, Col n2) -> compare n1 n2 + (CastWith n1 t1 _, CastWith n2 t2 _) -> compare n1 n2 <> compare t1 t2 + (CastExprWith t1 _ _, CastExprWith t2 _ _) -> compare t1 t2 (Lit v1, Lit v2) -> compare v1 v2 (If c1 t1 e1', If c2 t2 e2') -> compare c1 c2 <> exprComp t1 t2 <> exprComp e1' e2' @@ -307,6 +329,10 @@ -- Different constructors - compare by priority (Col _, _) -> LT (_, Col _) -> GT + (CastWith{}, _) -> LT + (_, CastWith{}) -> GT + (CastExprWith{}, _) -> LT + (_, CastExprWith{}) -> GT (Lit _, _) -> LT (_, Lit _) -> GT (Unary{}, _) -> LT @@ -334,6 +360,8 @@ where replace' = case expr of (Col _) -> expr + (CastWith{}) -> expr + (CastExprWith t f e) -> CastExprWith t f (replaceExpr new old e) (Lit _) -> expr (If cond l r) -> If (replaceExpr new old cond) (replaceExpr new old l) (replaceExpr new old r) @@ -343,6 +371,8 @@ eSize :: Expr a -> Int eSize (Col _) = 1 +eSize (CastWith{}) = 1 +eSize (CastExprWith _ _ e) = 1 + eSize e eSize (Lit _) = 1 eSize (If c l r) = 1 + eSize c + eSize l + eSize r eSize (Unary _ e) = 1 + eSize e @@ -351,6 +381,8 @@ getColumns :: Expr a -> [T.Text] getColumns (Col cName) = [cName] +getColumns (CastWith name _ _) = [name] +getColumns (CastExprWith _ _ e) = getColumns e getColumns expr@(Lit _) = [] getColumns (If cond l r) = getColumns cond <> getColumns l <> getColumns r getColumns (Unary op value) = getColumns value @@ -366,6 +398,8 @@ go :: Int -> Int -> Expr a -> String go depth prec expr = case expr of Col name -> T.unpack name + CastWith name _ _ -> T.unpack name + CastExprWith tag _ inner -> T.unpack tag ++ "(" ++ go depth 0 inner ++ ")" Lit value -> show value If cond t e -> let inner = diff --git a/src/DataFrame/Internal/Interpreter.hs b/src/DataFrame/Internal/Interpreter.hs --- a/src/DataFrame/Internal/Interpreter.hs +++ b/src/DataFrame/Internal/Interpreter.hs @@ -36,7 +36,10 @@ import DataFrame.Internal.DataFrame import DataFrame.Internal.Expression import DataFrame.Internal.Types -import Type.Reflection (Typeable, typeRep) +import Type.Reflection ( + Typeable, + typeRep, + ) ------------------------------------------------------------------------------- -- Value: the unified result type @@ -231,6 +234,237 @@ numGroups gdf = VU.length (offsets gdf) - 1 ------------------------------------------------------------------------------- +-- promoteColumnWith: unified numeric / text coercion for CastWith +------------------------------------------------------------------------------- + +{- | Apply a result-handler @onResult@ to each element of a column after +coercing it to type @a@. Covers three modes in one: + +* @onResult = either (const Nothing) Just@ → like @cast@ (returns @Maybe a@) +* @onResult = either (const def) id@ → like @castWithDefault@ (returns @a@) +* @onResult = either (Left . T.pack) Right@ → like @castEither@ (returns @Either T.Text a@) + +Numeric coercion handles Double, Float, and Int targets. Text columns +(String / T.Text) are parsed via 'reads'. Any other mismatch returns +'Left TypeMismatchException'. +-} +promoteColumnWith :: + forall a b. + (Columnable a, Columnable b) => + (Either String a -> b) -> Column -> Either DataFrameException Column +promoteColumnWith onResult col + | hasElemType @b col = Right col + | hasElemType @a col = mapColumn @a (onResult . Right) col + | Just result <- tryMaybeWrap @a @b onResult col = result + | otherwise = + case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> promoteToDoubleWith onResult col + Nothing -> + case testEquality (typeRep @a) (typeRep @Float) of + Just Refl -> promoteToFloatWith onResult col + Nothing -> + case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> promoteToIntWith onResult col + Nothing -> tryParseWith @a onResult col + +promoteToDoubleWith :: + forall b. + (Columnable b) => + (Either String Double -> b) -> Column -> Either DataFrameException Column +promoteToDoubleWith onResult col = case col of + UnboxedColumn (v :: VU.Vector c) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (realToFrac :: c -> Double)) (VG.convert v)) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (fromIntegral :: c -> Double)) (VG.convert v)) + SFalse -> castMismatch @c @b + OptionalColumn (v :: V.Vector (Maybe c)) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + ( V.map + (maybe (onResult (Left "null")) (onResult . Right . (realToFrac :: c -> Double))) + v + ) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + ( V.map + ( maybe + (onResult (Left "null")) + (onResult . Right . (fromIntegral :: c -> Double)) + ) + v + ) + SFalse -> tryParseWith @Double onResult col + BoxedColumn _ -> tryParseWith @Double onResult col + +promoteToFloatWith :: + forall b. + (Columnable b) => + (Either String Float -> b) -> Column -> Either DataFrameException Column +promoteToFloatWith onResult col = case col of + UnboxedColumn (v :: VU.Vector c) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (realToFrac :: c -> Float)) (VG.convert v)) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (fromIntegral :: c -> Float)) (VG.convert v)) + SFalse -> castMismatch @c @b + OptionalColumn (v :: V.Vector (Maybe c)) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + ( V.map + (maybe (onResult (Left "null")) (onResult . Right . (realToFrac :: c -> Float))) + v + ) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + ( V.map + (maybe (onResult (Left "null")) (onResult . Right . (fromIntegral :: c -> Float))) + v + ) + SFalse -> tryParseWith @Float onResult col + BoxedColumn _ -> tryParseWith @Float onResult col + +promoteToIntWith :: + forall b. + (Columnable b) => + (Either String Int -> b) -> Column -> Either DataFrameException Column +promoteToIntWith onResult col = case col of + UnboxedColumn (v :: VU.Vector c) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (round . (realToFrac :: c -> Double))) (VG.convert v)) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + (V.map (onResult . Right . (fromIntegral :: c -> Int)) (VG.convert v)) + SFalse -> castMismatch @c @b + OptionalColumn (v :: V.Vector (Maybe c)) -> + case sFloating @c of + STrue -> + Right $ + fromVector @b + ( V.map + ( maybe + (onResult (Left "null")) + (onResult . Right . (round . (realToFrac :: c -> Double))) + ) + v + ) + SFalse -> case sIntegral @c of + STrue -> + Right $ + fromVector @b + ( V.map + (maybe (onResult (Left "null")) (onResult . Right . (fromIntegral :: c -> Int))) + v + ) + SFalse -> tryParseWith @Int onResult col + BoxedColumn _ -> tryParseWith @Int onResult col + +-- | Single parse primitive: apply @onResult@ to the result of 'reads'. +parseWith :: (Read a) => (Either String a -> b) -> String -> b +parseWith f s = case reads s of + [(x, "")] -> f (Right x) + _ -> f (Left s) + +tryParseWith :: + forall a b. + (Columnable a, Columnable b) => + (Either String a -> b) -> Column -> Either DataFrameException Column +tryParseWith onResult col = case col of + BoxedColumn (v :: V.Vector c) -> + case testEquality (typeRep @c) (typeRep @String) of + Just Refl -> Right $ fromVector @b $ V.map (parseWith onResult) v + Nothing -> + case testEquality (typeRep @c) (typeRep @T.Text) of + Just Refl -> Right $ fromVector @b $ V.map (parseWith onResult . T.unpack) v + Nothing -> castMismatch @c @b + OptionalColumn (v :: V.Vector (Maybe c)) -> + case testEquality (typeRep @c) (typeRep @String) of + Just Refl -> + Right $ + fromVector @b $ + V.map (maybe (onResult (Left "null")) (parseWith onResult)) v + Nothing -> + case testEquality (typeRep @c) (typeRep @T.Text) of + Just Refl -> + Right $ + fromVector @b $ + V.map (maybe (onResult (Left "null")) (parseWith onResult . T.unpack)) v + Nothing -> castMismatch @c @b + UnboxedColumn (_ :: VU.Vector c) -> castMismatch @c @b + +{- | When the output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the +column stores plain @c@ values, wrap each element in 'Just'. +The @Maybe (Maybe c)@ case applies join semantics: instead of producing +a double-wrapped column, a @Maybe c@ column is returned, so +@castExpr \@(Maybe Double)@ on a @Double@ column yields @Maybe Double@ +rather than @Maybe (Maybe Double)@. +Returns 'Nothing' when neither condition holds. +-} +tryMaybeWrap :: + forall a b. + (Columnable a, Columnable b) => + (Either String a -> b) -> Column -> Maybe (Either DataFrameException Column) +tryMaybeWrap _onResult col = case col of + UnboxedColumn (v :: VU.Vector c) -> + let wrapped = V.map Just (VG.convert v) :: V.Vector (Maybe c) + in case testEquality (typeRep @b) (typeRep @(Maybe c)) of + Just Refl -> Just $ Right $ fromVector @b wrapped + Nothing -> + -- join: b = Maybe (Maybe c) → produce Maybe c column + case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of + Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped + Nothing -> Nothing + BoxedColumn (v :: V.Vector c) -> + let wrapped = V.map Just v :: V.Vector (Maybe c) + in case testEquality (typeRep @b) (typeRep @(Maybe c)) of + Just Refl -> Just $ Right $ fromVector @b wrapped + Nothing -> + case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of + Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped + Nothing -> Nothing + -- OptionalColumn and NullableColumn are already handled by the hasElemType guards above. + _ -> Nothing + +castMismatch :: + forall src tgt. + (Typeable src, Typeable tgt) => + Either DataFrameException Column +castMismatch = + Left $ + TypeMismatchException + MkTypeErrorContext + { userType = Right (typeRep @tgt) + , expectedType = Right (typeRep @src) + , callingFunctionName = Just "cast" + , errorColumnName = Nothing + } + +------------------------------------------------------------------------------- -- eval: the unified interpreter ------------------------------------------------------------------------------- @@ -248,10 +482,51 @@ eval (FlatCtx df) (Col name) = case getColumn name df of Nothing -> + Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) + Just c + | hasElemType @a c -> Right (Flat c) + | otherwise -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Left (columnTypeString c) + , errorColumnName = Just (T.unpack name) + , callingFunctionName = Just "col" + } :: + TypeErrorContext a () + ) +eval (GroupCtx gdf) (Col name) = + case getColumn name (fullDataframe gdf) of + Nothing -> Left $ + ColumnNotFoundException + name + "" + (M.keys $ columnIndices $ fullDataframe gdf) + Just c + | hasElemType @a c -> + Right (Group (sliceGroups c (offsets gdf) (valueIndices gdf))) + | otherwise -> + Left $ + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @a) + , expectedType = Left (columnTypeString c) + , errorColumnName = Just (T.unpack name) + , callingFunctionName = Just "col" + } :: + TypeErrorContext a () + ) +-- CastWith --------------------------------------------------------------- + +eval (FlatCtx df) (CastWith name _tag onResult) = + case getColumn name df of + Nothing -> + Left $ ColumnNotFoundException name "" (M.keys $ columnIndices df) - Just c -> Right (Flat c) -eval (GroupCtx gdf) (Col name) = + Just c -> Flat <$> promoteColumnWith onResult c +eval (GroupCtx gdf) (CastWith name _tag onResult) = case getColumn name (fullDataframe gdf) of Nothing -> Left $ @@ -259,11 +534,20 @@ name "" (M.keys $ columnIndices $ fullDataframe gdf) - Just c -> - Right - ( Group - (sliceGroups c (offsets gdf) (valueIndices gdf)) - ) + Just c -> do + promoted <- promoteColumnWith onResult c + Right $ Group (sliceGroups promoted (offsets gdf) (valueIndices gdf)) +-- CastExprWith ----------------------------------------------------------- + +eval ctx (CastExprWith _tag onResult (inner :: Expr src)) = do + v <- eval @src ctx inner + case v of + Scalar s -> + Flat <$> promoteColumnWith onResult (fromList @src [s]) + Flat col -> + Flat <$> promoteColumnWith onResult col + Group gs -> + Group <$> V.mapM (promoteColumnWith onResult) gs -- Unary ------------------------------------------------------------------ eval ctx expr@(Unary (op :: UnaryOp b a) inner) = addContext expr $ do @@ -320,8 +604,7 @@ "" (M.keys $ columnIndices $ fullDataframe gdf) Just col -> - Flat . fromVector - <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf) + Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf) -- Fast path: MergeAgg on a bare Col in GroupCtx. eval @@ -391,10 +674,10 @@ InternalException "Cannot apply a merge aggregation to a scalar" Flat col -> - Scalar . finalize <$> foldlColumnWith @b step seed col + Scalar . finalize <$> foldlColumn @b step seed col Group gs -> Flat . fromVector - <$> V.mapM (fmap finalize . foldlColumnWith @b step seed) gs + <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs -- Aggregation: FoldAgg without seed (fold1) ------------------------------ diff --git a/src/DataFrame/Internal/Nullable.hs b/src/DataFrame/Internal/Nullable.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Internal/Nullable.hs @@ -0,0 +1,477 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} + +{- | Nullable-aware binary operations for expressions. + +This module provides two type classes, 'NullableArithOp' and 'NullableCmpOp', +which enable operators like '.+', '.-', '.*', './', '.==' etc. to work +transparently across combinations of nullable (@Maybe a@) and non-nullable +(@a@) column types. + +The partial functional dependencies uniquely determine the result type from +the operand types, so GHC infers it without annotations. + +The four combinations covered for each class: + +* @(a, a)@ — non-nullable × non-nullable +* @(Maybe a, a)@ — nullable × non-nullable +* @(a, Maybe a)@ — non-nullable × nullable +* @(Maybe a, Maybe a)@ — both nullable + +== Usage + +@ +-- Mixing nullable and non-nullable columns: +F.col \@Int \"x\" '.+' F.col \@(Maybe Int) \"y\" -- :: Expr (Maybe Int) + +-- Both non-nullable (existing behaviour preserved): +F.col \@Int \"x\" '.+' F.col \@Int \"y\" -- :: Expr Int + +-- Comparison with three-valued logic: +F.col \@(Maybe Int) \"x\" '.==' F.col \@Int \"y\" -- :: Expr (Maybe Bool) +@ +-} +module DataFrame.Internal.Nullable ( + -- * Type family + BaseType, + + -- * Arithmetic class + NullableArithOp (..), + + -- * Comparison class + NullableCmpOp (..), + + -- * Generalized nullable lift classes + NullLift1Op (..), + NullLift2Op (..), + + -- * Result-type type families (drive inference in nullLift / nullLift2) + NullLift1Result, + NullLift2Result, + + -- * Result-type type family for comparison operators + NullCmpResult, + + -- * Numeric widening + NumericWidenOp (..), + widenArithOp, + WidenResult, + + -- * Division widening (integral × integral → Double) + DivWidenOp (..), + divArithOp, + WidenResultDiv, +) where + +import Data.Int (Int32, Int64) +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Types (Promote, PromoteDiv) + +{- | Strip one layer of 'Maybe'. + +@ +BaseType (Maybe a) = a +BaseType a = a -- for any non-Maybe type +@ +-} +type family BaseType a where + BaseType (Maybe a) = a + BaseType a = a + +{- | Class for arithmetic binary operations that work transparently over +nullable and non-nullable column types. + +The functional dependency @a b -> c@ ensures GHC can infer the result type @c@ +from the operand types. The 'OVERLAPPABLE' pragma on the non-nullable instance +ensures the more specific @(Maybe a, Maybe a)@ instance wins when both operands +are nullable. +-} +class + ( Columnable a + , Columnable b + , Columnable c + ) => + NullableArithOp a b c + | a b -> c + where + {- | Lift an arithmetic function over the inner (non-Maybe) values. + 'Nothing' short-circuits: any 'Nothing' operand produces 'Nothing'. + -} + nullArithOp :: + (BaseType a -> BaseType a -> BaseType a) -> + a -> + b -> + c + +{- | Compute the result type of a nullable comparison. + +@ +NullCmpResult (Maybe a) b = Maybe Bool +NullCmpResult a (Maybe b) = Maybe Bool -- when a is apart from Maybe +NullCmpResult a b = Bool +@ + +Used by the comparison operators ('.==', '.<', etc.) so GHC infers the +return type without an explicit annotation. +-} +type family NullCmpResult a b where + NullCmpResult (Maybe a) b = Maybe Bool + NullCmpResult a (Maybe b) = Maybe Bool + NullCmpResult a b = Bool + +{- | Class for comparison binary operations that work transparently over +nullable and non-nullable column types. + +No functional dependency on @e@: the 'OVERLAPPING'\/'OVERLAPPABLE' pragmas on +instances disambiguate at call sites without a FundDep (which would conflict +when both operands are @Maybe@). GHC selects the unique most-specific instance +from the concrete operand types. +-} +class + ( Columnable a + , Columnable b + , Columnable e + ) => + NullableCmpOp a b e + where + {- | Lift a comparison function over the inner values (three-valued logic). + Returns 'Nothing' when either operand is 'Nothing'. + -} + nullCmpOp :: + (BaseType a -> BaseType a -> Bool) -> + a -> + b -> + e + +{- | Non-nullable × Non-nullable: apply directly, no wrapping. +Arithmetic result is @a@; comparison result is @Bool@. +-} +instance + {-# OVERLAPPABLE #-} + (Columnable a, a ~ BaseType a) => + NullableArithOp a a a + where + nullArithOp f = f + +instance + {-# OVERLAPPABLE #-} + (Columnable a, Columnable Bool, a ~ BaseType a) => + NullableCmpOp a a Bool + where + nullCmpOp f = f + +-- | Nullable × Non-nullable: 'Nothing' short-circuits. +instance + (Columnable a, Columnable (Maybe a)) => + NullableArithOp (Maybe a) a (Maybe a) + where + nullArithOp f Nothing _ = Nothing + nullArithOp f (Just x) y = Just (f x y) + +instance + (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) => + NullableCmpOp (Maybe a) a (Maybe Bool) + where + nullCmpOp f Nothing _ = Nothing + nullCmpOp f (Just x) y = Just (f x y) + +-- | Non-nullable × Nullable: 'Nothing' short-circuits. +instance + ( Columnable a + , Columnable (Maybe a) + , a ~ BaseType a + ) => + NullableArithOp a (Maybe a) (Maybe a) + where + nullArithOp f _ Nothing = Nothing + nullArithOp f x (Just y) = Just (f x y) + +instance + ( Columnable a + , Columnable (Maybe a) + , Columnable (Maybe Bool) + , a ~ BaseType a + ) => + NullableCmpOp a (Maybe a) (Maybe Bool) + where + nullCmpOp f _ Nothing = Nothing + nullCmpOp f x (Just y) = Just (f x y) + +-- | Nullable × Nullable: either 'Nothing' short-circuits. +instance + {-# OVERLAPPING #-} + (Columnable a, Columnable (Maybe a)) => + NullableArithOp (Maybe a) (Maybe a) (Maybe a) + where + nullArithOp f Nothing _ = Nothing + nullArithOp f _ Nothing = Nothing + nullArithOp f (Just x) (Just y) = Just (f x y) + +instance + {-# OVERLAPPING #-} + (Columnable a, Columnable (Maybe a), Columnable (Maybe Bool)) => + NullableCmpOp (Maybe a) (Maybe a) (Maybe Bool) + where + nullCmpOp f Nothing _ = Nothing + nullCmpOp f _ Nothing = Nothing + nullCmpOp f (Just x) (Just y) = Just (f x y) + +-- --------------------------------------------------------------------------- +-- Generalized nullable lift (unary) +-- --------------------------------------------------------------------------- + +{- | Lift a unary function over a column expression, propagating 'Nothing'. + +When @a@ is non-nullable the function is applied directly; when @a = Maybe x@ +the function is applied under the 'Just' and 'Nothing' short-circuits. + +Use via 'DataFrame.Functions.nullLift'. +-} + +{- | Compute the result type of a nullable unary lift. + +@ +NullLift1Result (Maybe a) r = Maybe r +NullLift1Result a r = r -- for any non-Maybe a +@ + +Used by 'DataFrame.Functions.nullLift' so GHC can infer the return type +without an explicit annotation. +-} +type family NullLift1Result a r where + NullLift1Result (Maybe a) r = Maybe r + NullLift1Result a r = r + +class + ( Columnable a + , Columnable r + , Columnable c + ) => + NullLift1Op a r c + where + applyNull1 :: (BaseType a -> r) -> a -> c + +-- | Non-nullable: apply directly. +instance + {-# OVERLAPPABLE #-} + (Columnable a, Columnable r, a ~ BaseType a) => + NullLift1Op a r r + where + applyNull1 f = f + +-- | Nullable: propagate 'Nothing'. +instance + {-# OVERLAPPING #-} + (Columnable a, Columnable r, Columnable (Maybe r)) => + NullLift1Op (Maybe a) r (Maybe r) + where + applyNull1 _ Nothing = Nothing + applyNull1 f (Just x) = Just (f x) + +-- --------------------------------------------------------------------------- +-- Generalized nullable lift (binary) +-- --------------------------------------------------------------------------- + +{- | Lift a binary function over two column expressions, propagating 'Nothing'. + +The four combinations: + +* @(a, b)@ — both non-nullable: result is @r@ +* @(Maybe a, b)@ — left nullable: result is @Maybe r@ +* @(a, Maybe b)@ — right nullable: result is @Maybe r@ +* @(Maybe a, Maybe b)@ — both nullable: result is @Maybe r@ + +Use via 'DataFrame.Functions.nullLift2'. +-} + +{- | Compute the result type of a nullable binary lift. + +@ +NullLift2Result (Maybe a) b r = Maybe r +NullLift2Result a (Maybe b) r = Maybe r -- when a is apart from Maybe +NullLift2Result a b r = r +@ + +Used by 'DataFrame.Functions.nullLift2' so GHC can infer the return type. +-} +type family NullLift2Result a b r where + NullLift2Result (Maybe a) b r = Maybe r + NullLift2Result a (Maybe b) r = Maybe r + NullLift2Result a b r = r + +class + ( Columnable a + , Columnable b + , Columnable r + , Columnable c + ) => + NullLift2Op a b r c + where + applyNull2 :: (BaseType a -> BaseType b -> r) -> a -> b -> c + +-- | Both non-nullable: apply directly. +instance + {-# OVERLAPPABLE #-} + (Columnable a, Columnable b, Columnable r, a ~ BaseType a, b ~ BaseType b) => + NullLift2Op a b r r + where + applyNull2 f = f + +-- | Left nullable: 'Nothing' short-circuits. +instance + {-# OVERLAPPABLE #-} + (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), b ~ BaseType b) => + NullLift2Op (Maybe a) b r (Maybe r) + where + applyNull2 _ Nothing _ = Nothing + applyNull2 f (Just x) y = Just (f x y) + +-- | Right nullable: 'Nothing' short-circuits. +instance + {-# OVERLAPPABLE #-} + (Columnable a, Columnable b, Columnable r, Columnable (Maybe r), a ~ BaseType a) => + NullLift2Op a (Maybe b) r (Maybe r) + where + applyNull2 _ _ Nothing = Nothing + applyNull2 f x (Just y) = Just (f x y) + +-- | Both nullable: either 'Nothing' short-circuits. +instance + {-# OVERLAPPING #-} + (Columnable a, Columnable b, Columnable r, Columnable (Maybe r)) => + NullLift2Op (Maybe a) (Maybe b) r (Maybe r) + where + applyNull2 _ Nothing _ = Nothing + applyNull2 _ _ Nothing = Nothing + applyNull2 f (Just x) (Just y) = Just (f x y) + +-- --------------------------------------------------------------------------- +-- Numeric widening +-- --------------------------------------------------------------------------- + +{- | Widen two numeric base types to their promoted common type. + +When @a ~ b@ the coercions are identity; otherwise one operand is widened +(e.g. 'Int' → 'Double'). +-} +class (Columnable (Promote a b)) => NumericWidenOp a b where + widen1 :: a -> Promote a b + widen2 :: b -> Promote a b + +-- | Same type: identity coercions. +instance {-# OVERLAPPING #-} (Columnable a) => NumericWidenOp a a where + widen1 = id + widen2 = id + +instance NumericWidenOp Int Double where widen1 = fromIntegral; widen2 = id +instance NumericWidenOp Double Int where + widen1 = id + widen2 = fromIntegral +instance NumericWidenOp Float Double where widen1 = realToFrac; widen2 = id +instance NumericWidenOp Double Float where + widen1 = id + widen2 = realToFrac +instance NumericWidenOp Int Float where widen1 = fromIntegral; widen2 = id +instance NumericWidenOp Float Int where + widen1 = id + widen2 = fromIntegral + +-- | Apply an arithmetic function after widening both operands to their common type. +widenArithOp :: + forall a b. + (NumericWidenOp a b) => + (Promote a b -> Promote a b -> Promote a b) -> + a -> + b -> + Promote a b +widenArithOp 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)) + +-- --------------------------------------------------------------------------- +-- Division widening (integral × integral → Double) +-- --------------------------------------------------------------------------- + +{- | Like 'NumericWidenOp' but uses 'PromoteDiv': integral×integral → Double. +Floating types still dominate (Double > Float), and any two integral types +(same or mixed) are both widened to Double. +-} +class (Columnable (PromoteDiv a b)) => DivWidenOp a b where + divWiden1 :: a -> PromoteDiv a b + divWiden2 :: b -> PromoteDiv a b + +-- Floating same-type (identity) +instance DivWidenOp Double Double where divWiden1 = id; divWiden2 = id +instance DivWidenOp Float Float where divWiden1 = id; divWiden2 = id + +-- Mixed Double/Float +instance DivWidenOp Double Float where divWiden1 = id; divWiden2 = realToFrac +instance DivWidenOp Float Double where divWiden1 = realToFrac; divWiden2 = id + +-- Double beats integral +instance DivWidenOp Double Int where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int Double where divWiden1 = fromIntegral; divWiden2 = id +instance DivWidenOp Double Int32 where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int32 Double where divWiden1 = fromIntegral; divWiden2 = id +instance DivWidenOp Double Int64 where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int64 Double where divWiden1 = fromIntegral; divWiden2 = id + +-- Float beats integral +instance DivWidenOp Float Int where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int Float where divWiden1 = fromIntegral; divWiden2 = id +instance DivWidenOp Float Int32 where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int32 Float where divWiden1 = fromIntegral; divWiden2 = id +instance DivWidenOp Float Int64 where divWiden1 = id; divWiden2 = fromIntegral +instance DivWidenOp Int64 Float where divWiden1 = fromIntegral; divWiden2 = id + +-- Integral × integral → Double +instance DivWidenOp Int Int where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int32 Int32 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int64 Int64 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int Int32 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int32 Int where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int Int64 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int64 Int where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int32 Int64 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral +instance DivWidenOp Int64 Int32 where + divWiden1 = fromIntegral + divWiden2 = fromIntegral + +-- | Apply an arithmetic function after widening both operands via 'PromoteDiv'. +divArithOp :: + forall a b. + (DivWidenOp a b) => + (PromoteDiv a b -> PromoteDiv a b -> PromoteDiv a b) -> + a -> + b -> + PromoteDiv a b +divArithOp f x y = f (divWiden1 @a @b x) (divWiden2 @a @b y) + +-- | Result type of a division-widening binary operator, accounting for nullable wrappers. +type WidenResultDiv a b = + NullLift2Result a b (PromoteDiv (BaseType a) (BaseType b)) diff --git a/src/DataFrame/Internal/Schema.hs b/src/DataFrame/Internal/Schema.hs --- a/src/DataFrame/Internal/Schema.hs +++ b/src/DataFrame/Internal/Schema.hs @@ -96,3 +96,15 @@ -} } deriving (Show, Eq) + +{- | Construct a 'Schema' from a list of @(columnName, schemaType)@ pairs. + +==== __Example__ +>>> :set -XTypeApplications +>>> import qualified Data.Text as T +>>> let s = makeSchema [("name", schemaType @T.Text), ("age", schemaType @Int)] +>>> M.member "age" (elements s) +True +-} +makeSchema :: [(T.Text, SchemaType)] -> Schema +makeSchema = Schema . M.fromList diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs --- a/src/DataFrame/Internal/Types.hs +++ b/src/DataFrame/Internal/Types.hs @@ -127,3 +127,28 @@ sFloating = sbool @(FloatingTypes a) type FloatingIf a = When (FloatingTypes a) (Real a, Fractional a) + +{- | Numeric type promotion: resolves the common type for mixed arithmetic. +Double dominates over Float/Int; Float dominates over Int; same types stay unchanged. +-} +type family Promote (a :: Type) (b :: Type) :: Type where + Promote a a = a + Promote Double _ = Double + Promote _ Double = Double + Promote Float _ = Float + Promote _ Float = Float + Promote Int64 _ = Int64 + Promote _ Int64 = Int64 + Promote Int32 _ = Int32 + Promote _ Int32 = Int32 + Promote a _ = a + +{- | Like 'Promote', but integral × integral → Double for use with './' . +Double\/Float still dominate; any two integral types (same or mixed) become Double. +-} +type family PromoteDiv (a :: Type) (b :: Type) :: Type where + PromoteDiv Double _ = Double + PromoteDiv _ Double = Double + PromoteDiv Float _ = Float + PromoteDiv _ Float = Float + PromoteDiv _ _ = Double -- Int/Int32/Int64 in any combination diff --git a/src/DataFrame/Lazy/IO/Binary.hs b/src/DataFrame/Lazy/IO/Binary.hs --- a/src/DataFrame/Lazy/IO/Binary.hs +++ b/src/DataFrame/Lazy/IO/Binary.hs @@ -54,6 +54,7 @@ import Data.Maybe (fromMaybe, isJust) import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl)) import Data.Word (Word16, Word32, Word64, Word8) +import qualified DataFrame.Internal.Binary as Binary import DataFrame.Internal.Column (Column (..)) import DataFrame.Internal.DataFrame (DataFrame (..)) import Foreign (ForeignPtr, castForeignPtr, plusForeignPtr, sizeOf) @@ -362,37 +363,12 @@ readWord32LE :: BS.ByteString -> Int -> ParseResult Word32 readWord32LE bs off | off + 4 > BS.length bs = Left "unexpected end of input" - | otherwise = - let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word32 - b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word32 - b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word32 - b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word32 - in Right - (off + 4, b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24)) + | otherwise = Right (off + 4, Binary.littleEndianWord32 (BS.drop off bs)) readWord64LE :: BS.ByteString -> Int -> ParseResult Word64 readWord64LE bs off | off + 8 > BS.length bs = Left "unexpected end of input" - | otherwise = - let b0 = fromIntegral (BSU.unsafeIndex bs off) :: Word64 - b1 = fromIntegral (BSU.unsafeIndex bs (off + 1)) :: Word64 - b2 = fromIntegral (BSU.unsafeIndex bs (off + 2)) :: Word64 - b3 = fromIntegral (BSU.unsafeIndex bs (off + 3)) :: Word64 - b4 = fromIntegral (BSU.unsafeIndex bs (off + 4)) :: Word64 - b5 = fromIntegral (BSU.unsafeIndex bs (off + 5)) :: Word64 - b6 = fromIntegral (BSU.unsafeIndex bs (off + 6)) :: Word64 - b7 = fromIntegral (BSU.unsafeIndex bs (off + 7)) :: Word64 - in Right - ( off + 8 - , b0 - .|. (b1 `shiftL` 8) - .|. (b2 `shiftL` 16) - .|. (b3 `shiftL` 24) - .|. (b4 `shiftL` 32) - .|. (b5 `shiftL` 40) - .|. (b6 `shiftL` 48) - .|. (b7 `shiftL` 56) - ) + | otherwise = Right (off + 8, Binary.littleEndianWord64 (BS.drop off bs)) -- | Read @n@ consecutive Word32LE values starting at offset @off@. readWord32Array :: BS.ByteString -> Int -> Int -> Either String [Word32] @@ -401,11 +377,7 @@ | otherwise = Right [ let i = off + k * 4 - b0 = fromIntegral (BSU.unsafeIndex bs i) :: Word32 - b1 = fromIntegral (BSU.unsafeIndex bs (i + 1)) :: Word32 - b2 = fromIntegral (BSU.unsafeIndex bs (i + 2)) :: Word32 - b3 = fromIntegral (BSU.unsafeIndex bs (i + 3)) :: Word32 - in b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24) + in Binary.littleEndianWord32 (BS.drop i bs) | k <- [0 .. n - 1] ] diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs --- a/src/DataFrame/Lazy/IO/CSV.hs +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -339,17 +339,13 @@ case testEquality (typeRep @a) (typeRep @T.Text) of Just Refl -> let val = TextEncoding.decodeUtf8Lenient bs - in if isNullish val - then VM.unsafeWrite col i T.empty >> return (Left val) - else VM.unsafeWrite col i val >> return (Right True) + in VM.unsafeWrite col i val >> return (Right True) Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs)) writeColumnBs i bs (MOptionalColumn (col :: VM.IOVector (Maybe a))) = case testEquality (typeRep @a) (typeRep @T.Text) of Just Refl -> let val = TextEncoding.decodeUtf8Lenient bs - in if isNullish val - then VM.unsafeWrite col i Nothing >> return (Left val) - else VM.unsafeWrite col i (Just val) >> return (Right True) + in VM.unsafeWrite col i (Just val) >> return (Right True) Nothing -> return (Left (TextEncoding.decodeUtf8Lenient bs)) writeColumnBs i bs (MUnboxedColumn (col :: VUM.IOVector a)) = case testEquality (typeRep @a) (typeRep @Double) of 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 @@ -134,5 +134,5 @@ sortBy cols ldf = ldf{plan = Sort cols (plan ldf)} -- | Retain at most @n@ rows. -limit :: Int -> LazyDataFrame -> LazyDataFrame -limit n ldf = ldf{plan = Limit n (plan ldf)} +take :: Int -> LazyDataFrame -> LazyDataFrame +take n ldf = ldf{plan = Limit n (plan ldf)} 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 @@ -459,26 +459,65 @@ via 'ParquetReadOptions'. -} executeParquetScan :: FilePath -> ScanConfig -> IO Stream -executeParquetScan path cfg = do - isDir <- doesDirectoryExist path - let pat = if isDir then path "*" else path - matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches - when (null files) $ - error ("executeParquetScan: no parquet files found for " ++ path) +executeParquetScan path cfg + | Parquet.isHFUri path = executeHFParquetScan path cfg + | otherwise = do + isDir <- doesDirectoryExist path + let pat = if isDir then path "*" else path + matches <- glob pat + files <- filterM (fmap not . doesDirectoryExist) matches + when (null files) $ + error ("executeParquetScan: no parquet files found for " ++ path) + let opts = + Parquet.defaultParquetReadOptions + { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg))) + , Parquet.predicate = scanPushdownPredicate cfg + } + ref <- newIORef files + return . Stream $ do + fs <- readIORef ref + case fs of + [] -> return Nothing + (f : rest) -> do + writeIORef ref rest + Just <$> Parquet.readParquetWithOpts opts f + +{- | HuggingFace Parquet scan. Files are resolved once (API call or direct URL) +then downloaded one at a time as the stream is pulled — so only one file's worth +of data is in memory at a time, regardless of dataset size. +-} + +-- TODO: mchavinda - this should be a more general online file scanner. +executeHFParquetScan :: FilePath -> ScanConfig -> IO Stream +executeHFParquetScan path cfg = do + ref <- case Parquet.parseHFUri path of + Left err -> error err + Right r -> pure r + mToken <- Parquet.getHFToken + hfFiles <- + if Parquet.hasGlob (Parquet.hfGlob ref) + then Parquet.resolveHFUrls mToken ref + else do + let url = Parquet.directHFUrl ref + filename = last $ T.splitOn "/" (Parquet.hfGlob ref) + pure [Parquet.HFParquetFile url "" "" filename] + when (null hfFiles) $ + error ("executeParquetScan: no HF parquet files found for " ++ path) let opts = Parquet.defaultParquetReadOptions { Parquet.selectedColumns = Just (M.keys (elements (scanSchema cfg))) , Parquet.predicate = scanPushdownPredicate cfg } - ref <- newIORef files + filesRef <- newIORef hfFiles return . Stream $ do - fs <- readIORef ref + fs <- readIORef filesRef case fs of [] -> return Nothing (f : rest) -> do - writeIORef ref rest - Just <$> Parquet.readParquetWithOpts opts f + writeIORef filesRef rest + -- Download a single file, read it, then return the batch. + [localPath] <- Parquet.downloadHFFiles mToken [f] + Just <$> Parquet.readParquetWithOpts opts localPath -- --------------------------------------------------------------------------- -- CSV scan implementation diff --git a/src/DataFrame/Monad.hs b/src/DataFrame/Monad.hs --- a/src/DataFrame/Monad.hs +++ b/src/DataFrame/Monad.hs @@ -12,6 +12,8 @@ import qualified DataFrame as D import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Nullable (BaseType) +import DataFrame.Operations.Transformations (ImputeOp) import qualified Data.Text as T import System.Random @@ -75,7 +77,11 @@ filterJustM expr = error $ "Cannot filter on compound expression: " ++ show expr -imputeM :: (Columnable a) => Expr (Maybe a) -> a -> FrameM (Expr a) +imputeM :: + (ImputeOp a, Columnable (BaseType a)) => + Expr a -> + BaseType a -> + FrameM (Expr (BaseType a)) imputeM expr@(Col name) value = FrameM $ \df -> let df' = D.impute expr value df in (df', Col name) diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -41,6 +41,7 @@ import DataFrame.Internal.DataFrame ( DataFrame (..), columnIndices, + derivingExpressions, empty, getColumn, ) @@ -344,6 +345,7 @@ let (r, c) = dataframeDimensions d n = max (columnLength column) r + exprs = M.delete name (derivingExpressions d) in case M.lookup name (columnIndices d) of Just i -> @@ -351,13 +353,13 @@ (V.map (expandColumn n) (columns d V.// [(i, column)])) (columnIndices d) (n, c) - M.empty + exprs Nothing -> DataFrame (V.map (expandColumn n) (columns d `V.snoc` column)) (M.insert name c (columnIndices d)) (n, c + 1) - M.empty + exprs {- | /O(n)/ Clones a column and places it under a new name in the dataframe. @@ -945,3 +947,22 @@ -} columnAsList :: forall a. (Columnable a) => Expr a -> DataFrame -> [a] columnAsList expr df = either throw V.toList (columnAsVector expr df) + +{- | Returns the provenance of all columns in the DataFrame as a list of +@(name, expression)@ pairs. Derived columns show their expression; +raw columns show an identity @col \@type name@ expression. +-} +showDerivedExpressions :: DataFrame -> [NamedExpr] +showDerivedExpressions df = + let exprs = derivingExpressions df + names = columnNames df + toNamedExpr name = case M.lookup name exprs of + Just uexpr -> (name, uexpr) + Nothing -> (name, identityUExpr name) + in map toNamedExpr names + where + identityUExpr name = case getColumn name df of + Just (BoxedColumn (_ :: V.Vector a)) -> UExpr (Col @a name) + Just (UnboxedColumn (_ :: VU.Vector a)) -> UExpr (Col @a name) + Just (OptionalColumn (_ :: V.Vector (Maybe a))) -> UExpr (Col @(Maybe a) name) + Nothing -> error $ "showDerivedExpressions: column not found: " ++ T.unpack name diff --git a/src/DataFrame/Operations/Join.hs b/src/DataFrame/Operations/Join.hs --- a/src/DataFrame/Operations/Join.hs +++ b/src/DataFrame/Operations/Join.hs @@ -1,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE NumericUnderscores #-} @@ -13,7 +14,9 @@ import Control.Monad (forM_, when) import Control.Monad.ST (ST, runST) import qualified Data.HashMap.Strict as HM +#if !MIN_VERSION_base(4,20,0) import Data.List (foldl') +#endif import qualified Data.Map.Strict as M import Data.Maybe (fromMaybe) import Data.STRef (newSTRef, readSTRef, writeSTRef) diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs --- a/src/DataFrame/Operations/Merge.hs +++ b/src/DataFrame/Operations/Merge.hs @@ -18,11 +18,11 @@ (<>) a b = let addColumns a' b' df name - | fst (D.dimensions a') == 0 && fst (D.dimensions b') == 0 = df - | fst (D.dimensions a') == 0 = fromMaybe df $ do + | snd (D.dimensions a') == 0 && snd (D.dimensions b') == 0 = df + | snd (D.dimensions a') == 0 = fromMaybe df $ do col <- D.getColumn name b' pure $ D.insertColumn name col df - | fst (D.dimensions b') == 0 = fromMaybe df $ do + | snd (D.dimensions b') == 0 = fromMaybe df $ do col <- D.getColumn name a' pure $ D.insertColumn name col df | otherwise = @@ -49,8 +49,11 @@ Just a'' -> let concatedColumns = D.concatColumnsEither a'' b'' in D.insertColumn name concatedColumns df + result = L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) in - L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b) + result + { D.derivingExpressions = D.derivingExpressions a <> D.derivingExpressions b + } instance Monoid D.DataFrame where mempty = D.empty @@ -58,7 +61,12 @@ -- | Add two dataframes side by side/horizontally. (|||) :: D.DataFrame -> D.DataFrame -> D.DataFrame (|||) a b = - D.fold - (\name acc -> D.insertColumn name (D.unsafeGetColumn name b) acc) - (D.columnNames b) - a + let result = + D.fold + (\name acc -> D.insertColumn name (D.unsafeGetColumn name b) acc) + (D.columnNames b) + a + in result + { D.derivingExpressions = + D.derivingExpressions result <> D.derivingExpressions b + } diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -1,10 +1,12 @@ {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} module DataFrame.Operations.Statistics where @@ -30,12 +32,13 @@ ) import DataFrame.Internal.Expression import DataFrame.Internal.Interpreter +import DataFrame.Internal.Nullable (BaseType) import DataFrame.Internal.Row (showValue, toAny) import DataFrame.Internal.Statistics import DataFrame.Internal.Types import DataFrame.Operations.Core import DataFrame.Operations.Subset (filterJust) -import DataFrame.Operations.Transformations (impute) +import DataFrame.Operations.Transformations (ImputeOp (..), imputeCore) import Text.Printf (printf) import Type.Reflection (typeRep) @@ -287,22 +290,28 @@ -- 20 @ -} +instance {-# OVERLAPPING #-} (Columnable b) => ImputeOp (Maybe b) where + runImpute = imputeCore + + runImputeWith f col@(Col columnName) df = + case interpret @b (filterJust columnName df) (f (Col @b columnName)) of + Left e -> throw e + Right (TColumn value) -> case headColumn @b value of + Left e -> throw e + Right h -> + if all (== h) (toList @b value) + then imputeCore col h df + else error "Impute expression returned more than one value" + runImputeWith _ _ df = df + imputeWith :: - forall b. - (Columnable b) => - (Expr b -> Expr b) -> - Expr (Maybe b) -> + forall a. + (ImputeOp a, Columnable (BaseType a)) => + (Expr (BaseType a) -> Expr (BaseType a)) -> + Expr a -> DataFrame -> DataFrame -imputeWith f col@(Col columnName) df = case interpret @b (filterJust columnName df) (f (Col @b columnName)) of - Left e -> throw e - Right (TColumn value) -> case headColumn @b value of - Left e -> throw e - Right h -> - if all (== h) (toList @b value) - then impute col h df - else error "Impute expression returned more than one value" -imputeWith _ _ df = df +imputeWith = runImputeWith applyStatistic :: (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} @@ -33,12 +34,15 @@ import DataFrame.Internal.Column import DataFrame.Internal.DataFrame ( DataFrame (..), + derivingExpressions, empty, getColumn, + unsafeGetColumn, ) import DataFrame.Internal.Expression import DataFrame.Internal.Interpreter import DataFrame.Operations.Core +import DataFrame.Operations.Merge () import DataFrame.Operations.Transformations (apply) import System.Random import Type.Reflection @@ -242,7 +246,10 @@ (T.pack $ show $ cs L.\\ columnNames df) "select" (columnNames df) - | otherwise = L.foldl' addKeyValue empty cs + | otherwise = + let result = L.foldl' addKeyValue empty cs + filteredExprs = M.filterWithKey (\k _ -> k `L.elem` cs) (derivingExpressions df) + in result{derivingExpressions = filteredExprs} where addKeyValue d k = fromMaybe df $ do col <- getColumn k df @@ -467,3 +474,81 @@ (v, g') = uniformR (0 :: Double, 1 :: Double) g in v : go g' (n - 1) + +-- | Convert any Column to a vector of Text labels (one per row). +columnToTextVec :: Column -> V.Vector T.Text +columnToTextVec (BoxedColumn (col :: V.Vector a)) = + case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> col + Nothing -> V.map (T.pack . show) col +columnToTextVec (UnboxedColumn col) = V.map (T.pack . show) (V.convert col) +columnToTextVec (OptionalColumn col) = V.map (T.pack . show) col + +-- | Build a map from stringified label to row indices. +groupByIndices :: Column -> M.Map T.Text (VU.Vector Int) +groupByIndices col = + let textVec = columnToTextVec col + (grouped, _) = + V.foldl' + (\(!m, !i) key -> (M.insertWith (++) key [i] m, i + 1)) + (M.empty, 0) + textVec + in M.map (VU.fromList . L.reverse) grouped + +-- | Select rows at the given indices from all columns. +rowsAtIndices :: VU.Vector Int -> DataFrame -> DataFrame +rowsAtIndices ixs df = + df + { columns = V.map (atIndicesStable ixs) (columns df) + , dataframeDimensions = (VU.length ixs, snd (dataframeDimensions df)) + } + +{- | Sample a dataframe, preserving per-stratum proportions. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.stratifiedSample (mkStdGen 42) 0.8 "label" df +@ +-} +stratifiedSample :: + forall a g. + (SplitGen g, RandomGen g, Columnable a) => + g -> Double -> Expr a -> DataFrame -> DataFrame +stratifiedSample gen p strataCol df = + let col = case strataCol of + Col name -> unsafeGetColumn name df + _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + groups = M.elems (groupByIndices col) + go _ [] = mempty + go g (ixs : rest) = + let stratum = rowsAtIndices ixs df + (g1, g2) = splitGen g + in sample g1 p stratum <> go g2 rest + in go gen groups + +{- | Split a dataframe into two, preserving per-stratum proportions. + +==== __Example__ +@ +ghci> import System.Random +ghci> D.stratifiedSplit (mkStdGen 42) 0.8 "label" df +@ +-} +stratifiedSplit :: + forall a g. + (SplitGen g, RandomGen g, Columnable a) => + g -> Double -> Expr a -> DataFrame -> (DataFrame, DataFrame) +stratifiedSplit gen p strataCol df = + let col = case strataCol of + Col name -> unsafeGetColumn name df + _ -> unwrapTypedColumn (either throw id (interpret @a df strataCol)) + groups = M.elems (groupByIndices col) + go _ [] = (mempty, mempty) + go g (ixs : rest) = + let stratum = rowsAtIndices ixs df + (g1, g2) = splitGen g + (tr, va) = randomSplit g1 p stratum + (trAcc, vaAcc) = go g2 rest + in (tr <> trAcc, va <> vaAcc) + in go gen groups diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -1,10 +1,14 @@ +{-# LANGUAGE ConstrainedClassMethods #-} {-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} module DataFrame.Operations.Transformations where @@ -27,6 +31,7 @@ import DataFrame.Internal.DataFrame (DataFrame (..), getColumn) import DataFrame.Internal.Expression import DataFrame.Internal.Interpreter +import DataFrame.Internal.Nullable (BaseType) import DataFrame.Operations.Core -- | O(k) Apply a function to a given column in a dataframe. @@ -86,7 +91,12 @@ forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> (Expr a, DataFrame) deriveWithExpr name expr df = case interpret @a df (normalize expr) of Left e -> throw e - Right (TColumn value) -> (Col name, insertColumn name value df) + Right (TColumn value) -> + ( Col name + , (insertColumn name value df) + { derivingExpressions = M.insert name (UExpr expr) (derivingExpressions df) + } + ) deriveMany :: [NamedExpr] -> DataFrame -> DataFrame deriveMany exprs df = @@ -188,20 +198,45 @@ Left e -> throw e Right column' -> insertColumn columnName column' df --- | Replace all instances of `Nothing` in a column with the given value. -impute :: +-- | Core impute implementation for nullable columns. Silently no-ops on non-nullable columns. +imputeCore :: forall b. (Columnable b) => Expr (Maybe b) -> b -> DataFrame -> DataFrame -impute (Col columnName) value df = case getColumn columnName df of +imputeCore (Col columnName) value df = case getColumn columnName df of Nothing -> throw $ ColumnNotFoundException columnName "impute" (M.keys $ columnIndices df) Just (OptionalColumn _) -> case safeApply (fromMaybe value) columnName df of Left (TypeMismatchException context) -> throw $ TypeMismatchException (context{callingFunctionName = Just "impute"}) Left exception -> throw exception Right res -> res - _ -> error $ "Cannot impute to a non-Empty column: " ++ T.unpack columnName -impute _ _ df = df + _ -> df +imputeCore _ _ df = df + +class (Columnable a) => ImputeOp a where + runImpute :: Expr a -> BaseType a -> DataFrame -> DataFrame + runImputeWith :: + (Columnable (BaseType a)) => + (Expr (BaseType a) -> Expr (BaseType a)) -> + Expr a -> + DataFrame -> + DataFrame + +instance {-# OVERLAPPABLE #-} (Columnable a) => ImputeOp a where + runImpute _ _ df = df + runImputeWith _ _ df = df + +{- | Replace all instances of `Nothing` in a column with the given value. +When the column is already non-nullable, this is a silent no-op. +-} +impute :: + forall a. + (ImputeOp a) => + Expr a -> + BaseType a -> + DataFrame -> + DataFrame +impute = runImpute diff --git a/src/DataFrame/Operations/Typing.hs b/src/DataFrame/Operations/Typing.hs --- a/src/DataFrame/Operations/Typing.hs +++ b/src/DataFrame/Operations/Typing.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -6,61 +5,87 @@ module DataFrame.Operations.Typing where +import qualified Data.Map as M import qualified Data.Text as T import qualified Data.Vector as V +import Control.Applicative (asum) +import Control.Monad (join) import Data.Maybe (fromMaybe) import qualified Data.Proxy as P import Data.Time -import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import Data.Type.Equality (TestEquality (..)) import DataFrame.Internal.Column (Column (..), fromVector) -import DataFrame.Internal.DataFrame (DataFrame (..)) +import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn) import DataFrame.Internal.Parsing import DataFrame.Internal.Schema +import DataFrame.Operations.Core import Text.Read -import Type.Reflection (typeRep) +import Type.Reflection type DateFormat = String -parseDefaults :: [T.Text] -> Int -> Bool -> DateFormat -> DataFrame -> DataFrame -parseDefaults missing n safeRead dateFormat df = df{columns = V.map (parseDefault missing n safeRead dateFormat) (columns df)} +-- | Options controlling how text columns are parsed into typed values. +data ParseOptions = ParseOptions + { missingValues :: [T.Text] + -- ^ Values to treat as @Nothing@ when 'parseSafe' is @True@. + , sampleSize :: Int + -- ^ Number of rows to inspect when inferring a column's type (0 = all rows). + , parseSafe :: Bool + {- ^ When @True@, treat 'missingValues' and nullish strings as @Nothing@. + When @False@, only empty strings become @Nothing@. + -} + , parseDateFormat :: DateFormat + -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@). + } -parseDefault :: [T.Text] -> Int -> Bool -> DateFormat -> Column -> Column -parseDefault missing n safeRead dateFormat (BoxedColumn (c :: V.Vector a)) = +{- | Sensible out-of-the-box parse options: infer from the first 100 rows, + treat common nullish strings as missing, and expect ISO 8601 dates. +-} +defaultParseOptions :: ParseOptions +defaultParseOptions = + ParseOptions + { missingValues = [] + , sampleSize = 100 + , parseSafe = True + , parseDateFormat = "%Y-%m-%d" + } + +parseDefaults :: ParseOptions -> DataFrame -> DataFrame +parseDefaults opts df = df{columns = V.map (parseDefault opts) (columns df)} + +parseDefault :: ParseOptions -> Column -> Column +parseDefault opts (BoxedColumn (c :: V.Vector a)) = case (typeRep @a) `testEquality` (typeRep @T.Text) of Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of - Just Refl -> parseFromExamples missing n safeRead dateFormat (V.map T.pack c) + Just Refl -> parseFromExamples opts (V.map T.pack c) Nothing -> BoxedColumn c - Just Refl -> parseFromExamples missing n safeRead dateFormat c -parseDefault missing n safeRead dateFormat (OptionalColumn (c :: V.Vector (Maybe a))) = + Just Refl -> parseFromExamples opts c +parseDefault opts (OptionalColumn (c :: V.Vector (Maybe a))) = case (typeRep @a) `testEquality` (typeRep @T.Text) of Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of Just Refl -> - parseFromExamples - missing - n - safeRead - dateFormat - (V.map (T.pack . fromMaybe "") c) + parseFromExamples opts (V.map (T.pack . fromMaybe "") c) Nothing -> BoxedColumn c - Just Refl -> parseFromExamples missing n safeRead dateFormat (V.map (fromMaybe "") c) -parseDefault _ _ _ _ column = column + Just Refl -> parseFromExamples opts (V.map (fromMaybe "") c) +parseDefault _ column = column -parseFromExamples :: - [T.Text] -> Int -> Bool -> DateFormat -> V.Vector T.Text -> Column -parseFromExamples missing n safeRead dateFormat cols = +parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column +parseFromExamples opts cols = let - converter = if safeRead then convertNullish missing else convertOnlyEmpty - examples = V.map converter (V.take n cols) + converter = + if parseSafe opts then convertNullish (missingValues opts) else convertOnlyEmpty + examples = V.map converter (V.take (sampleSize opts) cols) asMaybeText = V.map converter cols + dfmt = parseDateFormat opts in - case makeParsingAssumption dateFormat examples of + case makeParsingAssumption dfmt examples of BoolAssumption -> handleBoolAssumption asMaybeText IntAssumption -> handleIntAssumption asMaybeText DoubleAssumption -> handleDoubleAssumption asMaybeText TextAssumption -> handleTextAssumption asMaybeText - DateAssumption -> handleDateAssumption dateFormat asMaybeText - NoAssumption -> handleNoAssumption dateFormat asMaybeText + DateAssumption -> handleDateAssumption dfmt asMaybeText + NoAssumption -> handleNoAssumption dfmt asMaybeText handleBoolAssumption :: V.Vector (Maybe T.Text) -> Column handleBoolAssumption asMaybeText @@ -197,19 +222,63 @@ | NoAssumption | TextAssumption -parseWithTypes :: [SchemaType] -> DataFrame -> DataFrame -parseWithTypes ts df = df{columns = go 0 ts (columns df)} +parseWithTypes :: Bool -> M.Map T.Text SchemaType -> DataFrame -> DataFrame +parseWithTypes safe ts df + | M.null ts = df + | otherwise = + M.foldrWithKey + (\k v d -> insertColumn k (asType v (unsafeGetColumn k d)) d) + df + ts where - go :: Int -> [SchemaType] -> V.Vector Column -> V.Vector Column - go n [] xs = xs - go n (t : rest) xs - | n >= V.length xs = xs - | otherwise = - go (n + 1) rest (V.update xs (V.fromList [(n, asType t (xs V.! n))])) asType :: SchemaType -> Column -> Column - asType (SType (_ :: P.Proxy a)) c@(BoxedColumn (col :: V.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> c - Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of - Just Refl -> fromVector (V.map ((readMaybe @a) . T.unpack) col) - Nothing -> fromVector (V.map ((readMaybe @a) . show) col) + asType (SType (_ :: P.Proxy a)) c@(BoxedColumn (col :: V.Vector b)) = case typeRep @a of + App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of + Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> fromVector (V.map (join . (readAsMaybe @a) . T.unpack) col) + Nothing -> fromVector (V.map (join . (readAsMaybe @a) . show) col) + Nothing -> case t1 of + App t1' t2' -> case eqTypeRep t1' (typeRep @Either) of + Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> fromVector (V.map ((readAsEither @a) . T.unpack) col) + Nothing -> fromVector (V.map ((readAsEither @a) . show) col) + Nothing -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> + if safe + then fromVector (V.map ((readMaybe @a) . T.unpack) col) + else fromVector (V.map ((read @a) . T.unpack) col) + Nothing -> + if safe + then fromVector (V.map ((readMaybe @a) . show) col) + else fromVector (V.map ((read @a) . show) col) + _ -> c + _ -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> c + Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of + Just Refl -> + if safe + then fromVector (V.map ((readMaybe @a) . T.unpack) col) + else fromVector (V.map ((read @a) . T.unpack) col) + Nothing -> + if safe + then fromVector (V.map ((readMaybe @a) . show) col) + else fromVector (V.map ((read @a) . show) col) asType _ c = c + +readAsMaybe :: (Read a) => String -> Maybe a +readAsMaybe s + | null s = Nothing + | otherwise = readMaybe $ "Just " <> s + +readAsEither :: (Read a) => String -> a +readAsEither v = case asum [readMaybe $ "Left " <> s, readMaybe $ "Right " <> s] of + Nothing -> error $ "Couldn't read value: " <> s + Just v -> v + where + s = if null v then "\"\"" else v diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs --- a/src/DataFrame/Operators.hs +++ b/src/DataFrame/Operators.hs @@ -1,5 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} module DataFrame.Operators where @@ -15,15 +17,31 @@ binaryPrecedence, binarySymbol ), - Expr (Binary, Col, If, Lit), + Expr (Binary, Col, If, Lit, Unary), NamedExpr, UExpr (UExpr), + UnaryOp (MkUnaryOp, unaryFn, unaryName, unarySymbol), ) +import DataFrame.Internal.Nullable ( + BaseType, + DivWidenOp, + NullCmpResult, + NullLift2Op (applyNull2), + NullableCmpOp (nullCmpOp), + NumericWidenOp, + WidenResult, + WidenResultDiv, + divArithOp, + widenArithOp, + ) +import DataFrame.Internal.Types (Promote, PromoteDiv) -infix 8 .^^ -infix 4 .==, .<, .<=, .>=, .>, ./= -infixr 3 .&& -infixr 2 .|| +infix 8 .^^, .^^., .^, .^. +infix 6 .+, .- +infix 7 .*, ./ +infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=. +infixr 3 .&&, .&&. +infixr 2 .||, .||. infixr 0 .= (|>) :: a -> (a -> b) -> b @@ -50,112 +68,238 @@ (.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr (.=) = flip as -(.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -(.==) = - Binary - ( MkBinaryOp - { binaryFn = (==) - , binaryName = "eq" - , binarySymbol = Just "==" - , binaryCommutative = True - , binaryPrecedence = 4 - } - ) +liftDecorated :: + (Columnable a, Columnable b) => + (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b +liftDecorated f name rep = Unary (MkUnaryOp{unaryFn = f, unaryName = name, unarySymbol = rep}) -(./=) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool -(./=) = +lift2Decorated :: + (Columnable c, Columnable b, Columnable a) => + (c -> b -> a) -> + T.Text -> + Maybe T.Text -> + Bool -> + Int -> + Expr c -> + Expr b -> + Expr a +lift2Decorated f name rep comm prec = Binary ( MkBinaryOp - { binaryFn = (/=) - , binaryName = "neq" - , binarySymbol = Just "/=" - , binaryCommutative = True - , binaryPrecedence = 4 + { binaryFn = f + , binaryName = name + , binarySymbol = rep + , binaryCommutative = comm + , binaryPrecedence = prec } ) -(.<) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.<) = - Binary - ( MkBinaryOp - { binaryFn = (<) - , binaryName = "lt" - , binarySymbol = Just "<" - , binaryCommutative = False - , binaryPrecedence = 4 - } - ) +(.==.) :: + (Columnable a, Eq a) => + Expr a -> + Expr a -> + Expr Bool +(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4 -(.>) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool -(.>) = - Binary - ( MkBinaryOp - { binaryFn = (>) - , binaryName = "gt" - , binarySymbol = Just ">" - , binaryCommutative = False - , binaryPrecedence = 4 - } - ) +(./=.) :: + (Columnable a, Eq a) => + Expr a -> + Expr a -> + Expr Bool +(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4 -(.<=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -(.<=) = - Binary - ( MkBinaryOp - { binaryFn = (<=) - , binaryName = "leq" - , binarySymbol = Just "<=" - , binaryCommutative = False - , binaryPrecedence = 4 - } - ) +(.<.) :: + (Columnable a, Ord a) => + Expr a -> + Expr a -> + Expr Bool +(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4 -(.>=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool -(.>=) = - Binary - ( MkBinaryOp - { binaryFn = (>=) - , binaryName = "geq" - , binarySymbol = Just ">=" - , binaryCommutative = False - , binaryPrecedence = 4 - } - ) +(.>.) :: + (Columnable a, Ord a) => + Expr a -> + Expr a -> + Expr Bool +(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4 -(.&&) :: Expr Bool -> Expr Bool -> Expr Bool -(.&&) = - Binary - ( MkBinaryOp - { binaryFn = (&&) - , binaryName = "and" - , binarySymbol = Just "&&" - , binaryCommutative = True - , binaryPrecedence = 3 - } - ) +(.<=.) :: + (Columnable a, Ord a) => + Expr a -> + Expr a -> + Expr Bool +(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4 -(.||) :: Expr Bool -> Expr Bool -> Expr Bool -(.||) = - Binary - ( MkBinaryOp - { binaryFn = (||) - , binaryName = "or" - , binarySymbol = Just "||" - , binaryCommutative = True - , binaryPrecedence = 2 - } - ) +(.>=.) :: + (Columnable a, Ord a) => + Expr a -> + Expr a -> + Expr Bool +(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4 -(.^^) :: (Columnable a, Num a) => Expr a -> Int -> Expr a -(.^^) expr i = - Binary - ( MkBinaryOp - { binaryFn = (^) - , binaryName = "pow" - , binarySymbol = Just "^" - , binaryCommutative = False - , binaryPrecedence = 8 - } - ) - expr - (Lit i) +(.+.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a +(.+.) = (+) + +(.-.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a +(.-.) = (-) + +(.*.) :: (Columnable a, Num a) => Expr a -> Expr a -> Expr a +(.*.) = (*) + +(./.) :: (Columnable a, Fractional a) => Expr a -> Expr a -> Expr a +(./.) = (/) + +-- Nullable-aware arithmetic operators + +{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands. +@col \@Int "x" .+ col \@(Maybe Int) "y" -- :: Expr (Maybe Int)@ +-} +(.+) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + Expr a -> + Expr b -> + Expr (WidenResult a b) +(.+) = lift2Decorated (applyNull2 (widenArithOp (+))) "nulladd" (Just ".+") True 6 + +-- | Nullable-aware subtraction. +(.-) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + Expr a -> + Expr b -> + Expr (WidenResult a b) +(.-) = lift2Decorated (applyNull2 (widenArithOp (-))) "nullsub" (Just ".-") False 6 + +-- | Nullable-aware multiplication. +(.*) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + Expr a -> + Expr b -> + Expr (WidenResult a b) +(.*) = lift2Decorated (applyNull2 (widenArithOp (*))) "nullmul" (Just ".*") True 7 + +-- | Nullable-aware division. Integral operands are promoted to Double. +(./) :: + ( DivWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b) + , Fractional (PromoteDiv (BaseType a) (BaseType b)) + ) => + Expr a -> + Expr b -> + Expr (WidenResultDiv a b) +(./) = lift2Decorated (applyNull2 (divArithOp (/))) "nulldiv" (Just "./") False 7 + +-- Nullable-aware comparison operators (three-valued logic: Nothing if either operand is Nothing) + +-- | Nullable-aware equality. Returns @Maybe Bool@ when either operand is nullable. +(.==) :: + (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.==) = lift2Decorated (nullCmpOp (==)) "eq" (Just ".==") True 4 + +-- | Nullable-aware inequality. +(./=) :: + (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(./=) = lift2Decorated (nullCmpOp (/=)) "neq" (Just "./=") True 4 + +-- | Nullable-aware less-than. +(.<) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.<) = lift2Decorated (nullCmpOp (<)) "lt" (Just ".<") False 4 + +-- | Nullable-aware greater-than. +(.>) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.>) = lift2Decorated (nullCmpOp (>)) "gt" (Just ".>") False 4 + +-- | Nullable-aware less-than-or-equal. +(.<=) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.<=) = lift2Decorated (nullCmpOp (<=)) "leq" (Just ".<=") False 4 + +-- | Nullable-aware greater-than-or-equal. +(.>=) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.>=) = lift2Decorated (nullCmpOp (>=)) "geq" (Just ".>=") False 4 + +(.&&.) :: Expr Bool -> Expr Bool -> Expr Bool +(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3 + +(.||.) :: Expr Bool -> Expr Bool -> Expr Bool +(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2 + +-- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable. +(.&&) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.&&) = lift2Decorated (nullCmpOp (&&)) "nulland" (Just ".&&") True 3 + +-- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable. +(.||) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + Expr a -> + Expr b -> + Expr (NullCmpResult a b) +(.||) = lift2Decorated (nullCmpOp (||)) "nullor" (Just ".||") True 2 + +(.^^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Fractional (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + Expr a -> Expr b -> Expr a +(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8 + +(.^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Num (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + Expr a -> Expr b -> Expr a +(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8 + +-- Same-type (non-nullable) exponentiation operators + +(.^^.) :: + (Columnable a, Columnable b, Fractional a, Integral b) => + Expr a -> Expr b -> Expr a +(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8 + +(.^.) :: + (Columnable a, Columnable b, Num a, Integral b) => + Expr a -> Expr b -> Expr a +(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8 diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -260,7 +260,7 @@ [] [] of Nothing -> Left "No programs found" - Just p -> Right (F.ifThenElse (p .> 0) 1 0) + Just p -> Right (F.ifThenElse (p .> (0 :: Expr Double)) 1 0) percentiles :: DataFrame -> [Expr Double] percentiles df = diff --git a/src/DataFrame/Typed.hs b/src/DataFrame/Typed.hs --- a/src/DataFrame/Typed.hs +++ b/src/DataFrame/Typed.hs @@ -70,14 +70,30 @@ ifThenElse, lift, lift2, + nullLift, + nullLift2, - -- * Comparison operators + -- * Same-type comparison operators (.==.), (./=.), (.<.), (.<=.), (.>=.), (.>.), + + -- * Nullable-aware arithmetic operators + (.+), + (.-), + (.*), + (./), + + -- * Nullable-aware comparison operators (three-valued logic) + (.==), + (./=), + (.<), + (.<=), + (.>=), + (.>), -- * Logical operators (.&&.), 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 @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -53,8 +54,10 @@ -- * Unary / binary lifting lift, lift2, + nullLift, + nullLift2, - -- * Comparison operators + -- * Same-type comparison operators (.==.), (./=.), (.<.), @@ -62,9 +65,39 @@ (.>=.), (.>.), + -- * Same-type arithmetic operators + (.+.), + (.-.), + (.*.), + (./.), + + -- * Same-type exponentiation operators + (.^^.), + (.^.), + + -- * Nullable-aware arithmetic operators + (.+), + (.-), + (.*), + (./), + + -- * Nullable-aware exponentiation operators + (.^^), + (.^), + + -- * Nullable-aware comparison operators (three-valued logic) + (.==), + (./=), + (.<), + (.<=), + (.>=), + (.>), + -- * Logical operators (.&&.), (.||.), + (.&&), + (.||), DataFrame.Typed.Expr.not, -- * Aggregation combinators @@ -75,6 +108,12 @@ maximum, collect, + -- * Cast / coercion expressions + castExpr, + castExprWithDefault, + castExprEither, + unsafeCastExpr, + -- * Named expression helper as, @@ -83,6 +122,7 @@ desc, ) where +import Data.Either (fromRight) import Data.Proxy (Proxy (..)) import Data.String (IsString (..)) import qualified Data.Text as T @@ -98,6 +138,22 @@ UExpr (..), UnaryOp (..), ) +import DataFrame.Internal.Nullable ( + BaseType, + DivWidenOp, + NullCmpResult, + NullLift1Op (applyNull1), + NullLift1Result, + NullLift2Op (applyNull2), + NullLift2Result, + NullableCmpOp (nullCmpOp), + NumericWidenOp, + WidenResult, + WidenResultDiv, + divArithOp, + widenArithOp, + ) +import DataFrame.Internal.Types (Promote, PromoteDiv) import DataFrame.Typed.Schema (AssertPresent, Lookup) import DataFrame.Typed.Types (TExpr (..), TSortOrder (..)) @@ -189,9 +245,36 @@ (a -> b -> c) -> TExpr cols a -> TExpr cols b -> TExpr cols c lift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp f "binaryUdf" Nothing False 0) a b) +{- | Typed 'nullLift': lift a unary function with nullable propagation. +When the input is @Maybe a@, 'Nothing' short-circuits; when plain @a@, applies directly. +The return type is inferred via 'NullLift1Result': no annotation needed. +-} +nullLift :: + (NullLift1Op a r (NullLift1Result a r), Columnable (NullLift1Result a r)) => + (BaseType a -> r) -> + TExpr cols a -> + TExpr cols (NullLift1Result a r) +nullLift f (TExpr e) = TExpr (Unary (MkUnaryOp (applyNull1 f) "nullLift" Nothing) e) + +{- | Typed 'nullLift2': lift a binary function with nullable propagation. +Any 'Nothing' operand short-circuits to 'Nothing' in the result. +The return type is inferred via 'NullLift2Result': no annotation needed. +-} +nullLift2 :: + (NullLift2Op a b r (NullLift2Result a b r), Columnable (NullLift2Result a b r)) => + (BaseType a -> BaseType b -> r) -> + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullLift2Result a b r) +nullLift2 f (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (applyNull2 f) "nullLift2" Nothing False 0) a b) + infixl 4 .==., ./=., .<., .<=., .>=., .>. -infixr 3 .&&. -infixr 2 .||. +infix 4 .==, ./=, .<, .<=, .>=, .> +infixr 3 .&&., .&& +infixr 2 .||., .|| +infixl 6 .+., .-. +infixl 7 .*., ./. +infix 8 .^^., .^^, .^., .^ (.==.) :: (Columnable a, Eq a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool @@ -217,12 +300,220 @@ (Columnable a, Ord a) => TExpr cols a -> TExpr cols a -> TExpr cols Bool (.>.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (>) "gt" (Just ">") False 4) a b) +-- Same-type arithmetic operators + +(.+.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.+.) = (+) + +(.-.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.-.) = (-) + +(.*.) :: (Columnable a, Num a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(.*.) = (*) + +(./.) :: + (Columnable a, Fractional a) => TExpr cols a -> TExpr cols a -> TExpr cols a +(./.) = (/) + +-- Same-type exponentiation operators + +(.^^.) :: + (Columnable a, Columnable b, Fractional a, Integral b) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^^) "pow" (Just ".^^.") False 8) a b) + +(.^.) :: + (Columnable a, Columnable b, Num a, Integral b) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (^) "pow" (Just ".^.") False 8) a b) + (.&&.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just "&&") True 3) a b) +(.&&.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (&&) "and" (Just ".&&.") True 3) a b) (.||.) :: TExpr cols Bool -> TExpr cols Bool -> TExpr cols Bool -(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just "||") True 2) a b) +(.||.) (TExpr a) (TExpr b) = TExpr (Binary (MkBinaryOp (||) "or" (Just ".||.") True 2) a b) +-- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable. +(.&&) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.&&) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (nullCmpOp (&&)) "nulland" (Just ".&&") True 3) a b) + +-- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable. +(.||) :: + (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (NullCmpResult a b) +(.||) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (nullCmpOp (||)) "nullor" (Just ".||") True 2) a b) + +------------------------------------------------------------------------------- +-- Nullable-aware arithmetic operators +------------------------------------------------------------------------------- + +infixl 6 .+, .- +infixl 7 .*, ./ + +{- | Nullable-aware addition. Works for all combinations of nullable\/non-nullable operands. +@col \@\"x\" '.+' col \@\"y\" -- :: TExpr cols (Maybe Int) when y :: Maybe Int@ +-} +(.+) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.+) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (+))) "nulladd" (Just "+") True 6) + a + b + ) + +-- | Nullable-aware subtraction. +(.-) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.-) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (-))) "nullsub" (Just "-") False 6) + a + b + ) + +-- | Nullable-aware multiplication. +(.*) :: + ( NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (Promote (BaseType a) (BaseType b)) (WidenResult a b) + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResult a b) +(.*) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (widenArithOp (*))) "nullmul" (Just "*") True 7) + a + b + ) + +-- | Nullable-aware division. Integral operands are promoted to Double. +(./) :: + ( DivWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b) + , Fractional (PromoteDiv (BaseType a) (BaseType b)) + ) => + TExpr cols a -> + TExpr cols b -> + TExpr cols (WidenResultDiv a b) +(./) (TExpr a) (TExpr b) = + TExpr + ( Binary + (MkBinaryOp (applyNull2 (divArithOp (/))) "nulldiv" (Just "/") False 7) + a + b + ) + +-- | Nullable-aware exponentiation (fractional base, integral exponent). +(.^^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Fractional (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^^) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (applyNull2 (^^)) "pow" (Just ".^^") False 8) a b) + +-- | Nullable-aware exponentiation (num base, integral exponent). +(.^) :: + ( Columnable (BaseType a) + , Columnable (BaseType b) + , Num (BaseType a) + , Integral (BaseType b) + , NumericWidenOp (BaseType a) (BaseType b) + , NullLift2Op a b (BaseType a) a + , Num (Promote (BaseType a) (BaseType b)) + ) => + TExpr cols a -> TExpr cols b -> TExpr cols a +(.^) (TExpr a) (TExpr b) = + TExpr (Binary (MkBinaryOp (applyNull2 (^)) "pow" (Just ".^") False 8) a b) + +------------------------------------------------------------------------------- +-- Nullable-aware comparison operators (three-valued logic) +------------------------------------------------------------------------------- + +-- | Nullable-aware equality. Returns @Maybe Bool@ when either operand is nullable. +(.==) :: + (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + 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) + +-- | Nullable-aware inequality. +(./=) :: + (NullableCmpOp a b (NullCmpResult a b), Eq (BaseType a)) => + 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) + +-- | Nullable-aware less-than. +(.<) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + 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) + +-- | Nullable-aware less-than-or-equal. +(.<=) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + 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) + +-- | Nullable-aware greater-than-or-equal. +(.>=) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + 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) + +-- | Nullable-aware greater-than. +(.>) :: + (NullableCmpOp a b (NullCmpResult a b), Ord (BaseType a)) => + 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) + not :: TExpr cols Bool -> TExpr cols Bool not (TExpr e) = TExpr (Unary (MkUnaryOp Prelude.not "not" (Just "!")) e) @@ -258,6 +549,50 @@ collect :: (Columnable a) => TExpr cols a -> TExpr cols [a] collect (TExpr e) = TExpr (Agg (FoldAgg "collect" (Just []) (flip (:))) e) + +------------------------------------------------------------------------------- +-- Cast / coercion expressions +------------------------------------------------------------------------------- + +castExpr :: + forall b cols src. + (Columnable b, Columnable src) => TExpr cols src -> TExpr cols (Maybe b) +castExpr (TExpr e) = + TExpr + (CastExprWith @b @(Maybe b) @src "castExpr" (either (const Nothing) Just) e) + +castExprWithDefault :: + forall b cols src. + (Columnable b, Columnable src) => b -> TExpr cols src -> TExpr cols b +castExprWithDefault def (TExpr e) = + TExpr + ( CastExprWith @b @b @src + ("castExprWithDefault:" <> T.pack (show def)) + (fromRight def) + e + ) + +castExprEither :: + forall b cols src. + (Columnable b, Columnable src) => TExpr cols src -> TExpr cols (Either T.Text b) +castExprEither (TExpr e) = + TExpr + ( CastExprWith @b @(Either T.Text b) @src + "castExprEither" + (either (Left . T.pack) Right) + e + ) + +unsafeCastExpr :: + forall b cols src. + (Columnable b, Columnable src) => TExpr cols src -> TExpr cols b +unsafeCastExpr (TExpr e) = + TExpr + ( CastExprWith @b @b @src + "unsafeCastExpr" + (fromRight (error "unsafeCastExpr: unexpected Nothing in column")) + e + ) ------------------------------------------------------------------------------- -- Named expression helper diff --git a/tests/DecisionTree.hs b/tests/DecisionTree.hs new file mode 100644 --- /dev/null +++ b/tests/DecisionTree.hs @@ -0,0 +1,764 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module DecisionTree where + +import qualified DataFrame as D +import DataFrame.DecisionTree +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Operators + +import Data.Function (on) +import Data.List (maximumBy, sort) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified Data.Vector as V +import Test.HUnit + +------------------------------------------------------------------------ +-- Shared fixtures +------------------------------------------------------------------------ + +-- 4 rows: label = ["A","B","A","C"], x = [1.0,2.0,3.0,4.0] +fixtureDF :: D.DataFrame +fixtureDF = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "B", "A", "C"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + +allIndices :: V.Vector Int +allIndices = V.fromList [0, 1, 2, 3] + +leftTree :: Tree T.Text +leftTree = Leaf "A" + +rightTree :: Tree T.Text +rightTree = Leaf "B" + +-- x <= 2.5: True for idx 0,1 (→ left); False for idx 2,3 (→ right) +splitCond :: Expr Bool +splitCond = F.col @Double "x" .<= F.lit (2.5 :: Double) + +-- Pre-computed care points for the full fixture +carePoints3 :: [CarePoint] +carePoints3 = + identifyCarePoints @T.Text "label" fixtureDF allIndices leftTree rightTree + +------------------------------------------------------------------------ +-- Unit tests: identifyCarePoints +------------------------------------------------------------------------ + +carePointsBothWrong :: Test +carePointsBothWrong = + TestCase $ + assertBool + "idx 3 (label=C, neither A nor B) should not be a care point" + (3 `notElem` map cpIndex carePoints3) + +carePointsLeftCorrect :: Test +carePointsLeftCorrect = TestCase $ do + let cp0 = filter ((== 0) . cpIndex) carePoints3 + assertBool "idx 0 should be a care point" (not (null cp0)) + assertEqual + "idx 0 (label=A matches left Leaf A) should route GoLeft" + GoLeft + (cpCorrectDir (head cp0)) + +carePointsRightCorrect :: Test +carePointsRightCorrect = TestCase $ do + let cp1 = filter ((== 1) . cpIndex) carePoints3 + assertBool "idx 1 should be a care point" (not (null cp1)) + assertEqual + "idx 1 (label=B matches right Leaf B) should route GoRight" + GoRight + (cpCorrectDir (head cp1)) + +carePointsMixed :: Test +carePointsMixed = TestCase $ do + assertEqual "exactly 3 care points" 3 (length carePoints3) + let idxs = map cpIndex carePoints3 + assertBool "idx 0 present" (0 `elem` idxs) + assertBool "idx 1 present" (1 `elem` idxs) + assertBool "idx 2 present" (2 `elem` idxs) + assertBool "idx 3 absent" (3 `notElem` idxs) + +carePointsBothCorrect :: Test +carePointsBothCorrect = TestCase $ do + let df2 = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0] :: [Double])) + ] + cps = + identifyCarePoints @T.Text + "label" + df2 + (V.fromList [0, 1]) + (Leaf "A") + (Leaf "A") + assertEqual "no care points when both subtrees agree" 0 (length cps) + +------------------------------------------------------------------------ +-- Unit tests: majorityValueFromIndices +------------------------------------------------------------------------ + +majorityVoteTest :: Test +majorityVoteTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["cat", "dog", "cat", "cat"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + assertEqual + "majority is cat (3 votes)" + "cat" + (majorityValueFromIndices @T.Text "label" df (V.fromList [0, 1, 2, 3])) + +majorityVoteSubset :: Test +majorityVoteSubset = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["cat", "dog", "cat", "cat"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + -- indices [0,1,3]: cat×2, dog×1 → "cat" wins clearly + result = majorityValueFromIndices @T.Text "label" df (V.fromList [0, 1, 3]) + assertEqual "majority from subset [0,1,3] is cat" "cat" result + +------------------------------------------------------------------------ +-- Unit tests: computeTreeLoss +------------------------------------------------------------------------ + +computeLossZero :: Test +computeLossZero = TestCase $ do + -- target = ["A","A","B","B"]: perfect split on x <= 2.5 + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + loss = computeTreeLoss @T.Text "label" df (V.fromList [0, 1, 2, 3]) stump + assertEqual "perfect stump has zero loss" 0.0 loss + +computeLossHalf :: Test +computeLossHalf = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + constTree = Leaf "A" :: Tree T.Text + loss = computeTreeLoss @T.Text "label" df (V.fromList [0, 1, 2, 3]) constTree + assertEqual "constant leaf misclassifies half of balanced data" 0.5 loss + +------------------------------------------------------------------------ +-- Unit tests: partitionIndices +------------------------------------------------------------------------ + +partitionDisjoint :: Test +partitionDisjoint = TestCase $ do + let (lft, rgt) = partitionIndices splitCond fixtureDF allIndices + leftSet = V.toList lft + rightSet = V.toList rgt + intersection = filter (`elem` rightSet) leftSet + assertEqual "left and right partitions are disjoint" [] intersection + +partitionUnion :: Test +partitionUnion = TestCase $ do + let (lft, rgt) = partitionIndices splitCond fixtureDF allIndices + combined = sort (V.toList lft ++ V.toList rgt) + assertEqual + "union of partitions equals the original index set" + [0, 1, 2, 3] + combined + +------------------------------------------------------------------------ +-- Unit tests: countCarePointErrors +------------------------------------------------------------------------ + +countErrorsAllCorrect :: Test +countErrorsAllCorrect = TestCase $ do + -- x <= 1.5: idx 0 goes left (True), idx 1 goes right (False) + -- CarePoint 0 GoLeft → goesLeft=True, shouldGoLeft=True → correct + -- CarePoint 1 GoRight → goesLeft=False, shouldGoLeft=False → correct + let cps = [CarePoint 0 GoLeft, CarePoint 1 GoRight] + cond = F.col @Double "x" .<= F.lit (1.5 :: Double) + errs = countCarePointErrors cond fixtureDF cps + assertEqual "condition routes all care points correctly" 0 errs + +countErrorsAllWrong :: Test +countErrorsAllWrong = TestCase $ do + -- x > 1.5: idx 0 goes right (False), idx 1 goes left (True) — both wrong + -- CarePoint 0 GoLeft → goesLeft=False, shouldGoLeft=True → wrong + -- CarePoint 1 GoRight → goesLeft=True, shouldGoLeft=False → wrong + let cps = [CarePoint 0 GoLeft, CarePoint 1 GoRight] + cond = F.col @Double "x" .> F.lit (1.5 :: Double) + errs = countCarePointErrors cond fixtureDF cps + assertEqual "reversed condition misroutes all care points" 2 errs + +------------------------------------------------------------------------ +-- Unit tests: predictWithTree +------------------------------------------------------------------------ + +predictLeaf :: Test +predictLeaf = + TestCase $ + assertEqual + "leaf prediction ignores row index" + "Z" + (predictWithTree @T.Text "label" fixtureDF 0 (Leaf "Z")) + +predictBranch :: Test +predictBranch = TestCase $ do + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + assertEqual + "idx 0 (x=1.0 <= 2.5) routes left -> A" + "A" + (predictWithTree @T.Text "label" fixtureDF 0 stump) + assertEqual + "idx 3 (x=4.0 > 2.5) routes right -> B" + "B" + (predictWithTree @T.Text "label" fixtureDF 3 stump) + +------------------------------------------------------------------------ +-- Integration tests +------------------------------------------------------------------------ + +-- 20-row, linearly separable: x in [1..10] -> "pos", x in [11..20] -> "neg" +sepDF :: D.DataFrame +sepDF = + let xs = map fromIntegral [1 .. 20 :: Int] :: [Double] + labels = map (\x -> if x <= 10.0 then "pos" else "neg") xs :: [T.Text] + in D.fromNamedColumns + [ ("label", DI.fromList labels) + , ("x", DI.fromList xs) + ] + +-- Candidate conditions that bracket the decision boundary +sepConds :: [Expr Bool] +sepConds = + [ F.col @Double "x" .<= F.lit (10.5 :: Double) + , F.col @Double "x" .> F.lit (10.5 :: Double) + ] + +testCfg :: TreeConfig +testCfg = + defaultTreeConfig + { taoIterations = 5 + , expressionPairs = 4 + , minLeafSize = 1 + } + +-- Initial tree deliberately wrong: routes "pos" rows to the "neg" leaf +wrongStump :: Tree T.Text +wrongStump = + Branch + (F.col @Double "x" .> F.lit (10.5 :: Double)) + (Leaf "pos") + (Leaf "neg") + +taoNoDegradation :: Test +taoNoDegradation = TestCase $ do + let indices = V.enumFromN 0 20 + initialLoss = computeTreeLoss @T.Text "label" sepDF indices wrongStump + optimized = + taoOptimize @T.Text testCfg "label" sepConds sepDF indices wrongStump + finalLoss = computeTreeLoss @T.Text "label" sepDF indices optimized + assertBool + "taoOptimize must not increase loss" + (finalLoss <= initialLoss + 1e-9) + +taoMonotone :: Test +taoMonotone = TestCase $ do + let indices = V.enumFromN 0 20 + initLoss = computeTreeLoss @T.Text "label" sepDF indices wrongStump + stepTree = taoIteration @T.Text testCfg "label" sepConds sepDF indices + -- Track (tree, loss) pairs: take 6 snapshots (initial + 5 steps) + step (tree, _) = + let tree' = stepTree tree + in (tree', computeTreeLoss @T.Text "label" sepDF indices tree') + snapshots = take 6 $ iterate step (wrongStump, initLoss) + losses = map snd snapshots + pairs = zip losses (tail losses) + assertBool + "loss must be non-increasing across taoIteration steps" + (all (\(a, b) -> b <= a + 1e-9) pairs) + +taoConvergesPureLabels :: Test +taoConvergesPureLabels = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (replicate 10 ("A" :: T.Text))) + , ("x", DI.fromList ([1.0 .. 10.0] :: [Double])) + ] + indices = V.enumFromN 0 10 + initTree = Leaf "A" :: Tree T.Text + initLoss = computeTreeLoss @T.Text "label" df indices initTree + result = + taoOptimize @T.Text testCfg "label" sepConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual "pure-label initial loss must be zero" 0.0 initLoss + assertEqual "pure-label final loss must still be zero" 0.0 finalLoss + +taoDeadBranchNoCrash :: Test +taoDeadBranchNoCrash = TestCase $ do + -- Threshold below all x values: x <= 0.5 is False for every row + -- → all indices route to the right child; left partition is always empty + let badCond = F.col @Double "x" .<= F.lit (0.5 :: Double) + indices = V.enumFromN 0 20 + initTree = Branch badCond (Leaf "pos") (Leaf "neg") :: Tree T.Text + result = + taoOptimize @T.Text testCfg "label" [badCond] sepDF indices initTree + finalLoss = computeTreeLoss @T.Text "label" sepDF indices result + assertBool + "dead-branch tree must produce a valid loss in [0,1]" + (finalLoss >= 0.0 && finalLoss <= 1.0) + +------------------------------------------------------------------------ +-- Shared fixtures: 4x4 grid +------------------------------------------------------------------------ + +gridPairs :: [(Double, Double)] +gridPairs = [(x, y) | y <- [1 .. 4], x <- [1 .. 4]] + +gridBaseDF :: D.DataFrame +gridBaseDF = + D.fromNamedColumns + [ ("x", DI.fromList (map fst gridPairs)) + , ("y", DI.fromList (map snd gridPairs)) + ] + +------------------------------------------------------------------------ +-- Oblique recovery tests +------------------------------------------------------------------------ + +taoRecoversSingleObliqueDerived :: Test +taoRecoversSingleObliqueDerived = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("pos" :: T.Text)) + (F.lit ("neg" :: T.Text)) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + initTree = + Branch + (F.col @Double "x" .<= F.lit (2.5 :: Double)) + (Leaf "pos") + (Leaf "neg") :: + Tree T.Text + conds = + [ (F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double) + , (F.col @Double "x" + F.col @Double "y") .> F.lit (4.5 :: Double) + ] + cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + result = taoOptimize @T.Text cfg "label" conds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "TAO recovers single oblique (x+y) split with zero loss" + 0.0 + finalLoss + +taoRecoversNestedObliqueDerived :: Test +taoRecoversNestedObliqueDerived = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("low" :: T.Text)) + ( F.ifThenElse + ((F.col @Double "x" - F.col @Double "y") .<= F.lit (0.5 :: Double)) + (F.lit "mid") + (F.lit "high") + ) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + initTree = + Branch + (F.col @Double "x" .<= F.lit (1.5 :: Double)) + (Leaf "low") + ( Branch + (F.col @Double "y" .<= F.lit (3.5 :: Double)) + (Leaf "mid") + (Leaf "high") + ) :: + Tree T.Text + conds = + [ (F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double) + , (F.col @Double "x" + F.col @Double "y") .> F.lit (4.5 :: Double) + , (F.col @Double "x" - F.col @Double "y") .<= F.lit (0.5 :: Double) + , (F.col @Double "x" - F.col @Double "y") .> F.lit (0.5 :: Double) + ] + cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + result = taoOptimize @T.Text cfg "label" conds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "TAO recovers nested oblique (x+y)/(x-y) tree with zero loss" + 0.0 + finalLoss + +taoAxisAlignedInsufficientForOblique :: Test +taoAxisAlignedInsufficientForOblique = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("pos" :: T.Text)) + (F.lit ("neg" :: T.Text)) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + axisConds = + [F.col @Double "x" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + ++ [F.col @Double "y" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + initTree = + Branch + (F.col @Double "x" .<= F.lit (2.5 :: Double)) + (Leaf "pos") + (Leaf "neg") :: + Tree T.Text + cfg = defaultTreeConfig{taoIterations = 10, expressionPairs = 6, minLeafSize = 1} + result = taoOptimize @T.Text cfg "label" axisConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertBool + "axis-aligned stump cannot recover oblique label (loss must remain > 0.1)" + (finalLoss > 0.1) + +------------------------------------------------------------------------ +-- Nullable numeric feature tests +------------------------------------------------------------------------ + +-- Cleanly separable: Just 1..6 -> "pos", Just 7..12 -> "neg", no nulls. +-- Uses OptionalColumn directly to exercise the new nullable numeric path. +nullableSepDF :: D.DataFrame +nullableSepDF = + D.fromNamedColumns + [ ("label", DI.fromList (replicate 6 "pos" ++ replicate 6 "neg" :: [T.Text])) + , + ( "x" + , DI.OptionalColumn + ( V.fromList $ + map (Just . fromIntegral) ([1 .. 6] :: [Int]) + ++ map (Just . fromIntegral) ([7 .. 12] :: [Int]) :: + V.Vector (Maybe Double) + ) + ) + ] + +-- DF with genuine nulls interspersed. +nullsMixedDF :: D.DataFrame +nullsMixedDF = + D.fromNamedColumns + [ ("label", DI.fromList (["pos", "pos", "pos", "neg", "neg", "neg"] :: [T.Text])) + , + ( "x" + , DI.OptionalColumn + ( V.fromList + [Just 1.0, Nothing, Just 3.0, Just 7.0, Nothing, Just 9.0] :: + V.Vector (Maybe Double) + ) + ) + ] + +-- numericCols picks up OptionalColumn (Maybe Double) as NMaybeDouble. +numericColsNullableDoubleTest :: Test +numericColsNullableDoubleTest = TestCase $ do + let exprs = numericCols nullableSepDF + hasMD = any (\case NMaybeDouble _ -> True; _ -> False) exprs + assertBool + "numericCols finds NMaybeDouble for OptionalColumn (Maybe Double)" + hasMD + +-- numericCols picks up OptionalColumn (Maybe Int) as NMaybeDouble (via whenPresent). +numericColsNullableIntTest :: Test +numericColsNullableIntTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["pos", "neg"] :: [T.Text])) + , + ( "n" + , DI.OptionalColumn (V.fromList [Just (1 :: Int), Just 2] :: V.Vector (Maybe Int)) + ) + ] + hasMD = any (\case NMaybeDouble _ -> True; _ -> False) (numericCols df) + assertBool "numericCols finds NMaybeDouble for OptionalColumn (Maybe Int)" hasMD + +-- generateNumericConds is non-empty for a DF with an OptionalColumn (Maybe Double). +numericCondsNullableNonEmptyTest :: Test +numericCondsNullableNonEmptyTest = + TestCase $ + assertBool + "generateNumericConds non-empty for OptionalColumn (Maybe Double)" + (not (null (generateNumericConds defaultTreeConfig nullableSepDF))) + +-- Null values evaluate to False for threshold conditions (null rows route right). +nullValueRoutesFalseTest :: Test +nullValueRoutesFalseTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "B"] :: [T.Text])) + , + ( "x" + , DI.OptionalColumn + (V.fromList [Nothing, Just (5.0 :: Double)] :: V.Vector (Maybe Double)) + ) + ] + -- Nothing <= 6.0 = Nothing -> fromMaybe False = False -> right + -- Just 5.0 <= 6.0 = Just True -> fromMaybe False = True -> left + cond = F.fromMaybe False (F.col @(Maybe Double) "x" .<= F.lit (6.0 :: Double)) + (lft, rgt) = partitionIndices cond df (V.fromList [0, 1]) + assertBool "null row (idx 0) routes to right (false) partition" (0 `V.elem` rgt) + assertBool "Just 5.0 <= 6.0 routes to left (true) partition" (1 `V.elem` lft) + +-- fitDecisionTree with an OptionalColumn nullable feature achieves zero loss +-- on cleanly separable data (no actual nulls in the column). +nullableFitZeroLossTest :: Test +nullableFitZeroLossTest = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + featureDf = D.exclude ["label"] nullableSepDF + conds = generateNumericConds cfg featureDf + initTree = buildGreedyTree @T.Text cfg (maxTreeDepth cfg) "label" conds nullableSepDF + indices = V.enumFromN 0 12 + result = taoOptimize @T.Text cfg "label" conds nullableSepDF indices initTree + loss = computeTreeLoss @T.Text "label" nullableSepDF indices result + assertEqual "zero loss on cleanly separable OptionalColumn data" 0.0 loss + +-- fitDecisionTree with genuine nulls: loss is a valid probability and no crash. +nullableFitWithNullsNoCrashTest :: Test +nullableFitWithNullsNoCrashTest = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 3, expressionPairs = 4, minLeafSize = 1} + featureDf = D.exclude ["label"] nullsMixedDF + conds = generateNumericConds cfg featureDf + initTree = buildGreedyTree @T.Text cfg (maxTreeDepth cfg) "label" conds nullsMixedDF + indices = V.enumFromN 0 6 + result = taoOptimize @T.Text cfg "label" conds nullsMixedDF indices initTree + loss = computeTreeLoss @T.Text "label" nullsMixedDF indices result + assertBool + "loss is in [0,1] with null values present" + (loss >= 0.0 && loss <= 1.0) + +-- numericExprsWithTerms produces cross-column combinations when one col is +-- OptionalColumn (Maybe Double) and another is a plain UnboxedColumn Double. +numericExprsWithTermsMixedTest :: Test +numericExprsWithTermsMixedTest = TestCase $ do + let df = + D.fromNamedColumns + [ + ( "x" + , DI.OptionalColumn + (V.fromList [Just 1.0, Just 2.0, Just 3.0] :: V.Vector (Maybe Double)) + ) + , ("y", DI.fromList ([4.0, 5.0, 6.0] :: [Double])) + ] + cfg = defaultSynthConfig{maxExprDepth = 2, enableArithOps = True} + exprs = numericExprsWithTerms cfg df + assertBool + "more than 2 expressions: base cols + combinations" + (length exprs > 2) + assertBool + "combined exprs include NMaybeDouble (nullable arithmetic)" + (any (\case NMaybeDouble _ -> True; _ -> False) exprs) + +------------------------------------------------------------------------ +-- Probability tree tests +------------------------------------------------------------------------ + +-- probsFromIndices: counts correct on a 3-row slice +probsFromIndicesBasic :: Test +probsFromIndicesBasic = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0] :: [Double])) + ] + probs = probsFromIndices @T.Text "label" df (V.fromList [0, 1, 2]) + assertBool "A prob ≈ 2/3" (abs (probs M.! "A" - 2 / 3) < 1e-9) + assertBool "B prob ≈ 1/3" (abs (probs M.! "B" - 1 / 3) < 1e-9) + +-- probsFromIndices: only a subset of rows counted +probsFromIndicesSubset :: Test +probsFromIndicesSubset = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + probs = probsFromIndices @T.Text "label" df (V.fromList [0, 1]) + assertEqual "only rows 0,1 → A:1.0" (M.fromList [("A", 1.0)]) probs + +-- probsFromIndices: single class → probability 1.0 +probsFromIndicesSingleClass :: Test +probsFromIndicesSingleClass = TestCase $ do + let probs = probsFromIndices @T.Text "label" fixtureDF (V.fromList [0, 2]) + assertEqual "rows 0,2 both A → A:1.0" (M.fromList [("A", 1.0)]) probs + +-- buildProbTree: Leaf preserves distribution +buildProbTreeLeaf :: Test +buildProbTreeLeaf = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "A"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0] :: [Double])) + ] + pt = buildProbTree @T.Text (Leaf "A") "label" df (V.fromList [0, 1, 2]) + case pt of + Leaf m -> assertEqual "pure-A leaf → {A:1.0}" (M.fromList [("A", 1.0)]) m + _ -> assertFailure "expected Leaf" + +-- buildProbTree: Branch distributes rows to left/right leaves correctly +buildProbTreeBranch :: Test +buildProbTreeBranch = TestCase $ do + -- splitCond: x <= 2.5 → idx 0,1 go left; idx 2,3 go right + -- left (idx 0,1): labels ["A","B"] → {A:0.5, B:0.5} + -- right (idx 2,3): labels ["A","C"] → {A:0.5, C:0.5} + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + pt = buildProbTree @T.Text stump "label" fixtureDF allIndices + case pt of + Branch _ (Leaf lm) (Leaf rm) -> do + assertBool "left leaf has A:0.5" (abs (M.findWithDefault 0 "A" lm - 0.5) < 1e-9) + assertBool "left leaf has B:0.5" (abs (M.findWithDefault 0 "B" lm - 0.5) < 1e-9) + assertBool + "right leaf has A:0.5" + (abs (M.findWithDefault 0 "A" rm - 0.5) < 1e-9) + assertBool + "right leaf has C:0.5" + (abs (M.findWithDefault 0 "C" rm - 0.5) < 1e-9) + _ -> assertFailure "expected Branch with two Leaves" + +-- probExprs: leaf tree produces Lit values +probExprsLeaf :: Test +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") + +-- probExprs: class absent from one leaf gets Lit 0.0 on that side +probExprsMissingClass :: Test +probExprsMissingClass = TestCase $ do + let pt = + Branch + splitCond + (Leaf (M.fromList [("A", 1.0)])) + (Leaf (M.fromList [("B", 1.0)])) :: + ProbTree T.Text + pe = probExprs pt + assertEqual + "A expr: If cond (Lit 1.0) (Lit 0.0)" + (F.ifThenElse splitCond (Lit 1.0) (Lit 0.0)) + (pe M.! "A") + assertEqual + "B expr: If cond (Lit 0.0) (Lit 1.0)" + (F.ifThenElse splitCond (Lit 0.0) (Lit 1.0)) + (pe M.! "B") + +-- probExprs: keys equal all classes that appear across any leaf +probExprsAllClasses :: Test +probExprsAllClasses = TestCase $ do + let pt = + Branch + splitCond + (Leaf (M.fromList [("A", 1.0)])) + (Leaf (M.fromList [("B", 0.6), ("C", 0.4)])) :: + ProbTree T.Text + pe = probExprs pt + assertEqual "three classes in result" (sort ["A", "B", "C"]) (sort (M.keys pe)) + +-- Probabilities sum to 1.0 at every row after applying probExprs +probsSumToOne :: Test +probsSumToOne = TestCase $ do + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + pt = buildProbTree @T.Text stump "label" fixtureDF allIndices + pe = probExprs pt + sumExpr = foldl1 (.+) (M.elems pe) + case interpret @Double fixtureDF sumExpr of + Left e -> assertFailure (show e) + Right (DI.TColumn col) -> + case DI.toVector @Double col of + Left e -> assertFailure (show e) + Right vals -> + mapM_ + (\v -> assertBool ("sum ≈ 1.0, got " ++ show v) (abs (v - 1.0) < 1e-9)) + (V.toList vals) + +-- argmax of probExprs agrees with fitDecisionTree on sepDF +probArgmaxMatchesClassifier :: Test +probArgmaxMatchesClassifier = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + hardExpr = fitDecisionTree @T.Text cfg (Col "label") sepDF + pe = fitProbTree @T.Text cfg (Col "label") sepDF + indices = [0 .. D.nRows sepDF - 1] + case interpret @T.Text sepDF hardExpr of + Left e -> assertFailure (show e) + Right (DI.TColumn hardCol) -> + case DI.toVector @T.Text hardCol of + Left e -> assertFailure (show e) + Right hardVals -> do + probCols <- + mapM + ( \(cls, expr) -> case interpret @Double sepDF expr of + Left e -> assertFailure (show e) >> return (cls, V.empty) + Right (DI.TColumn col) -> case DI.toVector @Double col of + Left e -> assertFailure (show e) >> return (cls, V.empty) + Right v -> return (cls, v) + ) + (M.toList pe) + mapM_ + ( \i -> + let argmax = fst $ maximumBy (compare `on` (V.! i) . snd) probCols + hard = hardVals V.! i + in assertEqual ("row " ++ show i) hard argmax + ) + indices + +------------------------------------------------------------------------ +-- Test list +------------------------------------------------------------------------ + +tests :: [Test] +tests = + [ TestLabel "carePointsBothWrong" carePointsBothWrong + , TestLabel "carePointsLeftCorrect" carePointsLeftCorrect + , TestLabel "carePointsRightCorrect" carePointsRightCorrect + , TestLabel "carePointsMixed" carePointsMixed + , TestLabel "carePointsBothCorrect" carePointsBothCorrect + , TestLabel "majorityVoteTest" majorityVoteTest + , TestLabel "majorityVoteSubset" majorityVoteSubset + , TestLabel "computeLossZero" computeLossZero + , TestLabel "computeLossHalf" computeLossHalf + , TestLabel "partitionDisjoint" partitionDisjoint + , TestLabel "partitionUnion" partitionUnion + , TestLabel "countErrorsAllCorrect" countErrorsAllCorrect + , TestLabel "countErrorsAllWrong" countErrorsAllWrong + , TestLabel "predictLeaf" predictLeaf + , TestLabel "predictBranch" predictBranch + , TestLabel "taoNoDegradation" taoNoDegradation + , TestLabel "taoMonotone" taoMonotone + , TestLabel "taoConvergesPureLabels" taoConvergesPureLabels + , TestLabel "taoDeadBranchNoCrash" taoDeadBranchNoCrash + , TestLabel "taoRecoversSingleObliqueDerived" taoRecoversSingleObliqueDerived + , TestLabel "taoRecoversNestedObliqueDerived" taoRecoversNestedObliqueDerived + , TestLabel + "taoAxisAlignedInsufficientForOblique" + taoAxisAlignedInsufficientForOblique + , TestLabel "numericColsNullableDouble" numericColsNullableDoubleTest + , TestLabel "numericColsNullableInt" numericColsNullableIntTest + , TestLabel "numericCondsNullableNonEmpty" numericCondsNullableNonEmptyTest + , TestLabel "nullValueRoutesFalse" nullValueRoutesFalseTest + , TestLabel "nullableFitZeroLoss" nullableFitZeroLossTest + , TestLabel "nullableFitWithNullsNoCrash" nullableFitWithNullsNoCrashTest + , TestLabel "numericExprsWithTermsMixed" numericExprsWithTermsMixedTest + , TestLabel "probsFromIndicesBasic" probsFromIndicesBasic + , TestLabel "probsFromIndicesSubset" probsFromIndicesSubset + , TestLabel "probsFromIndicesSingleClass" probsFromIndicesSingleClass + , TestLabel "buildProbTreeLeaf" buildProbTreeLeaf + , TestLabel "buildProbTreeBranch" buildProbTreeBranch + , TestLabel "probExprsLeaf" probExprsLeaf + , TestLabel "probExprsMissingClass" probExprsMissingClass + , TestLabel "probExprsAllClasses" probExprsAllClasses + , TestLabel "probsSumToOne" probsSumToOne + , TestLabel "probArgmaxMatchesClassifier" probArgmaxMatchesClassifier + ] diff --git a/tests/LazyParquet.hs b/tests/LazyParquet.hs --- a/tests/LazyParquet.hs +++ b/tests/LazyParquet.hs @@ -110,7 +110,7 @@ ( do actual <- L.runDataFrame - ( L.limit + ( L.take 3 ( L.sortBy [("id", Ascending)] diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -9,6 +9,7 @@ import Test.HUnit import Test.QuickCheck +import qualified DecisionTree import qualified Functions import qualified IO.JSON import qualified Internal.Parsing @@ -23,6 +24,8 @@ import qualified Operations.InsertColumn import qualified Operations.Join import qualified Operations.Merge +import qualified Operations.Nullable +import qualified Operations.Provenance import qualified Operations.ReadCsv import qualified Operations.Shuffle import qualified Operations.Sort @@ -36,7 +39,8 @@ tests :: Test tests = TestList $ - Internal.Parsing.tests + DecisionTree.tests + ++ Internal.Parsing.tests ++ Operations.Aggregations.tests ++ Operations.Apply.tests ++ Operations.Core.tests @@ -46,10 +50,13 @@ ++ Operations.InsertColumn.tests ++ Operations.Join.tests ++ Operations.Merge.tests + ++ Operations.Nullable.tests + ++ Operations.Provenance.tests ++ Operations.ReadCsv.tests ++ Operations.Shuffle.tests ++ Operations.Sort.tests ++ Operations.Statistics.tests + ++ Operations.Subset.hunitTests ++ Operations.Take.tests ++ Operations.Typing.tests ++ Functions.tests diff --git a/tests/Monad.hs b/tests/Monad.hs --- a/tests/Monad.hs +++ b/tests/Monad.hs @@ -22,6 +22,8 @@ let finalRowCount = D.nRows finalDf let realRate = roundToTwoPlaces $ fromIntegral finalRowCount / fromIntegral rowCount let diff = abs $ expectedRate - realRate - assert (diff <= 0.11) + -- calculates the 99.99% confidence interval (quickcheck runs 100 tests, aim for 1/10000) + let tolerance = 3.89 * sqrt (expectedRate * (1 - expectedRate) / fromIntegral rowCount) + assert (diff <= tolerance) tests = [prop_sampleM] diff --git a/tests/Operations/Apply.hs b/tests/Operations/Apply.hs --- a/tests/Operations/Apply.hs +++ b/tests/Operations/Apply.hs @@ -13,6 +13,7 @@ import qualified DataFrame.Internal.Column as DI import qualified DataFrame.Internal.DataFrame as D import qualified DataFrame.Internal.DataFrame as DI +import DataFrame.Operations.Statistics (imputeWith) import DataFrame.Operations.Transformations (impute) import Assertions @@ -238,12 +239,30 @@ imputeOnNonOptional :: Test imputeOnNonOptional = TestCase - ( assertExpectException - "[Error Case]" - "Cannot impute to a non-Empty column: plain" - (print $ impute (F.col @(Maybe Int) "plain") 0 imputeData) + ( assertEqual + "impute is a no-op on a non-nullable column" + imputeData + (impute (F.col @(Maybe Int) "plain") 0 imputeData) ) +imputePlainNoOp :: Test +imputePlainNoOp = + TestCase + ( assertEqual + "impute with non-Maybe expr is always a no-op" + imputeData + (impute (F.col @Int "plain") 0 imputeData) + ) + +imputeWithPlainNoOp :: Test +imputeWithPlainNoOp = + TestCase + ( assertEqual + "imputeWith with non-Maybe expr is always a no-op" + imputeData + (imputeWith id (F.col @Int "plain") imputeData) + ) + tests :: [Test] tests = [ TestLabel "applyBoxedToUnboxed" applyBoxedToUnboxed @@ -263,4 +282,6 @@ , TestLabel "imputeHappyPath" imputeHappyPath , TestLabel "imputeColumnNotFound" imputeColumnNotFound , TestLabel "imputeOnNonOptional" imputeOnNonOptional + , TestLabel "imputePlainNoOp" imputePlainNoOp + , TestLabel "imputeWithPlainNoOp" imputeWithPlainNoOp ] diff --git a/tests/Operations/Nullable.hs b/tests/Operations/Nullable.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Nullable.hs @@ -0,0 +1,773 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module Operations.Nullable where + +import qualified Data.Vector as V +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as DI +import DataFrame.Internal.Expression (Expr) +import DataFrame.Operators ((.*), (.+), (.-), (./), (.==)) +import qualified DataFrame.Typed as DT +import qualified DataFrame.Typed.Expr as TE +import DataFrame.Typed.Types (TExpr (..)) + +import Test.HUnit + +-- --------------------------------------------------------------------------- +-- Untyped Expr layer tests +-- --------------------------------------------------------------------------- + +-- DataFrame with one plain Int column and one Maybe Int column +testData :: D.DataFrame +testData = + D.fromNamedColumns + [ ("x", DI.fromList ([1, 2, 3] :: [Int])) + , ("y", DI.OptionalColumn (V.fromList [Just 10, Nothing, Just 30 :: Maybe Int])) + ] + +-- | col @Int .+ col @(Maybe Int) should give Maybe Int column +addIntMaybeInt :: Test +addIntMaybeInt = + TestCase + ( assertEqual + "Int .+ Maybe Int = Maybe Int" + (Just $ DI.OptionalColumn (V.fromList [Just 11, Nothing, Just 33 :: Maybe Int])) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @Int "x" .+ F.col @(Maybe Int) "y") + testData + ) + ) + +-- | col @(Maybe Int) .+ col @Int should give Maybe Int column +addMaybeIntInt :: Test +addMaybeIntInt = + TestCase + ( assertEqual + "Maybe Int .+ Int = Maybe Int" + (Just $ DI.OptionalColumn (V.fromList [Just 11, Nothing, Just 33 :: Maybe Int])) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @(Maybe Int) "y" .+ F.col @Int "x") + testData + ) + ) + +-- | col @Int .+ col @Int (same-type non-nullable) should give Int column +addIntInt :: Test +addIntInt = + TestCase + ( assertEqual + "Int .+ Int = Int (non-nullable preserved)" + (Just $ DI.fromList [2, 4, 6 :: Int]) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @Int "x" .+ F.col @Int "x") + testData + ) + ) + +-- | col @(Maybe Int) .+ col @(Maybe Int) should give Maybe Int column +addMaybeMaybe :: Test +addMaybeMaybe = + TestCase + ( assertEqual + "Maybe Int .+ Maybe Int = Maybe Int" + (Just $ DI.OptionalColumn (V.fromList [Just 20, Nothing, Just 60 :: Maybe Int])) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @(Maybe Int) "y" .+ F.col @(Maybe Int) "y") + testData + ) + ) + +-- | Subtraction: Int .- Maybe Int +subIntMaybeInt :: Test +subIntMaybeInt = + TestCase + ( assertEqual + "Int .- Maybe Int = Maybe Int" + ( Just $ + DI.OptionalColumn (V.fromList [Just (-9), Nothing, Just (-27) :: Maybe Int]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @Int "x" .- F.col @(Maybe Int) "y") + testData + ) + ) + +-- | Comparison: Int .== Maybe Int gives Maybe Bool +eqIntMaybeInt :: Test +eqIntMaybeInt = + TestCase + ( assertEqual + "Int .== Maybe Int = Maybe Bool" + ( Just $ + DI.OptionalColumn (V.fromList [Just False, Nothing, Just False :: Maybe Bool]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @Int "x" .== F.col @(Maybe Int) "y" :: Expr (Maybe Bool)) + testData + ) + ) + +-- | Comparison: Int .== Int gives plain Bool (backwards compatible) +eqIntInt :: Test +eqIntInt = + TestCase + ( assertEqual + "Int .== Int = Bool (non-nullable, backwards compatible)" + (Just $ DI.fromList [True, False, False :: Bool]) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @Int "x" .== F.lit (1 :: Int) :: Expr Bool) + testData + ) + ) + +-- --------------------------------------------------------------------------- +-- nullLift / nullLift2 — untyped layer +-- --------------------------------------------------------------------------- + +-- | nullLift negate on Maybe Int propagates Nothing +nullLiftMaybeInt :: Test +nullLiftMaybeInt = + TestCase + ( assertEqual + "nullLift negate (Maybe Int) propagates Nothing" + ( Just $ + DI.OptionalColumn (V.fromList [Just (-10), Nothing, Just (-30) :: Maybe Int]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.nullLift negate (F.col @(Maybe Int) "y") :: Expr (Maybe Int)) + testData + ) + ) + +-- | nullLift negate on plain Int — no Nothing, plain result +nullLiftInt :: Test +nullLiftInt = + TestCase + ( assertEqual + "nullLift negate (Int) gives Int column" + (Just $ DI.fromList [-1, -2, -3 :: Int]) + ( DI.getColumn "result" $ + D.derive + "result" + (F.nullLift negate (F.col @Int "x") :: Expr Int) + testData + ) + ) + +-- | nullLift2 (+) Int (Maybe Int) → Maybe Int +nullLift2IntMaybeInt :: Test +nullLift2IntMaybeInt = + TestCase + ( assertEqual + "nullLift2 (+) Int (Maybe Int) = Maybe Int" + (Just $ DI.OptionalColumn (V.fromList [Just 11, Nothing, Just 33 :: Maybe Int])) + ( DI.getColumn "result" $ + D.derive + "result" + (F.nullLift2 (+) (F.col @Int "x") (F.col @(Maybe Int) "y") :: Expr (Maybe Int)) + testData + ) + ) + +-- | nullLift2 (+) (Maybe Int) Int → Maybe Int +nullLift2MaybeIntInt :: Test +nullLift2MaybeIntInt = + TestCase + ( assertEqual + "nullLift2 (+) (Maybe Int) Int = Maybe Int" + (Just $ DI.OptionalColumn (V.fromList [Just 11, Nothing, Just 33 :: Maybe Int])) + ( DI.getColumn "result" $ + D.derive + "result" + (F.nullLift2 (+) (F.col @(Maybe Int) "y") (F.col @Int "x") :: Expr (Maybe Int)) + testData + ) + ) + +-- | nullLift2 (+) Int Int → plain Int +nullLift2IntInt :: Test +nullLift2IntInt = + TestCase + ( assertEqual + "nullLift2 (+) Int Int = Int (non-nullable)" + (Just $ DI.fromList [2, 4, 6 :: Int]) + ( DI.getColumn "result" $ + D.derive + "result" + (F.nullLift2 (+) (F.col @Int "x") (F.col @Int "x") :: Expr Int) + testData + ) + ) + +-- --------------------------------------------------------------------------- +-- Cross-type arithmetic (Int × Double promotion) +-- --------------------------------------------------------------------------- + +-- DataFrame with Int, Maybe Int, Double, Maybe Double columns +crossData :: D.DataFrame +crossData = + D.fromNamedColumns + [ ("x", DI.fromList ([1, 2, 3] :: [Int])) + , ("y", DI.OptionalColumn (V.fromList [Just 10, Nothing, Just 30 :: Maybe Int])) + , ("d", DI.fromList ([1.5, 2.5, 3.5] :: [Double])) + , + ( "md" + , DI.OptionalColumn (V.fromList [Just 10.5, Nothing, Just 30.5 :: Maybe Double]) + ) + ] + +-- | col @Int .+ col @Double → Double +addIntDouble :: Test +addIntDouble = + TestCase + ( assertEqual + "Int .+ Double = Double" + (Just $ DI.fromList [2.5, 4.5, 6.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .+ F.col @Double "d") crossData + ) + ) + +-- | col @Double .+ col @Int → Double +addDoubleInt :: Test +addDoubleInt = + TestCase + ( assertEqual + "Double .+ Int = Double" + (Just $ DI.fromList [2.5, 4.5, 6.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Double "d" .+ F.col @Int "x") crossData + ) + ) + +-- | col @(Maybe Int) .+ col @Double → Maybe Double +addMaybeIntDouble :: Test +addMaybeIntDouble = + TestCase + ( assertEqual + "Maybe Int .+ Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 11.5, Nothing, Just 33.5 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @(Maybe Int) "y" .+ F.col @Double "d") crossData + ) + ) + +-- | col @Int .+ col @(Maybe Double) → Maybe Double +addIntMaybeDouble :: Test +addIntMaybeDouble = + TestCase + ( assertEqual + "Int .+ Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 11.5, Nothing, Just 33.5 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .+ F.col @(Maybe Double) "md") crossData + ) + ) + +-- | col @(Maybe Int) .+ col @(Maybe Double) → Maybe Double +addMaybeIntMaybeDouble :: Test +addMaybeIntMaybeDouble = + TestCase + ( assertEqual + "Maybe Int .+ Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 20.5, Nothing, Just 60.5 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @(Maybe Int) "y" .+ F.col @(Maybe Double) "md") + crossData + ) + ) + +-- | col @Int .- col @Double → Double +subIntDouble :: Test +subIntDouble = + TestCase + ( assertEqual + "Int .- Double = Double" + (Just $ DI.fromList [-0.5, -0.5, -0.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .- F.col @Double "d") crossData + ) + ) + +-- | col @Int .* col @Double → Double +mulIntDouble :: Test +mulIntDouble = + TestCase + ( assertEqual + "Int .* Double = Double" + (Just $ DI.fromList [1.5, 5.0, 10.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .* F.col @Double "d") crossData + ) + ) + +-- | col @Double .- col @Int → Double +subDoubleInt :: Test +subDoubleInt = + TestCase + ( assertEqual + "Double .- Int = Double" + (Just $ DI.fromList [0.5, 0.5, 0.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Double "d" .- F.col @Int "x") crossData + ) + ) + +-- | col @(Maybe Int) .- col @Double → Maybe Double +subMaybeIntDouble :: Test +subMaybeIntDouble = + TestCase + ( assertEqual + "Maybe Int .- Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 8.5, Nothing, Just 26.5 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @(Maybe Int) "y" .- F.col @Double "d") crossData + ) + ) + +-- | col @Int .- col @(Maybe Double) → Maybe Double +subIntMaybeDouble :: Test +subIntMaybeDouble = + TestCase + ( assertEqual + "Int .- Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn + (V.fromList [Just (-9.5), Nothing, Just (-27.5) :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .- F.col @(Maybe Double) "md") crossData + ) + ) + +-- | col @(Maybe Int) .- col @(Maybe Double) → Maybe Double +subMaybeIntMaybeDouble :: Test +subMaybeIntMaybeDouble = + TestCase + ( assertEqual + "Maybe Int .- Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn + (V.fromList [Just (-0.5), Nothing, Just (-0.5) :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @(Maybe Int) "y" .- F.col @(Maybe Double) "md") + crossData + ) + ) + +-- | col @Double .* col @Int → Double +mulDoubleInt :: Test +mulDoubleInt = + TestCase + ( assertEqual + "Double .* Int = Double" + (Just $ DI.fromList [1.5, 5.0, 10.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Double "d" .* F.col @Int "x") crossData + ) + ) + +-- | col @(Maybe Int) .* col @Double → Maybe Double +mulMaybeIntDouble :: Test +mulMaybeIntDouble = + TestCase + ( assertEqual + "Maybe Int .* Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 15.0, Nothing, Just 105.0 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @(Maybe Int) "y" .* F.col @Double "d") crossData + ) + ) + +-- | col @Int .* col @(Maybe Double) → Maybe Double +mulIntMaybeDouble :: Test +mulIntMaybeDouble = + TestCase + ( assertEqual + "Int .* Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 10.5, Nothing, Just 91.5 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" .* F.col @(Maybe Double) "md") crossData + ) + ) + +-- | col @(Maybe Int) .* col @(Maybe Double) → Maybe Double +mulMaybeIntMaybeDouble :: Test +mulMaybeIntMaybeDouble = + TestCase + ( assertEqual + "Maybe Int .* Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 105.0, Nothing, Just 915.0 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive + "result" + (F.col @(Maybe Int) "y" .* F.col @(Maybe Double) "md") + crossData + ) + ) + +-- Division tests use clean-dividing data: x=[2,4,6], d=[1,2,3] +divData :: D.DataFrame +divData = + D.fromNamedColumns + [ ("x", DI.fromList ([2, 4, 6] :: [Int])) + , ("y", DI.OptionalColumn (V.fromList [Just 4, Nothing, Just 6 :: Maybe Int])) + , ("d", DI.fromList ([1.0, 2.0, 3.0] :: [Double])) + , + ( "md" + , DI.OptionalColumn (V.fromList [Just 1.0, Nothing, Just 3.0 :: Maybe Double]) + ) + ] + +-- | col @Int ./ col @Double → Double +divIntDouble :: Test +divIntDouble = + TestCase + ( assertEqual + "Int ./ Double = Double" + (Just $ DI.fromList [2.0, 2.0, 2.0 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" ./ F.col @Double "d") divData + ) + ) + +-- | col @Double ./ col @Int → Double +divDoubleInt :: Test +divDoubleInt = + TestCase + ( assertEqual + "Double ./ Int = Double" + (Just $ DI.fromList [0.5, 0.5, 0.5 :: Double]) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Double "d" ./ F.col @Int "x") divData + ) + ) + +-- | col @(Maybe Int) ./ col @Double → Maybe Double +divMaybeIntDouble :: Test +divMaybeIntDouble = + TestCase + ( assertEqual + "Maybe Int ./ Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 4.0, Nothing, Just 2.0 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @(Maybe Int) "y" ./ F.col @Double "d") divData + ) + ) + +-- | col @Int ./ col @(Maybe Double) → Maybe Double +divIntMaybeDouble :: Test +divIntMaybeDouble = + TestCase + ( assertEqual + "Int ./ Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 2.0, Nothing, Just 2.0 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @Int "x" ./ F.col @(Maybe Double) "md") divData + ) + ) + +-- | col @(Maybe Int) ./ col @(Maybe Double) → Maybe Double +divMaybeIntMaybeDouble :: Test +divMaybeIntMaybeDouble = + TestCase + ( assertEqual + "Maybe Int ./ Maybe Double = Maybe Double" + ( Just $ + DI.OptionalColumn (V.fromList [Just 4.0, Nothing, Just 2.0 :: Maybe Double]) + ) + ( DI.getColumn "result" $ + D.derive "result" (F.col @(Maybe Int) "y" ./ F.col @(Maybe Double) "md") divData + ) + ) + +-- --------------------------------------------------------------------------- +-- Typed TExpr layer — cross-type arithmetic +-- --------------------------------------------------------------------------- + +type CrossSchema = + '[ DT.Column "x" Int + , DT.Column "y" (Maybe Int) + , DT.Column "d" Double + , DT.Column "md" (Maybe Double) + ] + +typedCrossData :: DT.TypedDataFrame CrossSchema +typedCrossData = + either (error . show) id $ + DT.freezeWithError @CrossSchema crossData + +-- | Typed: col @"x" .+ col @"d" → Double +typedAddIntDouble :: Test +typedAddIntDouble = + TestCase + ( assertEqual + "Typed: Int .+ Double = Double" + [2.5, 4.5, 6.5 :: Double] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"x" TE..+ TE.col @"d") typedCrossData + ) + ) + +-- | Typed: col @"y" .+ col @"d" → Maybe Double +typedAddMaybeIntDouble :: Test +typedAddMaybeIntDouble = + TestCase + ( assertEqual + "Typed: Maybe Int .+ Double = Maybe Double" + [Just 11.5, Nothing, Just 33.5] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"y" TE..+ TE.col @"d") typedCrossData + ) + ) + +-- | Typed: col @"x" .+ col @"md" → Maybe Double +typedAddIntMaybeDouble :: Test +typedAddIntMaybeDouble = + TestCase + ( assertEqual + "Typed: Int .+ Maybe Double = Maybe Double" + [Just 11.5, Nothing, Just 33.5] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"x" TE..+ TE.col @"md") typedCrossData + ) + ) + +-- | Typed: col @"y" .+ col @"md" → Maybe Double +typedAddMaybeIntMaybeDouble :: Test +typedAddMaybeIntMaybeDouble = + TestCase + ( assertEqual + "Typed: Maybe Int .+ Maybe Double = Maybe Double" + [Just 20.5, Nothing, Just 60.5] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"y" TE..+ TE.col @"md") typedCrossData + ) + ) + +-- | Typed: col @"x" .- col @"d" → Double +typedSubIntDouble :: Test +typedSubIntDouble = + TestCase + ( assertEqual + "Typed: Int .- Double = Double" + [-0.5, -0.5, -0.5 :: Double] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"x" TE..- TE.col @"d") typedCrossData + ) + ) + +-- | Typed: col @"x" .* col @"d" → Double +typedMulIntDouble :: Test +typedMulIntDouble = + TestCase + ( assertEqual + "Typed: Int .* Double = Double" + [1.5, 5.0, 10.5 :: Double] + ( DT.columnAsList @"result" $ + DT.derive @"result" (TE.col @"x" TE..* TE.col @"d") typedCrossData + ) + ) + +-- --------------------------------------------------------------------------- +-- Typed TExpr layer tests +-- --------------------------------------------------------------------------- + +type TestSchema = '[DT.Column "x" Int, DT.Column "y" (Maybe Int)] + +typedTestData :: DT.TypedDataFrame TestSchema +typedTestData = + either (error . show) id $ + DT.freezeWithError @TestSchema testData + +-- | Typed: col @"x" .+ col @"y" should give Maybe Int column +typedAddIntMaybeInt :: Test +typedAddIntMaybeInt = + TestCase + ( assertEqual + "Typed: Int .+ Maybe Int = Maybe Int" + [Just 11, Nothing, Just 33] + ( DT.columnAsList @"result" $ + DT.derive + @"result" + (TE.col @"x" TE..+ TE.col @"y") + typedTestData + ) + ) + +-- | Typed: nullLift negate on Maybe Int propagates Nothing +typedNullLiftMaybeInt :: Test +typedNullLiftMaybeInt = + TestCase + ( assertEqual + "Typed nullLift negate (Maybe Int) propagates Nothing" + [Just (-10), Nothing, Just (-30 :: Int)] + ( DT.columnAsList @"result" $ + DT.derive + @"result" + (TE.nullLift negate (TE.col @"y") :: TExpr TestSchema (Maybe Int)) + typedTestData + ) + ) + +-- | Typed: nullLift negate on plain Int +typedNullLiftInt :: Test +typedNullLiftInt = + TestCase + ( assertEqual + "Typed nullLift negate (Int) gives Int column" + [-1, -2, -3 :: Int] + ( DT.columnAsList @"result" $ + DT.derive + @"result" + (TE.nullLift negate (TE.col @"x") :: TExpr TestSchema Int) + typedTestData + ) + ) + +-- | Typed: nullLift2 (+) col @"x" col @"y" → Maybe Int +typedNullLift2IntMaybeInt :: Test +typedNullLift2IntMaybeInt = + TestCase + ( assertEqual + "Typed nullLift2 (+) Int (Maybe Int) = Maybe Int" + [Just 11, Nothing, Just 33 :: Maybe Int] + ( DT.columnAsList @"result" $ + DT.derive + @"result" + (TE.nullLift2 (+) (TE.col @"x") (TE.col @"y") :: TExpr TestSchema (Maybe Int)) + typedTestData + ) + ) + +-- | Typed: col @"y" .== col @"y" should give Maybe Bool column +typedEqMaybeMaybe :: Test +typedEqMaybeMaybe = + TestCase + ( assertEqual + "Typed: Maybe Int .== Maybe Int = Maybe Bool" + [Just True, Nothing, Just True] + ( DT.columnAsList @"result" $ + DT.derive + @"result" + (TE.col @"y" TE..== TE.col @"y") + typedTestData + ) + ) + +-- | apply negate to a Maybe Int column — should fmap over inner values +applyFmapMaybeInt :: Test +applyFmapMaybeInt = + TestCase + ( assertEqual + "apply negate to Maybe Int column fmaps" + (Just $ DI.OptionalColumn (V.fromList [Just (-10), Nothing, Just (-30 :: Int)])) + ( DI.getColumn "y" $ + D.apply @Int negate "y" testData + ) + ) + +-- | apply to a plain Int column still works +applyPlainInt :: Test +applyPlainInt = + TestCase + ( assertEqual + "apply negate to plain Int column still works" + (Just $ DI.fromList [-1, -2, -3 :: Int]) + ( DI.getColumn "x" $ + D.apply @Int negate "x" testData + ) + ) + +tests :: [Test] +tests = + [ TestLabel "addIntMaybeInt" addIntMaybeInt + , TestLabel "addMaybeIntInt" addMaybeIntInt + , TestLabel "addIntInt" addIntInt + , TestLabel "addMaybeMaybe" addMaybeMaybe + , TestLabel "subIntMaybeInt" subIntMaybeInt + , TestLabel "eqIntMaybeInt" eqIntMaybeInt + , TestLabel "eqIntInt" eqIntInt + , TestLabel "nullLiftMaybeInt" nullLiftMaybeInt + , TestLabel "nullLiftInt" nullLiftInt + , TestLabel "nullLift2IntMaybeInt" nullLift2IntMaybeInt + , TestLabel "nullLift2MaybeIntInt" nullLift2MaybeIntInt + , TestLabel "nullLift2IntInt" nullLift2IntInt + , TestLabel "typedAddIntMaybeInt" typedAddIntMaybeInt + , TestLabel "typedEqMaybeMaybe" typedEqMaybeMaybe + , TestLabel "typedNullLiftMaybeInt" typedNullLiftMaybeInt + , TestLabel "typedNullLiftInt" typedNullLiftInt + , TestLabel "typedNullLift2IntMaybeInt" typedNullLift2IntMaybeInt + , TestLabel "applyFmapMaybeInt" applyFmapMaybeInt + , TestLabel "applyPlainInt" applyPlainInt + , TestLabel "addIntDouble" addIntDouble + , TestLabel "addDoubleInt" addDoubleInt + , TestLabel "addMaybeIntDouble" addMaybeIntDouble + , TestLabel "addIntMaybeDouble" addIntMaybeDouble + , TestLabel "addMaybeIntMaybeDouble" addMaybeIntMaybeDouble + , TestLabel "subIntDouble" subIntDouble + , TestLabel "subDoubleInt" subDoubleInt + , TestLabel "subMaybeIntDouble" subMaybeIntDouble + , TestLabel "subIntMaybeDouble" subIntMaybeDouble + , TestLabel "subMaybeIntMaybeDouble" subMaybeIntMaybeDouble + , TestLabel "mulIntDouble" mulIntDouble + , TestLabel "mulDoubleInt" mulDoubleInt + , TestLabel "mulMaybeIntDouble" mulMaybeIntDouble + , TestLabel "mulIntMaybeDouble" mulIntMaybeDouble + , TestLabel "mulMaybeIntMaybeDouble" mulMaybeIntMaybeDouble + , TestLabel "divIntDouble" divIntDouble + , TestLabel "divDoubleInt" divDoubleInt + , TestLabel "divMaybeIntDouble" divMaybeIntDouble + , TestLabel "divIntMaybeDouble" divIntMaybeDouble + , TestLabel "divMaybeIntMaybeDouble" divMaybeIntMaybeDouble + , TestLabel "typedAddIntDouble" typedAddIntDouble + , TestLabel "typedAddMaybeIntDouble" typedAddMaybeIntDouble + , TestLabel "typedAddIntMaybeDouble" typedAddIntMaybeDouble + , TestLabel "typedAddMaybeIntMaybeDouble" typedAddMaybeIntMaybeDouble + , TestLabel "typedSubIntDouble" typedSubIntDouble + , TestLabel "typedMulIntDouble" typedMulIntDouble + ] diff --git a/tests/Operations/Provenance.hs b/tests/Operations/Provenance.hs new file mode 100644 --- /dev/null +++ b/tests/Operations/Provenance.hs @@ -0,0 +1,237 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Operations.Provenance where + +import qualified Data.List as L +import qualified Data.Map as M +import qualified Data.Text as T +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import qualified DataFrame.Internal.DataFrame as DI +import DataFrame.Operations.Merge ((|||)) + +import Test.HUnit + +-- Base frame with no derived columns. +base :: D.DataFrame +base = + D.fromNamedColumns + [ ("x", DI.fromList [1 .. 5 :: Int]) + , ("y", DI.fromList [2 .. 6 :: Int]) + ] + +-- A frame with one derived column "z". +withZ :: D.DataFrame +withZ = D.derive "z" (F.col @Int "x" + F.col "y") base + +-- ── insertColumn ────────────────────────────────────────────────────────────── + +-- Inserting a new column must not wipe provenance of existing derived columns. +insertPreservesProvenance :: Test +insertPreservesProvenance = + TestCase + ( assertBool + "insertColumn should preserve existing derivingExpressions" + ( M.member + "z" + (DI.derivingExpressions (D.insertColumn "w" (DI.fromList [0 :: Int]) withZ)) + ) + ) + +-- Overwriting a derived column removes only *that* expression, not others. +insertOverwriteDropsOwnExpr :: Test +insertOverwriteDropsOwnExpr = + let + df2 = D.derive "w" (F.col @Int "x") withZ + df3 = D.insertColumn "z" (DI.fromList [99 :: Int]) df2 + in + TestCase $ do + assertBool + "overwritten column z should be removed from derivingExpressions" + (not $ M.member "z" (DI.derivingExpressions df3)) + assertBool + "sibling column w expression should be preserved" + (M.member "w" (DI.derivingExpressions df3)) + +-- ── derive ──────────────────────────────────────────────────────────────────── + +-- derive adds the expression to derivingExpressions. +deriveTracksExpression :: Test +deriveTracksExpression = + TestCase + ( assertBool + "derive should record z in derivingExpressions" + (M.member "z" (DI.derivingExpressions withZ)) + ) + +-- Multiple derives accumulate. +deriveManyTracksAll :: Test +deriveManyTracksAll = + let df = D.derive "w" (F.col @Int "x") withZ + in TestCase + ( assertEqual + "two derive calls should leave two expressions" + 2 + (M.size (DI.derivingExpressions df)) + ) + +-- Re-deriving a column replaces its expression and keeps the count stable. +deriveOverwriteReplacesExpression :: Test +deriveOverwriteReplacesExpression = + let df = D.derive "z" (F.col @Int "y") withZ -- overwrite z + in TestCase + ( assertEqual + "re-deriving z should not duplicate the entry" + 1 + (M.size (DI.derivingExpressions df)) + ) + +-- ── deriveWithExpr ──────────────────────────────────────────────────────────── + +-- deriveWithExpr should also track the expression. +deriveWithExprTracksExpression :: Test +deriveWithExprTracksExpression = + let (_, df) = D.deriveWithExpr @Int "z" (F.col @Int "x" + F.col "y") base + in TestCase + ( assertBool + "deriveWithExpr should record z in derivingExpressions" + (M.member "z" (DI.derivingExpressions df)) + ) + +-- ── showDerivedExpressions ──────────────────────────────────────────────────── + +-- A frame with no derived columns returns identity provenance for each column. +showDerivedEmpty :: Test +showDerivedEmpty = + TestCase + ( assertEqual + "raw-column frame should have identity provenance for each column" + ["x", "y"] + (map fst (D.showDerivedExpressions base)) + ) + +-- Frame with one derived column has an entry for the derived column. +showDerivedContainsName :: Test +showDerivedContainsName = + TestCase + ( assertBool + "showDerivedExpressions should include the derived column name" + ("z" `elem` map fst (D.showDerivedExpressions withZ)) + ) + +-- ── Semigroup (<>) provenance propagation ───────────────────────────────────── + +-- Vertical merge preserves expressions from both sides. +semiGroupPreservesLeft :: Test +semiGroupPreservesLeft = + let merged = withZ <> base + in TestCase + ( assertBool + "<> should preserve derivingExpressions from the left frame" + (M.member "z" (DI.derivingExpressions merged)) + ) + +semiGroupPreservesBoth :: Test +semiGroupPreservesBoth = + let dfW = D.derive "w" (F.col @Int "y") base + merged = withZ <> dfW + in TestCase + ( assertEqual + "<> should union derivingExpressions from both frames" + 2 + (M.size (DI.derivingExpressions merged)) + ) + +-- Left frame wins when both sides have an expression for the same column. +semiGroupLeftBias :: Test +semiGroupLeftBias = + let dfLeft = D.derive "z" (F.col @Int "x") base + dfRight = D.derive "z" (F.col @Int "y") base + merged = dfLeft <> dfRight + in TestCase + ( assertEqual + "<> should retain exactly one entry for a shared column name" + 1 + (M.size (DI.derivingExpressions merged)) + ) + +-- Merging with an empty frame must not lose provenance. +semiGroupWithEmpty :: Test +semiGroupWithEmpty = + TestCase + ( assertBool + "withZ <> empty should keep z's expression" + (M.member "z" (DI.derivingExpressions (withZ <> D.empty))) + ) + +emptyWithSemiGroup :: Test +emptyWithSemiGroup = + TestCase + ( assertBool + "empty <> withZ should keep z's expression" + (M.member "z" (DI.derivingExpressions (D.empty <> withZ))) + ) + +-- ── Horizontal merge (|||) provenance propagation ──────────────────────────── + +horizontalMergePreservesLeft :: Test +horizontalMergePreservesLeft = + let dfW = D.derive "w" (F.col @Int "y") base + extra = D.fromNamedColumns [("q", DI.fromList [0 :: Int, 0, 0, 0, 0])] + merged = dfW ||| extra + in TestCase + ( assertBool + "||| should preserve derivingExpressions from the left frame" + (M.member "w" (DI.derivingExpressions merged)) + ) + +horizontalMergePreservesRight :: Test +horizontalMergePreservesRight = + let extra = D.fromNamedColumns [("q", DI.fromList [0 :: Int, 0, 0, 0, 0])] + dfW = + D.derive + "w" + (F.col @Int "y") + (D.fromNamedColumns [("y", DI.fromList [2 .. 6 :: Int])]) + merged = extra ||| dfW + in TestCase + ( assertBool + "||| should bring in derivingExpressions from the right frame" + (M.member "w" (DI.derivingExpressions merged)) + ) + +horizontalMergePreservesBoth :: Test +horizontalMergePreservesBoth = + let dfZ = withZ + dfW = D.fromNamedColumns [("q", DI.fromList [0 :: Int, 0, 0, 0, 0])] + -- give dfW a derived column on a separate base + dfWD = D.derive "w" (F.col @Int "q") dfW + merged = dfZ ||| dfWD + in TestCase + ( assertEqual + "||| should union derivingExpressions" + 2 + (M.size (DI.derivingExpressions merged)) + ) + +tests :: [Test] +tests = + [ TestLabel "insertPreservesProvenance" insertPreservesProvenance + , TestLabel "insertOverwriteDropsOwnExpr" insertOverwriteDropsOwnExpr + , TestLabel "deriveTracksExpression" deriveTracksExpression + , TestLabel "deriveManyTracksAll" deriveManyTracksAll + , TestLabel "deriveOverwriteReplacesExpression" deriveOverwriteReplacesExpression + , TestLabel "deriveWithExprTracksExpression" deriveWithExprTracksExpression + , TestLabel "showDerivedEmpty" showDerivedEmpty + , TestLabel "showDerivedContainsName" showDerivedContainsName + , TestLabel "semiGroupPreservesLeft" semiGroupPreservesLeft + , TestLabel "semiGroupPreservesBoth" semiGroupPreservesBoth + , TestLabel "semiGroupLeftBias" semiGroupLeftBias + , TestLabel "semiGroupWithEmpty" semiGroupWithEmpty + , TestLabel "emptyWithSemiGroup" emptyWithSemiGroup + , TestLabel "horizontalMergePreservesLeft" horizontalMergePreservesLeft + , TestLabel "horizontalMergePreservesRight" horizontalMergePreservesRight + , TestLabel "horizontalMergePreservesBoth" horizontalMergePreservesBoth + ] diff --git a/tests/Operations/ReadCsv.hs b/tests/Operations/ReadCsv.hs --- a/tests/Operations/ReadCsv.hs +++ b/tests/Operations/ReadCsv.hs @@ -18,12 +18,14 @@ import Data.Function (on) import Data.Maybe (fromMaybe) import Data.Type.Equality (testEquality, (:~:) (Refl)) -import DataFrame.Internal.Column (Column (..)) +import DataFrame.Internal.Column (Column (..), columnTypeString) +import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame ( DataFrame (..), columnIndices, columns, dataframeDimensions, + getColumn, ) import System.Directory (removeFile) import System.IO (IOMode (..), withFile) @@ -36,6 +38,11 @@ tempDir :: FilePath tempDir = "./tests/data/unstable_csv/" +readCsvNoInfer :: FilePath -> IO DataFrame +readCsvNoInfer = + D.readCsvWithOpts + D.defaultReadOptions{D.typeSpec = D.NoInference} + -------------------------------------------------------------------------------- -- Pretty-printer -------------------------------------------------------------------------------- @@ -143,6 +150,219 @@ testJsonDataFast :: Test testJsonDataFast = testFastCsv "json_data" (fixtureDir <> "json_data.csv") +arbuthnotPath :: FilePath +arbuthnotPath = "./tests/data/arbuthnot.csv" + +-- SpecifyTypes with NoInference fallback: named column is typed, rest stay Text +specifyTypesNoInferenceFallback :: Test +specifyTypesNoInferenceFallback = + TestLabel "specifyTypes_noInference_fallback" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.typeSpec = + D.SpecifyTypes + [("year", D.schemaType @Int)] + D.NoInference + } + arbuthnotPath + -- "year" must be Int + case getColumn "year" df of + Just col@(UnboxedColumn _) -> assertEqual "year should be Int" "Int" (columnTypeString col) + _ -> assertFailure "expected UnboxedColumn for 'year'" + -- "boys" unspecified + NoInference → stays Text + case getColumn "boys" df of + Just col@(BoxedColumn _) -> assertEqual "boys should be Text" "Text" (columnTypeString col) + _ -> assertFailure "expected BoxedColumn for 'boys' with NoInference fallback" + +-- SpecifyTypes with InferFromSample fallback: named column typed, rest inferred +specifyTypesInferFallback :: Test +specifyTypesInferFallback = + TestLabel "specifyTypes_inferFromSample_fallback" $ TestCase $ do + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.typeSpec = + D.SpecifyTypes + [("year", D.schemaType @Int)] + (D.InferFromSample 100) + } + arbuthnotPath + -- "year" must be Int (explicitly specified) + case getColumn "year" df of + Just col@(UnboxedColumn _) -> assertEqual "year should be Int" "Int" (columnTypeString col) + _ -> assertFailure "expected UnboxedColumn for 'year'" + -- "boys" unspecified + InferFromSample → inferred as Int + case getColumn "boys" df of + Just col@(UnboxedColumn _) -> assertEqual "boys should be Int" "Int" (columnTypeString col) + _ -> + assertFailure "expected UnboxedColumn for 'boys' with InferFromSample fallback" + +-- SpecifyTypes: typeInferenceSampleSize delegates to fallback +specifyTypesSampleSize :: Test +specifyTypesSampleSize = + TestLabel "specifyTypes_sampleSize_from_fallback" $ TestCase $ do + -- Use a small sample size; all numeric columns should still be inferred + df <- + D.readCsvWithOpts + D.defaultReadOptions + { D.typeSpec = + D.SpecifyTypes + [] + (D.InferFromSample 10) + } + arbuthnotPath + case getColumn "girls" df of + Just col@(UnboxedColumn _) -> assertEqual "girls should be Int" "Int" (columnTypeString col) + _ -> + assertFailure + "expected UnboxedColumn for 'girls' via fallback InferFromSample 10" + +testCrlfCsv :: Test +testCrlfCsv = TestLabel "malformed_crlf_csv" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "crlf.csv") + let (rows, cols) = dataframeDimensions df + assertEqual "crlf.csv: 2 data rows" 2 rows + assertEqual "crlf.csv: 2 columns" 2 cols + case getColumn "name" df of + Nothing -> assertFailure "crlf.csv: column 'name' missing" + Just col -> + assertEqual + "crlf.csv: name has no \\r" + (DI.fromList @T.Text ["Alice", "Bob"]) + col + +testCrlfTsv :: Test +testCrlfTsv = TestLabel "malformed_crlf_tsv" $ TestCase $ do + df <- D.readTsvUnstable (fixtureDir <> "crlf.tsv") + let (rows, _) = dataframeDimensions df + assertEqual "crlf.tsv: 1 data row" 1 rows + case getColumn "name" df of + Nothing -> assertFailure "crlf.tsv: column 'name' missing" + Just col -> + assertEqual + "crlf.tsv: name has no \\r" + (DI.fromList @T.Text ["Alice"]) + col + +testHeaderOnly :: Test +testHeaderOnly = TestLabel "malformed_header_only" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "header_only.csv") + let (rows, cols) = dataframeDimensions df + assertEqual "header_only.csv: 0 data rows" 0 rows + assertEqual "header_only.csv: 3 columns" 3 cols + let names = map fst . L.sortBy (compare `on` snd) . M.toList $ columnIndices df + assertEqual "header_only.csv: column names" ["first", "second", "third"] names + +testTrailingBlankLine :: Test +testTrailingBlankLine = TestLabel "malformed_trailing_blank_line" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "trailing_blank_line.csv") + -- blank line contributes 1 extra delimiter; (3+3+1) div 3 = 2, numRow=1 + assertEqual + "trailing_blank_line.csv: 1 data row visible" + 1 + (fst (dataframeDimensions df)) + +testAllEmptyRow :: Test +testAllEmptyRow = TestLabel "malformed_all_empty_row" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "all_empty_row.csv") + assertEqual "all_empty_row.csv: 1 data row" 1 (fst (dataframeDimensions df)) + let checkEmpty colName = + case getColumn colName df of + Nothing -> assertFailure ("column '" <> T.unpack colName <> "' missing") + Just col -> + assertEqual + (T.unpack colName <> " is Nothing (empty field → null)") + (DI.fromList @(Maybe T.Text) [Nothing]) + col + mapM_ checkEmpty ["a", "b", "c"] + +testSingleCol :: Test +testSingleCol = TestLabel "malformed_single_col" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "single_col.csv") + let (rows, cols) = dataframeDimensions df + assertEqual "single_col.csv: 3 data rows" 3 rows + assertEqual "single_col.csv: 1 column" 1 cols + case getColumn "name" df of + Nothing -> assertFailure "single_col.csv: column 'name' missing" + Just col -> + assertEqual + "single_col.csv: correct values" + (DI.fromList @T.Text ["Alice", "Bob", "Carol"]) + col + +testWhitespaceFields :: Test +testWhitespaceFields = TestLabel "malformed_whitespace_fields" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "whitespace_fields.csv") + assertEqual + "whitespace_fields.csv: 2 data rows" + 2 + (fst (dataframeDimensions df)) + case getColumn "name" df of + Nothing -> assertFailure "whitespace_fields.csv: 'name' missing" + Just col -> assertEqual "name stripped" (DI.fromList @T.Text ["Alice", "Bob"]) col + case getColumn "city" df of + Nothing -> assertFailure "whitespace_fields.csv: 'city' missing" + Just col -> + assertEqual + "city stripped" + (DI.fromList @T.Text ["New York", "Los Angeles"]) + col + +-- File: a,b,c header; row "1,2" (short); row "X,Y,Z" (full) +-- Total delimiters: 3 (header) + 2 (short) + 3 (full) = 8 +-- numCol=3, totalRows=8 div 3=2, numRow=1 +-- Row-1 stride offsets 3,4,5 → fields "1","2","X" (X bleeds in from next row) +testMissingFields :: Test +testMissingFields = TestLabel "malformed_missing_fields" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "missing_fields.csv") + assertEqual + "missing_fields.csv: integer-division gives 1 visible row" + 1 + (fst (dataframeDimensions df)) + case getColumn "c" df of + Nothing -> assertFailure "missing_fields.csv: column 'c' missing" + Just col -> + -- "X" bleeds from the start of the next row into the missing slot + assertEqual + "missing_fields.csv: col 'c' bleeds 'X' from next row" + (DI.fromList @T.Text ["X"]) + col + +-- File: a,b,c header; row "1,2,3,EXTRA" +-- Total delimiters: 3 + 4 = 7; 7 div 3 = 2, numRow=1 +-- Row-1 strides 3,4,5 → "1","2","3" — EXTRA (stride 6) is never accessed +testExtraFields :: Test +testExtraFields = TestLabel "malformed_extra_fields" $ TestCase $ do + df <- readCsvNoInfer (fixtureDir <> "extra_fields.csv") + assertEqual "extra_fields.csv: 1 data row" 1 (fst (dataframeDimensions df)) + assertEqual "extra_fields.csv: 3 columns" 3 (snd (dataframeDimensions df)) + case getColumn "c" df of + Nothing -> assertFailure "extra_fields.csv: column 'c' missing" + Just col -> + assertEqual + "extra_fields.csv: 'c' = '3' (EXTRA ignored)" + (DI.fromList @T.Text ["3"]) + col + +testNoTrailingNewline :: Test +testNoTrailingNewline = TestLabel "malformed_no_trailing_newline" $ TestCase $ do + df <- D.readCsvUnstable (fixtureDir <> "no_trailing_newline.csv") + assertEqual + "no_trailing_newline.csv: 1 data row" + 1 + (fst (dataframeDimensions df)) + case getColumn "name" df of + Nothing -> assertFailure "no_trailing_newline.csv: 'name' missing" + Just col -> assertEqual "name = Alice" (DI.fromList @T.Text ["Alice"]) col + case getColumn "city" df of + Nothing -> assertFailure "no_trailing_newline.csv: 'city' missing" + Just col -> + assertEqual + "city = London (synthetic delimiter worked)" + (DI.fromList @T.Text ["London"]) + col + tests :: [Test] tests = [ testSimpleFast @@ -154,4 +374,17 @@ , testQuotesAndNewlinesFast , testEmptyValuesFast , testJsonDataFast + , specifyTypesNoInferenceFallback + , specifyTypesInferFallback + , specifyTypesSampleSize + , testCrlfCsv + , testCrlfTsv + , testHeaderOnly + , testTrailingBlankLine + , testAllEmptyRow + , testSingleCol + , testWhitespaceFields + , testMissingFields + , testExtraFields + , testNoTrailingNewline ] diff --git a/tests/Operations/Subset.hs b/tests/Operations/Subset.hs --- a/tests/Operations/Subset.hs +++ b/tests/Operations/Subset.hs @@ -1,11 +1,16 @@ {-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} module Operations.Subset where +import qualified Data.Text as T +import qualified DataFrame as D +import qualified DataFrame.Internal.Column as Col import DataFrame.Internal.DataFrame -import qualified DataFrame.Operations.Core as D -import qualified DataFrame.Operations.Subset as D +import DataFrame.Operations.Merge () import System.Random +import Test.HUnit prop_dropZero :: DataFrame -> Bool prop_dropZero df = D.drop 0 df == df @@ -75,6 +80,113 @@ sampled = D.sample gen 0.0 df in fst (dataframeDimensions sampled) == 0 +prop_stratifiedSplit_deterministic :: DataFrame -> Bool +prop_stratifiedSplit_deterministic _ = + let df = + D.fromNamedColumns + [ ("label", Col.fromList (replicate 50 ("A" :: T.Text) ++ replicate 50 "B")) + , ("val", Col.fromList ([1 .. 100] :: [Int])) + ] + (tr, va) = D.stratifiedSplit (mkStdGen 314) 0.7 (D.col @T.Text "label") df + in fst (dataframeDimensions tr) + fst (dataframeDimensions va) == 100 + +strataDf :: DataFrame +strataDf = + D.fromNamedColumns + [ ("label", Col.fromList (replicate 5 ("A" :: T.Text) ++ replicate 5 "B")) + , ("val", Col.fromList ([1 .. 10] :: [Int])) + ] + +unit_stratifiedSample_full :: Test +unit_stratifiedSample_full = + TestCase $ + let sampled = D.stratifiedSample (mkStdGen 42) 1.0 (D.col @T.Text "label") strataDf + in assertEqual + "p=1.0 preserves row count" + (fst $ dataframeDimensions strataDf) + (fst $ dataframeDimensions sampled) + +unit_stratifiedSplit_rowCount :: Test +unit_stratifiedSplit_rowCount = + TestCase $ + let (tr, va) = D.stratifiedSplit (mkStdGen 99) 0.8 (D.col @T.Text "label") strataDf + in assertEqual + "train+validation == total" + (fst $ dataframeDimensions strataDf) + (fst (dataframeDimensions tr) + fst (dataframeDimensions va)) + +unit_stratifiedSplit_singleRowStratum :: Test +unit_stratifiedSplit_singleRowStratum = + TestCase $ + let tinyDf = + D.fromNamedColumns + [ ("label", Col.fromList (["A", "A", "A", "A", "A", "B"] :: [T.Text])) + , ("val", Col.fromList ([1 .. 6] :: [Int])) + ] + (tr, va) = D.stratifiedSplit (mkStdGen 7) 0.8 (D.col @T.Text "label") tinyDf + in assertEqual + "single-row stratum: no rows lost" + (fst $ dataframeDimensions tinyDf) + (fst (dataframeDimensions tr) + fst (dataframeDimensions va)) + +-- | Count occurrences of a label in a column, expressed as a fraction of total rows. +labelProportion :: T.Text -> T.Text -> DataFrame -> Double +labelProportion col label df = + let total = fst (dataframeDimensions df) + vals = case getColumn col df of + Just c -> Col.toList @T.Text c + Nothing -> [] + n = length (filter (== label) vals) + in fromIntegral n / fromIntegral total + +unit_stratifiedSplit_proportions :: Test +unit_stratifiedSplit_proportions = + TestCase $ + let aCount = 100 + bCount = 50 + df = + D.fromNamedColumns + [ + ( "label" + , Col.fromList (replicate aCount ("A" :: T.Text) ++ replicate bCount "B") + ) + , ("val", Col.fromList ([1 .. aCount + bCount] :: [Int])) + ] + (tr, va) = D.stratifiedSplit (mkStdGen 42) 0.8 (D.col @T.Text "label") df + origProp = labelProportion "label" "A" df + trProp = labelProportion "label" "A" tr + vaProp = labelProportion "label" "A" va + tol = 0.05 :: Double + in do + assertBool + ( "train A-proportion " + ++ show trProp + ++ " differs from original " + ++ show origProp + ++ " by more than " + ++ show tol + ) + (abs (trProp - origProp) < tol) + assertBool + ( "validation A-proportion " + ++ show vaProp + ++ " differs from original " + ++ show origProp + ++ " by more than " + ++ show tol + ) + (abs (vaProp - origProp) < tol) + +hunitTests :: [Test] +hunitTests = + [ TestLabel "unit_stratifiedSample_full" unit_stratifiedSample_full + , TestLabel "unit_stratifiedSplit_rowCount" unit_stratifiedSplit_rowCount + , TestLabel + "unit_stratifiedSplit_singleRowStratum" + unit_stratifiedSplit_singleRowStratum + , TestLabel "unit_stratifiedSplit_proportions" unit_stratifiedSplit_proportions + ] + tests = [ prop_dropZero , prop_takeZero @@ -92,4 +204,5 @@ , prop_excludeAll , prop_cubePreservesSmall , prop_sampleEmptyApprox + , prop_stratifiedSplit_deterministic ] diff --git a/tests/Operations/Typing.hs b/tests/Operations/Typing.hs --- a/tests/Operations/Typing.hs +++ b/tests/Operations/Typing.hs @@ -42,7 +42,10 @@ beforeParse :: [T.Text] beforeParse = ["True", "true", "TRUE"] ++ ["False", "false", "FALSE"] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools without missing values as UnboxedColumn of Bools" @@ -57,7 +60,10 @@ beforeParse :: [T.Text] beforeParse = T.pack . show <$> [1 .. 50] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints without missing values as UnboxedColumn of Ints" @@ -74,7 +80,10 @@ T.pack . show <$> [1.0 .. 50.0] ++ [3.14, 2.22, 8.55, 23.3, 12.22222235049450945049504950] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Doubles without missing values as UnboxedColumn of Doubles" @@ -121,7 +130,10 @@ , "2020-02-26" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates without missing values as BoxedColumn of Days" @@ -206,7 +218,10 @@ , "Meowth, that's right!" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Text without missing values as BoxedColumn of Text" @@ -222,7 +237,10 @@ beforeParse :: [T.Text] beforeParse = ["True", "true", "TRUE"] ++ ["False", "false", "FALSE"] ++ ["1", "0", "1"] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses mixture of Bools and Ints as Text" @@ -449,7 +467,10 @@ , "12.22222235049451" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles" @@ -544,7 +565,10 @@ , "2020-02-20" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Dates as BoxedColumn of Texts" @@ -683,7 +707,10 @@ , "12.22222235049451" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Texts and Doubles as BoxedColumn of Texts" @@ -736,7 +763,10 @@ , "Meowth" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Dates and Texts as BoxedColumn of Texts" @@ -754,7 +784,10 @@ beforeParse = replicate 10 "true" ++ replicate 10 "false" expected = DI.UnboxedColumn $ VU.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools without missing values as UnboxedColumn of Bools, when safeRead is off" @@ -872,7 +905,10 @@ ] expected = DI.UnboxedColumn $ VU.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints without missing values as UnboxedColumn of Ints, when safeRead is off" @@ -1000,7 +1036,10 @@ ] expected = DI.UnboxedColumn $ VU.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Doubles without missing values as UnboxedColumn of Doubles, when safeRead is off" @@ -1048,7 +1087,10 @@ ] expected = DI.BoxedColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates without missing values as BoxedColumn of Days" @@ -1134,7 +1176,10 @@ ] expected = DI.BoxedColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Text without missing values as BoxedColumn of Text" @@ -1150,7 +1195,10 @@ beforeParse = replicate 10 "" ++ replicate 10 "true" ++ replicate 10 "false" expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools and empty Strings as OptionalColumn of Bools, when safeRead is off" @@ -1278,7 +1326,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints and empty strings as OptionalColumn of Ints, when safeRead is off" @@ -1416,7 +1467,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses combination of Ints, Doubles and empty strings as OptionalColumn of Doubles, when safeRead is off" @@ -1474,7 +1528,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates and Empty Strings as OptionalColumn of Dates, with safeRead off" @@ -1574,7 +1631,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Texts and Empty Strings as OptionalColumn of Text, with safeRead off" @@ -1590,7 +1650,10 @@ beforeParse = replicate 10 "N/A" ++ replicate 10 "True" ++ replicate 10 "False" expected = DI.BoxedColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools with nullish values as BoxedColumn of Texts, when safeRead is off" @@ -1718,7 +1781,10 @@ ] expected = DI.BoxedColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints with nullish values as BoxedColumn of Texts, when safeRead is off" @@ -1856,7 +1922,10 @@ ] expected = DI.BoxedColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses combination of Ints, Doubles and empty strings as OptionalColumn of Doubles, when safeRead is off" @@ -1996,7 +2065,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints with nullish values AND empty strings as OptionalColumn of Texts, when safeRead is off" @@ -2102,7 +2174,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 10 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 10, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Texts and Empty Strings as OptionalColumn of Text, with safeRead off" @@ -2118,7 +2193,10 @@ beforeParse :: [T.Text] beforeParse = replicate 10 "" ++ replicate 10 "true" expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools and empty strings as OptionalColumn of Bools, when safeRead is on" @@ -2245,7 +2323,10 @@ , "50" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints and empty strings as OptionalColumn of Ints, when safeRead is on" @@ -2382,7 +2463,10 @@ , "12.22222235049451" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses combination of Ints, Doubles and empty strings as OptionalColumn of Doubles, when safeRead is on" @@ -2439,7 +2523,10 @@ , "" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates and Empty Strings as OptionalColumn of Dates, with safeRead on" @@ -2538,7 +2625,10 @@ , "Meowth, that's right!" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Texts and Empty Strings as OptionalColumn of Text, with safeRead on" @@ -2665,7 +2755,10 @@ , "50" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints with nullish values as OptionalColumn of Ints, when safeRead is on" @@ -2802,7 +2895,10 @@ , "12.03" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses combination of Ints, Doubles and empty strings as OptionalColumn of Doubles, when safeRead is off" @@ -2941,7 +3037,10 @@ , "" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints with nullish values AND empty strings as OptionalColumn of Ints, when safeRead is on" @@ -3038,7 +3137,10 @@ , "3.14" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints and Doubles with nullish values AND empty strings as OptionalColumn of Doubles, when safeRead is on" @@ -3143,7 +3245,10 @@ , "N/A" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Texts and Empty Strings as OptionalColumn of Text, with safeRead off" @@ -3159,7 +3264,10 @@ beforeParse :: [T.Text] beforeParse = "false" : replicate 50 "true" expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 1 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools without missing values as UnboxedColumn of Ints with only one example" @@ -3174,7 +3282,10 @@ beforeParse :: [T.Text] beforeParse = "false" : replicate 50 "true" expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 49 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 49}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Bools without missing values as UnboxedColumn of Ints with only one example" @@ -3291,7 +3402,10 @@ , "50" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 1 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints without missing values as UnboxedColumn of Ints with only one example" @@ -3408,7 +3522,10 @@ , "50" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 25 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 25}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints without missing values as UnboxedColumn of Ints with some examples" @@ -3525,7 +3642,10 @@ , "50" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 49 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 49}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Ints without missing values as UnboxedColumn of Ints with many examples" @@ -3572,7 +3692,10 @@ , "2020-02-26" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 1 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates without missing values as BoxedColumn of Days with only one example" @@ -3619,7 +3742,10 @@ , "2020-02-26" ] expected = DI.BoxedColumn $ V.fromList afterParse - actual = D.parseDefault [] 15 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 15}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses Dates without missing values as BoxedColumn of Days with many examples" @@ -3846,7 +3972,10 @@ , "12.22222235049451" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 1 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles with just one example" @@ -4073,7 +4202,10 @@ , "12.22222235049451" ] expected = DI.UnboxedColumn $ VU.fromList afterParse - actual = D.parseDefault [] 50 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 50}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles with just one example" @@ -4320,7 +4452,10 @@ , "12.22222235049451" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 1 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = False}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles with just one example" @@ -4569,7 +4704,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 30 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints and Doubles as UnboxedColumn of Doubles with just one example" @@ -4737,7 +4875,10 @@ , "12.22222235049451" ] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 1 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 1, D.parseSafe = False}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints, Doubles, empty strings, nullish as OptionalColumn of Text with just one example, when safeRead is off" @@ -4906,7 +5047,10 @@ ] expected = DI.OptionalColumn $ V.fromList afterParse actual = - D.parseDefault [] 30 False "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + D.parseDefault + (D.defaultParseOptions{D.sampleSize = 30, D.parseSafe = False}) + $ DI.fromVector + $ V.fromList beforeParse in TestCase ( assertEqual "Correctly parses a mixture of Ints, Doubles, empty strings, nullish as OptionalColumn of Text with many examples, when safeRead is off" @@ -4923,7 +5067,10 @@ beforeParse :: [T.Text] beforeParse = replicate 100 "NaN" ++ ["100000"] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase (assertEqual "Correctly parses many Nulls followed by one Int" expected actual) @@ -4934,7 +5081,10 @@ beforeParse :: [T.Text] beforeParse = replicate 100 "NaN" ++ ["3.14"] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase (assertEqual "Correctly parses many Nulls followed by one Int" expected actual) @@ -4945,7 +5095,10 @@ beforeParse :: [T.Text] beforeParse = replicate 100 "NaN" ++ ["2024-12-25"] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase (assertEqual "Correctly parses many Nulls followed by one Int" expected actual) @@ -4956,7 +5109,10 @@ beforeParse :: [T.Text] beforeParse = replicate 100 "NaN" ++ ["2024-12-25", "2024-12-w6"] expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase (assertEqual "Correctly parses many Nulls followed by one Int" expected actual) @@ -4967,7 +5123,10 @@ beforeParse :: [T.Text] beforeParse = replicate 100 "NaN" expected = DI.OptionalColumn $ V.fromList afterParse - actual = D.parseDefault [] 10 True "%Y-%m-%d" $ DI.fromVector $ V.fromList beforeParse + actual = + D.parseDefault (D.defaultParseOptions{D.sampleSize = 10}) $ + DI.fromVector $ + V.fromList beforeParse in TestCase (assertEqual "Correctly parses many Nulls followed by one Int" expected actual) diff --git a/tests/Parquet.hs b/tests/Parquet.hs --- a/tests/Parquet.hs +++ b/tests/Parquet.hs @@ -6,10 +6,32 @@ import Assertions (assertExpectException) import qualified DataFrame as D import qualified DataFrame.Functions as F +import qualified DataFrame.IO.Parquet as DP +import qualified Data.ByteString as BS import Data.Int +import qualified Data.Set as S import Data.Text (Text) +import qualified Data.Text as T import Data.Time +import Data.Word +import DataFrame.IO.Parquet.Thrift ( + columnMetaData, + columnPathInSchema, + columnStatistics, + rowGroupColumns, + rowGroups, + schema, + ) +import DataFrame.IO.Parquet.Types (columnNullCount) +import DataFrame.Internal.Binary ( + littleEndianWord32, + littleEndianWord64, + word32ToLittleEndian, + word64ToLittleEndian, + ) +import DataFrame.Internal.Column (hasMissing) +import DataFrame.Internal.DataFrame (unsafeGetColumn) import GHC.IO (unsafePerformIO) import Test.HUnit @@ -53,13 +75,20 @@ ) ] +testBothReadParquetPaths :: ((FilePath -> IO D.DataFrame) -> Test) -> Test +testBothReadParquetPaths test = + TestList + [ test D.readParquet + , test (DP._readParquetWithOpts (Just True) D.defaultParquetReadOptions) + ] + allTypesPlain :: Test -allTypesPlain = +allTypesPlain = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesPlain" allTypes - (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.parquet")) + (unsafePerformIO (readParquet "./tests/data/alltypes_plain.parquet")) ) allTypesTinyPagesDimensions :: Test @@ -155,7 +184,7 @@ ] allTypesTinyPagesLastFew :: Test -allTypesTinyPagesLastFew = +allTypesTinyPagesLastFew = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesTinyPages dimensions" @@ -164,27 +193,27 @@ -- Excluding doubles because they are weird to compare. ( fmap (D.takeLast 10 . D.exclude ["double_col"]) - (D.readParquet "./tests/data/alltypes_tiny_pages.parquet") + (readParquet "./tests/data/alltypes_tiny_pages.parquet") ) ) ) allTypesPlainSnappy :: Test -allTypesPlainSnappy = +allTypesPlainSnappy = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesPlainSnappy" (D.filter (F.col @Int32 "id") (`elem` [6, 7]) allTypes) - (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.snappy.parquet")) + (unsafePerformIO (readParquet "./tests/data/alltypes_plain.snappy.parquet")) ) allTypesDictionary :: Test -allTypesDictionary = +allTypesDictionary = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "allTypesPlainSnappy" (D.filter (F.col @Int32 "id") (`elem` [0, 1]) allTypes) - (unsafePerformIO (D.readParquet "./tests/data/alltypes_dictionary.parquet")) + (unsafePerformIO (readParquet "./tests/data/alltypes_dictionary.parquet")) ) selectedColumnsWithOpts :: Test @@ -460,12 +489,12 @@ ] transactionsTest :: Test -transactionsTest = +transactionsTest = testBothReadParquetPaths $ \readParquet -> TestCase ( assertEqual "transactions" transactions - (unsafePerformIO (D.readParquet "./tests/data/transactions.parquet")) + (unsafePerformIO (readParquet "./tests/data/transactions.parquet")) ) mtCarsDataset :: D.DataFrame @@ -926,6 +955,66 @@ (unsafePerformIO (D.readParquet "./tests/data/mtcars.parquet")) ) +littleEndianWord64KnownPattern :: Test +littleEndianWord64KnownPattern = + TestCase + ( assertEqual + "littleEndianWord64KnownPattern" + (0x1122334455667788 :: Word64) + ( littleEndianWord64 + (BS.pack [0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11]) + ) + ) + +littleEndianWord32KnownPattern :: Test +littleEndianWord32KnownPattern = + TestCase + ( assertEqual + "littleEndianWord32KnownPattern" + (0x11223344 :: Word32) + (littleEndianWord32 (BS.pack [0x44, 0x33, 0x22, 0x11])) + ) + +littleEndianWord64ShortInputPadsZeroes :: Test +littleEndianWord64ShortInputPadsZeroes = + TestCase + ( assertEqual + "littleEndianWord64ShortInputPadsZeroes" + (0x00CCBBAA :: Word64) + (littleEndianWord64 (BS.pack [0xAA, 0xBB, 0xCC])) + ) + +littleEndianWord32ShortInputPadsZeroes :: Test +littleEndianWord32ShortInputPadsZeroes = + TestCase + ( assertEqual + "littleEndianWord32ShortInputPadsZeroes" + (0x0000BEEF :: Word32) + (littleEndianWord32 (BS.pack [0xEF, 0xBE])) + ) + +littleEndianWord64RoundTrip :: Test +littleEndianWord64RoundTrip = + TestCase + ( assertEqual + "littleEndianWord64RoundTrip" + value + (littleEndianWord64 (word64ToLittleEndian value)) + ) + where + value = 0x1122334455667788 :: Word64 + +littleEndianWord32RoundTrip :: Test +littleEndianWord32RoundTrip = + TestCase + ( assertEqual + "littleEndianWord32RoundTrip" + value + (littleEndianWord32 (word32ToLittleEndian value)) + ) + where + value = 0xA1B2C3D4 :: Word32 + -- --------------------------------------------------------------------------- -- Group 1: Plain variant -- --------------------------------------------------------------------------- @@ -1611,6 +1700,58 @@ (D.readParquet "./tests/data/nation.dict-malformed.parquet") ) +shardedNullableSchema :: Test +shardedNullableSchema = + TestCase $ do + metas <- + mapM + (fmap fst . DP.readMetadataFromPath) + ["data/sharded/part-0.parquet", "data/sharded/part-1.parquet"] + let nullableCols = + S.fromList + [ last (map T.pack colPath) + | meta <- metas + , rg <- rowGroups meta + , cc <- rowGroupColumns rg + , let cm = columnMetaData cc + colPath = columnPathInSchema cm + , not (null colPath) + , columnNullCount (columnStatistics cm) > 0 + ] + df = + foldl + (\acc meta -> acc <> F.schemaToEmptyDataFrame nullableCols (schema meta)) + D.empty + metas + assertBool "id should be nullable" (hasMissing (unsafeGetColumn "id" df)) + assertBool "name should be nullable" (hasMissing (unsafeGetColumn "name" df)) + assertBool "score should be nullable" (hasMissing (unsafeGetColumn "score" df)) + +singleShardNoNulls :: Test +singleShardNoNulls = + TestCase $ do + (meta, _) <- DP.readMetadataFromPath "data/sharded/part-0.parquet" + let nullableCols = + S.fromList + [ last (map T.pack colPath) + | rg <- rowGroups meta + , cc <- rowGroupColumns rg + , let cm = columnMetaData cc + colPath = columnPathInSchema cm + , not (null colPath) + , columnNullCount (columnStatistics cm) > 0 + ] + df = F.schemaToEmptyDataFrame nullableCols (schema meta) + assertBool + "id should NOT be nullable" + (not (hasMissing (unsafeGetColumn "id" df))) + assertBool + "name should NOT be nullable" + (not (hasMissing (unsafeGetColumn "name" df))) + assertBool + "score should NOT be nullable" + (not (hasMissing (unsafeGetColumn "score" df))) + tests :: [Test] tests = [ allTypesPlain @@ -1626,6 +1767,12 @@ , allTypesTinyPagesLastFew , allTypesTinyPagesDimensions , transactionsTest + , littleEndianWord64KnownPattern + , littleEndianWord32KnownPattern + , littleEndianWord64ShortInputPadsZeroes + , littleEndianWord32ShortInputPadsZeroes + , littleEndianWord64RoundTrip + , littleEndianWord32RoundTrip , -- Group 1 allTypesTinyPagesPlain , -- Group 2: compression codecs @@ -1698,4 +1845,7 @@ , unknownLogicalType , -- Group 12: malformed files nationDictMalformed + , -- Group 13: metadata-based null detection + shardedNullableSchema + , singleShardNoNulls ] diff --git a/tests/data/unstable_csv/all_empty_row.csv b/tests/data/unstable_csv/all_empty_row.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/all_empty_row.csv @@ -0,0 +1,2 @@ +a,b,c +,, diff --git a/tests/data/unstable_csv/crlf.csv b/tests/data/unstable_csv/crlf.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/crlf.csv @@ -0,0 +1,3 @@ +name,city +Alice,New York +Bob,Los Angeles diff --git a/tests/data/unstable_csv/crlf.tsv b/tests/data/unstable_csv/crlf.tsv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/crlf.tsv @@ -0,0 +1,2 @@ +name city +Alice New York diff --git a/tests/data/unstable_csv/extra_fields.csv b/tests/data/unstable_csv/extra_fields.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/extra_fields.csv @@ -0,0 +1,2 @@ +a,b,c +1,2,3,EXTRA diff --git a/tests/data/unstable_csv/header_only.csv b/tests/data/unstable_csv/header_only.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/header_only.csv @@ -0,0 +1,1 @@ +first,second,third diff --git a/tests/data/unstable_csv/missing_fields.csv b/tests/data/unstable_csv/missing_fields.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/missing_fields.csv @@ -0,0 +1,3 @@ +a,b,c +1,2 +X,Y,Z diff --git a/tests/data/unstable_csv/no_trailing_newline.csv b/tests/data/unstable_csv/no_trailing_newline.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/no_trailing_newline.csv @@ -0,0 +1,2 @@ +name,city +Alice,London diff --git a/tests/data/unstable_csv/single_col.csv b/tests/data/unstable_csv/single_col.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/single_col.csv @@ -0,0 +1,4 @@ +name +Alice +Bob +Carol diff --git a/tests/data/unstable_csv/trailing_blank_line.csv b/tests/data/unstable_csv/trailing_blank_line.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/trailing_blank_line.csv @@ -0,0 +1,3 @@ +a,b,c +1,2,3 + diff --git a/tests/data/unstable_csv/whitespace_fields.csv b/tests/data/unstable_csv/whitespace_fields.csv new file mode 100644 --- /dev/null +++ b/tests/data/unstable_csv/whitespace_fields.csv @@ -0,0 +1,3 @@ +name,city + Alice , New York + Bob , Los Angeles