dataframe 3.5.0.0 → 3.6.0.0
raw patch · 30 files changed
+452/−1188 lines, 30 filesdep ~aesondep ~dataframedep ~dataframe-core
Dependency ranges changed: aeson, dataframe, dataframe-core, dataframe-expr-serializer, dataframe-lazy, dataframe-learn, dataframe-operations, dataframe-th, dataframe-viz
Files
- CHANGELOG.md +47/−0
- README.md +1/−1
- app/LazyBenchmark.hs +9/−8
- app/Synthesis.hs +1/−1
- benchmark/Main.hs +1/−1
- cbits/arrow_abi.h +0/−44
- dataframe.cabal +32/−65
- ffi/DataFrame/IO/Arrow.hs +0/−552
- ffi/DataFrame/IR.hs +0/−481
- src/DataFrame.hs +1/−1
- tests/IR/ExprJsonRoundtrip.hs +1/−1
- tests/Internal/ColumnBuilder.hs +41/−9
- tests/Internal/DictEncode.hs +3/−3
- tests/Internal/PackedText.hs +39/−1
- tests/LazyParity.hs +1/−1
- tests/Main.hs +5/−0
- tests/Operations/Aggregations.hs +1/−1
- tests/Operations/Join.hs +1/−1
- tests/Operations/Nullable.hs +70/−2
- tests/Operations/ParallelGroupBy.hs +1/−1
- tests/Operations/Record.hs +5/−3
- tests/Operations/Shuffle.hs +76/−2
- tests/Operations/Statistics.hs +54/−0
- tests/Operations/Subset.hs +51/−0
- tests/Operations/VectorKernel.hs +1/−1
- tests/Operations/Window.hs +1/−1
- tests/Operations/WriteCsv.hs +4/−2
- tests/PackedTextMigration.hs +1/−1
- tests/PrettyPrint.hs +1/−1
- tests/Simplify.hs +3/−3
CHANGELOG.md view
@@ -1,5 +1,52 @@ # Revision history for dataframe +## 3.6.0.0++### Breaking changes+* `DataFrame.Operators` has been renamed to `DataFrame.Expression.Operators`.+ The module's contents are unchanged; update the import if you were using it+ directly (`import DataFrame` is unaffected).+* `dataframe-core`'s exposed `DataFrame.Internal.*` modules were reorganised+ into namespaced groups, e.g. `DataFrame.Internal.Hash` →+ `DataFrame.Internal.Algorithms.Hash`, `DataFrame.Internal.ColumnBuilder` →+ `DataFrame.Internal.Column.Builder`, `DataFrame.Internal.GroupingPar` →+ `DataFrame.Internal.Grouping.Partitioned`. Likewise in+ `dataframe-operations`: `DataFrame.Operations.AggregateScatter` →+ `DataFrame.Operations.Aggregation.Run` and+ `DataFrame.Operations.JoinPar` → `DataFrame.Operations.Join.Parallel`.+* The `arrow-bridge` public sublibrary of this package is gone; it now ships+ as the standalone package `dataframe-arrow-bridge` (`1.0.0.0`), with the+ same modules (`DataFrame.IO.Arrow`, `DataFrame.IR`, and the re-exported+ `DataFrame.IR.ExprJson`). A sublibrary of a separately uploaded package+ cannot be resolved by dependents on Hackage — `dataframe-arrow` failed to+ configure with "missing or private dependencies: dataframe:arrow-bridge" —+ so consumers should now depend on `dataframe-arrow-bridge` directly. The+ meta-package no longer ships `cbits/arrow_abi.h`, which only the sublibrary+ referenced (`dataframe-arrow` carries its own copy).+* Coordinated bumps: `dataframe-core` and `dataframe-operations` →+ `2.5.0.0`, `dataframe-learn` → `2.4.2.0`, `dataframe-lazy` → `2.4.1.0`,+ `dataframe-th` → `2.2.1.0`, `dataframe-viz` → `1.3.2.0`,+ `dataframe-expr-serializer` → `1.2.1.0`, `dataframe-arrow` → `1.0.3.0`,+ with inter-package bounds raised to match.++### Improvements+* Faster `groupBy` and joins through better parallelism.++### Bug fixes+* Validity bits are now counted only up to the column's length, so statistics+ on a nullable column no longer include trailing bits from the bitmap's+ padding.+* `range` clamps both endpoints before subtracting, fixing an underflow when+ the range starts past the end of the frame.+* `sliceGroups` permutes the group bitmap alongside the data, so nulls stay+ attached to their rows.+* `shuffledIndices` uses Fisher-Yates, making the shuffle uniform.+* Packed text columns are sliced through their selection layer.++## 3.5.0.0++* Reduced memory pressure in CSV reads, `groupBy`, and joins.+ ## 3.4.0.0 * `impute` on a non-nullable expression is now fails and throws when given a column with the wrong type.
README.md view
@@ -89,7 +89,7 @@ import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Typed as DT-import DataFrame.Operators+import DataFrame.Expression.Operators import Data.Text (Text) import Data.Int (Int64)
app/LazyBenchmark.hs view
@@ -8,7 +8,7 @@ Usage: cabal run lazy-bench [-- [OPTIONS]] (+RTS -s -RTS for heap stats) --rows N rows to generate (default 1_000_000_000)- --file PATH output CSV path (default /tmp/lazy_1b.csv)+ --file PATH output CSV path (default: lazy_1b.csv in the system temp dir) --skip-gen reuse the file if it already exists -} module Main where@@ -19,10 +19,10 @@ import qualified Data.Text as T import Data.Time (UTCTime, diffUTCTime, getCurrentTime) import qualified DataFrame as D+import DataFrame.Expression.Operators import qualified DataFrame.Lazy as L-import DataFrame.Operators import DataFrame.Schema (Schema (..), schemaType)-import System.Directory (doesFileExist, getFileSize)+import System.Directory (doesFileExist, getFileSize, getTemporaryDirectory) import System.Environment (getArgs) import System.Exit (exitFailure) import System.IO (@@ -42,8 +42,8 @@ defaultRows :: Int defaultRows = 1_000_000_000 -defaultFile :: FilePath-defaultFile = "/tmp/lazy_1b.csv"+defaultFile :: IO FilePath+defaultFile = (<> "/lazy_1b.csv") <$> getTemporaryDirectory -- Rows written per Builder flush to disk. chunkSize :: Int@@ -59,8 +59,8 @@ , optSkipGen :: Bool } -parseArgs :: [String] -> Either String Opts-parseArgs = go (Opts defaultRows defaultFile False)+parseArgs :: FilePath -> [String] -> Either String Opts+parseArgs defFile = go (Opts defaultRows defFile False) where go opts [] = Right opts go opts ("--rows" : n : rest) = case reads n of@@ -170,7 +170,8 @@ main = do hSetBuffering stdout LineBuffering args <- getArgs- opts <- case parseArgs args of+ defFile <- defaultFile+ opts <- case parseArgs defFile args of Left err -> putStrLn ("Error: " ++ err) >> exitFailure Right o -> return o
app/Synthesis.hs view
@@ -10,8 +10,8 @@ import qualified Data.Text as T import qualified DataFrame as D import DataFrame.DecisionTree+import DataFrame.Expression.Operators import qualified DataFrame.Functions as F-import DataFrame.Operators import qualified DataFrame.Typed as DT import System.Random
benchmark/Main.hs view
@@ -9,9 +9,9 @@ import Control.DeepSeq (NFData (..)) import Control.Monad (void) import Criterion.Main+import DataFrame.Expression.Operators import DataFrame.Internal.DataFrame (forceDataFrame) import DataFrame.Operations.Join-import DataFrame.Operators import System.Process hiding (env) import System.Random.Stateful
− cbits/arrow_abi.h
@@ -1,44 +0,0 @@-/* Arrow C Data Interface structs (verbatim from specification).- See https://arrow.apache.org/docs/format/CDataInterface.html */--#pragma once--#include <stdint.h>--#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
dataframe.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-version: 3.5.0.0+version: 3.6.0.0 synopsis: A fast, safe, and intuitive DataFrame library. description: A fast, safe, and intuitive DataFrame library for exploratory data analysis.@@ -15,8 +15,7 @@ category: Data tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 extra-doc-files: CHANGELOG.md README.md-extra-source-files: cbits/arrow_abi.h- tests/data/typing/texts.txt+extra-source-files: tests/data/typing/texts.txt tests/data/typing/texts_with_empties.txt tests/data/typing/texts_with_empties_and_nullish.txt data/titanic/*.csv@@ -88,7 +87,7 @@ DataFrame.Operations.Core, DataFrame.Operations.Join, DataFrame.Operations.Merge,- DataFrame.Operators,+ DataFrame.Expression.Operators, DataFrame.Operations.Permutation, DataFrame.Operations.Subset, DataFrame.Operations.Statistics,@@ -128,13 +127,13 @@ DataFrame.Typed.Record, DataFrame.Typed.Generic build-depends: base >= 4 && <5,- dataframe-core >= 2.4 && < 2.5,+ dataframe-core >= 2.5 && < 2.6, dataframe-json >= 1.2.0.1 && < 1.3,- dataframe-expr-serializer >= 1.2.0.1 && < 1.3,- dataframe-operations >= 2.4 && < 2.5,+ dataframe-expr-serializer >= 1.2.1 && < 1.3,+ dataframe-operations >= 2.5 && < 2.6, dataframe-parsing >= 2.2 && < 2.3,- dataframe-viz >= 1.3 && < 1.4,- dataframe-learn >= 2.4 && < 2.5+ dataframe-viz >= 1.3.2 && < 1.4,+ dataframe-learn >= 2.4.2 && < 2.5 if !flag(no-csv) reexported-modules: DataFrame.IO.CSV,@@ -167,11 +166,11 @@ DataFrame.Lazy.IO.Binary, DataFrame.Lazy.IO.CSV, DataFrame.Typed.Lazy- build-depends: dataframe-lazy >= 2.4 && < 2.5+ build-depends: dataframe-lazy >= 2.4.1 && < 2.5 cpp-options: -DWITH_LAZY if !flag(no-th)- build-depends: dataframe-th >= 2.2 && < 2.3+ build-depends: dataframe-th >= 2.2.1 && < 2.3 cpp-options: -DWITH_TH exposed-modules: DataFrame.TH, DataFrame.Typed.TH@@ -187,44 +186,12 @@ hs-source-dirs: src default-language: Haskell2010 -library arrow-bridge- import: warnings- visibility: public- hs-source-dirs: ffi- exposed-modules: DataFrame.IO.Arrow- DataFrame.IR- -- The expr/pipeline JSON codec now lives in its own lightweight package;- -- re-export it so the Python FFI and Haskell consumers keep importing- -- @DataFrame.IR.ExprJson@ unchanged.- reexported-modules: DataFrame.IR.ExprJson- -- The Arrow bridge IR loads CSV + Parquet readers, so it requires- -- both backends. If either flag is off, skip building it.- if flag(no-csv) || flag(no-parquet)- buildable: False- build-depends:- base >= 4 && < 5,- dataframe-core >= 2.4 && < 2.5,- dataframe-expr-serializer >= 1.2.0.1 && < 1.3,- dataframe-csv >= 2.3 && < 2.4,- dataframe-json >= 1.2.0.1 && < 1.3,- dataframe-lazy >= 2.4 && < 2.5,- dataframe-operations >= 2.4 && < 2.5,- dataframe-parquet >= 1.5 && < 1.6,- dataframe-parsing >= 2.2 && < 2.3,- text >= 2.1 && < 3,- aeson >= 0.11 && < 3,- bytestring >= 0.11 && < 0.14,- containers >= 0.6.7 && < 0.10,- vector >= 0.13 && < 0.15- include-dirs: cbits- default-language: Haskell2010- executable dataframe-benchmark-example import: warnings main-is: Benchmark.hs build-depends: base >= 4 && < 5,- dataframe >= 3.5 && < 3.6,- dataframe-operations >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-operations >= 2.5 && < 2.6, random >= 1 && < 2, time >= 1.12 && < 2, vector >= 0.13 && < 0.15,@@ -236,10 +203,10 @@ import: warnings main-is: Synthesis.hs build-depends: base >= 4 && < 5,- dataframe >= 3.5 && < 3.6,- dataframe-core >= 2.4 && < 2.5,- dataframe-learn >= 2.4 && < 2.5,- dataframe-operations >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-core >= 2.5 && < 2.6,+ dataframe-learn >= 2.4.2 && < 2.5,+ dataframe-operations >= 2.5 && < 2.6, random >= 1 && < 2, text >= 2.1 && < 3 hs-source-dirs: app@@ -270,9 +237,9 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10,- dataframe >= 3.5 && < 3.6,- dataframe-core >= 2.4 && < 2.5,- dataframe-lazy >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-core >= 2.5 && < 2.6,+ dataframe-lazy >= 2.4.1 && < 2.5, dataframe-parsing >= 2.2 && < 2.3, directory >= 1.3.0.0 && < 2, random >= 1 && < 2,@@ -291,9 +258,9 @@ criterion >= 1 && < 2, deepseq >= 1.4 && < 2, process >= 1.6 && < 2,- dataframe >= 3.5 && < 3.6,- dataframe-core >= 2.4 && < 2.5,- dataframe-operations >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-core >= 2.5 && < 2.6,+ dataframe-operations >= 2.5 && < 2.6, random >= 1 && < 2, default-language: Haskell2010 ghc-options:@@ -377,16 +344,16 @@ build-depends: base >= 4 && < 5, aeson >= 0.11.0.0 && < 3, bytestring >= 0.11 && < 0.14,- dataframe >= 3.5 && < 3.6,- dataframe-core >= 2.4 && < 2.5,- dataframe-core >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-core >= 2.5 && < 2.6,+ dataframe-core >= 2.5 && < 2.6, dataframe-csv >= 2.3 && < 2.4,- dataframe-expr-serializer >= 1.2.0.1 && < 1.3,+ dataframe-expr-serializer >= 1.2.1 && < 1.3, dataframe-fastcsv >= 1.4.0.1 && < 1.5, dataframe-json >= 1.2.0.1 && < 1.3,- dataframe-lazy >= 2.4 && < 2.5,- dataframe-learn >= 2.4 && < 2.5,- dataframe-operations >= 2.4 && < 2.5,+ dataframe-lazy >= 2.4.1 && < 2.5,+ dataframe-learn >= 2.4.2 && < 2.5,+ dataframe-operations >= 2.5 && < 2.6, dataframe-parquet >= 1.5 && < 1.6, dataframe-parsing >= 2.2 && < 2.3, HUnit >= 1.6 && < 1.8,@@ -413,9 +380,9 @@ other-modules: Internal.PackedText build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14,- dataframe >= 3.5 && < 3.6,- dataframe-core >= 2.4 && < 2.5,- dataframe-operations >= 2.4 && < 2.5,+ dataframe >= 3.6 && < 3.7,+ dataframe-core >= 2.5 && < 2.6,+ dataframe-operations >= 2.5 && < 2.6, HUnit >= 1.6 && < 1.8, text >= 2.1 && < 3, vector >= 0.13 && < 0.15
− ffi/DataFrame/IO/Arrow.hs
@@ -1,552 +0,0 @@-{-# 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,- 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, when, zipWithM_)-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))-import Foreign-import Foreign.C.String (CString, newCString, peekCString)-import Type.Reflection (typeRep)--import DataFrame.Internal.Column (Column (..))-import DataFrame.Internal.DataFrame (DataFrame (..), 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--{- | Allocate an Arrow-format validity bitmap from a 'DI.Bitmap'.-Returns (ptr, nullCount). Caller must 'free' the pointer.--}-bitmapToPtr :: Int -> DI.Bitmap -> IO (Ptr Word8, Int)-bitmapToPtr n bm = do- let numBytes = max 1 ((n + 7) `div` 8)- validCount = VU.foldl' (\acc b -> acc + popCount b) 0 bm- nullCount = n - validCount- bitmapPtr <- mallocBytes numBytes :: IO (Ptr Word8)- VU.imapM_ (pokeElemOff bitmapPtr) bm- when (VU.length bm < numBytes) $- mapM_- (\i -> pokeElemOff bitmapPtr i (0 :: Word8))- [VU.length bm .. numBytes - 1]- return (bitmapPtr, nullCount)---- | Read an Arrow validity bitmap into a 'DI.Bitmap'.-readArrowBitmap :: Ptr Word8 -> Int -> IO DI.Bitmap-readArrowBitmap bitmapPtr n = VU.generateM ((n + 7) `div` 8) (peekElemOff bitmapPtr)--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 Nothing (vec :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = do- let n = V.length vec- bss = map TE.encodeUtf8 (V.toList vec)- 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 Nothing (vec :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do- let n = V.length vec- dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double)- V.imapM_ (pokeElemOff dataPtr) vec- sPtr <- makeLeafSchema "g" colName- aPtr <- makeLeafArray n 0 [nullPtr, castPtr dataPtr] (free dataPtr)- return (sPtr, aPtr)-columnToArrow colName (BoxedColumn Nothing (vec :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do- let n = V.length vec- dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64)- 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)--- Nullable Int (UnboxedColumn with bitmap)-columnToArrow colName (UnboxedColumn (Just bm) (vec :: VU.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) = do- let n = VU.length vec- (bitmapPtr, nullCount) <- bitmapToPtr n bm- dataPtr <- mallocArray (max 1 n) :: IO (Ptr Int64)- VU.imapM_ (\i v -> pokeElemOff dataPtr i (fromIntegral v :: Int64)) vec- sPtr <- makeLeafSchema "l" colName- aPtr <-- makeLeafArray- n- (fromIntegral nullCount)- [castPtr bitmapPtr, castPtr dataPtr]- (free bitmapPtr >> free dataPtr)- return (sPtr, aPtr)--- Nullable Double (UnboxedColumn with bitmap)-columnToArrow colName (UnboxedColumn (Just bm) (vec :: VU.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) = do- let n = VU.length vec- (bitmapPtr, nullCount) <- bitmapToPtr n bm- dataPtr <- mallocArray (max 1 n) :: IO (Ptr Double)- VU.imapM_ (\i v -> pokeElemOff dataPtr i (realToFrac v :: Double)) vec- sPtr <- makeLeafSchema "g" colName- aPtr <-- makeLeafArray- n- (fromIntegral nullCount)- [castPtr bitmapPtr, castPtr dataPtr]- (free bitmapPtr >> free dataPtr)- return (sPtr, aPtr)--- Nullable Text (BoxedColumn with bitmap)-columnToArrow colName (BoxedColumn (Just bm) (vec :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) = do- let n = V.length vec- -- For null positions, use empty BS (null placeholder in vec is never evaluated)- bss =- map- (\i -> if DI.bitmapTestBit bm i then TE.encodeUtf8 (vec V.! i) else BS.empty)- [0 .. n - 1]- cumOff = scanl (+) 0 (map BS.length bss)- total = last cumOff- (bitmapPtr, nullCount) <- bitmapToPtr n bm- 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- (_, col) : _ -> DI.columnLength col- 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)--{- | 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- bm <- readArrowBitmap bitmapPtr n- vec <- VU.generateM n $ \i -> fmap fromIntegral (peekElemOff dataPtr i :: IO Int64)- return $ UnboxedColumn (Just bm) (vec :: VU.Vector Int)- else do- vec <- VU.generateM n $ \i -> fmap fromIntegral (peekElemOff dataPtr i :: IO Int64)- return $ UnboxedColumn Nothing (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- bm <- readArrowBitmap bitmapPtr n- vec <- VU.generateM n $ \i -> fmap fromIntegral (peekElemOff dataPtr i :: IO Int32)- return $ UnboxedColumn (Just bm) (vec :: VU.Vector Int)- else do- vec <- VU.generateM n $ \i -> fmap fromIntegral (peekElemOff dataPtr i :: IO Int32)- return $ UnboxedColumn Nothing (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- bm <- readArrowBitmap bitmapPtr n- vec <- VU.generateM n (peekElemOff dataPtr)- return $ UnboxedColumn (Just bm) (vec :: VU.Vector Double)- else do- vec <- VU.generateM n (peekElemOff dataPtr)- return $ UnboxedColumn Nothing (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- bm <- readArrowBitmap bitmapPtr n- vec <- VU.generateM n $ \i -> fmap (realToFrac :: Float -> Double) (peekElemOff dataPtr i)- return $ UnboxedColumn (Just bm) (vec :: VU.Vector Double)- else do- vec <- VU.generateM n $ \i -> fmap (realToFrac :: Float -> Double) (peekElemOff dataPtr i)- return $ UnboxedColumn Nothing (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- let readText i = do- start <- fromIntegral <$> peekElemOff offsetPtr i- end <- fromIntegral <$> peekElemOff offsetPtr (i + 1)- TE.decodeUtf8- <$> BS.packCStringLen (castPtr (charPtr `plusPtr` start), end - start)- if nullCnt > 0- then do- let bitmapPtr = castPtr bitmapVoid :: Ptr Word8- bm <- readArrowBitmap bitmapPtr n- vec <- V.generateM n readText- return $ BoxedColumn (Just bm) vec- else do- vec <- V.generateM n readText- return $ BoxedColumn Nothing 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- readText i = do- start <- fromIntegral <$> peekElemOff offsetPtr i- end <- fromIntegral <$> peekElemOff offsetPtr (i + 1)- TE.decodeUtf8- <$> BS.packCStringLen (castPtr (charPtr `plusPtr` start), end - start)- if nullCnt > 0- then do- let bitmapPtr = castPtr bitmapVoid :: Ptr Word8- bm <- readArrowBitmap bitmapPtr n- vec <- V.generateM n readText- return $ BoxedColumn (Just bm) vec- else do- vec <- V.generateM n readText- return $ BoxedColumn Nothing vec
− ffi/DataFrame/IR.hs
@@ -1,481 +0,0 @@-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE ExplicitNamespaces #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE OverloadedStrings #-}-{-# LANGUAGE RankNTypes #-}-{-# 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 qualified Data.Aeson as Aeson-import Data.Aeson.Types (Parser)-import qualified Data.ByteString as BS-import Data.Int (Int16, Int32, Int64, Int8)-import qualified Data.Text as T-import Data.Type.Equality (- TestEquality (testEquality),- type (:~:) (Refl),- type (:~~:) (HRefl),- )-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as VU-import Data.Word (Word16, Word32, Word64, Word8)-import Foreign (wordPtrToPtr)-import Type.Reflection (SomeTypeRep (..), eqTypeRep, typeRep)--import DataFrame.Functions (count, mean, meanMaybe, sumMaybe)-import qualified DataFrame.Functions as Functions-import DataFrame.IO.Arrow (arrowToDataframe)-import DataFrame.IO.CSV (- CsvReader,- defaultReadOptions,- readSeparated,- readTsv,- writeCsv,- )-import DataFrame.IO.JSON (readJSON)-import qualified DataFrame.IO.Parquet as Parquet-import DataFrame.IR.ExprJson (SomeExpr (..), decodeExprAny, decodeExprAt)-import DataFrame.Internal.Column (- Column (..),- Columnable,- mergedHead,- )-import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn)-import DataFrame.Internal.Expression (Expr (..), NamedExpr)-import qualified DataFrame.Lazy as Lazy-import DataFrame.Operations.Aggregation (aggregate, distinct, groupBy)-import DataFrame.Operations.Core (insertVector, renameMany)-import DataFrame.Operations.Join (JoinType (..), join)-import DataFrame.Operations.Permutation (SortOrder (..), sortBy)-import qualified DataFrame.Operations.Statistics as Stats-import DataFrame.Operations.Subset (exclude, filterWhere, range, select)-import qualified DataFrame.Operations.Subset as Subset-import DataFrame.Operations.Transformations (derive)-import DataFrame.Operators ((.=))-import DataFrame.Schema (Schema, makeSchema, schemaType)---- ------------------------------------------------------------------------------ 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- | -- | predicate JSON, child plan- Filter Aeson.Value PlanNode- | -- | column name, expr JSON, child plan- Derive T.Text Aeson.Value PlanNode- | Exclude [T.Text] PlanNode- | Rename [(T.Text, T.Text)] PlanNode- | Distinct PlanNode- | TakeLast Int PlanNode- | Drop Int PlanNode- | DropLast Int PlanNode- | Range Int Int PlanNode- | -- | joinType ("inner"|"left"|"right"|"outer"), shared key columns, left, right- Join T.Text [T.Text] PlanNode PlanNode- | Describe PlanNode- | -- | first column, second column, child plan- Correlation T.Text T.Text PlanNode- | Frequencies T.Text PlanNode- | ReadParquet FilePath- | ReadJson FilePath- | -- | path, separator (single character), child plan; runs as a terminal op- WriteCsv FilePath PlanNode- | {- | path, schema (column name → type-tag map). Reads via the lazy- engine with predicate / projection pushdown; subsequent ops- currently still run eagerly on the materialized result.- -}- ScanCsv FilePath [(T.Text, T.Text)]- | ScanParquet FilePath [(T.Text, T.Text)]- deriving (Show)---- ------------------------------------------------------------------------------ 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"- "Filter" -> Filter <$> o .: "predicate" <*> o .: "input"- "Derive" -> Derive <$> o .: "name" <*> o .: "expr" <*> o .: "input"- "Exclude" -> Exclude <$> o .: "cols" <*> o .: "input"- "Rename" -> Rename <$> o .: "pairs" <*> o .: "input"- "Distinct" -> Distinct <$> o .: "input"- "TakeLast" -> TakeLast <$> o .: "n" <*> o .: "input"- "Drop" -> Drop <$> o .: "n" <*> o .: "input"- "DropLast" -> DropLast <$> o .: "n" <*> o .: "input"- "Range" -> Range <$> o .: "start" <*> o .: "end" <*> o .: "input"- "Join" ->- Join- <$> o .: "how"- <*> o .: "on"- <*> o .: "left"- <*> o .: "right"- "Describe" -> Describe <$> o .: "input"- "Correlation" ->- Correlation- <$> o .: "first"- <*> o .: "second"- <*> o .: "input"- "Frequencies" -> Frequencies <$> o .: "col" <*> o .: "input"- "ReadParquet" -> ReadParquet <$> o .: "path"- "ReadJson" -> ReadJson <$> o .: "path"- "WriteCsv" -> WriteCsv <$> o .: "path" <*> o .: "input"- "ScanCsv" -> ScanCsv <$> o .: "path" <*> o .: "schema"- "ScanParquet" -> ScanParquet <$> o .: "path" <*> o .: "schema"- _ -> fail $ "DataFrame.IR: unknown op: " ++ T.unpack op--executePlan :: CsvReader -> PlanNode -> IO DataFrame-executePlan _reader (ReadCsv path) =- readSeparated defaultReadOptions path-executePlan _reader (ReadTsv path) =- readTsv path-executePlan _reader (FromArrow schemaAddr arrayAddr) =- arrowToDataframe- (wordPtrToPtr (fromIntegral schemaAddr))- (wordPtrToPtr (fromIntegral arrayAddr))-executePlan reader (Select cols node) =- select cols <$> executePlan reader node-executePlan reader (GroupBy keys aggs node) = do- df <- executePlan reader node- nes <- mapM (buildNamedExpr df) aggs- return $ aggregate nes (groupBy keys df)-executePlan reader (Sort cols ascending node) = do- df <- executePlan reader node- let orders = map (\c -> mkSortOrder ascending c (unsafeGetColumn c df)) cols- return $ sortBy orders df-executePlan reader (Limit k node) =- Subset.take k <$> executePlan reader node-executePlan reader (Filter predJson node) = do- df <- executePlan reader node- case decodeExprAt @Bool predJson of- Right pred_ -> return $ filterWhere pred_ df- Left err -> ioError $ userError $ "DataFrame.IR.Filter: " <> err-executePlan reader (Derive name exprJson node) = do- df <- executePlan reader node- case decodeExprAny exprJson of- Right (SomeExpr _trep expr) -> return $ derive name expr df- Left err -> ioError $ userError $ "DataFrame.IR.Derive: " <> err-executePlan reader (Exclude cols node) =- exclude cols <$> executePlan reader node-executePlan reader (Rename pairs node) =- renameMany pairs <$> executePlan reader node-executePlan reader (Distinct node) =- distinct <$> executePlan reader node-executePlan reader (TakeLast n node) =- Subset.takeLast n <$> executePlan reader node-executePlan reader (Drop n node) =- Subset.drop n <$> executePlan reader node-executePlan reader (DropLast n node) =- Subset.dropLast n <$> executePlan reader node-executePlan reader (Range start end node) =- range (start, end) <$> executePlan reader node-executePlan reader (Join how on leftPlan rightPlan) = do- left <- executePlan reader leftPlan- right <- executePlan reader rightPlan- jt <- case how of- "inner" -> return INNER- "left" -> return LEFT- "right" -> return RIGHT- "outer" -> return FULL_OUTER- "full_outer" -> return FULL_OUTER- other ->- ioError . userError $- "DataFrame.IR.Join: unknown join type " <> T.unpack other- return $ join jt on left right-executePlan reader (Describe node) = Stats.summarize <$> executePlan reader node-executePlan reader (Correlation a b node) = do- df <- executePlan reader node- let r = Stats.correlation a b df- valueCol = case r of- Just d -> V.singleton d- Nothing -> V.singleton (0 / 0 :: Double)- return $- insertVector "first" (V.singleton a) $- insertVector "second" (V.singleton b) $- insertVector "correlation" valueCol mempty-executePlan reader (Frequencies colName node) = do- df <- executePlan reader node- runFrequencies colName df-executePlan _reader (ReadParquet path) = Parquet.readParquet path-executePlan _reader (ReadJson path) = readJSON path-executePlan reader (WriteCsv path node) = do- df <- executePlan reader node- writeCsv path df- return df-executePlan reader (ScanCsv path schemaPairs) = do- schema <- buildSchema schemaPairs- Lazy.runDataFrame (Lazy.scanCsvWith reader schema (T.pack path))-executePlan _reader (ScanParquet path schemaPairs) = do- schema <- buildSchema schemaPairs- Lazy.runDataFrame (Lazy.scanParquet schema (T.pack path))--{- | Build a SortOrder from a column's runtime type.-Uses type dispatch to recover Ord for known column types.--}-mkSortOrder :: Bool -> T.Text -> Column -> SortOrder-mkSortOrder isAsc name col = dispatchType (columnTypeRep col)- where- columnTypeRep :: Column -> SomeTypeRep- columnTypeRep (UnboxedColumn _ (_ :: VU.Vector a)) = SomeTypeRep (typeRep @a)- columnTypeRep (BoxedColumn _ (_ :: V.Vector a)) = SomeTypeRep (typeRep @a)- columnTypeRep (PackedText _ _) = SomeTypeRep (typeRep @T.Text)- columnTypeRep c@(MergedColumn _ _) = columnTypeRep (mergedHead c)- mk :: (Columnable a, Ord a) => Expr a -> SortOrder- mk = if isAsc then Asc else Desc- dispatchType (SomeTypeRep tr)- | Just HRefl <- eqTypeRep tr (typeRep @Int) = mk (Col @Int name)- | Just HRefl <- eqTypeRep tr (typeRep @Int8) = mk (Col @Int8 name)- | Just HRefl <- eqTypeRep tr (typeRep @Int16) = mk (Col @Int16 name)- | Just HRefl <- eqTypeRep tr (typeRep @Int32) = mk (Col @Int32 name)- | Just HRefl <- eqTypeRep tr (typeRep @Int64) = mk (Col @Int64 name)- | Just HRefl <- eqTypeRep tr (typeRep @Word) = mk (Col @Word name)- | Just HRefl <- eqTypeRep tr (typeRep @Word8) = mk (Col @Word8 name)- | Just HRefl <- eqTypeRep tr (typeRep @Word16) = mk (Col @Word16 name)- | Just HRefl <- eqTypeRep tr (typeRep @Word32) = mk (Col @Word32 name)- | Just HRefl <- eqTypeRep tr (typeRep @Word64) = mk (Col @Word64 name)- | Just HRefl <- eqTypeRep tr (typeRep @Integer) = mk (Col @Integer name)- | Just HRefl <- eqTypeRep tr (typeRep @Double) = mk (Col @Double name)- | Just HRefl <- eqTypeRep tr (typeRep @Float) = mk (Col @Float name)- | Just HRefl <- eqTypeRep tr (typeRep @Bool) = mk (Col @Bool name)- | Just HRefl <- eqTypeRep tr (typeRep @Char) = mk (Col @Char name)- | Just HRefl <- eqTypeRep tr (typeRep @T.Text) = mk (Col @T.Text name)- | Just HRefl <- eqTypeRep tr (typeRep @String) = mk (Col @String name)- | Just HRefl <- eqTypeRep tr (typeRep @BS.ByteString) =- mk (Col @BS.ByteString name)- | otherwise = error $ "mkSortOrder: unsupported column type: " ++ show tr---- | 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)- "min" -> minMaxExpr Functions.minimum name colName (unsafeGetColumn colName df)- "max" -> minMaxExpr Functions.maximum name colName (unsafeGetColumn colName df)- "median" -> doubleStatExpr Functions.median name colName (unsafeGetColumn colName df)- "variance" -> doubleStatExpr Functions.variance name colName (unsafeGetColumn colName df)- "std" -> doubleStatExpr stdDevExpr name colName (unsafeGetColumn colName df)- other ->- ioError $- userError $- "DataFrame.IR: unknown aggregation '" ++ T.unpack other ++ "'"---- | Variance → standard deviation; sqrt of the underlying variance Expr.-stdDevExpr :: (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double-stdDevExpr e = sqrt (Functions.variance e)---- | Build a 'Schema' from a list of (col, type-tag) pairs sent over the wire.-buildSchema :: [(T.Text, T.Text)] -> IO Schema-buildSchema pairs = do- sch <- mapM resolve pairs- return (makeSchema sch)- where- resolve (name, tag) = case tag of- "int" -> return (name, schemaType @Int)- "int8" -> return (name, schemaType @Int8)- "int16" -> return (name, schemaType @Int16)- "int32" -> return (name, schemaType @Int32)- "int64" -> return (name, schemaType @Int64)- "double" -> return (name, schemaType @Double)- "float" -> return (name, schemaType @Float)- "bool" -> return (name, schemaType @Bool)- "text" -> return (name, schemaType @T.Text)- "string" -> return (name, schemaType @String)- other ->- ioError . userError $- "DataFrame.IR.buildSchema: unsupported schema type tag '"- ++ T.unpack other- ++ "' for column '"- ++ T.unpack name- ++ "'"---- | Dispatch 'frequencies' on the column's runtime element type.-runFrequencies :: T.Text -> DataFrame -> IO DataFrame-runFrequencies colName df = dispatchType (columnTypeRep (unsafeGetColumn colName df))- where- columnTypeRep :: Column -> SomeTypeRep- columnTypeRep (UnboxedColumn _ (_ :: VU.Vector a)) = SomeTypeRep (typeRep @a)- columnTypeRep (BoxedColumn _ (_ :: V.Vector a)) = SomeTypeRep (typeRep @a)- columnTypeRep (PackedText _ _) = SomeTypeRep (typeRep @T.Text)- columnTypeRep c@(MergedColumn _ _) = columnTypeRep (mergedHead c)-- fr :: forall a. (Columnable a, Ord a) => IO DataFrame- fr = return $ Stats.frequencies (Col @a colName) df-- dispatchType :: SomeTypeRep -> IO DataFrame- dispatchType (SomeTypeRep tr)- | Just HRefl <- eqTypeRep tr (typeRep @Int) = fr @Int- | Just HRefl <- eqTypeRep tr (typeRep @Int8) = fr @Int8- | Just HRefl <- eqTypeRep tr (typeRep @Int16) = fr @Int16- | Just HRefl <- eqTypeRep tr (typeRep @Int32) = fr @Int32- | Just HRefl <- eqTypeRep tr (typeRep @Int64) = fr @Int64- | Just HRefl <- eqTypeRep tr (typeRep @Word) = fr @Word- | Just HRefl <- eqTypeRep tr (typeRep @Word8) = fr @Word8- | Just HRefl <- eqTypeRep tr (typeRep @Word16) = fr @Word16- | Just HRefl <- eqTypeRep tr (typeRep @Word32) = fr @Word32- | Just HRefl <- eqTypeRep tr (typeRep @Word64) = fr @Word64- | Just HRefl <- eqTypeRep tr (typeRep @Integer) = fr @Integer- | Just HRefl <- eqTypeRep tr (typeRep @Double) = fr @Double- | Just HRefl <- eqTypeRep tr (typeRep @Float) = fr @Float- | Just HRefl <- eqTypeRep tr (typeRep @Bool) = fr @Bool- | Just HRefl <- eqTypeRep tr (typeRep @Char) = fr @Char- | Just HRefl <- eqTypeRep tr (typeRep @T.Text) = fr @T.Text- | Just HRefl <- eqTypeRep tr (typeRep @String) = fr @String- | otherwise =- ioError . userError $- "DataFrame.IR.Frequencies: unsupported column type for '"- ++ T.unpack colName- ++ "'"--countExpr :: T.Text -> T.Text -> Column -> IO NamedExpr-countExpr name colName (UnboxedColumn Nothing (_ :: VU.Vector a)) = return $ name .= count (Col @a colName)-countExpr name colName (UnboxedColumn (Just _) (_ :: VU.Vector a)) = return $ name .= count (Col @(Maybe a) colName)-countExpr name colName (BoxedColumn Nothing (_ :: V.Vector a)) = return $ name .= count (Col @a colName)-countExpr name colName (BoxedColumn (Just _) (_ :: V.Vector a)) = return $ name .= count (Col @(Maybe a) colName)-countExpr name colName (PackedText Nothing _) = return $ name .= count (Col @T.Text colName)-countExpr name colName (PackedText (Just _) _) = return $ name .= count (Col @(Maybe T.Text) colName)-countExpr name colName c@(MergedColumn _ _) = countExpr name colName (mergedHead c)--sumExpr :: T.Text -> T.Text -> Column -> IO NamedExpr-sumExpr name colName (UnboxedColumn Nothing (_ :: 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 (UnboxedColumn (Just _) (_ :: VU.Vector 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 (BoxedColumn Nothing (_ :: 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 (BoxedColumn (Just _) (_ :: V.Vector 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 _ 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 Nothing (_ :: 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 (UnboxedColumn (Just _) (_ :: VU.Vector 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 (BoxedColumn Nothing (_ :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =- return $ name .= mean (Col @Double colName)-meanExpr name colName (BoxedColumn (Just _) (_ :: V.Vector 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 _ colName _ =- ioError $- userError $- "DataFrame.IR: mean: unsupported column type for '" ++ T.unpack colName ++ "'"---- | min / max — preserve column type, require Ord.-minMaxExpr ::- (forall a. (Columnable a, Ord a) => Expr a -> Expr a) ->- T.Text ->- T.Text ->- Column ->- IO NamedExpr-minMaxExpr op name colName (UnboxedColumn Nothing (_ :: VU.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) =- return $ name .= op (Col @Int colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =- return $ name .= op (Col @Double colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Float) =- return $ name .= op (Col @Float colName)-minMaxExpr op name colName (BoxedColumn Nothing (_ :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) =- return $ name .= op (Col @T.Text colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) =- return $ name .= op (Col @Int colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =- return $ name .= op (Col @Double colName)-minMaxExpr _ _ colName _ =- ioError . userError $- "DataFrame.IR: min/max: unsupported column type for '"- ++ T.unpack colName- ++ "'"---- | median / variance / std — return Double, require Real + Unbox.-doubleStatExpr ::- (forall a. (Columnable a, Real a, VU.Unbox a) => Expr a -> Expr Double) ->- T.Text ->- T.Text ->- Column ->- IO NamedExpr-doubleStatExpr op name colName (UnboxedColumn Nothing (_ :: VU.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) =- return $ name .= op (Col @Int colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =- return $ name .= op (Col @Double colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Float) =- return $ name .= op (Col @Float colName)-doubleStatExpr op name colName (BoxedColumn Nothing (_ :: V.Vector a))- | Just Refl <- testEquality (typeRep @a) (typeRep @Int) =- return $ name .= op (Col @Int colName)- | Just Refl <- testEquality (typeRep @a) (typeRep @Double) =- return $ name .= op (Col @Double colName)-doubleStatExpr _ _ colName _ =- ioError . userError $- "DataFrame.IR: median/variance/std: unsupported column type for '"- ++ T.unpack colName- ++ "'"
src/DataFrame.hs view
@@ -284,7 +284,7 @@ effectiveSafeRead, parseDefaults, )-import DataFrame.Operators as Operators+import DataFrame.Expression.Operators as Operators #ifdef WITH_TH import DataFrame.TH as TH ( declareColumns,
tests/IR/ExprJsonRoundtrip.hs view
@@ -15,12 +15,12 @@ import qualified Data.Vector.Unboxed as VU import qualified DataFrame as D+import DataFrame.Expression.Operators (ifThenElse, (.>.)) import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.Expression (Expr, UExpr (..)) import DataFrame.Internal.Interpreter (interpret)-import DataFrame.Operators (ifThenElse, (.>.)) import DataFrame.LinearModel (defaultLinearConfig) import DataFrame.Model (fit, predict)
tests/Internal/ColumnBuilder.hs view
@@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-}@@ -18,11 +19,42 @@ import Data.Text.Encoding (decodeUtf8Lenient, encodeUtf8) import Data.Type.Equality (testEquality, (:~:) (Refl)) import Data.Word (Word8)-import DataFrame.Internal.Column hiding (mergeColumns)-import DataFrame.Internal.ColumnBuilder+import DataFrame.Internal.Column (+ Column (BoxedColumn, PackedText, UnboxedColumn),+ Columnable,+ columnElemIsNull,+ columnVersionString,+ fromVector,+ materializePacked,+ )+import DataFrame.Internal.Column.Bitmap (bitmapTestBit)+import DataFrame.Internal.Column.Builder (+ ColumnBuilder (appendNull, builderLength, freezeBuilder),+ appendDouble,+ appendInt,+ appendText,+ appendTextSlice,+ appendTextSliceFromPtr,+ concatColumns,+ newDoubleBuilder,+ newIntBuilder,+ newTextBuilder,+ ) import Foreign.Ptr (castPtr)-import Test.HUnit-import Test.QuickCheck+import Test.HUnit (+ Test (TestCase, TestLabel),+ assertEqual,+ assertFailure,+ )+import Test.QuickCheck (+ Gen,+ Property,+ Testable (property),+ chooseInt,+ forAll,+ (.&&.),+ (===),+ ) import Type.Reflection (typeRep) -- Build a column by appending each list element (Nothing => appendNull).@@ -211,7 +243,7 @@ let nullIdx = [1, 4, 9, 10, 14] :: [Int] xs = [if i `elem` nullIdx then Nothing else Just i | i <- [0 .. 14]] chunks = [take 3 xs, take 5 (drop 3 xs), drop 8 xs]- merged = mergeColumns (map buildIntColumn chunks)+ merged = concatColumns (map buildIntColumn chunks) assertEqual "merged equals unsplit" (buildIntColumn xs) merged forM_ [0 .. 14] $ \i -> assertEqual@@ -261,27 +293,27 @@ prop_mergeIntMatchesUnsplit :: [Maybe Int] -> Property prop_mergeIntMatchesUnsplit xs = forAll (splitsOf xs) $ \chunks ->- let merged = mergeColumns (map buildIntColumn chunks)+ let merged = concatColumns (map buildIntColumn chunks) unsplit = buildIntColumn xs in merged === unsplit .&&. columnVersionString merged === columnVersionString unsplit prop_mergeDoubleMatchesUnsplit :: [Maybe Double] -> Property prop_mergeDoubleMatchesUnsplit xs = forAll (splitsOf xs) $ \chunks ->- mergeColumns (map buildDoubleColumn chunks) === buildDoubleColumn xs+ concatColumns (map buildDoubleColumn chunks) === buildDoubleColumn xs prop_mergeTextMatchesUnsplit :: [Maybe String] -> Property prop_mergeTextMatchesUnsplit ss = let xs = map (fmap T.pack) ss in forAll (splitsOf xs) $ \chunks ->- let merged = mergeColumns (map buildTextColumn chunks)+ let merged = concatColumns (map buildTextColumn chunks) unsplit = buildTextColumn xs in merged === unsplit .&&. columnTexts merged === columnTexts unsplit prop_mergeNullsLandAtRightRows :: [Maybe Int] -> Property prop_mergeNullsLandAtRightRows xs = forAll (splitsOf xs) $ \chunks ->- let merged = mergeColumns (map buildIntColumn chunks)+ let merged = concatColumns (map buildIntColumn chunks) in map isNothing xs === [columnElemIsNull merged i | i <- [0 .. length xs - 1]]
tests/Internal/DictEncode.hs view
@@ -2,7 +2,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {- | Correctness oracle for the dictionary-encode building block-('DataFrame.Internal.DictEncode'). A text column encodes to dense+('DataFrame.Internal.Column.Encode'). A text column encodes to dense first-appearance @Int@ codes: two rows share a code iff their text is equal, the codes are a contiguous @0 .. card-1@ range in first-appearance order, packed and boxed Text encode identically, and the cap parameter bails past a cardinality.@@ -21,8 +21,8 @@ import Data.Text.Encoding (encodeUtf8) import Data.Word (Word8) import qualified DataFrame.Internal.Column as DI-import DataFrame.Internal.DictEncode (dictEncodeColumn, dictEncodeColumnUpTo)-import DataFrame.Internal.PackedText (mkPackedContiguous)+import DataFrame.Internal.Column.Encode (dictEncodeColumn, dictEncodeColumnUpTo)+import DataFrame.Internal.Data.PackedText (mkPackedContiguous) import Test.HUnit sampleRows :: [T.Text]
tests/Internal/PackedText.hs view
@@ -14,6 +14,7 @@ import qualified Data.Text.Array as A import qualified Data.Vector.Unboxed as VU +import Control.Exception (SomeException, evaluate, try) import Control.Monad (zipWithM_) import qualified Data.ByteString as B import Data.Text.Encoding (encodeUtf8)@@ -21,8 +22,8 @@ import qualified DataFrame as D import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI+import DataFrame.Internal.Data.PackedText (mkPackedContiguous) import DataFrame.Internal.DataFrame (unsafeGetColumn)-import DataFrame.Internal.PackedText (mkPackedContiguous) import qualified DataFrame.Operations.Aggregation as Agg import Test.HUnit @@ -169,6 +170,40 @@ (DI.isPackedText (unsafeGetColumn "v" joinedP)) assertBool "left join packed == boxed" (joinedP == joinedB) +-- A slice of a packed column stays packed: decoding to Text would make a+-- ten-row slice cost the whole frame.+slicePreservesPacked :: Test+slicePreservesPacked = TestCase $ do+ let sp = DI.sliceColumn 3 4 packedCol+ sb = DI.sliceColumn 3 4 boxedCol+ assertBool "sliced packed stays PackedText" (DI.isPackedText sp)+ assertBool "sliced packed == sliced boxed" (sp == sb)+ assertEqual+ "sliced packed toList == boxed"+ (DI.toList @T.Text sb)+ (DI.toList @T.Text sp)++-- takeLastColumn is sliceColumn from an offset, so it must stay packed too.+takeLastPreservesPacked :: Test+takeLastPreservesPacked = TestCase $ do+ let tp = DI.takeLastColumn 3 packedCol+ tb = DI.takeLastColumn 3 boxedCol+ assertBool "takeLast packed stays PackedText" (DI.isPackedText tp)+ assertEqual+ "takeLast packed toList == boxed"+ (DI.toList @T.Text tb)+ (DI.toList @T.Text tp)++-- Every other representation rejects a slice running past the end; the packed+-- arm must not quietly decode the overrun as empty strings.+sliceRejectsInvalidBounds :: Test+sliceRejectsInvalidBounds = TestCase $ do+ r <- try (evaluate (DI.columnLength (DI.sliceColumn 8 5 packedCol)))+ case r :: Either SomeException Int of+ Left _ -> pure ()+ Right len ->+ assertFailure ("expected an invalid slice, got length " ++ show len)+ tests :: [Test] tests = [ TestLabel "PackedText display parity" displayParity@@ -183,4 +218,7 @@ , TestLabel "PackedText left-join sentinel preserves packed" leftJoinSentinelPreservesPacked+ , TestLabel "PackedText slice preserves packed" slicePreservesPacked+ , TestLabel "PackedText takeLast preserves packed" takeLastPreservesPacked+ , TestLabel "PackedText slice rejects bad bounds" sliceRejectsInvalidBounds ]
tests/LazyParity.hs view
@@ -13,6 +13,7 @@ import Data.Text (Text) import qualified Data.Text as T import qualified DataFrame as D+import DataFrame.Expression.Operators (as, (|>)) import qualified DataFrame.Functions as F import qualified DataFrame.IO.CSV as Csv import qualified DataFrame.Internal.Column as DI@@ -23,7 +24,6 @@ import DataFrame.Operations.Join (JoinType (LEFT)) import qualified DataFrame.Operations.Join as Join import qualified DataFrame.Operations.Permutation as Perm-import DataFrame.Operators (as, (|>)) import DataFrame.Schema (Schema (..), schemaType) import System.Directory (removeFile) import System.IO.Temp (emptySystemTempFile)
tests/Main.hs view
@@ -143,6 +143,10 @@ mapM (quickCheckWithResult stdArgs) Operations.Subset.tests+ subsetPropRes <-+ mapM+ (quickCheckWithResult stdArgs)+ Operations.Subset.properties monadRes <- mapM (quickCheckWithResult stdArgs) Monad.tests cbRes <- mapM@@ -151,6 +155,7 @@ propsRes <- mapM (quickCheckWithResult stdArgs) Properties.tests catRes <- mapM (quickCheckWithResult stdArgs) Properties.Categorical.tests if not (all isSuccessful propRes)+ || not (all isSuccessful subsetPropRes) || not (all isSuccessful cbRes) || not (all isSuccessful monadRes) || not (all isSuccessful propsRes)
tests/Operations/Aggregations.hs view
@@ -11,7 +11,7 @@ import qualified DataFrame.Typed as DT import Data.Function-import DataFrame.Operators+import DataFrame.Expression.Operators import Test.HUnit values :: [(T.Text, DI.Column)]
tests/Operations/Join.hs view
@@ -9,7 +9,7 @@ import Data.Text (Text, unpack) import qualified DataFrame as D import qualified DataFrame.Functions as F-import DataFrame.Internal.Types (These (..))+import DataFrame.Internal.Column.Types (These (..)) import DataFrame.Operations.Join import qualified DataFrame.Typed as DT import Test.HUnit
tests/Operations/Nullable.hs view
@@ -5,13 +5,15 @@ module Operations.Nullable where +import Data.Function ((&))+import qualified Data.Text as T import qualified Data.Vector as V import qualified DataFrame as D+import DataFrame.Expression.Operators (as, (.*), (.+), (.-), (./), (.==)) 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 (..))@@ -723,9 +725,75 @@ ) ) +-- ---------------------------------------------------------------------------+-- Grouped aggregation over nullable columns+-- ---------------------------------------------------------------------------++-- | Interleaved, so grouping's row permutation is not the identity.+interleavedKeys :: [T.Text]+interleavedKeys = ["a", "b", "a", "b", "a", "b", "a", "b"]++-- | Group @a@ holds 1,2,3,4; group @b@ is entirely null.+nullsInOneGroupValues :: [Maybe Double]+nullsInOneGroupValues =+ [Just 1, Nothing, Just 2, Nothing, Just 3, Nothing, Just 4, Nothing]++nullsInOneGroup :: D.DataFrame+nullsInOneGroup =+ D.fromNamedColumns+ [ ("k", DI.fromList interleavedKeys)+ , ("v", DI.fromVector (V.fromList nullsInOneGroupValues))+ ]++-- | Group @a@ sums to 7, group @b@ to 40.+nullsInBothGroupsValues :: [Maybe Double]+nullsInBothGroupsValues =+ [Just 1, Just 10, Just 2, Nothing, Nothing, Just 30, Just 4, Nothing]++nullsInBothGroups :: D.DataFrame+nullsInBothGroups =+ D.fromNamedColumns+ [ ("k", DI.fromList interleavedKeys)+ , ("v", DI.fromVector (V.fromList nullsInBothGroupsValues))+ ]++sumGroupedNullable :: D.DataFrame -> D.DataFrame+sumGroupedNullable df =+ df+ & D.groupBy ["k"]+ & D.aggregate [F.sumMaybe (F.col @(Maybe Double) "v") `as` "s"]+ & D.sortBy [D.Asc (F.col @T.Text "k")]++expectedGroupSums :: [Double] -> D.DataFrame+expectedGroupSums sums =+ D.fromNamedColumns+ [ ("k", DI.fromList (["a", "b"] :: [T.Text]))+ , ("s", DI.fromList sums)+ ]++sumMaybeOverInterleavedGroups :: Test+sumMaybeOverInterleavedGroups =+ TestCase+ ( assertEqual+ "sumMaybe sums each group's own non-null values"+ (expectedGroupSums [10.0, 0.0])+ (sumGroupedNullable nullsInOneGroup)+ )++sumMaybeWithNullsInEveryGroup :: Test+sumMaybeWithNullsInEveryGroup =+ TestCase+ ( assertEqual+ "sumMaybe skips nulls in every group, not just one"+ (expectedGroupSums [7.0, 40.0])+ (sumGroupedNullable nullsInBothGroups)+ )+ tests :: [Test] tests =- [ TestLabel "addIntMaybeInt" addIntMaybeInt+ [ TestLabel "sumMaybeOverInterleavedGroups" sumMaybeOverInterleavedGroups+ , TestLabel "sumMaybeWithNullsInEveryGroup" sumMaybeWithNullsInEveryGroup+ , TestLabel "addIntMaybeInt" addIntMaybeInt , TestLabel "addMaybeIntInt" addMaybeIntInt , TestLabel "addIntInt" addIntInt , TestLabel "addMaybeMaybe" addMaybeMaybe
tests/Operations/ParallelGroupBy.hs view
@@ -184,7 +184,7 @@ assertEqual ("oracle group-row count n=" ++ show n) refN outN {- | The low-cardinality DIRECT-INDEXED grouping fast path-('DataFrame.Internal.GroupingDirect') fires from 'D.groupBy' on a single clean+('DataFrame.Internal.Grouping.Direct') fires from 'D.groupBy' on a single clean small-range Int key. It emits groups in ascending key-value order rather than the hash path's order, so we cannot compare index structures directly; instead we assert it produces the SAME per-key aggregate values as the hash 'groupBySeq'
tests/Operations/Record.hs view
@@ -17,16 +17,17 @@ import Data.Int (Int64) import qualified Data.Map.Strict as M import qualified Data.Text as T-import qualified Data.Text.IO as TIO+import qualified Data.Text.IO.Utf8 as TIO import GHC.Generics (Generic) import qualified DataFrame as D+import DataFrame.Expression.Operators import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI-import DataFrame.Operators import qualified DataFrame.Schema as IS import DataFrame.Typed (Schema) import qualified DataFrame.Typed as DT+import System.Directory (getTemporaryDirectory) import Test.HUnit @@ -289,7 +290,8 @@ , "2,eu,20.5" , "3,ap,30.0" ]- tmp = "/tmp/dataframe_test_deriveSchema.csv"+ tmpDir <- getTemporaryDirectory+ let tmp = tmpDir <> "/dataframe_test_deriveSchema.csv" TIO.writeFile tmp csv df <- D.readCsvWithSchema orderSchema tmp assertEqual
tests/Operations/Shuffle.hs view
@@ -5,11 +5,13 @@ import qualified DataFrame as D +import Data.List (permutations)+import qualified Data.Map.Strict as M import qualified Data.Set as Set import qualified Data.Vector.Unboxed as VU import DataFrame.Operations.Permutation (shuffle, shuffledIndices) import System.Random (mkStdGen)-import Test.HUnit (Test (..), assertEqual)+import Test.HUnit (Test (..), assertBool, assertEqual) testDataFrame :: D.DataFrame testDataFrame =@@ -91,9 +93,81 @@ , TestCase (assertEqual "There are no repeated indecis" computed actual) ] +-- A one-row frame has exactly one permutation.+shuffleSingleRow :: Test+shuffleSingleRow =+ TestCase+ ( assertEqual+ "shuffling one index yields that index"+ (VU.fromList [0 :: Int])+ (shuffledIndices (mkStdGen 7) 1)+ )++{- | Chi-squared statistic of observed counts against a flat expectation:+sum over cells of (observed - expected)^2 / expected.+-}+chiSquared :: [Int] -> Double+chiSquared counts =+ let expected = fromIntegral (sum counts) / fromIntegral (length counts)+ in sum [(fromIntegral o - expected) ^ (2 :: Int) / expected | o <- counts]++{- | Every permutation of n items is equally likely under a uniform shuffle,+so the counts over all n! outcomes are chi-squared with n! - 1 degrees of+freedom. Testing the whole permutation, rather than one position at a time,+also catches a shuffle whose positions are individually uniform but+correlated. Seeds are fixed, so the sample -- and the verdict -- is+deterministic.++n = 5 gives 120 outcomes; 12000 draws puts 100 in each on average. The bound+is the 0.999 quantile of chi-squared with 119 degrees of freedom.+-}+shufflePermutationsAreUniform :: Test+shufflePermutationsAreUniform =+ let n = 5+ trials = 12000+ observed =+ M.fromListWith+ (+)+ [(VU.toList (shuffledIndices (mkStdGen s) n), 1 :: Int) | s <- [1 .. trials]]+ counts = [M.findWithDefault 0 p observed | p <- permutations [0 .. n - 1]]+ stat = chiSquared counts+ in TestCase+ ( assertBool+ ("chi-squared over all permutations is " ++ show stat ++ ", above 172.4")+ (stat < 172.4)+ )++{- | The frequency test from Knuth 3.3.2: each item lands in each position with+probability 1/n, so the n x n position-by-item table is chi-squared with+(n - 1)^2 degrees of freedom. A larger n than the permutation test can afford,+to catch bias that only shows at scale, such as a shuffle that leaves a+suffix untouched or never leaves an item in place.++n = 10 and 5000 draws put 500 in each cell. The bound is the 0.999 quantile of+chi-squared with 81 degrees of freedom.+-}+shufflePositionsAreUniform :: Test+shufflePositionsAreUniform =+ let n = 10+ trials = 5000+ samples = [VU.toList (shuffledIndices (mkStdGen s) n) | s <- [1 .. trials]]+ cell p i = length [() | xs <- samples, xs !! p == i]+ stat = chiSquared [cell p i | p <- [0 .. n - 1], i <- [0 .. n - 1]]+ in TestCase+ ( assertBool+ ( "chi-squared over the position-by-item table is "+ ++ show stat+ ++ ", above 126.1"+ )+ (stat < 126.1)+ )+ tests :: [Test] tests =- [ TestLabel "shuffleShuffles" shuffleShuffles+ [ TestLabel "shuffleSingleRow" shuffleSingleRow+ , TestLabel "shufflePermutationsAreUniform" shufflePermutationsAreUniform+ , TestLabel "shufflePositionsAreUniform" shufflePositionsAreUniform+ , TestLabel "shuffleShuffles" shuffleShuffles , TestLabel "shufflePreservesData" shufflePreservesData , TestLabel "shufflePreservesColumnNames" shufflePreservesColumnNames , TestLabel "shuffleSameSeedIsSameShuffle" shuffleSameSeedIsSameShuffle
tests/Operations/Statistics.hs view
@@ -7,6 +7,7 @@ import qualified Data.Vector.Unboxed as VU import qualified DataFrame as D import qualified DataFrame.Internal.Column as DI+import DataFrame.Internal.DataFrame (getColumn) import qualified DataFrame.Internal.Statistics as D import Assertions@@ -198,6 +199,56 @@ ) ) +{- | A sliced column's bitmap keeps whole bytes, so its trailing bits still+describe rows past the end of the slice. Counting non-null rows has to stop+at the column's length rather than fold the whole byte vector.+-}+sixteenNullableRows :: D.DataFrame+sixteenNullableRows =+ D.fromNamedColumns+ [ ("x", D.fromList (map (Just . fromIntegral) [1 .. 16 :: Int] :: [Maybe Double]))+ ]++countAfterTake :: Test+countAfterTake =+ TestCase+ ( assertEqual+ "summarize counts the rows a take actually kept"+ (Just (DI.fromList ([3.0] :: [Double])))+ ( getColumn "x" $+ D.take 1 $+ D.summarize (D.take 3 sixteenNullableRows)+ )+ )++countAfterRange :: Test+countAfterRange =+ TestCase+ ( assertEqual+ "summarize counts the rows a range actually kept"+ (Just (DI.fromList ([5.0] :: [Double])))+ ( getColumn "x" $+ D.take 1 $+ D.summarize (D.range (8, 13) sixteenNullableRows)+ )+ )++-- allMissing folds the same bitmap, so it needs the same length cutoff: the+-- padding bits of a sliced all-null column would otherwise read as present.+allMissingAfterSlice :: Test+allMissingAfterSlice =+ TestCase+ ( assertEqual+ "an all-null slice is still all-null"+ (Just True)+ (DI.allMissing <$> getColumn "x" (D.take 3 allNullRows))+ )++allNullRows :: D.DataFrame+allNullRows =+ D.fromNamedColumns+ [("x", DI.fromList (replicate 16 (Nothing :: Maybe Double)))]+ -- correlation correlationDf :: D.DataFrame@@ -277,6 +328,9 @@ , TestLabel "wrongQuantileNumber" wrongQuantileNumber , TestLabel "wrongQuantileIndex" wrongQuantileIndex , TestLabel "summarizeOptional" summarizeOptional+ , TestLabel "countAfterTake" countAfterTake+ , TestLabel "countAfterRange" countAfterRange+ , TestLabel "allMissingAfterSlice" allMissingAfterSlice , TestLabel "correlationPerfectPositive" correlationPerfectPositive , TestLabel "correlationPerfectNegative" correlationPerfectNegative , TestLabel "correlationSelfIdentity" correlationSelfIdentity
tests/Operations/Subset.hs view
@@ -9,8 +9,10 @@ import qualified DataFrame.Internal.Column as Col import DataFrame.Internal.DataFrame import DataFrame.Operations.Merge ()+import GenDataFrame () import System.Random import Test.HUnit+import Test.QuickCheck (Property, property) prop_dropZero :: DataFrame -> Bool prop_dropZero df = D.drop 0 df == df@@ -53,6 +55,17 @@ let rows = fst (dataframeDimensions df) in D.range (0, rows) df == df +prop_rangeClampsToBounds :: DataFrame -> Int -> Int -> Bool+prop_rangeClampsToBounds df a b =+ fst (dataframeDimensions (D.range (a, b) df)) == expected+ where+ rows = fst (dataframeDimensions df)+ -- Rows in [a, b) that actually exist. Clamps before subtracting so the+ -- oracle itself cannot overflow on extreme endpoints.+ lo = min (max a 0) rows+ hi = min (max b lo) rows+ expected = hi - lo+ prop_selectAll :: DataFrame -> Bool prop_selectAll df = D.select (D.columnNames df) df == df @@ -177,6 +190,37 @@ ) (abs (vaProp - origProp) < tol) +tenRows :: DataFrame+tenRows = fromNamedColumns [("x", Col.fromList ([0 .. 9] :: [Int]))]++-- Endpoints that overflow Int if the length is computed before clamping.+unit_rangeExtremeEndpoints :: Test+unit_rangeExtremeEndpoints =+ TestCase+ ( assertEqual+ "range (1, minBound) is empty, not a wrapped-around full range"+ 0+ (fst (dataframeDimensions (D.range (1, minBound) tenRows)))+ )++unit_rangeEndPastEnd :: Test+unit_rangeEndPastEnd =+ TestCase+ ( assertEqual+ "range (8, 20) on a 10-row frame yields rows 8 and 9"+ (Just (Col.fromList ([8, 9] :: [Int])))+ (getColumn "x" (D.range (8, 20) tenRows))+ )++unit_rangeStartBeforeZero :: Test+unit_rangeStartBeforeZero =+ TestCase+ ( assertEqual+ "range (-5, 3) on a 10-row frame yields rows 0 to 2"+ (Just (Col.fromList ([0, 1, 2] :: [Int])))+ (getColumn "x" (D.range (-5, 3) tenRows))+ )+ hunitTests :: [Test] hunitTests = [ TestLabel "unit_stratifiedSample_full" unit_stratifiedSample_full@@ -185,7 +229,14 @@ "unit_stratifiedSplit_singleRowStratum" unit_stratifiedSplit_singleRowStratum , TestLabel "unit_stratifiedSplit_proportions" unit_stratifiedSplit_proportions+ , TestLabel "unit_rangeEndPastEnd" unit_rangeEndPastEnd+ , TestLabel "unit_rangeExtremeEndpoints" unit_rangeExtremeEndpoints+ , TestLabel "unit_rangeStartBeforeZero" unit_rangeStartBeforeZero ]++-- Properties whose shape does not fit [DataFrame -> Bool].+properties :: [Property]+properties = [property prop_rangeClampsToBounds] tests :: [DataFrame -> Bool] tests =
tests/Operations/VectorKernel.hs view
@@ -13,7 +13,7 @@ Floating-point sums computed by the scatter follow the same left-to-right group-order fold as the interpreter, so the two paths agree bit-for-bit here.-The PARALLEL kernel ('DataFrame.Internal.AggKernelPar') splits the work by+The PARALLEL path ('DataFrame.Internal.Aggregation.Kernel.Scatter') splits the work by disjoint group-id range, and because each group's rows stay in their original @valueIndices@ order within one worker's range, the per-group fold order is unchanged from the sequential scatter — so the parallel path is also
tests/Operations/Window.hs view
@@ -10,7 +10,7 @@ import qualified DataFrame.Internal.DataFrame as DI import Data.Function ((&))-import DataFrame.Operators+import DataFrame.Expression.Operators import Test.HUnit -- Similar to example discussed in https://www.sumsar.net/blog/pandas-feels-clunky-when-coming-from-r/
tests/Operations/WriteCsv.hs view
@@ -4,10 +4,11 @@ module Operations.WriteCsv where import qualified Data.Text as T-import qualified Data.Text.IO as TIO+import qualified Data.Text.IO.Utf8 as TIO import qualified DataFrame as D import qualified DataFrame.Internal.Column as DI import DataFrame.Internal.DataFrame (DataFrame (..), toCsv, toSeparated)+import System.Directory (getTemporaryDirectory) import Test.HUnit -- Basic test: Int and Text columns produce correct CSV@@ -81,7 +82,8 @@ , ("b", DI.fromList @T.Text ["hello", "world", "test"]) ] let csvText = toCsv df- let tmpPath = "/tmp/dataframe_test_toCsv_roundtrip.csv"+ tmpDir <- getTemporaryDirectory+ let tmpPath = tmpDir <> "/dataframe_test_toCsv_roundtrip.csv" TIO.writeFile tmpPath csvText df' <- D.readCsv tmpPath assertEqual
tests/PackedTextMigration.hs view
@@ -20,7 +20,7 @@ import qualified DataFrame as D import qualified DataFrame.Internal.Column as DI-import DataFrame.Internal.PackedText (mkPackedContiguous)+import DataFrame.Internal.Data.PackedText (mkPackedContiguous) import DataFrame.DecisionTree (defaultTreeConfig) import DataFrame.DecisionTree.Model ()
tests/PrettyPrint.hs view
@@ -8,8 +8,8 @@ module PrettyPrint (tests) where import qualified Data.Text as T+import DataFrame.Expression.Operators import DataFrame.Internal.Expression (Expr, prettyPrint, prettyPrintWidth)-import DataFrame.Operators import Test.HUnit a, b, c :: Expr Double
tests/Simplify.hs view
@@ -2,16 +2,16 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} -{- | Specification for 'DataFrame.Internal.Simplify.simplify': each case is the+{- | Specification for 'DataFrame.Internal.Expression.Simplify.simplify': each case is the full predicate expression, compared with 'eqExpr'. -} module Simplify (tests) where +import DataFrame.Expression.Operators import qualified DataFrame.Functions as F import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression (Expr, eqExpr)-import DataFrame.Internal.Simplify (simplify)-import DataFrame.Operators+import DataFrame.Internal.Expression.Simplify (simplify) import Test.HUnit