diff --git a/dataframe-core.cabal b/dataframe-core.cabal
--- a/dataframe-core.cabal
+++ b/dataframe-core.cabal
@@ -1,6 +1,6 @@
-cabal-version: 2.4
+cabal-version: 3.4
name: dataframe-core
-version: 1.1.0.4
+version: 2.0.0.0
synopsis: Core data structures for the dataframe library.
description:
Minimal interchange-format types for the @dataframe@ ecosystem:
@@ -28,8 +28,9 @@
-Wunused-local-binds
-Wunused-packages
-library
+library internal
import: warnings
+ visibility: public
exposed-modules:
DataFrame.Errors
DataFrame.Operators
@@ -54,23 +55,37 @@
DataFrame.Internal.Nullable
DataFrame.Internal.PackedText
DataFrame.Internal.ParRadixSort
+ DataFrame.Internal.Pretty
DataFrame.Internal.RadixRank
DataFrame.Internal.RowHash
DataFrame.Internal.Row
DataFrame.Internal.Simplify
DataFrame.Internal.Types
DataFrame.Internal.Utf8
+ build-depends: base >= 4 && < 5,
+ containers >= 0.6.7 && < 0.10,
+ primitive >= 0.7 && < 0.11,
+ random >= 1 && < 2,
+ text >= 2.1 && < 3,
+ vector >= 0.13 && < 0.15
+ hs-source-dirs: src-internal
+ default-language: Haskell2010
+
+library
+ import: warnings
+ exposed-modules:
+ DataFrame.Core
DataFrame.Typed.Freeze
DataFrame.Typed.Generic
DataFrame.Typed.Record
DataFrame.Typed.Schema
DataFrame.Typed.Types
DataFrame.Typed.Util
+ reexported-modules: DataFrame.Errors,
+ DataFrame.Operators
build-depends: base >= 4 && < 5,
- containers >= 0.6.7 && < 0.9,
- primitive >= 0.7 && < 0.10,
- random >= 1 && < 2,
text >= 2.1 && < 3,
- vector ^>= 0.13
+ vector >= 0.13 && < 0.15,
+ dataframe-core:internal
hs-source-dirs: src
default-language: Haskell2010
diff --git a/src-internal/DataFrame/Display/Terminal/Colours.hs b/src-internal/DataFrame/Display/Terminal/Colours.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Display/Terminal/Colours.hs
@@ -0,0 +1,7 @@
+module DataFrame.Display.Terminal.Colours where
+
+red, green, brightGreen, brightBlue :: String -> String
+red = id
+green = id
+brightGreen = id
+brightBlue = id
diff --git a/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs b/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Display/Terminal/PrettyPrint.hs
@@ -0,0 +1,114 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module DataFrame.Display.Terminal.PrettyPrint where
+
+import qualified Data.Text as T
+import qualified Data.Vector as V
+
+{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with
+ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for
+notebooks.
+-}
+data RenderFormat = Plain | Markdown
+ deriving (Show, Eq)
+
+-- Utility functions to show a DataFrame as a Markdown-ish table.
+
+-- Fill functions, adapted from:
+-- https://stackoverflow.com/questions/5929377/format-list-output-in-haskell
+type Filler = Int -> T.Text -> T.Text
+
+-- A description of one table column.
+data ColDesc t = ColDesc
+ { colTitleFill :: Filler
+ , colTitle :: T.Text
+ , colValueFill :: Filler
+ }
+
+-- Fill text to a given width with a pad character, aligning left, right, or center.
+fillLeft :: Char -> Int -> T.Text -> T.Text
+fillLeft c n s = s <> T.replicate (n - T.length s) (T.singleton c)
+
+fillRight :: Char -> Int -> T.Text -> T.Text
+fillRight c n s = T.replicate (n - T.length s) (T.singleton c) <> s
+
+fillCenter :: Char -> Int -> T.Text -> T.Text
+fillCenter c n s =
+ T.replicate l (T.singleton c) <> s <> T.replicate r (T.singleton c)
+ where
+ x = n - T.length s
+ l = x `div` 2
+ r = x - l
+
+-- Fill with spaces.
+left :: Int -> T.Text -> T.Text
+left = fillLeft ' '
+
+right :: Int -> T.Text -> T.Text
+right = fillRight ' '
+
+center :: Int -> T.Text -> T.Text
+center = fillCenter ' '
+
+{- | Render a table from column-major data. @columns@ has one 'V.Vector' per
+column; widths are computed in one pass per column (no row-major transpose),
+and row lines are built by indexing each column at row @i@.
+-}
+showTable ::
+ RenderFormat ->
+ [T.Text] ->
+ [T.Text] ->
+ [V.Vector T.Text] ->
+ T.Text
+showTable fmt header types columns =
+ let isMarkdown = fmt == Markdown
+ esc = if isMarkdown then escapeMarkdownCell else id
+ hdr = map esc header
+ tys = map esc types
+ cols = map (V.map esc) columns
+ consolidatedHeader =
+ if isMarkdown
+ then zipWith (\h t -> h <> "
" <> t) hdr tys
+ else hdr
+ cs = map (\h -> ColDesc center h left) consolidatedHeader
+ nRows = case cols of
+ (c : _) -> V.length c
+ [] -> 0
+ columnMaxWidth col
+ | V.null col = 0
+ | otherwise = V.foldl' (\acc x -> max acc (T.length x)) 0 col
+ widths =
+ zipWith3
+ (\h t col -> T.length h `max` T.length t `max` columnMaxWidth col)
+ consolidatedHeader
+ tys
+ cols
+ dashesOf w = T.replicate w "-"
+ border = T.intercalate "---" (map dashesOf widths)
+ separator = T.intercalate "-|-" (map dashesOf widths)
+ fillCells fill cells =
+ T.intercalate " | " (zipWith3 fill cs widths cells)
+ rowCells i = map (V.! i) cols
+ rowLines = [fillCells colValueFill (rowCells i) | i <- [0 .. nRows - 1]]
+ wrapMd t = T.concat ["| ", t, " |"]
+ outputLines =
+ if isMarkdown
+ then
+ wrapMd (fillCells colTitleFill consolidatedHeader)
+ : wrapMd separator
+ : map wrapMd rowLines
+ else
+ border
+ : fillCells colTitleFill consolidatedHeader
+ : separator
+ : fillCells colTitleFill tys
+ : separator
+ : rowLines
+ in T.unlines outputLines
+
+{- | Escape a value for a GitHub-flavoured Markdown table cell: a bare @|@ would
+end the cell and a newline would end the row, so both are neutralised (the pipe
+backslash-escaped, the newline turned into a @\@).
+-}
+escapeMarkdownCell :: T.Text -> T.Text
+escapeMarkdownCell = T.replace "\n" "
" . T.replace "|" "\\|"
diff --git a/src-internal/DataFrame/Errors.hs b/src-internal/DataFrame/Errors.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Errors.hs
@@ -0,0 +1,187 @@
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
+
+module DataFrame.Errors where
+
+import qualified Data.Map.Lazy as ML
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+
+import Control.Exception
+import qualified Data.List as L
+import Data.Typeable (Typeable)
+import DataFrame.Display.Terminal.Colours
+import Type.Reflection (TypeRep)
+
+data TypeErrorContext a b = MkTypeErrorContext
+ { userType :: Either String (TypeRep a)
+ , expectedType :: Either String (TypeRep b)
+ , errorColumnName :: Maybe String
+ , callingFunctionName :: Maybe String
+ }
+
+data DataFrameException where
+ TypeMismatchException ::
+ forall a b.
+ (Typeable a, Typeable b) =>
+ TypeErrorContext a b ->
+ DataFrameException
+ AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException
+ ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException
+ EmptyDataSetException :: T.Text -> DataFrameException
+ InternalException :: T.Text -> DataFrameException
+ NonColumnReferenceException :: T.Text -> DataFrameException
+ UnaggregatedException :: T.Text -> DataFrameException
+ WrongQuantileNumberException :: Int -> DataFrameException
+ WrongQuantileIndexException :: VU.Vector Int -> Int -> DataFrameException
+ deriving (Exception)
+
+instance Show DataFrameException where
+ show :: DataFrameException -> String
+ show (TypeMismatchException context) =
+ let
+ errorString =
+ typeMismatchError
+ (either id show (userType context))
+ (either id show (expectedType context))
+ in
+ addCallPointInfo
+ (errorColumnName context)
+ (callingFunctionName context)
+ errorString
+ show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns
+ show (EmptyDataSetException callPoint) = emptyDataSetError callPoint
+ show (WrongQuantileNumberException q) = wrongQuantileNumberError q
+ show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q
+ show (InternalException msg) = "Internal error: " ++ T.unpack msg
+ show (NonColumnReferenceException msg) = "Expression must be a column reference in: " ++ T.unpack msg
+ show (UnaggregatedException expr) = "Expression is not fully aggregated: " ++ T.unpack expr
+ show (AggregatedAndNonAggregatedException expr1 expr2) =
+ "Cannot combine aggregated and non-aggregated expressions: \n"
+ ++ T.unpack expr1
+ ++ "\n"
+ ++ T.unpack expr2
+
+columnNotFound :: T.Text -> T.Text -> [T.Text] -> String
+columnNotFound missingColumn = columnsNotFound [missingColumn]
+
+columnsNotFound :: [T.Text] -> T.Text -> [T.Text] -> String
+columnsNotFound missingColumns callPoint availableColumns =
+ red "\n\n[ERROR] "
+ ++ missingColumnsLabel missingColumns
+ ++ ": "
+ ++ T.unpack (T.intercalate ", " missingColumns)
+ ++ " for operation "
+ ++ T.unpack callPoint
+ ++ formatSuggestions missingColumns availableColumns
+ ++ "\n\n"
+ where
+ missingColumnsLabel [_] = "Column not found"
+ missingColumnsLabel _ = "Columns not found"
+
+ formatSuggestions [missingColumn] columns =
+ case guessColumnName missingColumn columns of
+ "" -> ""
+ guessed ->
+ "\n\tDid you mean "
+ ++ T.unpack guessed
+ ++ "?"
+ formatSuggestions names columns =
+ case traverse (`suggestColumnName` columns) names of
+ Just guessedColumns
+ | not (null guessedColumns) ->
+ "\n\tDid you mean "
+ ++ formatColumnSuggestions guessedColumns
+ ++ "?"
+ _ -> ""
+
+ suggestColumnName missingColumn columns = case guessColumnName missingColumn columns of
+ "" -> Nothing
+ guessed -> Just guessed
+
+ formatColumnSuggestions guessedColumns =
+ "["
+ ++ L.intercalate ", " (map (show . T.unpack) guessedColumns)
+ ++ "]"
+
+typeMismatchError :: String -> String -> String
+typeMismatchError givenType expType =
+ red $
+ red "\n\n[Error]: Type Mismatch"
+ ++ "\n\tWhile running your code I tried to "
+ ++ "get a column of type: "
+ ++ red (show givenType)
+ ++ " but the column in the dataframe was actually of type: "
+ ++ green (show expType)
+
+emptyDataSetError :: T.Text -> String
+emptyDataSetError callPoint =
+ red "\n\n[ERROR] "
+ ++ T.unpack callPoint
+ ++ " cannot be called on empty data sets"
+
+wrongQuantileNumberError :: Int -> String
+wrongQuantileNumberError q =
+ red "\n\n[ERROR] "
+ ++ "Quantile number q should satisfy "
+ ++ "q >= 2, but here q is "
+ ++ show q
+
+wrongQuantileIndexError :: VU.Vector Int -> Int -> String
+wrongQuantileIndexError qs q =
+ red "\n\n[ERROR] "
+ ++ "For quantile number q, "
+ ++ "each quantile index i "
+ ++ "should satisfy 0 <= i <= q, "
+ ++ "but here q is "
+ ++ show q
+ ++ " and indexes are "
+ ++ show qs
+
+addCallPointInfo :: Maybe String -> Maybe String -> String -> String
+addCallPointInfo (Just name) (Just cp) err =
+ err
+ ++ ( "\n\tThis happened when calling function "
+ ++ brightGreen cp
+ ++ " on "
+ ++ brightGreen name
+ )
+addCallPointInfo Nothing (Just cp) err =
+ err
+ ++ ( "\n\tThis happened when calling function "
+ ++ brightGreen cp
+ )
+addCallPointInfo (Just name) Nothing err =
+ err
+ ++ ( "\n\tOn "
+ ++ name
+ ++ "\n\n"
+ )
+addCallPointInfo Nothing Nothing err = err
+
+guessColumnName :: T.Text -> [T.Text] -> T.Text
+guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of
+ [] -> ""
+ res -> (snd . minimum) res
+
+editDistance :: T.Text -> T.Text -> Int
+editDistance xs ys = table ML.! (m, n)
+ where
+ (m, n) = (T.length xs, T.length ys)
+ xv = V.fromList (T.unpack xs)
+ yv = V.fromList (T.unpack ys)
+ table :: ML.Map (Int, Int) Int
+ table = ML.fromList [((i, j), dist i j) | i <- [0 .. m], j <- [0 .. n]]
+ dist 0 j = j
+ dist i 0 = i
+ dist i j =
+ minimum
+ [ table ML.! (i - 1, j) + 1
+ , table ML.! (i, j - 1) + 1
+ , (if xv V.! (i - 1) == yv V.! (j - 1) then 0 else 1)
+ + table ML.! (i - 1, j - 1)
+ ]
diff --git a/src-internal/DataFrame/Internal/AggKernel.hs b/src-internal/DataFrame/Internal/AggKernel.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/AggKernel.hs
@@ -0,0 +1,260 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | Vectorized scatter-accumulate aggregation kernel.
+module DataFrame.Internal.AggKernel (
+ Reduction (..),
+ scatterReduce,
+ scatterColumnToDouble,
+) where
+
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Monad (when)
+import Control.Monad.ST (ST, runST)
+import DataFrame.Internal.Column (
+ Column (..),
+ Columnable,
+ fromUnboxedVector,
+ materializePacked,
+ )
+import Type.Reflection (typeRep)
+
+{- | A recognised fast-path reduction over a single value column. The element
+type (Int vs Double) is resolved at scatter time; sum/min/max preserve the
+column's element type, everything else produces a Double column.
+-}
+data Reduction
+ = RSum
+ | RCount
+ | RMin
+ | RMax
+ | RMean
+ | RStd
+ | RVar
+ | RTop2Sum
+ deriving (Eq, Show)
+
+{- | Coerce an unboxed Int or Double column to an unboxed Double vector for the
+moment/mean/sd/median family. Returns 'Nothing' for boxed, nullable, or other
+element types (the caller then falls back to the interpreter).
+-}
+scatterColumnToDouble :: Column -> Maybe (VU.Vector Double)
+scatterColumnToDouble = \case
+ UnboxedColumn Nothing (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> Just v
+ Nothing -> case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> Just (VU.map fromIntegral v)
+ Nothing -> Nothing
+ p@(PackedText _ _) -> scatterColumnToDouble (materializePacked p)
+ _ -> Nothing
+
+scatterReduce ::
+ Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column
+scatterReduce red g nGroups col = case col of
+ UnboxedColumn Nothing (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> Just (reduceTyped red g nGroups v intIdent)
+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> Just (reduceTyped red g nGroups v dblIdent)
+ Nothing -> Nothing
+ p@(PackedText _ _) -> scatterReduce red g nGroups (materializePacked p)
+ _ -> Nothing
+{-# INLINEABLE scatterReduce #-}
+
+-- | Per-type seed identities for the order-preserving reductions.
+data Idents a = Idents {minSeed :: !a, maxSeed :: !a}
+
+intIdent :: Idents Int
+intIdent = Idents maxBound minBound
+
+dblIdent :: Idents Double
+dblIdent = Idents (1 / 0) (negate (1 / 0))
+
+reduceTyped ::
+ forall a.
+ (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>
+ Reduction -> VU.Vector Int -> Int -> VU.Vector a -> Idents a -> Column
+reduceTyped red g nGroups v idents = case red of
+ RCount -> fromUnboxedVector (countScatter g nGroups)
+ RSum -> fromUnboxedVector (sumScatter g nGroups v)
+ RMin -> fromUnboxedVector (extremaScatter min (minSeed idents) g nGroups v)
+ RMax -> fromUnboxedVector (extremaScatter max (maxSeed idents) g nGroups v)
+ RMean -> fromUnboxedVector (meanScatter g nGroups v)
+ RVar -> fromUnboxedVector (varScatter False g nGroups v)
+ RStd -> fromUnboxedVector (varScatter True g nGroups v)
+ RTop2Sum -> fromUnboxedVector (top2Scatter g nGroups v)
+{-# INLINE reduceTyped #-}
+
+countScatter :: VU.Vector Int -> Int -> VU.Vector Int
+countScatter g nGroups = runST $ do
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ let n = VU.length g
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ c <- VUM.unsafeRead cnt k
+ VUM.unsafeWrite cnt k (c + 1)
+ go (i + 1)
+ go 0
+ VU.unsafeFreeze cnt
+
+sumScatter ::
+ (VU.Unbox a, Num a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a
+sumScatter g nGroups v = runST $ do
+ s <- VUM.replicate nGroups 0
+ let n = VU.length v
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ cur <- VUM.unsafeRead s k
+ VUM.unsafeWrite s k (cur + VU.unsafeIndex v i)
+ go (i + 1)
+ go 0
+ VU.unsafeFreeze s
+{-# INLINE sumScatter #-}
+
+extremaScatter ::
+ (VU.Unbox a) =>
+ (a -> a -> a) -> a -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a
+extremaScatter combine seed g nGroups v = runST $ do
+ m <- VUM.replicate nGroups seed
+ let n = VU.length v
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ cur <- VUM.unsafeRead m k
+ VUM.unsafeWrite m k (combine cur (VU.unsafeIndex v i))
+ go (i + 1)
+ go 0
+ VU.unsafeFreeze m
+{-# INLINE extremaScatter #-}
+
+meanScatter ::
+ (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
+meanScatter g nGroups v = runST $ do
+ s <- VUM.replicate nGroups (0 :: Double)
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ scatterSumCount g v s cnt
+ finalizeMean nGroups s cnt
+{-# INLINE meanScatter #-}
+
+scatterSumCount ::
+ (VU.Unbox a, Real a) =>
+ VU.Vector Int ->
+ VU.Vector a ->
+ VUM.MVector s Double ->
+ VUM.MVector s Int ->
+ ST s ()
+scatterSumCount g v s cnt = go 0
+ where
+ n = VU.length v
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ !x = realToFrac (VU.unsafeIndex v i)
+ curS <- VUM.unsafeRead s k
+ VUM.unsafeWrite s k (curS + x)
+ curC <- VUM.unsafeRead cnt k
+ VUM.unsafeWrite cnt k (curC + 1)
+ go (i + 1)
+{-# INLINE scatterSumCount #-}
+
+finalizeMean ::
+ Int -> VUM.MVector s Double -> VUM.MVector s Int -> ST s (VU.Vector Double)
+finalizeMean nGroups s cnt = do
+ out <- VUM.new nGroups
+ let go !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ sv <- VUM.unsafeRead s k
+ c <- VUM.unsafeRead cnt k
+ VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)
+ go (k + 1)
+ go 0
+ VU.unsafeFreeze out
+
+varScatter ::
+ (VU.Unbox a, Real a) =>
+ Bool -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
+varScatter takeSqrt g nGroups v = runST $ do
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ meanV <- VUM.replicate nGroups (0 :: Double)
+ m2 <- VUM.replicate nGroups (0 :: Double)
+ let n = VU.length v
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ !x = realToFrac (VU.unsafeIndex v i)
+ c <- VUM.unsafeRead cnt k
+ mu <- VUM.unsafeRead meanV k
+ mm <- VUM.unsafeRead m2 k
+ let !c' = c + 1
+ !delta = x - mu
+ !mu' = mu + delta / fromIntegral c'
+ !mm' = mm + delta * (x - mu')
+ VUM.unsafeWrite cnt k c'
+ VUM.unsafeWrite meanV k mu'
+ VUM.unsafeWrite m2 k mm'
+ go (i + 1)
+ go 0
+ out <- VUM.new nGroups
+ let fin !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ c <- VUM.unsafeRead cnt k
+ mm <- VUM.unsafeRead m2 k
+ let var = if c < 2 then 0 else mm / fromIntegral (c - 1)
+ VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)
+ fin (k + 1)
+ fin 0
+ VU.unsafeFreeze out
+{-# INLINE varScatter #-}
+
+top2Scatter ::
+ (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
+top2Scatter g nGroups v = runST $ do
+ let ninf = negate (1 / 0) :: Double
+ m1 <- VUM.replicate nGroups ninf
+ m2 <- VUM.replicate nGroups ninf
+ let n = VU.length v
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ !x = realToFrac (VU.unsafeIndex v i)
+ a1 <- VUM.unsafeRead m1 k
+ if x > a1
+ then do
+ VUM.unsafeWrite m1 k x
+ VUM.unsafeWrite m2 k a1
+ else do
+ a2 <- VUM.unsafeRead m2 k
+ when (x > a2) (VUM.unsafeWrite m2 k x)
+ go (i + 1)
+ go 0
+ out <- VUM.new nGroups
+ let fin !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ a1 <- VUM.unsafeRead m1 k
+ a2 <- VUM.unsafeRead m2 k
+ let s = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)
+ VUM.unsafeWrite out k s
+ fin (k + 1)
+ fin 0
+ VU.unsafeFreeze out
+{-# INLINE top2Scatter #-}
diff --git a/src-internal/DataFrame/Internal/AggKernelDirect.hs b/src-internal/DataFrame/Internal/AggKernelDirect.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/AggKernelDirect.hs
@@ -0,0 +1,338 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+module DataFrame.Internal.AggKernelDirect (
+ directThreshold,
+ directReduce,
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import System.IO.Unsafe (unsafePerformIO)
+import Type.Reflection (typeRep)
+
+import DataFrame.Internal.AggKernel (Reduction (..))
+import DataFrame.Internal.Column (
+ Column (..),
+ fromUnboxedVector,
+ materializePacked,
+ )
+
+{- | Group-domain size at or below which the direct-indexed accumulator path is
+taken; wider domains keep the group-range kernel. The admitted reductions are
+order-independent, so the per-worker accumulator merge is exact.
+-}
+directThreshold :: Int
+directThreshold = 262144
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+{- | Below this many rows the parallel fan-out is not worth it; a single
+sequential direct pass runs instead (tiny accumulator, one tight loop). Matches
+the grouping/scatter parallel threshold.
+-}
+parThreshold :: Int
+parThreshold = 200000
+
+{- | Run a recognised reduction through the direct-indexed path. 'Nothing' (so
+the caller falls back to the order-preserving kernel) unless the reduction is
+order-independent at this element type AND the column is a clean unboxed Int/Double.
+-}
+directReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column
+directReduce red g nGroups col = case col of
+ UnboxedColumn Nothing (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> directInt red g nGroups v
+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> directDouble red g nGroups v
+ Nothing -> Nothing
+ p@(PackedText _ _) -> directReduce red g nGroups (materializePacked p)
+ _ -> Nothing
+{-# INLINEABLE directReduce #-}
+
+-- | The order-independent reductions over an Int column.
+directInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column
+directInt red g nGroups v = case red of
+ RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))
+ RSum -> Just (fromUnboxedVector (sumIntDirect g nGroups v))
+ RMin -> Just (fromUnboxedVector (extremaIntDirect True g nGroups v))
+ RMax -> Just (fromUnboxedVector (extremaIntDirect False g nGroups v))
+ RMean -> Just (fromUnboxedVector (meanIntDirect g nGroups v))
+ _ -> Nothing
+
+{- | Over a Double column only @count@ is order-independent; the float
+sum/mean/variance reductions must keep the order-preserving kernel.
+-}
+directDouble ::
+ Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column
+directDouble red g nGroups v = case red of
+ RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))
+ _ -> Nothing
+
+-- | Whether to fan out at this row count.
+shouldPar :: Int -> Bool
+shouldPar n = n >= parThreshold && capabilities > 1
+
+{- | Fork @caps@ workers over disjoint contiguous row ranges of @[0, n)@, each
+producing its own private accumulator (no shared array, no sync). Returns the
+partials in worker order for the caller's merge; rethrows the first failure.
+-}
+runPartialsOver ::
+ Int -> Int -> (Int -> Int -> IO (VUM.IOVector Int)) -> IO [VUM.IOVector Int]
+runPartialsOver n caps fill = do
+ let !per = (n + caps - 1) `div` caps
+ spawn w = do
+ var <- newEmptyMVar
+ let !lo = min n (w * per)
+ !hi = min n (lo + per)
+ _ <- forkIO (try (fill lo hi) >>= putMVar var)
+ pure var
+ vars <- mapM spawn [0 .. caps - 1]
+ results <- mapM takeMVar vars
+ mapM (either (throwIO @SomeException) pure) results
+
+{- | As 'runPartialsOver' but each worker produces a PAIR of accumulators (e.g.
+sum and count for the fused integer mean).
+-}
+runPartialsPairOver ::
+ Int ->
+ Int ->
+ (Int -> Int -> IO (VUM.IOVector Int, VUM.IOVector Int)) ->
+ IO [(VUM.IOVector Int, VUM.IOVector Int)]
+runPartialsPairOver n caps fill = do
+ let !per = (n + caps - 1) `div` caps
+ spawn w = do
+ var <- newEmptyMVar
+ let !lo = min n (w * per)
+ !hi = min n (lo + per)
+ _ <- forkIO (try (fill lo hi) >>= putMVar var)
+ pure var
+ vars <- mapM spawn [0 .. caps - 1]
+ results <- mapM takeMVar vars
+ mapM (either (throwIO @SomeException) pure) results
+
+-------------------------------------------------------------------------------
+-- Count (order-independent: per-group row count)
+-------------------------------------------------------------------------------
+
+countDirect :: VU.Vector Int -> Int -> Int -> VU.Vector Int
+countDirect g nGroups n
+ | not (shouldPar n) =
+ unsafePerformIO (countChunk g nGroups 0 n >>= VU.unsafeFreeze)
+ | otherwise = unsafePerformIO $ do
+ parts <- runPartialsOver n capabilities (countChunk g nGroups)
+ mergeIntSum nGroups parts
+{-# NOINLINE countDirect #-}
+
+countChunk :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
+countChunk g nGroups lo hi = do
+ acc <- VUM.replicate nGroups (0 :: Int)
+ let go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ c <- VUM.unsafeRead acc k
+ VUM.unsafeWrite acc k (c + 1)
+ go (i + 1)
+ go lo
+ pure acc
+
+-------------------------------------------------------------------------------
+-- Integer sum (exact: merge order irrelevant)
+-------------------------------------------------------------------------------
+
+sumIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
+sumIntDirect g nGroups v
+ | not (shouldPar n) =
+ unsafePerformIO (sumIntChunk g v nGroups 0 n >>= VU.unsafeFreeze)
+ | otherwise = unsafePerformIO $ do
+ parts <- runPartialsOver n capabilities (sumIntChunk g v nGroups)
+ mergeIntSum nGroups parts
+ where
+ !n = VU.length v
+{-# NOINLINE sumIntDirect #-}
+
+sumIntChunk ::
+ VU.Vector Int -> VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
+sumIntChunk g v nGroups lo hi = do
+ acc <- VUM.replicate nGroups (0 :: Int)
+ let go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ c <- VUM.unsafeRead acc k
+ VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)
+ go (i + 1)
+ go lo
+ pure acc
+
+-------------------------------------------------------------------------------
+-- Integer min / max (order-independent)
+-------------------------------------------------------------------------------
+
+extremaIntDirect ::
+ Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
+extremaIntDirect isMin g nGroups v
+ | not (shouldPar n) =
+ unsafePerformIO (extremaIntChunk isMin g v nGroups 0 n >>= VU.unsafeFreeze)
+ | otherwise = unsafePerformIO $ do
+ parts <- runPartialsOver n capabilities (extremaIntChunk isMin g v nGroups)
+ mergeExtremaInt isMin nGroups parts
+ where
+ !n = VU.length v
+{-# NOINLINE extremaIntDirect #-}
+
+extremaIntChunk ::
+ Bool ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ Int ->
+ Int ->
+ IO (VUM.IOVector Int)
+extremaIntChunk isMin g v nGroups lo hi = do
+ let !seed = if isMin then maxBound else minBound
+ combine a b = if isMin then min a b else max a b
+ acc <- VUM.replicate nGroups seed
+ let go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ c <- VUM.unsafeRead acc k
+ VUM.unsafeWrite acc k (combine c (VU.unsafeIndex v i))
+ go (i + 1)
+ go lo
+ pure acc
+
+-------------------------------------------------------------------------------
+-- Integer mean (exact integer sum + count, divided once -> order-independent)
+-------------------------------------------------------------------------------
+
+{- | Integer mean in ONE fused pass: a running integer sum and count per group,
+divided once at finalize. The integer sum is exact, so the parallel partial
+merge is byte-identical to the sequential single pass at any @-N@.
+-}
+meanIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double
+meanIntDirect g nGroups v
+ | not (shouldPar n) = unsafePerformIO $ do
+ (s, c) <- meanIntChunk g v nGroups 0 n
+ finalizeMeanInt nGroups s c
+ | otherwise = unsafePerformIO $ do
+ parts <- runPartialsPairOver n capabilities (meanIntChunk g v nGroups)
+ (s, c) <- mergePair nGroups parts
+ finalizeMeanInt nGroups s c
+ where
+ !n = VU.length v
+{-# NOINLINE meanIntDirect #-}
+
+meanIntChunk ::
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ Int ->
+ Int ->
+ IO (VUM.IOVector Int, VUM.IOVector Int)
+meanIntChunk g v nGroups lo hi = do
+ s <- VUM.replicate nGroups (0 :: Int)
+ c <- VUM.replicate nGroups (0 :: Int)
+ let go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ sv <- VUM.unsafeRead s k
+ VUM.unsafeWrite s k (sv + VU.unsafeIndex v i)
+ cv <- VUM.unsafeRead c k
+ VUM.unsafeWrite c k (cv + 1)
+ go (i + 1)
+ go lo
+ pure (s, c)
+
+finalizeMeanInt ::
+ Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO (VU.Vector Double)
+finalizeMeanInt nGroups s c = do
+ out <- VUM.new nGroups
+ let go !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ sv <- VUM.unsafeRead s k
+ cv <- VUM.unsafeRead c k
+ VUM.unsafeWrite
+ out
+ k
+ (if cv == 0 then 0 / 0 else fromIntegral sv / fromIntegral cv)
+ go (k + 1)
+ go 0
+ VU.unsafeFreeze out
+
+-------------------------------------------------------------------------------
+-- Partial accumulation + merge
+-------------------------------------------------------------------------------
+
+mergeIntSum :: Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
+mergeIntSum nGroups parts = case parts of
+ [] -> VU.unsafeFreeze =<< VUM.replicate nGroups 0
+ (p0 : rest) -> do
+ let add !p = do
+ let go !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ a <- VUM.unsafeRead p0 k
+ b <- VUM.unsafeRead p k
+ VUM.unsafeWrite p0 k (a + b)
+ go (k + 1)
+ go 0
+ mapM_ add rest
+ VU.unsafeFreeze p0
+
+{- | Merge per-worker (sum, count) partials into the first worker's pair by
+exact integer addition; returns the accumulated pair for finalize.
+-}
+mergePair ::
+ Int ->
+ [(VUM.IOVector Int, VUM.IOVector Int)] ->
+ IO (VUM.IOVector Int, VUM.IOVector Int)
+mergePair nGroups parts = case parts of
+ [] -> (,) <$> VUM.replicate nGroups 0 <*> VUM.replicate nGroups 0
+ ((s0, c0) : rest) -> do
+ let add (s, c) = do
+ let go !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ sa <- VUM.unsafeRead s0 k
+ sb <- VUM.unsafeRead s k
+ VUM.unsafeWrite s0 k (sa + sb)
+ ca <- VUM.unsafeRead c0 k
+ cb <- VUM.unsafeRead c k
+ VUM.unsafeWrite c0 k (ca + cb)
+ go (k + 1)
+ go 0
+ mapM_ add rest
+ pure (s0, c0)
+
+mergeExtremaInt :: Bool -> Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
+mergeExtremaInt isMin nGroups parts = case parts of
+ [] ->
+ VU.unsafeFreeze =<< VUM.replicate nGroups (if isMin then maxBound else minBound)
+ (p0 : rest) -> do
+ let combine a b = if isMin then min a b else max a b
+ add !p = do
+ let go !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ a <- VUM.unsafeRead p0 k
+ b <- VUM.unsafeRead p k
+ VUM.unsafeWrite p0 k (combine a b)
+ go (k + 1)
+ go 0
+ mapM_ add rest
+ VU.unsafeFreeze p0
diff --git a/src-internal/DataFrame/Internal/AggKernelPar.hs b/src-internal/DataFrame/Internal/AggKernelPar.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/AggKernelPar.hs
@@ -0,0 +1,391 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | Parallel scatter-accumulate aggregation kernel.
+module DataFrame.Internal.AggKernelPar (
+ scatterReducePar,
+ momentScatterPar,
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import Control.Monad (when)
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import System.IO.Unsafe (unsafePerformIO)
+import Type.Reflection (typeRep)
+
+import DataFrame.Internal.AggKernel (
+ Reduction (..),
+ scatterColumnToDouble,
+ scatterReduce,
+ )
+import DataFrame.Internal.AggPlan (Moments (..), momentScatter)
+import DataFrame.Internal.Column (
+ Column (..),
+ Columnable,
+ fromUnboxedVector,
+ materializePacked,
+ )
+
+parThreshold :: Int
+parThreshold = 200000
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+-- | Whether to take the parallel path at this row count.
+shouldPar :: Int -> Bool
+shouldPar n = n >= parThreshold && capabilities > 1
+
+groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int
+groupRangeBounds offs nGroups caps = VU.create $ do
+ b <- VUM.new (caps + 1)
+ let !nRows = VU.unsafeIndex offs nGroups
+ !per = max 1 ((nRows + caps - 1) `div` caps)
+ adv !target !gg
+ | gg >= nGroups = nGroups
+ | VU.unsafeIndex offs gg >= target = gg
+ | otherwise = adv target (gg + 1)
+ go !w !prev
+ | w >= caps = VUM.unsafeWrite b caps nGroups
+ | otherwise = do
+ let !target = min nRows (w * per)
+ !g = adv target prev
+ VUM.unsafeWrite b w g
+ go (w + 1) g
+ VUM.unsafeWrite b 0 0
+ go 1 0
+ pure b
+
+forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
+forEachRange bounds caps act
+ | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps)
+ | otherwise = do
+ vars <- mapM spawn [0 .. caps - 1]
+ results <- mapM takeMVar vars
+ mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
+ where
+ spawn w = do
+ var <- newEmptyMVar
+ let !s = VU.unsafeIndex bounds w
+ !e = VU.unsafeIndex bounds (w + 1)
+ _ <- forkIO (try (act s e) >>= putMVar var)
+ pure var
+
+scatterReducePar ::
+ Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> Maybe Column
+scatterReducePar red vis offs nGroups col
+ | not (shouldPar (VU.length vis)) || nGroups <= 1 =
+ scatterReduce red (rtgFromVis vis offs nGroups) nGroups col
+ | otherwise = case col of
+ UnboxedColumn Nothing (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> Just (reduceParTyped red vis offs nGroups v intIdent)
+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> Just (reduceParTyped red vis offs nGroups v dblIdent)
+ Nothing -> Nothing
+ p@(PackedText _ _) -> scatterReducePar red vis offs nGroups (materializePacked p)
+ _ -> Nothing
+{-# NOINLINE scatterReducePar #-}
+
+rtgFromVis :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int
+rtgFromVis vis offs nGroups = VU.create $ do
+ let n = VU.length vis
+ rtg <- VUM.new (max 1 n)
+ let go !g
+ | g >= nGroups = pure ()
+ | otherwise = do
+ let !e = VU.unsafeIndex offs (g + 1)
+ inner !pos
+ | pos >= e = pure ()
+ | otherwise = do
+ VUM.unsafeWrite rtg (VU.unsafeIndex vis pos) g
+ inner (pos + 1)
+ inner (VU.unsafeIndex offs g)
+ go (g + 1)
+ go 0
+ pure rtg
+
+data Idents a = Idents {minSeed :: !a, maxSeed :: !a}
+
+intIdent :: Idents Int
+intIdent = Idents maxBound minBound
+
+dblIdent :: Idents Double
+dblIdent = Idents (1 / 0) (negate (1 / 0))
+
+reduceParTyped ::
+ forall a.
+ (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>
+ Reduction ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Idents a ->
+ Column
+reduceParTyped red vis offs nGroups v idents =
+ let !caps = capabilities
+ !bounds = groupRangeBounds offs nGroups caps
+ in case red of
+ RCount -> fromUnboxedVector (unsafePerformIO (countPar vis offs nGroups caps bounds))
+ RSum -> fromUnboxedVector (unsafePerformIO (sumPar vis offs nGroups v caps bounds))
+ RMin ->
+ fromUnboxedVector
+ (unsafePerformIO (extremaPar min (minSeed idents) vis offs nGroups v caps bounds))
+ RMax ->
+ fromUnboxedVector
+ (unsafePerformIO (extremaPar max (maxSeed idents) vis offs nGroups v caps bounds))
+ RMean -> fromUnboxedVector (unsafePerformIO (meanPar vis offs nGroups v caps bounds))
+ RVar ->
+ fromUnboxedVector
+ (unsafePerformIO (varPar False vis offs nGroups v caps bounds))
+ RStd ->
+ fromUnboxedVector (unsafePerformIO (varPar True vis offs nGroups v caps bounds))
+ RTop2Sum -> fromUnboxedVector (unsafePerformIO (top2Par vis offs nGroups v caps bounds))
+{-# INLINE reduceParTyped #-}
+
+-- | Iterate the rows of groups @[gs, ge)@ in @valueIndices@/group order.
+overGroups ::
+ VU.Vector Int -> VU.Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
+overGroups vis offs gs ge step = grp gs
+ where
+ grp !g
+ | g >= ge = pure ()
+ | otherwise = do
+ let !e = VU.unsafeIndex offs (g + 1)
+ inner !pos
+ | pos >= e = pure ()
+ | otherwise = step g (VU.unsafeIndex vis pos) >> inner (pos + 1)
+ inner (VU.unsafeIndex offs g)
+ grp (g + 1)
+{-# INLINE overGroups #-}
+
+countPar ::
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector Int)
+countPar _vis offs nGroups caps bounds = do
+ out <- VUM.replicate nGroups (0 :: Int)
+ forEachRange bounds caps $ \gs ge ->
+ let grp !g
+ | g >= ge = pure ()
+ | otherwise = do
+ let !c = VU.unsafeIndex offs (g + 1) - VU.unsafeIndex offs g
+ VUM.unsafeWrite out g c
+ grp (g + 1)
+ in grp gs
+ VU.unsafeFreeze out
+
+sumPar ::
+ (VU.Unbox a, Num a) =>
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector a)
+sumPar vis offs nGroups v caps bounds = do
+ out <- VUM.replicate nGroups 0
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ cur <- VUM.unsafeRead out g
+ VUM.unsafeWrite out g (cur + VU.unsafeIndex v row)
+ VU.unsafeFreeze out
+{-# INLINE sumPar #-}
+
+extremaPar ::
+ (VU.Unbox a) =>
+ (a -> a -> a) ->
+ a ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector a)
+extremaPar combine seed vis offs nGroups v caps bounds = do
+ out <- VUM.replicate nGroups seed
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ cur <- VUM.unsafeRead out g
+ VUM.unsafeWrite out g (combine cur (VU.unsafeIndex v row))
+ VU.unsafeFreeze out
+{-# INLINE extremaPar #-}
+
+meanPar ::
+ (VU.Unbox a, Real a) =>
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector Double)
+meanPar vis offs nGroups v caps bounds = do
+ s <- VUM.replicate nGroups (0 :: Double)
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ let !x = realToFrac (VU.unsafeIndex v row)
+ cs <- VUM.unsafeRead s g
+ VUM.unsafeWrite s g (cs + x)
+ cc <- VUM.unsafeRead cnt g
+ VUM.unsafeWrite cnt g (cc + 1)
+ out <- VUM.new nGroups
+ let fin !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ sv <- VUM.unsafeRead s k
+ c <- VUM.unsafeRead cnt k
+ VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)
+ fin (k + 1)
+ fin 0
+ VU.unsafeFreeze out
+{-# INLINE meanPar #-}
+
+varPar ::
+ (VU.Unbox a, Real a) =>
+ Bool ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector Double)
+varPar takeSqrt vis offs nGroups v caps bounds = do
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ meanV <- VUM.replicate nGroups (0 :: Double)
+ m2 <- VUM.replicate nGroups (0 :: Double)
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ let !x = realToFrac (VU.unsafeIndex v row)
+ c <- VUM.unsafeRead cnt g
+ mu <- VUM.unsafeRead meanV g
+ mm <- VUM.unsafeRead m2 g
+ let !c' = c + 1
+ !delta = x - mu
+ !mu' = mu + delta / fromIntegral c'
+ !mm' = mm + delta * (x - mu')
+ VUM.unsafeWrite cnt g c'
+ VUM.unsafeWrite meanV g mu'
+ VUM.unsafeWrite m2 g mm'
+ out <- VUM.new nGroups
+ let fin !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ c <- VUM.unsafeRead cnt k
+ mm <- VUM.unsafeRead m2 k
+ let var = if c < 2 then 0 else mm / fromIntegral (c - 1)
+ VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)
+ fin (k + 1)
+ fin 0
+ VU.unsafeFreeze out
+{-# INLINE varPar #-}
+
+top2Par ::
+ (VU.Unbox a, Real a) =>
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector a ->
+ Int ->
+ VU.Vector Int ->
+ IO (VU.Vector Double)
+top2Par vis offs nGroups v caps bounds = do
+ let ninf = negate (1 / 0) :: Double
+ m1 <- VUM.replicate nGroups ninf
+ m2 <- VUM.replicate nGroups ninf
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ let !x = realToFrac (VU.unsafeIndex v row)
+ a1 <- VUM.unsafeRead m1 g
+ if x > a1
+ then do
+ VUM.unsafeWrite m1 g x
+ VUM.unsafeWrite m2 g a1
+ else do
+ a2 <- VUM.unsafeRead m2 g
+ when (x > a2) (VUM.unsafeWrite m2 g x)
+ out <- VUM.new nGroups
+ let fin !k
+ | k >= nGroups = pure ()
+ | otherwise = do
+ a1 <- VUM.unsafeRead m1 k
+ a2 <- VUM.unsafeRead m2 k
+ let sm = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)
+ VUM.unsafeWrite out k sm
+ fin (k + 1)
+ fin 0
+ VU.unsafeFreeze out
+{-# INLINE top2Par #-}
+
+-------------------------------------------------------------------------------
+-- Parallel fused two-column moments (Q9)
+-------------------------------------------------------------------------------
+
+{- | Parallel counterpart of 'momentScatter': one fused pass over both columns,
+each group's six sums accumulated within one worker's range. Byte-identical to
+'momentScatter'. 'Nothing' unless both columns are non-null unboxed Int/Double.
+-}
+momentScatterPar ::
+ VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Moments
+momentScatterPar vis offs nGroups colX colY
+ | not (shouldPar (VU.length vis)) || nGroups <= 1 =
+ momentScatter (rtgFromVis vis offs nGroups) nGroups colX colY
+ | otherwise = do
+ xs <- scatterColumnToDouble colX
+ ys <- scatterColumnToDouble colY
+ let !caps = capabilities
+ !bounds = groupRangeBounds offs nGroups caps
+ pure (unsafePerformIO (momentPar vis offs nGroups xs ys caps bounds))
+{-# NOINLINE momentScatterPar #-}
+
+momentPar ::
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Int ->
+ VU.Vector Double ->
+ VU.Vector Double ->
+ Int ->
+ VU.Vector Int ->
+ IO Moments
+momentPar vis offs nGroups xs ys caps bounds = do
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ sx <- VUM.replicate nGroups (0 :: Double)
+ sy <- VUM.replicate nGroups (0 :: Double)
+ sxx <- VUM.replicate nGroups (0 :: Double)
+ syy <- VUM.replicate nGroups (0 :: Double)
+ sxy <- VUM.replicate nGroups (0 :: Double)
+ let bump arr g d = VUM.unsafeRead arr g >>= \c -> VUM.unsafeWrite arr g (c + d)
+ forEachRange bounds caps $ \gs ge ->
+ overGroups vis offs gs ge $ \g row -> do
+ let !x = VU.unsafeIndex xs row
+ !y = VU.unsafeIndex ys row
+ VUM.unsafeRead cnt g >>= \c -> VUM.unsafeWrite cnt g (c + 1)
+ bump sx g x
+ bump sy g y
+ bump sxx g (x * x)
+ bump syy g (y * y)
+ bump sxy g (x * y)
+ Moments . fromUnboxedVector
+ <$> VU.unsafeFreeze cnt
+ <*> (fromUnboxedVector <$> VU.unsafeFreeze sx)
+ <*> (fromUnboxedVector <$> VU.unsafeFreeze sy)
+ <*> (fromUnboxedVector <$> VU.unsafeFreeze sxx)
+ <*> (fromUnboxedVector <$> VU.unsafeFreeze syy)
+ <*> (fromUnboxedVector <$> VU.unsafeFreeze sxy)
diff --git a/src-internal/DataFrame/Internal/AggPlan.hs b/src-internal/DataFrame/Internal/AggPlan.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/AggPlan.hs
@@ -0,0 +1,299 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | The aggregation fast-path planner and the two-column moment scatter.
+'planAgg' recognises a supported aggregate shape over a clean unboxed Int/Double
+column and returns an 'AggPlan'; 'momentScatter' fuses the six regression sums.
+-}
+module DataFrame.Internal.AggPlan (
+ AggPlan (..),
+ planAgg,
+ Moments (..),
+ momentScatter,
+ MomentPlan (..),
+ planMoments,
+) where
+
+import qualified Data.Map.Strict as M
+import qualified Data.Text as T
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Monad.ST (runST)
+import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble)
+import DataFrame.Internal.Column (Column (..), fromUnboxedVector)
+import DataFrame.Internal.DataFrame (
+ DataFrame (derivingExpressions),
+ GroupedDataFrame (..),
+ getColumn,
+ )
+import DataFrame.Internal.Expression (
+ AggStrategy (..),
+ BinaryOp (binaryCommutative, binaryName),
+ Expr (..),
+ UExpr (..),
+ )
+import Type.Reflection (Typeable, typeRep)
+
+{- | The plan 'planAgg' produces for a recognised output expression. The median
+plan carries only the column name (the holistic grouped sort lives in the
+operations layer, where @vector-algorithms@ is available).
+-}
+data AggPlan
+ = -- | A single scatter reduction over one named column.
+ PlanScatter Reduction T.Text
+ | -- | @max a - min b@ (Q7): two scatters then a vectorized combine.
+ PlanMaxMinusMin T.Text T.Text
+ | -- | Holistic median over one named column.
+ PlanMedian T.Text
+
+{- | Inspect a named output expression; return @Just plan@ on a recognised shape
+over a present clean column, else 'Nothing'. Nullable or non-Int/Double columns
+are rejected here so the scatter only sees a clean unboxed vector.
+-}
+planAgg :: GroupedDataFrame -> UExpr -> Maybe AggPlan
+planAgg gdf (UExpr (expr :: Expr a)) = case expr of
+ Agg (FoldAgg tag _ _) (Col name) -> foldPlan tag name
+ Agg (MergeAgg tag _ _ _ _) (Col name) -> mergePlan tag name
+ Agg (CollectAgg tag _) (Col name) -> collectPlan tag name
+ Binary
+ op
+ (Agg (FoldAgg lt Nothing _) (Col a))
+ (Agg (FoldAgg rt Nothing _) (Col b)) ->
+ if binaryName op == "sub" && lt == "maximum" && rt == "minimum"
+ then requireBoth a b (PlanMaxMinusMin a b)
+ else Nothing
+ _ -> Nothing
+ where
+ foldPlan tag name = case tag of
+ "sum" -> require name (PlanScatter RSum name)
+ "minimum" -> require name (PlanScatter RMin name)
+ "maximum" -> require name (PlanScatter RMax name)
+ _ -> Nothing
+ mergePlan tag name = case tag of
+ "mean" -> outputType @Double >> require name (PlanScatter RMean name)
+ "count" -> outputType @Int >> require name (PlanScatter RCount name)
+ _ -> Nothing
+ outputType :: forall t. (Typeable t) => Maybe ()
+ outputType = case testEquality (typeRep @a) (typeRep @t) of
+ Just Refl -> Just ()
+ Nothing -> Nothing
+ collectPlan tag name = case tag of
+ "stddev" -> require name (PlanScatter RStd name)
+ "variance" -> require name (PlanScatter RVar name)
+ "top2Sum" -> require name (PlanScatter RTop2Sum name)
+ "median" -> require name (PlanMedian name)
+ _ -> Nothing
+ require name plan = colUnboxedNumeric name >> Just plan
+ requireBoth a b plan = colUnboxedNumeric a >> colUnboxedNumeric b >> Just plan
+ colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of
+ Just c | isUnboxedNumeric c -> Just ()
+ _ -> Nothing
+
+-- | The matcher only fires on non-null unboxed Int/Double columns.
+isUnboxedNumeric :: Column -> Bool
+isUnboxedNumeric = \case
+ UnboxedColumn Nothing (_ :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> True
+ Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> True
+ Nothing -> False
+ _ -> False
+
+{- | A recognised moment (Q9 regression) aggregate group: six output columns that
+form the sufficient statistics of two base columns @x@ and @y@. The caller runs
+'momentScatter' once and binds each output name to a field of the result.
+-}
+data MomentPlan = MomentPlan
+ { mpColX :: T.Text
+ , mpColY :: T.Text
+ , mpNName :: T.Text
+ , mpSxName :: T.Text
+ , mpSyName :: T.Text
+ , mpSxxName :: T.Text
+ , mpSyyName :: T.Text
+ , mpSxyName :: T.Text
+ }
+
+{- | The shape of a sum's argument once unary coercions are peeled and derived
+columns are resolved through @derivingExpressions@: either linear in one base
+column or the product of two base columns (sorted).
+-}
+data Term
+ = Lin T.Text
+ | Prod T.Text T.Text
+ deriving (Eq, Ord, Show)
+
+{- | Recognise the moment shape across a whole @aggregate@ list: exactly
+@count@, @sum(x)@, @sum(y)@, @sum(x*x)@, @sum(y*y)@, @sum(x*y)@ over two distinct
+clean unboxed base columns. 'Nothing' on any other set.
+-}
+planMoments :: GroupedDataFrame -> [(T.Text, UExpr)] -> Maybe MomentPlan
+planMoments gdf aggs
+ | length aggs /= 6 = Nothing
+ | otherwise = do
+ let exprs = derivingExpressions (fullDataframe gdf)
+ roles <- traverse (classify exprs) aggs
+ let names = M.fromList [(r, nm) | (nm, r) <- roles]
+ nName <- M.lookup RoleN names
+ (x, y) <- pickBaseColumns roles
+ sxName <- M.lookup (RoleLin x) names
+ syName <- M.lookup (RoleLin y) names
+ sxxName <- M.lookup (RoleProd x x) names
+ syyName <- M.lookup (RoleProd y y) names
+ sxyName <- M.lookup (RoleProd x y) names
+ _ <- if x /= y then Just () else Nothing
+ _ <- colUnboxedNumeric x
+ _ <- colUnboxedNumeric y
+ pure
+ MomentPlan
+ { mpColX = x
+ , mpColY = y
+ , mpNName = nName
+ , mpSxName = sxName
+ , mpSyName = syName
+ , mpSxxName = sxxName
+ , mpSyyName = syyName
+ , mpSxyName = sxyName
+ }
+ where
+ colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of
+ Just c | isUnboxedNumeric c -> Just ()
+ _ -> Nothing
+
+-- | The output role each named aggregation plays in the moment shape.
+data Role
+ = RoleN
+ | RoleLin T.Text
+ | RoleProd T.Text T.Text
+ deriving (Eq, Ord, Show)
+
+-- | Tag a single named aggregation with its moment role, or reject the group.
+classify :: M.Map T.Text UExpr -> (T.Text, UExpr) -> Maybe (T.Text, Role)
+classify exprs (name, UExpr expr) = case expr of
+ Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN)
+ Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (UExpr arg)
+ _ -> Nothing
+
+termRole :: Term -> Role
+termRole (Lin a) = RoleLin a
+termRole (Prod a b) = RoleProd a b
+
+{- | Resolve a (sum-argument) expression to its 'Term'. Peels @toDouble@-style
+unary coercions, follows a derived column to its stored expression, and
+recognises a commutative product of two linear terms.
+-}
+resolveTerm :: M.Map T.Text UExpr -> UExpr -> Maybe Term
+resolveTerm exprs = go (8 :: Int)
+ where
+ go 0 _ = Nothing
+ go fuel (UExpr e) = case e of
+ Col nm -> case M.lookup nm exprs of
+ Just ue -> go (fuel - 1) ue
+ Nothing -> Just (Lin nm)
+ Unary _ inner -> go (fuel - 1) (UExpr inner)
+ Binary op l r
+ | binaryName op == "mult" && binaryCommutative op -> do
+ Lin a <- go (fuel - 1) (UExpr l)
+ Lin b <- go (fuel - 1) (UExpr r)
+ Just (sortProd a b)
+ _ -> Nothing
+
+-- | Products are unordered: store the pair sorted so @x*y@ and @y*x@ unify.
+sortProd :: T.Text -> T.Text -> Term
+sortProd a b
+ | a <= b = Prod a b
+ | otherwise = Prod b a
+
+{- | From the classified roles, find the unordered pair of base columns that the
+linear sums name. There must be exactly two distinct linear-sum columns.
+-}
+pickBaseColumns :: [(T.Text, Role)] -> Maybe (T.Text, T.Text)
+pickBaseColumns roles =
+ case lins of
+ [a, b] | a /= b -> Just (a, b)
+ _ -> Nothing
+ where
+ lins = M.keys (M.fromList [(c, ()) | (_, RoleLin c) <- roles])
+
+{- | The additive moment sums of two columns, each an @nGroups@-length column:
+@(n, Sx, Sy, Sxx, Syy, Sxy)@.
+-}
+data Moments = Moments
+ { mN :: Column
+ , mSx :: Column
+ , mSy :: Column
+ , mSxx :: Column
+ , mSyy :: Column
+ , mSxy :: Column
+ }
+
+{- | One pass over two Double-coercible columns @x@ and @y@ filling the count and
+five sums, collapsing the Q9 regression family's six folds into a single pass.
+'Nothing' unless both columns are non-null unboxed Int/Double.
+-}
+momentScatter :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments
+momentScatter g nGroups colX colY = do
+ xs <- scatterColumnToDouble colX
+ ys <- scatterColumnToDouble colY
+ let (cnt, sx, sy, sxx, syy, sxy) = momentPass g nGroups xs ys
+ pure
+ Moments
+ { mN = fromUnboxedVector cnt
+ , mSx = fromUnboxedVector sx
+ , mSy = fromUnboxedVector sy
+ , mSxx = fromUnboxedVector sxx
+ , mSyy = fromUnboxedVector syy
+ , mSxy = fromUnboxedVector sxy
+ }
+
+momentPass ::
+ VU.Vector Int ->
+ Int ->
+ VU.Vector Double ->
+ VU.Vector Double ->
+ ( VU.Vector Int
+ , VU.Vector Double
+ , VU.Vector Double
+ , VU.Vector Double
+ , VU.Vector Double
+ , VU.Vector Double
+ )
+momentPass g nGroups xs ys = runST $ do
+ cnt <- VUM.replicate nGroups (0 :: Int)
+ sx <- VUM.replicate nGroups (0 :: Double)
+ sy <- VUM.replicate nGroups (0 :: Double)
+ sxx <- VUM.replicate nGroups (0 :: Double)
+ syy <- VUM.replicate nGroups (0 :: Double)
+ sxy <- VUM.replicate nGroups (0 :: Double)
+ let n = VU.length xs
+ bump arr k d = VUM.unsafeRead arr k >>= \c -> VUM.unsafeWrite arr k (c + d)
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex g i
+ !x = VU.unsafeIndex xs i
+ !y = VU.unsafeIndex ys i
+ VUM.unsafeRead cnt k >>= \c -> VUM.unsafeWrite cnt k (c + 1)
+ bump sx k x
+ bump sy k y
+ bump sxx k (x * x)
+ bump syy k (y * y)
+ bump sxy k (x * y)
+ go (i + 1)
+ go 0
+ (,,,,,)
+ <$> VU.unsafeFreeze cnt
+ <*> VU.unsafeFreeze sx
+ <*> VU.unsafeFreeze sy
+ <*> VU.unsafeFreeze sxx
+ <*> VU.unsafeFreeze syy
+ <*> VU.unsafeFreeze sxy
diff --git a/src-internal/DataFrame/Internal/Column.hs b/src-internal/DataFrame/Internal/Column.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Column.hs
@@ -0,0 +1,1744 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+module DataFrame.Internal.Column where
+
+import qualified Data.Text as T
+import qualified Data.Vector as VB
+import qualified Data.Vector.Generic as VG
+import qualified Data.Vector.Mutable as VBM
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Exception (throw)
+import Control.Monad (forM_, when)
+import Control.Monad.ST (ST, runST)
+import Data.Bits (
+ complement,
+ popCount,
+ setBit,
+ shiftL,
+ shiftR,
+ testBit,
+ (.&.),
+ )
+import Data.Kind (Type)
+import Data.Maybe
+import Data.Type.Equality (TestEquality (..))
+import Data.Word (Word8)
+import DataFrame.Errors
+import DataFrame.Internal.PackedText (
+ PackedTextData (..),
+ packedGather,
+ packedIndexText,
+ packedLength,
+ packedRowOffsetVec,
+ packedSlice,
+ packedTake,
+ sliceEqBytes,
+ )
+import DataFrame.Internal.Types
+import DataFrame.Internal.Utf8 (sliceTextVector)
+import System.IO.Unsafe (unsafePerformIO)
+import System.Random
+import Type.Reflection
+
+-- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).
+type Bitmap = VU.Vector Word8
+
+{- | Type-erased column GADT. Pattern-matching on the constructor recovers the
+representation; nullability is an optional bit-packed 'Bitmap' (@Nothing@ = no
+nulls, @Just bm@ = bit @i@ set iff row @i@ is valid).
+-}
+data Column where
+ BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column
+ UnboxedColumn ::
+ (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column
+ -- Bit-packed Text: shared UTF-8 byte buffer + row offsets + optional bitmap;
+ -- Text is materialized on demand. Only CSV ingest emits this; user-built
+ -- Text columns stay 'BoxedColumn'.
+ PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> 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
+
+-- ---------------------------------------------------------------------------
+-- Bitmap helpers
+-- ---------------------------------------------------------------------------
+
+-- | Test whether row @i@ is valid (not null) in a bitmap.
+bitmapTestBit :: Bitmap -> Int -> Bool
+bitmapTestBit bm i = testBit (VU.unsafeIndex bm (i `shiftR` 3)) (i .&. 7)
+{-# INLINE bitmapTestBit #-}
+
+-- | Build a fully-valid bitmap for @n@ rows (all bits set).
+allValidBitmap :: Int -> Bitmap
+allValidBitmap n =
+ let bytes = (n + 7) `shiftR` 3
+ lastBits = n .&. 7
+ full = VU.replicate (bytes - 1) 0xFF
+ lastByte = if lastBits == 0 then 0xFF else (1 `shiftL` lastBits) - 1
+ in if bytes == 0 then VU.empty else VU.snoc full lastByte
+{-# INLINE allValidBitmap #-}
+
+{- | Build a bitmap from a @VU.Vector Word8@ validity vector
+(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.
+-}
+buildBitmapFromValid :: VU.Vector Word8 -> Bitmap
+buildBitmapFromValid valid =
+ let n = VU.length valid
+ bytes = (n + 7) `shiftR` 3
+ in VU.generate bytes $ \b ->
+ let base = b `shiftL` 3
+ setBitIf acc bit =
+ let idx = base + bit
+ in if idx < n && VU.unsafeIndex valid idx /= 0
+ then setBit acc bit
+ else acc
+ in foldl setBitIf (0 :: Word8) [0 .. 7]
+
+{- | Build a bitmap from a list of null-row indices.
+@nullIdxs@ are the positions that are NULL.
+-}
+buildBitmapFromNulls :: Int -> [Int] -> Bitmap
+buildBitmapFromNulls n nullIdxs =
+ let base = allValidBitmap n
+ in VU.modify
+ ( \mv ->
+ forM_ nullIdxs $ \i -> do
+ let byteIdx = i `shiftR` 3
+ bitIdx = i .&. 7
+ v <- VUM.unsafeRead mv byteIdx
+ VUM.unsafeWrite mv byteIdx (clearBit8 v bitIdx)
+ )
+ base
+ where
+ clearBit8 :: Word8 -> Int -> Word8
+ clearBit8 b bit = b .&. complement (1 `shiftL` bit)
+
+-- | Slice a bitmap for rows @[start .. start+len-1]@.
+bitmapSlice :: Int -> Int -> Bitmap -> Bitmap
+bitmapSlice start len bm
+ | start .&. 7 == 0 =
+ let startByte = start `shiftR` 3
+ bytes = min ((len + 7) `shiftR` 3) (VU.length bm - startByte)
+ in VU.slice startByte bytes bm
+ | otherwise =
+ let n = min len (VU.length bm `shiftL` 3 - start)
+ in buildBitmapFromValid $
+ VU.generate n $
+ \i -> if bitmapTestBit bm (start + i) then 1 else 0
+
+-- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.
+bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap
+bitmapConcat n1 bm1 n2 bm2 =
+ buildBitmapFromValid $
+ VU.generate (n1 + n2) $ \i ->
+ if i < n1
+ then if bitmapTestBit bm1 i then 1 else 0
+ else if bitmapTestBit bm2 (i - n1) then 1 else 0
+
+-- | Combine two bitmaps with AND (both must be valid for result to be valid).
+mergeBitmaps :: Bitmap -> Bitmap -> Bitmap
+mergeBitmaps = VU.zipWith (.&.)
+
+{- | Materialize a nullable column from @VB.Vector (Maybe a)@; picks 'UnboxedColumn'
+when @a@ is unboxable, else 'BoxedColumn'. Always attaches a bitmap so the column
+reads as nullable even with no 'Nothing' values.
+-}
+fromMaybeVec :: forall a. (Columnable a) => VB.Vector (Maybe a) -> Column
+fromMaybeVec v = case sUnbox @a of
+ STrue -> fromMaybeVecUnboxed v
+ SFalse ->
+ let n = VB.length v
+ nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]
+ bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs
+ dat = VB.map (fromMaybe (errorWithoutStackTrace "fromMaybeVec: Nothing slot")) v
+ in BoxedColumn (Just bm) dat
+
+{- | Materialize a nullable 'UnboxedColumn' to @VB.Vector (Maybe a)@ using runST.
+Always attaches a bitmap so the column is recognized as nullable even when
+no 'Nothing' values are present (preserves the Maybe type marker).
+-}
+fromMaybeVecUnboxed ::
+ forall a. (Columnable a, VU.Unbox a) => VB.Vector (Maybe a) -> Column
+fromMaybeVecUnboxed v =
+ let n = VB.length v
+ nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]
+ bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs
+ dat = runST $ do
+ mv <- VUM.new n
+ VG.iforM_ v $ \i mx -> forM_ mx (VUM.unsafeWrite mv i)
+ VU.unsafeFreeze mv
+ in UnboxedColumn (Just bm) dat
+
+-- | Whether row @i@ is null, respecting the bitmap.
+columnElemIsNull :: Column -> Int -> Bool
+columnElemIsNull (BoxedColumn (Just bm) _) i = not (bitmapTestBit bm i)
+columnElemIsNull (UnboxedColumn (Just bm) _) i = not (bitmapTestBit bm i)
+columnElemIsNull (PackedText (Just bm) _) i = not (bitmapTestBit bm i)
+columnElemIsNull _ _ = False
+
+-- | Return the 'Maybe Bitmap' from a column.
+columnBitmap :: Column -> Maybe Bitmap
+columnBitmap (BoxedColumn bm _) = bm
+columnBitmap (UnboxedColumn bm _) = bm
+columnBitmap (PackedText bm _) = bm
+
+{- | Decode a 'PackedText' into a @BoxedColumn Text@ (bit-identical to
+materializing at freeze). Identity on every other column.
+-}
+materializePacked :: Column -> Column
+materializePacked (PackedText bm p) = case packedRowOffsetVec p of
+ Just (arr, offs) -> BoxedColumn bm (sliceTextVector arr offs)
+ Nothing -> BoxedColumn bm (VB.generate (packedLength p) (packedIndexText p))
+materializePacked c = c
+{-# INLINE materializePacked #-}
+
+-- | Whether a column is a 'PackedText'.
+isPackedText :: Column -> Bool
+isPackedText (PackedText _ _) = True
+isPackedText _ = False
+{-# INLINE isPackedText #-}
+
+-- ---------------------------------------------------------------------------
+-- End bitmap helpers
+-- ---------------------------------------------------------------------------
+
+{- | A wrapper around the type-erased 'Column' carrying a phantom element type,
+used to type-check expressions. The phantom is not guaranteed to match the
+underlying vector's type.
+-}
+data TypedColumn a where
+ TColumn :: (Columnable a) => Column -> TypedColumn a
+
+instance (Eq a) => Eq (TypedColumn a) where
+ (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool
+ (==) (TColumn a) (TColumn b) = a == b
+
+-- | Gets the underlying value from a TypedColumn.
+unwrapTypedColumn :: TypedColumn a -> Column
+unwrapTypedColumn (TColumn value) = value
+
+-- | Gets the underlying vector from a TypedColumn.
+vectorFromTypedColumn :: TypedColumn a -> VB.Vector a
+vectorFromTypedColumn (TColumn value) = either throw id (toVector value)
+
+-- | Checks if a column contains missing values (has a bitmap).
+hasMissing :: Column -> Bool
+hasMissing (BoxedColumn (Just _) _) = True
+hasMissing (UnboxedColumn (Just _) _) = True
+hasMissing (PackedText (Just _) _) = True
+hasMissing _ = False
+
+-- | Checks if a column contains only missing values.
+allMissing :: Column -> Bool
+allMissing (BoxedColumn (Just bm) col) = VU.all (== 0) bm && not (VB.null col)
+allMissing (UnboxedColumn (Just bm) col) = VU.all (== 0) bm && not (VU.null col)
+allMissing (PackedText (Just bm) p) = VU.all (== 0) bm && packedLength p > 0
+allMissing _ = False
+
+-- | Checks if a column contains numeric values.
+isNumeric :: Column -> Bool
+isNumeric (UnboxedColumn _ (_vec :: VU.Vector a)) = case sNumeric @a of
+ STrue -> True
+ _ -> False
+isNumeric (BoxedColumn _ (_vec :: VB.Vector a)) = case testEquality (typeRep @a) (typeRep @Integer) of
+ Nothing -> False
+ Just Refl -> True
+isNumeric (PackedText _ _) = False
+
+{- | Whether the column stores element type @a@. For nullable columns, also
+'True' when @a = Maybe b@ and the column stores @b@ internally.
+-}
+hasElemType :: forall a. (Columnable a) => Column -> Bool
+hasElemType = \case
+ BoxedColumn bm (_column :: VB.Vector b) -> checkBoxed bm (typeRep @b)
+ UnboxedColumn bm (_column :: VU.Vector b) -> checkUnboxed bm (typeRep @b)
+ PackedText bm _ -> checkBoxed bm (typeRep @T.Text)
+ where
+ directMatch :: forall (b :: Type). TypeRep b -> Bool
+ directMatch = isJust . testEquality (typeRep @a)
+ checkMaybe :: forall (b :: Type). TypeRep b -> Bool
+ checkMaybe tb = case typeRep @a of
+ App tMaybe tInner -> case eqTypeRep tMaybe (typeRep @Maybe) of
+ Just HRefl -> isJust (testEquality tInner tb)
+ Nothing -> False
+ _ -> False
+ checkBoxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
+ checkBoxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)
+ checkUnboxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
+ checkUnboxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)
+
+-- | An internal/debugging function to get the column type of a column.
+columnVersionString :: Column -> String
+columnVersionString column = case column of
+ BoxedColumn Nothing _ -> "Boxed"
+ BoxedColumn (Just _) _ -> "NullableBoxed"
+ UnboxedColumn Nothing _ -> "Unboxed"
+ UnboxedColumn (Just _) _ -> "NullableUnboxed"
+ PackedText Nothing _ -> "Boxed"
+ PackedText (Just _) _ -> "NullableBoxed"
+
+{- | An internal/debugging function to get the type stored in the outermost vector
+of a column.
+-}
+columnTypeString :: Column -> String
+columnTypeString column = case column of
+ BoxedColumn Nothing (_ :: VB.Vector a) -> show (typeRep @a)
+ BoxedColumn (Just _) (_ :: VB.Vector a) -> showMaybeType @a
+ UnboxedColumn Nothing (_ :: VU.Vector a) -> show (typeRep @a)
+ UnboxedColumn (Just _) (_ :: VU.Vector a) -> showMaybeType @a
+ PackedText Nothing _ -> show (typeRep @T.Text)
+ PackedText (Just _) _ -> showMaybeType @T.Text
+ where
+ showMaybeType :: forall a. (Typeable a) => String
+ showMaybeType =
+ let s = show (typeRep @a)
+ in "Maybe " ++ if ' ' `elem` s then "(" ++ s ++ ")" else s
+
+instance (Show a) => Show (TypedColumn a) where
+ show :: (Show a) => TypedColumn a -> String
+ show (TColumn col) = show col
+
+{- | Force evaluation of all elements in a column. Replacement for the removed
+@instance NFData Column@; used by the IO and lazy-executor strict paths.
+-}
+forceColumn :: Column -> ()
+forceColumn (BoxedColumn Nothing (v :: VB.Vector a)) = VB.foldl' (const (`seq` ())) () v
+forceColumn (BoxedColumn (Just bm) (v :: VB.Vector a)) =
+ let n = VB.length v
+ go !i
+ | i >= n = ()
+ | bitmapTestBit bm i = VB.unsafeIndex v i `seq` go (i + 1)
+ | otherwise = go (i + 1)
+ in go 0
+forceColumn (UnboxedColumn _ v) = v `seq` ()
+forceColumn (PackedText _ (PackedTextData arr offs sel)) = arr `seq` offs `seq` sel `seq` ()
+
+instance Show Column where
+ show :: Column -> String
+ show (BoxedColumn Nothing column) = show column
+ show (BoxedColumn (Just bm) column) =
+ let n = VB.length column
+ elems =
+ [ if bitmapTestBit bm i then show (VB.unsafeIndex column i) else "null"
+ | i <- [0 .. n - 1]
+ ]
+ in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"
+ show (UnboxedColumn Nothing column) = show column
+ show (UnboxedColumn (Just bm) column) =
+ let n = VU.length column
+ elems =
+ [ if bitmapTestBit bm i then show (VU.unsafeIndex column i) else "null"
+ | i <- [0 .. n - 1]
+ ]
+ in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"
+ show c@(PackedText _ _) = show (materializePacked c)
+
+{- | Compare two nullable boxed columns element by element, skipping null slots.
+Uses a manual loop to avoid stream fusion forcing null-slot error thunks.
+-}
+eqBoxedCols ::
+ (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool
+eqBoxedCols bm1 a bm2 b
+ | VB.length a /= VB.length b = False
+ | otherwise = go 0
+ where
+ !n = VB.length a
+ go !i
+ | i >= n = True
+ | nullA || nullB = (nullA == nullB) && go (i + 1)
+ | VB.unsafeIndex a i == VB.unsafeIndex b i = go (i + 1)
+ | otherwise = False
+ where
+ nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
+ nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
+{-# INLINE eqBoxedCols #-}
+
+instance Eq Column where
+ (==) :: Column -> Column -> Bool
+ (==) (BoxedColumn bm1 (a :: VB.Vector t1)) (BoxedColumn bm2 (b :: VB.Vector t2)) =
+ case testEquality (typeRep @t1) (typeRep @t2) of
+ Nothing -> False
+ Just Refl -> eqBoxedCols bm1 a bm2 b
+ (==) (UnboxedColumn bm1 (a :: VU.Vector t1)) (UnboxedColumn bm2 (b :: VU.Vector t2)) =
+ case testEquality (typeRep @t1) (typeRep @t2) of
+ Nothing -> False
+ Just Refl ->
+ VU.length a == VU.length b
+ && VU.and
+ ( VU.imap
+ ( \i x ->
+ let nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
+ nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
+ in if nullA || nullB then nullA == nullB else x == VU.unsafeIndex b i
+ )
+ a
+ )
+ (==) (PackedText bm1 p1) (PackedText bm2 p2) = eqPackedCols bm1 p1 bm2 p2
+ (==) lhs@(PackedText _ _) rhs = materializePacked lhs == rhs
+ (==) lhs rhs@(PackedText _ _) = lhs == materializePacked rhs
+ (==) _ _ = False
+
+{- | Byte-slice equality of two packed-text columns, skipping null slots
+(a null compares equal only to a null), mirroring 'eqBoxedCols'.
+-}
+eqPackedCols ::
+ Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
+eqPackedCols bm1 p1 bm2 p2
+ | packedLength p1 /= packedLength p2 = False
+ | otherwise = go 0
+ where
+ !n = packedLength p1
+ go !i
+ | i >= n = True
+ | nullA || nullB = (nullA == nullB) && go (i + 1)
+ | otherwise =
+ let (a1, o1, l1) = packedSlice p1 i
+ (a2, o2, l2) = packedSlice p2 i
+ in sliceEqBytes a1 o1 l1 a2 o2 l2 && go (i + 1)
+ where
+ nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
+ nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
+{-# INLINE eqPackedCols #-}
+
+{- | A class for converting a vector to a column of the appropriate type.
+Given each Rep we tell the `toColumnRep` function which Column type to pick.
+-}
+class ColumnifyRep (r :: Rep) a where
+ toColumnRep :: VB.Vector a -> Column
+
+-- | Constraint synonym for what we can put into columns.
+type Columnable a =
+ ( Columnable' a
+ , ColumnifyRep (KindOf a) a
+ , UnboxIf a
+ , IntegralIf a
+ , FloatingIf a
+ , SBoolI (Unboxable a)
+ , SBoolI (Numeric a)
+ , SBoolI (IntegralTypes a)
+ , SBoolI (FloatingTypes a)
+ )
+
+instance
+ (Columnable a, VU.Unbox a) =>
+ ColumnifyRep 'RUnboxed a
+ where
+ toColumnRep :: (Columnable a, VUM.Unbox a) => VB.Vector a -> Column
+ toColumnRep v = UnboxedColumn Nothing (VU.convert v)
+
+instance
+ (Columnable a) =>
+ ColumnifyRep 'RBoxed a
+ where
+ toColumnRep :: (Columnable a) => VB.Vector a -> Column
+ toColumnRep = BoxedColumn Nothing
+
+instance
+ (Columnable a) =>
+ ColumnifyRep 'RNullableBoxed (Maybe a)
+ where
+ toColumnRep :: (Columnable a) => VB.Vector (Maybe a) -> Column
+ toColumnRep = fromMaybeVec
+
+{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.
+
+__Examples:__
+
+@
+> import qualified Data.Vector as V
+> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])
+[1,2,3,4]
+@
+-}
+fromVector ::
+ forall a.
+ (Columnable a, ColumnifyRep (KindOf a) a) =>
+ VB.Vector a -> Column
+fromVector = toColumnRep @(KindOf a)
+
+{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.
+
+__Examples:__
+
+@
+> import qualified Data.Vector.Unboxed as V
+> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])
+[1,2,3,4]
+@
+-}
+fromUnboxedVector ::
+ forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column
+fromUnboxedVector = UnboxedColumn Nothing
+
+{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.
+
+__Examples:__
+
+@
+> fromList [(1 :: Int), 2, 3, 4]
+[1,2,3,4]
+@
+-}
+fromList ::
+ forall a.
+ (Columnable a, ColumnifyRep (KindOf a) a) =>
+ [a] -> Column
+fromList = toColumnRep @(KindOf a) . VB.fromList
+
+{- | O(n) Create a column of random elements within a range.
+
+Takes a random number generator, a length, and a lower and upper bound for the random values.
+
+__Examples:__
+
+@
+> import System.Random (mkStdGen)
+> mkRandom (mkStdGen 42) 4 0 10
+[4,2,6,5]
+@
+-}
+mkRandom ::
+ (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>
+ g -> Int -> a -> a -> Column
+mkRandom pureGen k lo hi = fromList $ go pureGen k
+ where
+ go _g 0 = []
+ go g n =
+ let
+ (!v, !g') = uniformR (lo, hi) g
+ in
+ v : go g' (n - 1)
+
+-- An internal helper for type errors
+throwTypeMismatch ::
+ forall (a :: Type) (b :: Type).
+ (Typeable a, Typeable b) => Either DataFrameException Column
+throwTypeMismatch =
+ Left $
+ TypeMismatchException
+ MkTypeErrorContext
+ { userType = Right (typeRep @b)
+ , expectedType = Right (typeRep @a)
+ , callingFunctionName = Nothing
+ , errorColumnName = Nothing
+ }
+
+-- | An internal function to map a function over the values of a column.
+mapColumn ::
+ forall b c.
+ (Columnable b, Columnable c) =>
+ (b -> c) -> Column -> Either DataFrameException Column
+mapColumn f = \case
+ BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col
+ UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col
+ c@(PackedText _ _) -> mapColumn f (materializePacked c)
+ where
+ runBoxed ::
+ forall a.
+ (Columnable a) =>
+ Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
+ runBoxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of
+ Just Refl ->
+ let !n = VB.length col
+ in Right $ case sUnbox @c of
+ STrue -> UnboxedColumn Nothing $
+ VU.generate n $ \i ->
+ f
+ ( if maybe True (`bitmapTestBit` i) bm
+ then Just (VB.unsafeIndex col i)
+ else Nothing
+ )
+ SFalse -> fromVector @c $
+ VB.generate n $ \i ->
+ f
+ ( if maybe True (`bitmapTestBit` i) bm
+ then Just (VB.unsafeIndex col i)
+ else Nothing
+ )
+ Nothing -> case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl ->
+ Right $ case sUnbox @c of
+ STrue -> UnboxedColumn bm (VU.generate (VB.length col) (f . VB.unsafeIndex col))
+ SFalse -> case bm of
+ Nothing -> fromVector @c (VB.map f col)
+ Just _ -> BoxedColumn bm (VB.map f col)
+ Nothing -> throwTypeMismatch @a @b
+
+ runUnboxed ::
+ forall a.
+ (Columnable a, VU.Unbox a) =>
+ Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
+ runUnboxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of
+ Just Refl ->
+ let !n = VU.length col
+ in Right $ case sUnbox @c of
+ STrue -> UnboxedColumn Nothing $
+ VU.generate n $ \i ->
+ f
+ ( if maybe True (`bitmapTestBit` i) bm
+ then Just (VU.unsafeIndex col i)
+ else Nothing
+ )
+ SFalse -> fromVector @c $
+ VB.generate n $ \i ->
+ f
+ ( if maybe True (`bitmapTestBit` i) bm
+ then Just (VU.unsafeIndex col i)
+ else Nothing
+ )
+ Nothing -> case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Right $ case sUnbox @c of
+ STrue -> UnboxedColumn bm (VU.map f col)
+ SFalse -> case bm of
+ Nothing -> fromVector @c (VB.generate (VU.length col) (f . VU.unsafeIndex col))
+ Just _ -> BoxedColumn bm (VB.generate (VU.length col) (f . VU.unsafeIndex col))
+ Nothing -> throwTypeMismatch @a @b
+{-# INLINEABLE mapColumn #-}
+
+-- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.
+imapColumn ::
+ forall b c.
+ (Columnable b, Columnable c) =>
+ (Int -> b -> c) -> Column -> Either DataFrameException Column
+imapColumn f = \case
+ BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col
+ UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col
+ c@(PackedText _ _) -> imapColumn f (materializePacked c)
+ where
+ runBoxed ::
+ forall a.
+ (Columnable a) =>
+ Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
+ runBoxed bm col = case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Right $ case sUnbox @c of
+ STrue ->
+ UnboxedColumn
+ bm
+ (VU.generate (VB.length col) (\i -> f i (VB.unsafeIndex col i)))
+ SFalse -> BoxedColumn bm (VB.imap f col)
+ Nothing -> throwTypeMismatch @a @b
+
+ runUnboxed ::
+ forall a.
+ (Columnable a, VU.Unbox a) =>
+ Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
+ runUnboxed bm col = case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Right $ case sUnbox @c of
+ STrue -> UnboxedColumn bm (VU.imap f col)
+ SFalse -> BoxedColumn bm (VB.imap f (VG.convert col))
+ Nothing -> throwTypeMismatch @a @b
+
+-- | O(1) Gets the number of elements in the column.
+columnLength :: Column -> Int
+columnLength (BoxedColumn _ xs) = VB.length xs
+columnLength (UnboxedColumn _ xs) = VU.length xs
+columnLength (PackedText _ p) = packedLength p
+{-# INLINE columnLength #-}
+
+-- | O(n) Gets the number of non-null elements in the column.
+numElements :: Column -> Int
+numElements (BoxedColumn Nothing xs) = VB.length xs
+numElements (BoxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+numElements (UnboxedColumn Nothing xs) = VU.length xs
+numElements (UnboxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+numElements (PackedText Nothing p) = packedLength p
+numElements (PackedText (Just bm) _p) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+{-# INLINE numElements #-}
+
+-- | O(n) Takes the first n values of a column.
+takeColumn :: Int -> Column -> Column
+takeColumn n (BoxedColumn bm xs) =
+ BoxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)
+takeColumn n (UnboxedColumn bm xs) =
+ UnboxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)
+takeColumn n (PackedText bm p) =
+ PackedText (fmap (bitmapSlice 0 n) bm) (packedTake n p)
+{-# INLINE takeColumn #-}
+
+-- | O(n) Takes the last n values of a column.
+takeLastColumn :: Int -> Column -> Column
+takeLastColumn n column = sliceColumn (columnLength column - n) n column
+{-# INLINE takeLastColumn #-}
+
+-- | O(n) Takes n values after a given column index.
+sliceColumn :: Int -> Int -> Column -> Column
+sliceColumn start n (BoxedColumn bm xs) =
+ BoxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)
+sliceColumn start n (UnboxedColumn bm xs) =
+ UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)
+sliceColumn start n c@(PackedText _ _) = sliceColumn start n (materializePacked c)
+{-# INLINE sliceColumn #-}
+
+-- | 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 bm column) =
+ BoxedColumn
+ ( fmap
+ ( \bm0 ->
+ buildBitmapFromValid $
+ VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
+ )
+ bm
+ )
+ ( VB.generate
+ (VU.length indexes)
+ ((column `VB.unsafeIndex`) . (indexes `VU.unsafeIndex`))
+ )
+atIndicesStable indexes (UnboxedColumn bm column) =
+ UnboxedColumn
+ ( fmap
+ ( \bm0 ->
+ buildBitmapFromValid $
+ VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
+ )
+ bm
+ )
+ (VU.unsafeBackpermute column indexes)
+atIndicesStable indexes (PackedText bm p) =
+ PackedText
+ ( fmap
+ ( \bm0 ->
+ buildBitmapFromValid $
+ VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
+ )
+ bm
+ )
+ (packedGather indexes p)
+{-# INLINE atIndicesStable #-}
+
+{- | Like 'atIndicesStable' but treats negative indices as null.
+Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).
+-}
+gatherWithSentinel :: VU.Vector Int -> Column -> Column
+gatherWithSentinel indices col =
+ let !n = VU.length indices
+ newBm = buildBitmapFromValid $ VU.generate n $ \i ->
+ if VU.unsafeIndex indices i < 0 then 0 else 1
+ in case col of
+ PackedText srcBm p ->
+ let bm = case srcBm of
+ Nothing -> Just newBm
+ Just sb ->
+ Just
+ ( mergeBitmaps
+ newBm
+ ( buildBitmapFromValid $ VU.generate n $ \i ->
+ let idx = VU.unsafeIndex indices i
+ in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
+ )
+ )
+ in PackedText bm (packedGather indices p)
+ BoxedColumn srcBm v ->
+ let dat = VB.generate n $ \i ->
+ let !idx = VU.unsafeIndex indices i
+ in if idx < 0 then VB.unsafeIndex v 0 else VB.unsafeIndex v idx
+ bm = case srcBm of
+ Nothing -> Just newBm
+ Just sb ->
+ Just
+ ( mergeBitmaps
+ newBm
+ ( buildBitmapFromValid $ VU.generate n $ \i ->
+ let idx = VU.unsafeIndex indices i
+ in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
+ )
+ )
+ in BoxedColumn bm dat
+ UnboxedColumn srcBm v ->
+ let dat = runST $ do
+ mv <- VUM.new n
+ VG.iforM_ indices $ \i idx ->
+ when (idx >= 0) $ VUM.unsafeWrite mv i (VU.unsafeIndex v idx)
+ VU.unsafeFreeze mv
+ bm = case srcBm of
+ Nothing -> Just newBm
+ Just sb ->
+ Just
+ ( mergeBitmaps
+ newBm
+ ( buildBitmapFromValid $ VU.generate n $ \i ->
+ let idx = VU.unsafeIndex indices i
+ in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
+ )
+ )
+ in UnboxedColumn bm dat
+{-# INLINE gatherWithSentinel #-}
+
+-- | 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))
+{-# INLINE getIndices #-}
+
+-- | Internal helper to get indices in an unboxed vector.
+getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a
+getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i))
+{-# INLINE getIndicesUnboxed #-}
+
+findIndices ::
+ forall a.
+ (Columnable a) =>
+ (a -> Bool) ->
+ Column ->
+ Either DataFrameException (VU.Vector Int)
+findIndices predicate = \case
+ BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert
+ UnboxedColumn _ (v :: VU.Vector b) -> run v id
+ c@(PackedText _ _) -> findIndices predicate (materializePacked c)
+ where
+ run ::
+ forall b v.
+ (Typeable b, VG.Vector v b, VG.Vector v Int) =>
+ v b ->
+ (v Int -> VU.Vector Int) ->
+ Either DataFrameException (VU.Vector Int)
+ run column finalize = case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Right . finalize $ VG.findIndices predicate column
+ Nothing ->
+ Left $
+ TypeMismatchException
+ MkTypeErrorContext
+ { userType = Right (typeRep @a)
+ , expectedType = Right (typeRep @b)
+ , callingFunctionName = Just "findIndices"
+ , errorColumnName = Nothing
+ }
+
+-- | Fold (right) column with index.
+ifoldrColumn ::
+ forall a b.
+ (Columnable a, Columnable b) =>
+ (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
+ifoldrColumn f acc = \case
+ BoxedColumn _ column -> foldrWorker column
+ UnboxedColumn _ column -> foldrWorker column
+ c@(PackedText _ _) -> ifoldrColumn f acc (materializePacked c)
+ 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 @c)
+ , callingFunctionName = Just "ifoldrColumn"
+ , errorColumnName = Nothing
+ }
+ )
+
+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
+ c@(PackedText _ _) -> foldlColumn f acc (materializePacked c)
+ 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 @c)
+ , callingFunctionName = Just "ifoldrColumn"
+ , errorColumnName = Nothing
+ }
+ )
+
+foldl1Column ::
+ forall a.
+ (Columnable a) =>
+ (a -> a -> a) -> Column -> Either DataFrameException a
+foldl1Column f = \case
+ BoxedColumn _ column -> foldl1Worker column
+ UnboxedColumn _ column -> foldl1Worker column
+ c@(PackedText _ _) -> foldl1Column f (materializePacked c)
+ 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 @c)
+ , callingFunctionName = Just "foldl1Column"
+ , errorColumnName = Nothing
+ }
+ )
+
+{- | 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.
+-}
+foldl1DirectGroups ::
+ forall a.
+ (Columnable a) =>
+ (a -> a -> a) ->
+ Column ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ Either DataFrameException Column
+foldl1DirectGroups f col valueIndices offsets
+ | VU.length offsets <= 1 = pure $ fromVector @a VB.empty
+ | otherwise = case col of
+ UnboxedColumn _ (vec :: VU.Vector d) -> UnboxedColumn Nothing <$> foldl1Worker vec
+ BoxedColumn _ (vec :: VB.Vector d) -> BoxedColumn Nothing <$> foldl1Worker vec
+ PackedText _ _ -> foldl1DirectGroups f (materializePacked col) valueIndices offsets
+ 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 (rowToGroup[i] = group
+of row i). Random writes hit the small per-group accumulator array; when @acc@ is
+unboxable that array is unboxed, avoiding pointer indirection.
+-}
+foldLinearGroups ::
+ forall b acc.
+ (Columnable b, Columnable acc) =>
+ (acc -> b -> acc) ->
+ acc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+foldLinearGroups f seed col rowToGroup nGroups
+ | nGroups == 0 = Right (fromVector @acc VB.empty)
+ | otherwise = case col of
+ UnboxedColumn _ (vec :: VU.Vector d) -> foldLinearWorker vec
+ BoxedColumn _ (vec :: VB.Vector d) -> foldLinearWorker vec
+ PackedText _ _ ->
+ foldLinearGroups f seed (materializePacked col) rowToGroup nGroups
+ 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
+ }
+
+ runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column
+ runWith body = case sUnbox @acc of
+ STrue -> do
+ accs <- VUM.replicate nGroups seed
+ body (VUM.unsafeRead accs) (VUM.unsafeWrite accs)
+ UnboxedColumn Nothing <$> VU.unsafeFreeze accs
+ SFalse -> do
+ accs <- VBM.replicate nGroups seed
+ body (VBM.unsafeRead accs) (VBM.unsafeWrite accs)
+ fromVector @acc <$> VB.unsafeFreeze accs
+ {-# INLINE runWith #-}
+{-# INLINEABLE foldLinearGroups #-}
+
+headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a
+headColumn = \case
+ BoxedColumn _ col -> headWorker col
+ UnboxedColumn _ col -> headWorker col
+ c@(PackedText _ _) -> headColumn (materializePacked c)
+ 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
+zipColumns l@(PackedText _ _) r = zipColumns (materializePacked l) r
+zipColumns l r@(PackedText _ _) = zipColumns l (materializePacked r)
+zipColumns (BoxedColumn _ column) (BoxedColumn _ other) = BoxedColumn Nothing (VG.zip column other)
+zipColumns (BoxedColumn _ column) (UnboxedColumn _ other) =
+ BoxedColumn
+ Nothing
+ ( VB.generate
+ (min (VG.length column) (VG.length other))
+ (\i -> (column VG.! i, other VG.! i))
+ )
+zipColumns (UnboxedColumn _ column) (BoxedColumn _ other) =
+ BoxedColumn
+ Nothing
+ ( VB.generate
+ (min (VG.length column) (VG.length other))
+ (\i -> (column VG.! i, other VG.! i))
+ )
+zipColumns (UnboxedColumn _ column) (UnboxedColumn _ other) = UnboxedColumn Nothing (VG.zip column other)
+{-# INLINE zipColumns #-}
+
+-- | Merge two columns using `These`.
+mergeColumns :: Column -> Column -> Column
+mergeColumns colA colB = case (colA, colB) of
+ (PackedText _ _, _) -> mergeColumns (materializePacked colA) colB
+ (_, PackedText _ _) -> mergeColumns colA (materializePacked colB)
+ (BoxedColumn bmA c1, BoxedColumn bmB c2) -> case (bmA, bmB) of
+ (Just ba, Just bb) ->
+ BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
+ let nullA = not (bitmapTestBit ba i)
+ nullB = not (bitmapTestBit bb i)
+ in case (nullA, nullB) of
+ (True, True) -> error "mergeColumns: both null"
+ (False, True) -> This v1
+ (True, False) -> That v2
+ (False, False) -> These v1 v2
+ (Just ba, Nothing) ->
+ BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
+ if not (bitmapTestBit ba i) then That v2 else These v1 v2
+ (Nothing, Just bb) ->
+ BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
+ if not (bitmapTestBit bb i) then This v1 else These v1 v2
+ (Nothing, Nothing) ->
+ BoxedColumn Nothing $ mkVecSimple c1 c2 These
+ (BoxedColumn _ c1, UnboxedColumn _ c2) ->
+ BoxedColumn Nothing $ mkVecSimple c1 c2 These
+ (UnboxedColumn _ c1, BoxedColumn _ c2) ->
+ BoxedColumn Nothing $ mkVecSimple c1 c2 These
+ (UnboxedColumn _ c1, UnboxedColumn _ c2) ->
+ BoxedColumn Nothing $ mkVecSimple c1 c2 These
+ where
+ mkVec c1 c2 combineElements =
+ VB.generate
+ (min (VG.length c1) (VG.length c2))
+ (\i -> combineElements i (c1 VG.! i) (c2 VG.! i))
+ {-# INLINE mkVec #-}
+
+ mkVecSimple c1 c2 f =
+ VB.generate
+ (min (VG.length c1) (VG.length c2))
+ (\i -> f (c1 VG.! i) (c2 VG.! i))
+ {-# INLINE mkVecSimple #-}
+{-# INLINE mergeColumns #-}
+
+-- | An internal, column version of zipWith.
+zipWithColumns ::
+ forall a b c.
+ (Columnable a, Columnable b, Columnable c) =>
+ (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
+zipWithColumns f (UnboxedColumn bmL (column :: VU.Vector d)) (UnboxedColumn bmR (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of
+ Just Refl -> case testEquality (typeRep @b) (typeRep @e) of
+ Just Refl
+ | isNothing bmL
+ , isNothing bmR ->
+ pure $ case sUnbox @c of
+ STrue -> UnboxedColumn Nothing (VU.zipWith f column other)
+ SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)
+ _ -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)
+ Nothing -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)
+-- TODO: mchavinda - reuse pattern from interpret where we augment the
+-- error at the end.
+zipWithColumns f left right = zipWithColumnsGeneral f left right
+
+zipWithColumnsGeneral ::
+ forall a b c.
+ (Columnable a, Columnable b, Columnable c) =>
+ (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
+zipWithColumnsGeneral f left right = case toVector @a left of
+ Left (TypeMismatchException context) ->
+ Left $
+ TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})
+ Left e -> Left e
+ Right left' -> case toVector @b right of
+ Left (TypeMismatchException context) ->
+ Left $
+ TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})
+ Left e -> Left e
+ Right right' -> pure $ fromVector $ VB.zipWith f left' right'
+{-# INLINE zipWithColumnsGeneral #-}
+{-# INLINE zipWithColumns #-}
+
+-- writeColumn and freezeColumn' (CSV-ingest helpers) moved to
+-- DataFrame.IO.Internal.MutableColumn so the core column module does not
+-- need to depend on DataFrame.Internal.Parsing.
+
+{- | Freeze a mutable column into an @Either Text a@ column: every recorded
+null position becomes @Left rawText@ (preserving the original input), every
+other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.
+-}
+freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column
+freezeColumnEither nulls (MBoxedColumn col) = do
+ frozen <- VB.unsafeFreeze col
+ let nullMap = nulls
+ pure $
+ BoxedColumn Nothing $
+ VB.imap
+ ( \i v -> case lookup i nullMap of
+ Just t -> Left t
+ Nothing -> Right v
+ )
+ frozen
+freezeColumnEither nulls (MUnboxedColumn col) = do
+ c <- VU.unsafeFreeze col
+ let nullMap = nulls
+ pure $
+ BoxedColumn Nothing $
+ VB.generate (VU.length c) $ \i ->
+ case lookup i nullMap of
+ Just t -> Left t
+ Nothing -> Right (c VU.! i)
+{-# INLINE freezeColumnEither #-}
+
+{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).
+No-op when already nullable.
+-}
+ensureOptional :: Column -> Column
+ensureOptional c@(BoxedColumn (Just _) _) = c
+ensureOptional (BoxedColumn Nothing col) =
+ BoxedColumn (Just (allValidBitmap (VB.length col))) col
+ensureOptional c@(UnboxedColumn (Just _) _) = c
+ensureOptional (UnboxedColumn Nothing col) =
+ UnboxedColumn (Just (allValidBitmap (VU.length col))) col
+ensureOptional c@(PackedText (Just _) _) = c
+ensureOptional (PackedText Nothing p) =
+ PackedText (Just (allValidBitmap (packedLength p))) p
+
+-- | Fills the end of a column, up to n, with null rows. Does nothing if column has length >= n.
+expandColumn :: Int -> Column -> Column
+expandColumn n c@(PackedText _ p)
+ | n <= packedLength p = c
+ | otherwise = expandColumn n (materializePacked c)
+expandColumn n column@(BoxedColumn bm col)
+ | n <= VG.length col = column
+ | otherwise =
+ let extra = n - VG.length col
+ newBm = case bm of
+ Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])
+ Just b ->
+ Just
+ (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))
+ newCol = col <> VB.replicate extra (errorWithoutStackTrace "expandColumn: null slot")
+ in BoxedColumn newBm newCol
+expandColumn n column@(UnboxedColumn bm col)
+ | n <= VG.length col = column
+ | otherwise =
+ let extra = n - VG.length col
+ newBm = case bm of
+ Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])
+ Just b ->
+ Just
+ (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))
+ newCol = runST $ do
+ mv <- VUM.new n
+ VU.imapM_ (VUM.unsafeWrite mv) col
+ VU.unsafeFreeze mv
+ in UnboxedColumn newBm newCol
+
+-- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.
+leftExpandColumn :: Int -> Column -> Column
+leftExpandColumn n c@(PackedText _ p)
+ | n <= packedLength p = c
+ | otherwise = leftExpandColumn n (materializePacked c)
+leftExpandColumn n column@(BoxedColumn bm col)
+ | n <= VG.length col = column
+ | otherwise =
+ let extra = n - VG.length col
+ origLen = VG.length col
+ newBm = case bm of
+ Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])
+ Just b ->
+ let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0
+ in Just (bitmapConcat extra nullPart origLen b)
+ newCol =
+ VB.replicate extra (errorWithoutStackTrace "leftExpandColumn: null slot") <> col
+ in BoxedColumn newBm newCol
+leftExpandColumn n column@(UnboxedColumn bm col)
+ | n <= VG.length col = column
+ | otherwise =
+ let extra = n - VG.length col
+ origLen = VG.length col
+ newBm = case bm of
+ Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])
+ Just b ->
+ let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0
+ in Just (bitmapConcat extra nullPart origLen b)
+ newCol = runST $ do
+ mv <- VUM.new n
+ VU.imapM_ (\i x -> VUM.unsafeWrite mv (extra + i) x) col
+ VU.unsafeFreeze mv
+ in UnboxedColumn newBm newCol
+
+{- | Concatenates two columns.
+Returns Nothing if the columns are of different types.
+-}
+concatColumns :: Column -> Column -> Either DataFrameException Column
+concatColumns left right = case (left, right) of
+ (PackedText _ _, _) -> concatColumns (materializePacked left) right
+ (_, PackedText _ _) -> concatColumns left (materializePacked right)
+ (BoxedColumn bmL l, BoxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of
+ Just Refl ->
+ let newBm = case (bmL, bmR) of
+ (Nothing, Nothing) -> Nothing
+ (Just bl, Nothing) ->
+ Just
+ (bitmapConcat (VB.length l) bl (VB.length r) (allValidBitmap (VB.length r)))
+ (Nothing, Just br) ->
+ Just
+ (bitmapConcat (VB.length l) (allValidBitmap (VB.length l)) (VB.length r) br)
+ (Just bl, Just br) -> Just (bitmapConcat (VB.length l) bl (VB.length r) br)
+ in pure (BoxedColumn newBm (l <> r))
+ Nothing -> Left (mismatchErr (typeOf r) (typeOf l))
+ (UnboxedColumn bmL l, UnboxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of
+ Just Refl ->
+ let newBm = case (bmL, bmR) of
+ (Nothing, Nothing) -> Nothing
+ (Just bl, Nothing) ->
+ Just
+ (bitmapConcat (VU.length l) bl (VU.length r) (allValidBitmap (VU.length r)))
+ (Nothing, Just br) ->
+ Just
+ (bitmapConcat (VU.length l) (allValidBitmap (VU.length l)) (VU.length r) br)
+ (Just bl, Just br) -> Just (bitmapConcat (VU.length l) bl (VU.length r) br)
+ in pure (UnboxedColumn newBm (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
+ }
+ )
+
+{- | Like 'concatColumns' but also combines columns of different types by wrapping
+values in 'Either' (e.g. @[1,2]@ and @["a","b"]@ become
+@[Left 1, Left 2, Right "a", Right "b"]@).
+-}
+
+{- | O(n) Concatenate a list of same-type columns in a single allocation.
+All columns must have the same constructor and element type (as they will
+within a single Parquet column). Calls 'error' on mismatch.
+-}
+concatManyColumns :: [Column] -> Column
+concatManyColumns [] = fromList ([] :: [Maybe Int])
+concatManyColumns [c] = c
+concatManyColumns all'
+ | any isPackedText all' =
+ concatManyColumns (map materializePacked all')
+concatManyColumns (c0 : cs) = case c0 of
+ BoxedColumn bm0 v0 ->
+ let getCol (BoxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of
+ Just Refl -> (bm, v)
+ Nothing -> error "concatManyColumns: BoxedColumn type mismatch"
+ getCol _ = error "concatManyColumns: column constructor mismatch"
+ rest = map getCol cs
+ allVecs = v0 : map snd rest
+ allBms = bm0 : map fst rest
+ newBm
+ | all isNothing allBms = Nothing
+ | otherwise =
+ let pairs = zip allVecs allBms
+ expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VB.length v)) mb) pairs
+ go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
+ concatBms [] = VU.empty
+ concatBms [(b, _v)] = b
+ concatBms ((b1, v1) : (b2, v2) : rest') =
+ let merged = go b1 (VB.length v1) b2 (VB.length v2)
+ in concatBms ((merged, v1 <> v2) : rest')
+ in Just $ concatBms (zip expandedBms allVecs)
+ in BoxedColumn newBm (VB.concat allVecs)
+ UnboxedColumn bm0 v0 ->
+ let getCol (UnboxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of
+ Just Refl -> (bm, v)
+ Nothing -> error "concatManyColumns: UnboxedColumn type mismatch"
+ getCol _ = error "concatManyColumns: column constructor mismatch"
+ rest = map getCol cs
+ allVecs = v0 : map snd rest
+ allBms = bm0 : map fst rest
+ newBm
+ | all isNothing allBms = Nothing
+ | otherwise =
+ let pairs = zip allVecs allBms
+ expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VU.length v)) mb) pairs
+ go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
+ concatBms [] = VU.empty
+ concatBms [(b, _)] = b
+ concatBms ((b1, v1) : (b2, v2) : rest') =
+ let merged = go b1 (VU.length v1) b2 (VU.length v2)
+ in concatBms ((merged, v1 <> v2) : rest')
+ in Just $ concatBms (zip expandedBms allVecs)
+ in UnboxedColumn newBm (VU.concat allVecs)
+ PackedText _ _ -> concatManyColumns (map materializePacked (c0 : cs))
+
+concatColumnsEither :: Column -> Column -> Column
+concatColumnsEither l@(PackedText _ _) r = concatColumnsEither (materializePacked l) r
+concatColumnsEither l r@(PackedText _ _) = concatColumnsEither l (materializePacked r)
+concatColumnsEither (BoxedColumn bmL left) (BoxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of
+ Nothing ->
+ BoxedColumn Nothing $ fmap Left left <> fmap Right right
+ Just Refl ->
+ let newBm = case (bmL, bmR) of
+ (Nothing, Nothing) -> Nothing
+ (Just bl, Nothing) ->
+ Just
+ ( bitmapConcat
+ (VB.length left)
+ bl
+ (VB.length right)
+ (allValidBitmap (VB.length right))
+ )
+ (Nothing, Just br) ->
+ Just
+ ( bitmapConcat
+ (VB.length left)
+ (allValidBitmap (VB.length left))
+ (VB.length right)
+ br
+ )
+ (Just bl, Just br) -> Just (bitmapConcat (VB.length left) bl (VB.length right) br)
+ in BoxedColumn newBm $ left <> right
+concatColumnsEither (UnboxedColumn bmL left) (UnboxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of
+ Nothing ->
+ BoxedColumn Nothing $
+ fmap Left (VG.convert left) <> fmap Right (VG.convert right)
+ Just Refl ->
+ let newBm = case (bmL, bmR) of
+ (Nothing, Nothing) -> Nothing
+ (Just bl, Nothing) ->
+ Just
+ ( bitmapConcat
+ (VU.length left)
+ bl
+ (VU.length right)
+ (allValidBitmap (VU.length right))
+ )
+ (Nothing, Just br) ->
+ Just
+ ( bitmapConcat
+ (VU.length left)
+ (allValidBitmap (VU.length left))
+ (VU.length right)
+ br
+ )
+ (Just bl, Just br) -> Just (bitmapConcat (VU.length left) bl (VU.length right) br)
+ in UnboxedColumn newBm $ left <> right
+concatColumnsEither (BoxedColumn _ left) (UnboxedColumn _ right) =
+ BoxedColumn Nothing $ fmap Left left <> fmap Right (VG.convert right)
+concatColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =
+ BoxedColumn Nothing $ fmap Left (VG.convert left) <> fmap Right right
+
+-- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.
+newMutableColumn :: Int -> Column -> IO MutableColumn
+newMutableColumn n (BoxedColumn _ (_ :: VB.Vector a)) =
+ MBoxedColumn <$> (VBM.new n :: IO (VBM.IOVector a))
+newMutableColumn n (UnboxedColumn _ (_ :: VU.Vector a)) =
+ MUnboxedColumn <$> (VUM.new n :: IO (VUM.IOVector a))
+newMutableColumn n c@(PackedText _ _) = newMutableColumn n (materializePacked c)
+
+-- | Copy a column chunk into a mutable column starting at offset @off@.
+copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()
+copyIntoMutableColumn (MBoxedColumn (mv :: VBM.IOVector b)) off (BoxedColumn _ (v :: VB.Vector a)) =
+ case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> VG.imapM_ (\i x -> VBM.unsafeWrite mv (off + i) x) v
+ Nothing -> error "copyIntoMutableColumn: Boxed type mismatch"
+copyIntoMutableColumn (MUnboxedColumn (mv :: VUM.IOVector b)) off (UnboxedColumn _ (v :: VU.Vector a)) =
+ case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> VG.imapM_ (\i x -> VUM.unsafeWrite mv (off + i) x) v
+ Nothing -> error "copyIntoMutableColumn: Unboxed type mismatch"
+copyIntoMutableColumn mc off c@(PackedText _ _) =
+ copyIntoMutableColumn mc off (materializePacked c)
+copyIntoMutableColumn _ _ _ =
+ error "copyIntoMutableColumn: constructor mismatch"
+
+-- | Freeze a mutable column into an immutable column.
+freezeMutableColumn :: MutableColumn -> IO Column
+freezeMutableColumn (MBoxedColumn mv) = BoxedColumn Nothing <$> VB.unsafeFreeze mv
+freezeMutableColumn (MUnboxedColumn mv) = UnboxedColumn Nothing <$> VU.unsafeFreeze mv
+
+{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.
+
+__Examples:__
+
+@
+> column = fromList [(1 :: Int), 2, 3, 4]
+> toList @Int column
+[1,2,3,4]
+> toList @Double column
+exception: ...
+@
+-}
+toList :: forall a. (Columnable a) => Column -> [a]
+toList xs = case toVector @a xs of
+ Left err -> throw err
+ Right val -> VB.toList val
+
+{- | Type-safe conversion of a column to a vector of element type @a@ (specify via
+type application); 'Left' 'TypeMismatchException' when the column's type differs.
+
+>>> toVector @Int @VU.Vector column
+Right (unboxed vector of Ints)
+
+>>> toVector @Text @VB.Vector column
+Right (boxed vector of Text)
+-}
+toVector ::
+ forall a v.
+ (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)
+toVector col = case col of
+ PackedText _ _ -> toVector (materializePacked col)
+ BoxedColumn bm (inner :: VB.Vector c) ->
+ -- Check if user wants Maybe c (nullable) or c directly
+ case testEquality (typeRep @a) (typeRep @c) of
+ Just Refl -> Right $ VG.convert inner
+ Nothing ->
+ -- Try: a = Maybe c
+ case testEquality (typeRep @a) (typeRep @(Maybe c)) of
+ Just Refl ->
+ -- Use VB.generate to avoid fusion forcing null slots
+ let !n = VB.length inner
+ maybeVec = case bm of
+ Nothing -> VB.generate n (Just . VB.unsafeIndex inner)
+ Just bitmap -> VB.generate n $ \i ->
+ if bitmapTestBit bitmap i then Just (VB.unsafeIndex inner i) else Nothing
+ in Right $ VG.convert maybeVec
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @a)
+ , expectedType = Right (typeRep @c)
+ , callingFunctionName = Just "toVector"
+ , errorColumnName = Nothing
+ }
+ )
+ UnboxedColumn bm (inner :: VU.Vector c) ->
+ case testEquality (typeRep @a) (typeRep @c) of
+ Just Refl -> Right $ VG.convert inner
+ Nothing ->
+ case testEquality (typeRep @a) (typeRep @(Maybe c)) of
+ Just Refl ->
+ let maybeVec = case bm of
+ Nothing -> VB.generate (VU.length inner) (Just . VU.unsafeIndex inner)
+ Just bitmap -> VB.generate (VU.length inner) $ \i ->
+ if bitmapTestBit bitmap i then Just (VU.unsafeIndex inner i) else Nothing
+ in Right $ VG.convert maybeVec
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @a)
+ , expectedType = Right (typeRep @c)
+ , callingFunctionName = Just "toVector"
+ , errorColumnName = Nothing
+ }
+ )
+
+-- Some common types we will use for numerical computing.
+
+{- | Convert a column to an unboxed 'Double' vector, coercing numeric types
+('realToFrac' for floats, 'fromIntegral' for integrals; nulls become @NaN@).
+'Left' 'TypeMismatchException' when the column is not numeric.
+-}
+toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)
+toDoubleVector column =
+ case column of
+ PackedText _ _ -> toDoubleVector (materializePacked column)
+ UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> case bm of
+ Nothing -> Right f
+ Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f
+ Nothing -> case sFloating @a of
+ STrue ->
+ Right
+ ( VU.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> realToFrac x
+ )
+ f
+ )
+ SFalse -> case sIntegral @a of
+ STrue ->
+ Right
+ ( VU.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> fromIntegral x
+ )
+ f
+ )
+ SFalse ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Double)
+ , expectedType = Right (typeRep @a)
+ , callingFunctionName = Just "toDoubleVector"
+ , errorColumnName = Nothing
+ }
+ )
+ BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
+ Just Refl ->
+ Right
+ ( VB.convert $
+ VB.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> fromIntegral x
+ )
+ f
+ )
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Double)
+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
+ , callingFunctionName = Just "toDoubleVector"
+ , errorColumnName = Nothing
+ }
+ )
+
+{- | Convert a column to an unboxed 'Float' vector, coercing numeric types (nulls
+become @NaN@); 'Left' 'TypeMismatchException' when not numeric. Converting from
+'Double' may lose precision.
+-}
+toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)
+toFloatVector column =
+ case column of
+ PackedText _ _ -> toFloatVector (materializePacked column)
+ UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Float) of
+ Just Refl -> case bm of
+ Nothing -> Right f
+ Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f
+ Nothing -> case sFloating @a of
+ STrue ->
+ Right
+ ( VU.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> realToFrac x
+ )
+ f
+ )
+ SFalse -> case sIntegral @a of
+ STrue ->
+ Right
+ ( VU.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> fromIntegral x
+ )
+ f
+ )
+ SFalse ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Float)
+ , expectedType = Right (typeRep @a)
+ , callingFunctionName = Just "toFloatVector"
+ , errorColumnName = Nothing
+ }
+ )
+ BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
+ Just Refl ->
+ Right
+ ( VB.convert $
+ VB.imap
+ ( \i x -> case bm of
+ Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
+ _ -> fromIntegral x
+ )
+ f
+ )
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Float)
+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
+ , callingFunctionName = Just "toFloatVector"
+ , errorColumnName = Nothing
+ }
+ )
+
+{- | Convert a column to an unboxed 'Int' vector, coercing numeric types
+(floats are 'round'ed via banker's rounding); 'Left' 'TypeMismatchException'
+when the column is not numeric. Does not support nullable columns.
+-}
+toIntVector :: Column -> Either DataFrameException (VU.Vector Int)
+toIntVector column =
+ case column of
+ PackedText _ _ -> toIntVector (materializePacked column)
+ UnboxedColumn _ (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> Right f
+ Nothing -> case sFloating @a of
+ STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)
+ SFalse -> case sIntegral @a of
+ STrue -> Right (VU.map fromIntegral f)
+ SFalse ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Int)
+ , expectedType = Right (typeRep @a)
+ , callingFunctionName = Just "toIntVector"
+ , errorColumnName = Nothing
+ }
+ )
+ BoxedColumn _ (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
+ Just Refl -> Right (VB.convert $ VB.map fromIntegral f)
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @Int)
+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
+ , callingFunctionName = Just "toIntVector"
+ , errorColumnName = Nothing
+ }
+ )
+
+toUnboxedVector ::
+ forall a.
+ (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)
+toUnboxedVector column =
+ case column of
+ UnboxedColumn _ (f :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Right f
+ Nothing ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @a)
+ , expectedType = Right (typeRep @b)
+ , callingFunctionName = Just "toUnboxedVector"
+ , errorColumnName = Nothing
+ }
+ )
+ _ ->
+ Left $
+ TypeMismatchException
+ ( MkTypeErrorContext
+ { userType = Right (typeRep @a)
+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
+ , callingFunctionName = Just "toUnboxedVector"
+ , errorColumnName = Nothing
+ }
+ )
+{-# INLINE toUnboxedVector #-}
+
+-- Shared finaliser for the two parseUnboxedColumn* helpers. Freezes
+-- the mutable data vector, and only materialises the bitmap when the
+-- column actually had nulls.
+{-# INLINE finalizeParseResult #-}
+finalizeParseResult ::
+ (VU.Unbox a) =>
+ VUM.STVector s a ->
+ VUM.STVector s Word8 ->
+ Bool ->
+ ST s (Maybe (Maybe Bitmap, VU.Vector a))
+finalizeParseResult values vmask anyNull
+ | anyNull = do
+ vs <- VU.unsafeFreeze values
+ vm <- VU.unsafeFreeze vmask
+ return (Just (Just (buildBitmapFromValid vm), vs))
+ | otherwise = do
+ vs <- VU.unsafeFreeze values
+ return (Just (Nothing, vs))
diff --git a/src-internal/DataFrame/Internal/ColumnBuilder.hs b/src-internal/DataFrame/Internal/ColumnBuilder.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/ColumnBuilder.hs
@@ -0,0 +1,299 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+{- | Mutable, growable column builders for high-throughput ingest. No
+per-append @IORef@ traffic: hot counters live in an unboxed vector, payloads
+double on demand, and validity is only materialized once a null is seen.
+-}
+module DataFrame.Internal.ColumnBuilder (
+ ColumnBuilder (..),
+ NumBuilder,
+ IntBuilder,
+ DoubleBuilder,
+ TextBuilder,
+ TextChunk (..),
+ newIntBuilder,
+ newDoubleBuilder,
+ newNumBuilder,
+ newTextBuilder,
+ appendInt,
+ appendDouble,
+ appendNum,
+ appendText,
+ appendTextSlice,
+ appendTextSliceFromPtr,
+ freezeTextChunk,
+ mergeColumns,
+ mergeTextChunks,
+) where
+
+import qualified Data.Text as T
+import qualified Data.Text.Array as A
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Monad (when)
+import Control.Monad.ST (ST)
+import Data.Bits (shiftR)
+import Data.STRef
+import Data.Text.Internal (Text (..))
+import Data.Word (Word8)
+import DataFrame.Internal.Column hiding (mergeColumns)
+import DataFrame.Internal.ColumnMerge (
+ TextChunk (..),
+ mergeColumns,
+ mergeTextChunks,
+ packValidity,
+ )
+import Foreign.Ptr (Ptr)
+
+{- | Operations shared by all column builders. A builder must not be used
+again after 'freezeBuilder' (its storage is frozen in place, not copied).
+-}
+class ColumnBuilder b where
+ -- | Append a null row (sentinel payload + invalid bit).
+ appendNull :: b s -> ST s ()
+
+ -- | Rows appended so far.
+ builderLength :: b s -> ST s Int
+
+ -- | Freeze into a fully-forced 'Column'; bitmap only when a null was seen.
+ freezeBuilder :: b s -> ST s Column
+
+-- Counter slots shared by the builders: rows, any-null flag, text bytes used.
+cRows, cAnyNull, cBytes :: Int
+cRows = 0
+cAnyNull = 1
+cBytes = 2
+
+{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'
+is the sentinel written into null slots (protected by the bitmap).
+-}
+data NumBuilder a s = NumBuilder
+ { nbNull :: !a
+ , nbCounters :: !(VUM.MVector s Int)
+ , nbArrays :: !(STRef s (NumArrays a s))
+ }
+
+data NumArrays a s = NumArrays
+ { naData :: !(VUM.MVector s a)
+ , naValid :: !(VUM.MVector s Word8)
+ }
+
+type IntBuilder = NumBuilder Int
+
+type DoubleBuilder = NumBuilder Double
+
+-- | New numeric builder with a row-capacity hint and a null sentinel.
+newNumBuilder :: (VU.Unbox a) => a -> Int -> ST s (NumBuilder a s)
+newNumBuilder nullValue hint = do
+ let cap = max 16 hint
+ counters <- VUM.replicate 2 0
+ dat <- VUM.unsafeNew cap
+ val <- VUM.unsafeNew cap
+ NumBuilder nullValue counters <$> newSTRef (NumArrays dat val)
+
+newIntBuilder :: Int -> ST s (IntBuilder s)
+newIntBuilder = newNumBuilder 0
+
+newDoubleBuilder :: Int -> ST s (DoubleBuilder s)
+newDoubleBuilder = newNumBuilder 0
+
+appendNum :: (VU.Unbox a) => NumBuilder a s -> a -> ST s ()
+appendNum b !x = do
+ n <- VUM.unsafeRead (nbCounters b) cRows
+ anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
+ NumArrays dat val <- reserveNum b n
+ VUM.unsafeWrite dat n x
+ when (anyNull /= 0) $ VUM.unsafeWrite val n 1
+ VUM.unsafeWrite (nbCounters b) cRows (n + 1)
+{-# INLINE appendNum #-}
+
+appendInt :: IntBuilder s -> Int -> ST s ()
+appendInt = appendNum
+{-# INLINE appendInt #-}
+
+appendDouble :: DoubleBuilder s -> Double -> ST s ()
+appendDouble = appendNum
+{-# INLINE appendDouble #-}
+
+-- Fetch the arrays, growing (doubling) first if row @n@ would not fit.
+reserveNum :: (VU.Unbox a) => NumBuilder a s -> Int -> ST s (NumArrays a s)
+reserveNum b n = do
+ arrs <- readSTRef (nbArrays b)
+ if n < VUM.length (naData arrs) then pure arrs else growNum b arrs
+{-# INLINE reserveNum #-}
+
+growNum ::
+ (VU.Unbox a) => NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
+growNum b (NumArrays dat val) = do
+ let cap = VUM.length dat
+ dat' <- VUM.unsafeGrow dat cap
+ val' <- VUM.unsafeGrow val cap
+ let arrs = NumArrays dat' val'
+ writeSTRef (nbArrays b) arrs
+ pure arrs
+
+instance (Columnable a, VU.Unbox a) => ColumnBuilder (NumBuilder a) where
+ appendNull b = do
+ n <- VUM.unsafeRead (nbCounters b) cRows
+ anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
+ NumArrays dat val <- reserveNum b n
+ VUM.unsafeWrite dat n (nbNull b)
+ when (anyNull == 0) $ do
+ VUM.set (VUM.slice 0 n val) 1
+ VUM.unsafeWrite (nbCounters b) cAnyNull 1
+ VUM.unsafeWrite val n 0
+ VUM.unsafeWrite (nbCounters b) cRows (n + 1)
+ {-# INLINE appendNull #-}
+
+ builderLength b = VUM.unsafeRead (nbCounters b) cRows
+
+ freezeBuilder b = do
+ n <- VUM.unsafeRead (nbCounters b) cRows
+ anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
+ NumArrays dat val <- readSTRef (nbArrays b)
+ !vs <- freezeTrimmed n dat
+ if anyNull /= 0
+ then do
+ !bm <- packValidity n val
+ pure $! UnboxedColumn (Just bm) vs
+ else pure $! UnboxedColumn Nothing vs
+
+-- Zero-copy freeze; copies to exact size when slack exceeds a quarter of n.
+freezeTrimmed :: (VU.Unbox a) => Int -> VUM.MVector s a -> ST s (VU.Vector a)
+freezeTrimmed n mv
+ | VUM.length mv - n <= n `shiftR` 2 = VU.unsafeFreeze (VUM.slice 0 n mv)
+ | otherwise = VU.freeze (VUM.slice 0 n mv)
+
+{- | Builder for 'Text' columns. All field bytes go into one exponentially
+grown byte array; rows are recorded as offsets, so an append is a memcpy
+and freezing slices 'Text' values off the shared array without copying.
+-}
+data TextBuilder s = TextBuilder
+ { tbCounters :: !(VUM.MVector s Int)
+ , tbArrays :: !(STRef s (TextArrays s))
+ }
+
+data TextArrays s = TextArrays
+ { taBytes :: !(A.MArray s)
+ , taByteCap :: !Int
+ , taOffsets :: !(VUM.MVector s Int)
+ -- ^ Row @i@ spans bytes @[offsets!i, offsets!(i+1))@.
+ , taValid :: !(VUM.MVector s Word8)
+ }
+
+-- | New text builder with row-count and total-byte capacity hints.
+newTextBuilder :: Int -> Int -> ST s (TextBuilder s)
+newTextBuilder rowHint byteHint = do
+ let rcap = max 16 rowHint
+ bcap = max 64 byteHint
+ counters <- VUM.replicate 3 0
+ bytes <- A.new bcap
+ offsets <- VUM.unsafeNew (rcap + 1)
+ VUM.unsafeWrite offsets 0 0
+ val <- VUM.unsafeNew rcap
+ TextBuilder counters <$> newSTRef (TextArrays bytes bcap offsets val)
+
+-- | Append @len@ raw bytes at @off@ in @src@ as one field (one memcpy).
+appendTextSlice :: TextBuilder s -> A.Array -> Int -> Int -> ST s ()
+appendTextSlice b src off len = do
+ (n, pos, arrs) <- reserveText b len
+ A.copyI len (taBytes arrs) pos src off
+ finishTextAppend b arrs n (pos + len)
+{-# INLINE appendTextSlice #-}
+
+-- | 'appendTextSlice' from foreign memory (e.g. an mmapped file buffer).
+appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()
+appendTextSliceFromPtr b ptr len = do
+ (n, pos, arrs) <- reserveText b len
+ A.copyFromPointer (taBytes arrs) pos ptr len
+ finishTextAppend b arrs n (pos + len)
+{-# INLINE appendTextSliceFromPtr #-}
+
+-- | Append an already-decoded 'Text' (its bytes are UTF-8 already).
+appendText :: TextBuilder s -> T.Text -> ST s ()
+appendText b (Text src off len) = appendTextSlice b src off len
+{-# INLINE appendText #-}
+
+finishTextAppend :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
+finishTextAppend b arrs n endPos = do
+ anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
+ when (anyNull /= 0) $ VUM.unsafeWrite (taValid arrs) n 1
+ VUM.unsafeWrite (taOffsets arrs) (n + 1) endPos
+ VUM.unsafeWrite (tbCounters b) cRows (n + 1)
+ VUM.unsafeWrite (tbCounters b) cBytes endPos
+{-# INLINE finishTextAppend #-}
+
+reserveText :: TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
+reserveText b extra = do
+ n <- VUM.unsafeRead (tbCounters b) cRows
+ pos <- VUM.unsafeRead (tbCounters b) cBytes
+ arrs <- readSTRef (tbArrays b)
+ arrs' <-
+ if n < VUM.length (taValid arrs) && pos + extra <= taByteCap arrs
+ then pure arrs
+ else growText b arrs (n + 1) (pos + extra)
+ pure (n, pos, arrs')
+{-# INLINE reserveText #-}
+
+growText :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
+growText b (TextArrays bytes bcap offsets val) needRows needBytes = do
+ let rcap = VUM.length val
+ (offsets', val') <-
+ if needRows > rcap
+ then do
+ let rcap' = max (2 * rcap) needRows
+ o <- VUM.unsafeGrow offsets (rcap' - rcap)
+ v <- VUM.unsafeGrow val (rcap' - rcap)
+ pure (o, v)
+ else pure (offsets, val)
+ (bytes', bcap') <-
+ if needBytes > bcap
+ then do
+ let cap' = max (2 * bcap) needBytes
+ bs <- A.resizeM bytes cap'
+ pure (bs, cap')
+ else pure (bytes, bcap)
+ let arrs = TextArrays bytes' bcap' offsets' val'
+ writeSTRef (tbArrays b) arrs
+ pure arrs
+
+{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for byte-level merging
+('mergeTextChunks'): no 'T.Text' values are created until chunks merge.
+-}
+freezeTextChunk :: TextBuilder s -> ST s TextChunk
+freezeTextChunk b = do
+ n <- VUM.unsafeRead (tbCounters b) cRows
+ anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
+ used <- VUM.unsafeRead (tbCounters b) cBytes
+ TextArrays bytes bcap offsets val <- readSTRef (tbArrays b)
+ when (used < bcap) (A.shrinkM bytes used)
+ arr <- A.unsafeFreeze bytes
+ offs <- VU.unsafeFreeze (VUM.slice 0 (n + 1) offsets)
+ bm <-
+ if anyNull /= 0
+ then Just <$> packValidity n val
+ else pure Nothing
+ pure (TextChunk arr used offs bm)
+
+instance ColumnBuilder TextBuilder where
+ appendNull b = do
+ (n, pos, arrs) <- reserveText b 0
+ anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
+ when (anyNull == 0) $ do
+ VUM.set (VUM.slice 0 n (taValid arrs)) 1
+ VUM.unsafeWrite (tbCounters b) cAnyNull 1
+ VUM.unsafeWrite (taValid arrs) n 0
+ VUM.unsafeWrite (taOffsets arrs) (n + 1) pos
+ VUM.unsafeWrite (tbCounters b) cRows (n + 1)
+ VUM.unsafeWrite (tbCounters b) cBytes pos
+ {-# INLINE appendNull #-}
+
+ builderLength b = VUM.unsafeRead (tbCounters b) cRows
+
+ freezeBuilder b = do
+ chunk <- freezeTextChunk b
+ pure $! mergeTextChunks [chunk]
diff --git a/src-internal/DataFrame/Internal/ColumnMerge.hs b/src-internal/DataFrame/Internal/ColumnMerge.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/ColumnMerge.hs
@@ -0,0 +1,176 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks). Text
+columns merge at the byte level via 'TextChunk' \/ 'mergeTextChunks', so no
+per-chunk 'Data.Text.Text' values are ever materialized.
+-}
+module DataFrame.Internal.ColumnMerge (
+ TextChunk (..),
+ mergeColumns,
+ mergeTextChunks,
+ packedFromTextChunk,
+ packValidity,
+ spliceBitmaps,
+ tcRows,
+) where
+
+import qualified Data.Text.Array as A
+import qualified Data.Vector as VB
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Monad (foldM_, forM_, when)
+import Control.Monad.ST (ST, runST)
+import Data.Bits (shiftL, shiftR, (.&.), (.|.))
+import Data.Maybe (fromMaybe, isNothing)
+import Data.Type.Equality (testEquality, (:~:) (Refl))
+import Data.Word (Word8)
+import DataFrame.Internal.Column (
+ Bitmap,
+ Column (..),
+ Columnable,
+ allValidBitmap,
+ materializePacked,
+ )
+import DataFrame.Internal.PackedText (mkPackedContiguous)
+import Type.Reflection (typeRep)
+
+{- | A frozen text-builder chunk: raw UTF-8 bytes plus row offsets (row @i@
+spans bytes @[offsets!i, offsets!(i+1))@) and an optional validity bitmap.
+'Data.Text.Text' values are only created when chunks merge into a 'Column'.
+-}
+data TextChunk = TextChunk
+ { tcBytes :: !A.Array
+ , tcUsed :: !Int
+ , tcOffsets :: !(VU.Vector Int)
+ , tcBitmap :: !(Maybe Bitmap)
+ }
+
+tcRows :: TextChunk -> Int
+tcRows c = VU.length (tcOffsets c) - 1
+
+{- | Freeze a builder chunk directly into a packed-text column: no
+'Data.Text.Text' materialization, no UTF-8 validation pass (deferred to decode).
+Not yet called by any reader.
+-}
+packedFromTextChunk :: TextChunk -> Column
+packedFromTextChunk (TextChunk arr _used offs bm) =
+ PackedText bm (mkPackedContiguous arr offs)
+
+{- | Merge text chunks into one packed-text 'Column': one byte-array copy per
+chunk, one offset rebase, then wrap the shared buffer + offsets as 'PackedText'
+(no per-row header, decode deferred).
+-}
+mergeTextChunks :: [TextChunk] -> Column
+mergeTextChunks [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunks: empty list"
+mergeTextChunks [c] = packedFromTextChunk c
+mergeTextChunks cs = runST $ do
+ let totalBytes = sum (map tcUsed cs)
+ totalRows = sum (map tcRows cs)
+ arr <- A.new (max 1 totalBytes)
+ offs <- VUM.unsafeNew (totalRows + 1)
+ VUM.unsafeWrite offs 0 0
+ let splice !byteBase !rowBase c = do
+ let n = tcRows c
+ co = tcOffsets c
+ A.copyI (tcUsed c) arr byteBase (tcBytes c) 0
+ forM_ [1 .. n] $ \i ->
+ VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex co i)
+ pure (byteBase + tcUsed c, rowBase + n)
+ foldM_ (\(b, r) c -> splice b r c) (0, 0) cs
+ farr <- A.unsafeFreeze arr
+ foffs <- VU.unsafeFreeze offs
+ let !bm = spliceBitmaps [(tcBitmap c, tcRows c) | c <- cs]
+ pure (PackedText bm (mkPackedContiguous farr foffs))
+
+{- | Merge per-chunk columns into one column: one allocation + memcpy per
+payload, with bitmaps spliced across non-byte-aligned chunk boundaries.
+All chunks must have the same element type.
+-}
+mergeColumns :: [Column] -> Column
+mergeColumns [] = error "DataFrame.Internal.ColumnBuilder.mergeColumns: empty list"
+mergeColumns [c] = c
+mergeColumns cols@(c0 : _) = case c0 of
+ PackedText _ _ -> mergeColumns (map materializePacked cols)
+ UnboxedColumn _ (_ :: VU.Vector a) ->
+ let parts = map (unboxedPart @a) cols
+ !merged = VU.concat (map snd parts)
+ !bm = spliceBitmaps [(mb, VU.length v) | (mb, v) <- parts]
+ in UnboxedColumn bm merged
+ BoxedColumn _ (_ :: VB.Vector a) ->
+ let parts = map (boxedPart @a) cols
+ !merged = VB.concat (map snd parts)
+ !bm = spliceBitmaps [(mb, VB.length v) | (mb, v) <- parts]
+ in BoxedColumn bm merged
+
+unboxedPart ::
+ forall a. (Columnable a, VU.Unbox a) => Column -> (Maybe Bitmap, VU.Vector a)
+unboxedPart (UnboxedColumn mb (v :: VU.Vector b)) =
+ case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> (mb, v)
+ Nothing -> mergeMismatch
+unboxedPart _ = mergeMismatch
+
+boxedPart ::
+ forall a. (Columnable a) => Column -> (Maybe Bitmap, VB.Vector a)
+boxedPart (BoxedColumn mb (v :: VB.Vector b)) =
+ case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> (mb, v)
+ Nothing -> mergeMismatch
+boxedPart _ = mergeMismatch
+
+mergeMismatch :: a
+mergeMismatch =
+ error "DataFrame.Internal.ColumnBuilder.mergeColumns: chunk column types differ"
+
+{- | Splice chunk bitmaps end to end at the bit level. 'Nothing' if no chunk
+carries a bitmap; chunks without one count as all-valid otherwise.
+-}
+spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap
+spliceBitmaps parts
+ | all (isNothing . fst) parts = Nothing
+ | otherwise = Just $ VU.create $ do
+ let total = sum (map snd parts)
+ outBytes = (total + 7) `shiftR` 3
+ mv <- VUM.replicate outBytes 0
+ let orInto i w =
+ when (i < outBytes && w /= 0) $ do
+ old <- VUM.unsafeRead mv i
+ VUM.unsafeWrite mv i (old .|. w)
+ splice !bitPos (mb, len) = do
+ let bm = fromMaybe (allValidBitmap len) mb
+ sh = bitPos .&. 7
+ byte0 = bitPos `shiftR` 3
+ lastIdx = ((len + 7) `shiftR` 3) - 1
+ tailBits = len .&. 7
+ lastMask =
+ if tailBits == 0 then 0xFF else (1 `shiftL` tailBits) - 1
+ forM_ [0 .. lastIdx] $ \k -> do
+ let raw = VU.unsafeIndex bm k
+ masked = if k == lastIdx then raw .&. lastMask else raw
+ w = fromIntegral masked :: Word
+ orInto (byte0 + k) (fromIntegral (w `shiftL` sh))
+ when (sh /= 0) $
+ orInto (byte0 + k + 1) (fromIntegral (w `shiftR` (8 - sh)))
+ pure (bitPos + len)
+ foldM_ splice 0 parts
+ pure mv
+
+-- | Pack a 0\/1 byte-per-row validity prefix into a bit-packed 'Bitmap'.
+packValidity :: Int -> VUM.MVector s Word8 -> ST s Bitmap
+packValidity n val = do
+ bytes <- VU.unsafeFreeze (VUM.slice 0 n val)
+ let assemble b =
+ let base = b `shiftL` 3
+ m = min 8 (n - base)
+ go !acc !k
+ | k >= m = acc
+ | VU.unsafeIndex bytes (base + k) /= 0 =
+ go (acc .|. (1 `shiftL` k)) (k + 1)
+ | otherwise = go acc (k + 1)
+ in go (0 :: Word8) 0
+ pure $! VU.generate ((n + 7) `shiftR` 3) assemble
diff --git a/src-internal/DataFrame/Internal/DataFrame.hs b/src-internal/DataFrame/Internal/DataFrame.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/DataFrame.hs
@@ -0,0 +1,362 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+module DataFrame.Internal.DataFrame where
+
+import qualified Data.Map as M
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+
+import Control.Exception (throw)
+import Data.Function (on)
+import Data.List (sortBy, (\\))
+import Data.Maybe (fromMaybe)
+import Data.Type.Equality (
+ TestEquality (testEquality),
+ type (:~:) (Refl),
+ type (:~~:) (HRefl),
+ )
+import DataFrame.Display.Terminal.PrettyPrint
+import DataFrame.Errors
+import DataFrame.Internal.Column
+import DataFrame.Internal.Expression
+import DataFrame.Internal.PackedText (packedIndexText)
+import Text.Printf
+import Type.Reflection (Typeable, eqTypeRep, typeRep, pattern App)
+import Prelude hiding (null)
+
+data DataFrame = DataFrame
+ { columns :: V.Vector Column
+ -- ^ Column-oriented storage: the frame is a vector of columns.
+ , columnIndices :: M.Map T.Text Int
+ -- ^ Keeps the column names in the order they were inserted in.
+ , dataframeDimensions :: (Int, Int)
+ -- ^ (rows, columns)
+ , derivingExpressions :: M.Map T.Text UExpr
+ }
+
+{- | Force evaluation of all columns in a DataFrame. Replacement for the removed
+@instance NFData DataFrame@; used by the IO and lazy-executor strict paths.
+-}
+forceDataFrame :: DataFrame -> DataFrame
+forceDataFrame df@(DataFrame cols idx dims _exprs) =
+ V.foldl' (\() c -> forceColumn c) () cols `seq` idx `seq` dims `seq` df
+
+{- | A record that contains information about how and what
+rows are grouped in the dataframe. This can only be used with
+`aggregate`.
+-}
+data GroupedDataFrame = Grouped
+ { fullDataframe :: DataFrame
+ , groupedColumns :: [T.Text]
+ , valueIndices :: VU.Vector Int
+ , offsets :: VU.Vector Int
+ , rowToGroup :: VU.Vector Int
+ {- ^ rowToGroup[i] = group index for row i. Length n (one per row).
+ Built once in 'groupBy'; reused by every aggregation.
+ -}
+ }
+
+instance Show GroupedDataFrame where
+ show (Grouped df cols _indices _os _rtg) =
+ printf
+ "{ keyColumns: %s groupedColumns: %s }"
+ (show cols)
+ (show (M.keys (columnIndices df) \\ cols))
+
+instance Eq GroupedDataFrame where
+ (==) (Grouped df cols _indices _os _rtg) (Grouped df' cols' _indices' _os' _rtg') = (df == df') && (cols == cols')
+
+instance Eq DataFrame where
+ (==) :: DataFrame -> DataFrame -> Bool
+ a == b =
+ M.keys (columnIndices a) == M.keys (columnIndices b)
+ && foldr
+ ( \(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))
+ )
+ True
+ (M.toList $ columnIndices a)
+
+instance Show DataFrame where
+ show :: DataFrame -> String
+ show d =
+ let (r, _) = dataframeDimensions d
+ cfg = defaultTruncateConfig
+ shown = if maxRows cfg > 0 then min (maxRows cfg) r else r
+ body = asTextWith Plain (Just cfg) d
+ footer
+ | shown < r =
+ "\nShowing "
+ <> T.pack (show shown)
+ <> " rows out of "
+ <> T.pack (show r)
+ | otherwise = T.empty
+ in T.unpack (body <> footer)
+
+{- | Configures how a 'DataFrame' is rendered as text: 'maxRows' caps rendered
+rows, 'maxColumns' collapses middle columns past the limit into an ellipsis, and
+'maxCellWidth' truncates long cells. A non-positive field means \"no limit\".
+-}
+data TruncateConfig = TruncateConfig
+ { maxRows :: Int
+ , maxColumns :: Int
+ , maxCellWidth :: Int
+ }
+ deriving (Show, Eq)
+
+-- | Sensible defaults for GHCi: 20 rows, 10 columns, 30 characters per cell.
+defaultTruncateConfig :: TruncateConfig
+defaultTruncateConfig =
+ TruncateConfig{maxRows = 20, maxColumns = 10, maxCellWidth = 30}
+
+-- | Ellipsis character used to mark elided columns and clipped cells.
+ellipsisText :: T.Text
+ellipsisText = "\x2026"
+
+-- | For showing the dataframe as markdown in notebooks.
+toMarkdown :: DataFrame -> T.Text
+toMarkdown = asText Markdown
+
+-- | For showing the dataframe as a string markdown in notebooks.
+toMarkdown' :: DataFrame -> String
+toMarkdown' = T.unpack . toMarkdown
+
+asText :: RenderFormat -> DataFrame -> T.Text
+asText fmt = asTextWith fmt Nothing
+
+asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> T.Text
+asTextWith fmt mTrunc d =
+ let allHeaders =
+ map fst (sortBy (compare `on` snd) (M.toList (columnIndices d)))
+ nCols = length allHeaders
+ (totalRows, _) = dataframeDimensions d
+
+ rowCap = case mTrunc of
+ Just cfg | maxRows cfg > 0 -> min totalRows (maxRows cfg)
+ _ -> totalRows
+
+ (visibleHeaders, ellipsisAt) = pickColumns mTrunc nCols allHeaders
+
+ lookupCol name =
+ fmap
+ (takeColumn rowCap)
+ ((V.!?) (columns d) ((M.!) (columnIndices d) name))
+ survivingCols = map lookupCol visibleHeaders
+ survivingTypes = map (maybe "" getType) survivingCols
+ survivingData = map get survivingCols
+
+ clipCell = case mTrunc of
+ Just cfg | maxCellWidth cfg > 0 -> truncateCell (maxCellWidth cfg)
+ _ -> id
+
+ (finalHeaders, finalTypes, finalCols) = case ellipsisAt of
+ Nothing -> (visibleHeaders, survivingTypes, survivingData)
+ Just i ->
+ let ellipsisCol = V.replicate rowCap ellipsisText
+ in ( insertAt i ellipsisText visibleHeaders
+ , insertAt i ellipsisText survivingTypes
+ , insertAt i ellipsisCol survivingData
+ )
+
+ getType :: Column -> T.Text
+ showMaybeType :: forall a. (Typeable a) => String
+ showMaybeType =
+ let s = show (typeRep @a)
+ in "Maybe " <> if ' ' `elem` s then "(" <> s <> ")" else s
+ getType (BoxedColumn Nothing (_ :: V.Vector a)) = T.pack $ show (typeRep @a)
+ getType (BoxedColumn (Just _) (_ :: V.Vector a)) = T.pack $ showMaybeType @a
+ getType (UnboxedColumn Nothing (_ :: VU.Vector a)) = T.pack $ show (typeRep @a)
+ getType (UnboxedColumn (Just _) (_ :: VU.Vector a)) = T.pack $ showMaybeType @a
+ getType (PackedText Nothing _) = T.pack $ show (typeRep @T.Text)
+ getType (PackedText (Just _) _) = T.pack $ showMaybeType @T.Text
+
+ get :: Maybe Column -> V.Vector T.Text
+ get (Just (BoxedColumn (Just bm) (column :: V.Vector a))) =
+ V.generate (V.length column) $ \i ->
+ if bitmapTestBit bm i
+ then T.pack (show (Just (V.unsafeIndex column i)))
+ else "Nothing"
+ get (Just (BoxedColumn Nothing (column :: V.Vector a))) =
+ case testEquality (typeRep @a) (typeRep @T.Text) of
+ Just Refl -> column
+ Nothing -> case testEquality (typeRep @a) (typeRep @String) of
+ Just Refl -> V.map T.pack column
+ Nothing -> V.map (T.pack . show) column
+ get (Just (UnboxedColumn (Just bm) column)) =
+ V.generate (VU.length column) $ \i ->
+ if bitmapTestBit bm i
+ then T.pack (show (Just (VU.unsafeIndex column i)))
+ else "Nothing"
+ get (Just (UnboxedColumn Nothing column)) =
+ V.generate (VU.length column) (T.pack . show . VU.unsafeIndex column)
+ get (Just c@(PackedText _ _)) = get (Just (materializePacked c))
+ get Nothing = V.empty
+ in showTable
+ fmt
+ (map clipCell finalHeaders)
+ (map clipCell finalTypes)
+ (map (V.map clipCell) finalCols)
+
+{- | Decide which columns survive horizontal truncation and where to splice the
+ellipsis column. Splits with the extra column on the left for odd 'maxColumns';
+inserts the ellipsis only when it actually saves space.
+-}
+pickColumns ::
+ Maybe TruncateConfig ->
+ Int ->
+ [a] ->
+ ([a], Maybe Int)
+pickColumns mTrunc nCols xs = case mTrunc of
+ Just cfg
+ | let c = maxColumns cfg
+ , c > 0
+ , nCols > c + 1 ->
+ let leftN = (c + 1) `div` 2
+ rightN = c - leftN
+ in ( Prelude.take leftN xs ++ Prelude.drop (nCols - rightN) xs
+ , Just leftN
+ )
+ _ -> (xs, Nothing)
+
+-- | Splice @x@ into @xs@ at index @i@ (0-based), shifting later elements right.
+insertAt :: Int -> a -> [a] -> [a]
+insertAt i x xs = let (l, r) = splitAt i xs in l ++ x : r
+
+-- | Cap a single cell's rendered length, appending an ellipsis when shortened.
+truncateCell :: Int -> T.Text -> T.Text
+truncateCell n t
+ | n <= 0 = t
+ | T.compareLength t n /= GT = t
+ | n == 1 = ellipsisText
+ | otherwise = T.take (n - 1) t <> ellipsisText
+
+-- | O(1) Creates an empty dataframe
+empty :: DataFrame
+empty =
+ DataFrame
+ { columns = V.empty
+ , columnIndices = M.empty
+ , dataframeDimensions = (0, 0)
+ , derivingExpressions = M.empty
+ }
+
+-- | O(k) Get column names of the DataFrame in order of insertion.
+columnNames :: DataFrame -> [T.Text]
+columnNames = map fst . sortBy (compare `on` snd) . M.toList . columnIndices
+{-# INLINE columnNames #-}
+
+{- | Insert a column into a DataFrame. If a column with the same name already
+exists it is replaced in-place; otherwise the column is appended at the end.
+Other columns are expanded (padded with nulls) to match the new row count.
+-}
+insertColumn :: T.Text -> Column -> DataFrame -> DataFrame
+insertColumn name column d =
+ 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 ->
+ DataFrame
+ (V.map (expandColumn n) (columns d V.// [(i, column)]))
+ (columnIndices d)
+ (n, c)
+ exprs
+ Nothing ->
+ DataFrame
+ (V.map (expandColumn n) (columns d `V.snoc` column))
+ (M.insert name c (columnIndices d))
+ (n, c + 1)
+ exprs
+
+-- | Build a DataFrame from a list of @(name, column)@ pairs using 'insertColumn'.
+fromNamedColumns :: [(T.Text, Column)] -> DataFrame
+fromNamedColumns = foldl (\df (name, column) -> insertColumn name column df) empty
+
+{- | Safely retrieves a column by name from the dataframe.
+
+Returns 'Nothing' if the column does not exist.
+
+==== __Examples__
+
+>>> getColumn "age" df
+Just (UnboxedColumn ...)
+
+>>> getColumn "nonexistent" df
+Nothing
+-}
+getColumn :: T.Text -> DataFrame -> Maybe Column
+getColumn name df
+ | null df = Nothing
+ | otherwise = do
+ i <- columnIndices df M.!? name
+ columns df V.!? i
+
+{- | Retrieve a column by name, throwing 'ColumnsNotFoundException' if it does not
+exist. Unsafe version of 'getColumn'; use when the column is certain to exist.
+-}
+unsafeGetColumn :: T.Text -> DataFrame -> Column
+unsafeGetColumn name df = case getColumn name df of
+ Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
+ Just col -> col
+
+{- | Checks if the dataframe is empty (has no columns).
+
+Returns 'True' if the dataframe has no columns, 'False' otherwise.
+Note that a dataframe with columns but no rows is not considered null.
+-}
+null :: DataFrame -> Bool
+null df = V.null (columns df)
+
+-- | Convert a DataFrame to a CSV (comma-separated) text.
+toCsv :: DataFrame -> T.Text
+toCsv = toSeparated ','
+
+-- | Convert a DataFrame to a CSV (comma-separated) string.
+toCsv' :: DataFrame -> String
+toCsv' = T.unpack . toSeparated ','
+
+-- | Convert a DataFrame to a text representation with a custom separator.
+toSeparated :: Char -> DataFrame -> T.Text
+toSeparated sep df
+ | null df = T.empty
+ | otherwise =
+ let (rows, _) = dataframeDimensions df
+ headers = map fst (sortBy (compare `on` snd) (M.toList (columnIndices df)))
+ sepText = T.singleton sep
+ headerLine = T.intercalate sepText headers
+ dataLines = map (T.intercalate sepText . getRowAsText df) [0 .. rows - 1]
+ in T.unlines (headerLine : dataLines)
+
+getRowAsText :: DataFrame -> Int -> [T.Text]
+getRowAsText df i = map (`showElement` i) (V.toList (columns df))
+
+showElement :: Column -> Int -> T.Text
+showElement (BoxedColumn _ (c :: V.Vector a)) i = case c V.!? i of
+ Nothing -> error $ "Column index out of bounds at row " ++ show i
+ Just e
+ | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) -> e
+ | App t1 t2 <- typeRep @a
+ , Just HRefl <- eqTypeRep t1 (typeRep @Maybe) ->
+ case testEquality t2 (typeRep @T.Text) of
+ Just Refl -> fromMaybe "null" e
+ Nothing -> stripJust (T.pack (show e))
+ | otherwise -> T.pack (show e)
+showElement (UnboxedColumn _ c) i = case c VU.!? i of
+ Nothing -> error $ "Column index out of bounds at row " ++ show i
+ Just e -> T.pack (show e)
+showElement (PackedText bm p) i = case bm of
+ Just b | not (bitmapTestBit b i) -> "null"
+ _ -> packedIndexText p i
+
+stripJust :: T.Text -> T.Text
+stripJust = fromMaybe "null" . T.stripPrefix "Just "
diff --git a/src-internal/DataFrame/Internal/DictEncode.hs b/src-internal/DataFrame/Internal/DictEncode.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/DictEncode.hs
@@ -0,0 +1,129 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes: each row
+gets a first-appearance code @0..card-1@ (NULL reserved) plus the cardinality. A
+tested building block; profiled slower than the hash group-by, so unused for now.
+-}
+module DataFrame.Internal.DictEncode (
+ dictEncodeColumn,
+ dictEncodeColumnUpTo,
+ dictMaxCardinality,
+) where
+
+import Control.Monad.ST (runST)
+import qualified Data.Text as T
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import Type.Reflection (typeRep)
+
+import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)
+import DataFrame.Internal.Hash (fnvOffset, mixBytes, mixText, nullSalt)
+import DataFrame.Internal.HashTable (htInsert, newHashTable)
+import DataFrame.Internal.PackedText (
+ PackedTextData,
+ packedLength,
+ packedSlice,
+ sliceEqBytes,
+ )
+
+{- | Largest distinct-value count we will dictionary-encode. Above this the codes
+no longer index a reasonable direct accumulator and the encode pass is pure
+overhead, so the caller keeps the plain hash group-by.
+-}
+dictMaxCardinality :: Int
+dictMaxCardinality = 1048576
+
+{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes,
+returning @Just (codes, cardinality)@ (a NULL row gets its own reserved code).
+'Nothing' for non-text columns or cardinality above 'dictMaxCardinality'.
+-}
+dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int)
+dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality
+
+{- | Dictionary-encode like 'dictEncodeColumn' but bail to 'Nothing' as soon as
+the distinct count would exceed @maxCard@, letting a low-cardinality probe avoid
+a full high-cardinality pass.
+-}
+dictEncodeColumnUpTo :: Int -> Column -> Maybe (VU.Vector Int, Int)
+dictEncodeColumnUpTo maxCard (PackedText bm p) = encodePacked maxCard bm p
+dictEncodeColumnUpTo maxCard (BoxedColumn bm (v :: V.Vector a)) =
+ case testEquality (typeRep @a) (typeRep @T.Text) of
+ Just Refl -> encodeBoxedText maxCard bm v
+ Nothing -> Nothing
+dictEncodeColumnUpTo _ _ = Nothing
+
+{- | Encode a packed-text column: hash each row's raw UTF-8 bytes (the grouping
+'mixBytes'), re-verify byte equality on collisions, assign dense codes in
+first-appearance order. A null row hashes 'nullSalt'.
+-}
+encodePacked ::
+ Int -> Maybe Bitmap -> PackedTextData -> Maybe (VU.Vector Int, Int)
+encodePacked maxCard bm p =
+ let !n = packedLength p
+ valid i = case bm of
+ Just b -> bitmapTestBit b i
+ Nothing -> True
+ hashAt i =
+ if valid i
+ then let (arr, o, l) = packedSlice p i in mixBytes fnvOffset arr o l
+ else nullSalt
+ eqAt a b =
+ case (valid a, valid b) of
+ (True, True) ->
+ let (arrA, oA, lA) = packedSlice p a
+ (arrB, oB, lB) = packedSlice p b
+ in sliceEqBytes arrA oA lA arrB oB lB
+ (False, False) -> True
+ _ -> False
+ in buildCodes maxCard n hashAt eqAt
+
+{- | Encode a boxed 'Data.Text.Text' column, mirroring 'encodePacked' but over
+boxed values (used when a user-built Text column is grouped).
+-}
+encodeBoxedText ::
+ Int -> Maybe Bitmap -> V.Vector T.Text -> Maybe (VU.Vector Int, Int)
+encodeBoxedText maxCard bm v =
+ let !n = V.length v
+ valid i = case bm of
+ Just b -> bitmapTestBit b i
+ Nothing -> True
+ hashAt i =
+ if valid i then mixText fnvOffset (V.unsafeIndex v i) else nullSalt
+ eqAt a b =
+ case (valid a, valid b) of
+ (True, True) -> V.unsafeIndex v a == V.unsafeIndex v b
+ (False, False) -> True
+ _ -> False
+ in buildCodes maxCard n hashAt eqAt
+
+{- | The shared code-assignment loop: bucket every row through an open-addressing
+table on its precomputed hash, re-verify with @eqAt@ on a hit, assign dense
+first-appearance codes. Bails to 'Nothing' once the distinct count exceeds @maxCard@.
+-}
+buildCodes ::
+ Int -> Int -> (Int -> Int) -> (Int -> Int -> Bool) -> Maybe (VU.Vector Int, Int)
+buildCodes maxCard n hashAt eqAt
+ | n == 0 = Just (VU.empty, 0)
+ | otherwise = runST $ do
+ ht <- newHashTable (min n (maxCard + 1))
+ codes <- VUM.new n
+ let go !i !next
+ | i >= n = pure (Just next)
+ | next > maxCard = pure Nothing
+ | otherwise = do
+ let !h = hashAt i
+ (code, isNew) <- htInsert ht eqAt next i h
+ VUM.unsafeWrite codes i code
+ go (i + 1) (if isNew then next + 1 else next)
+ mres <- go 0 0
+ case mres of
+ Nothing -> pure Nothing
+ Just card -> do
+ frozen <- VU.unsafeFreeze codes
+ pure (Just (frozen, card))
diff --git a/src-internal/DataFrame/Internal/Expression.hs b/src-internal/DataFrame/Internal/Expression.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Expression.hs
@@ -0,0 +1,511 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE DisambiguateRecordFields #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE NoFieldSelectors #-}
+
+module DataFrame.Internal.Expression where
+
+import qualified Data.Map.Strict as M
+import Data.Maybe (fromMaybe)
+import Data.String
+import qualified Data.Text as T
+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
+import qualified Data.Vector.Generic as VG
+import DataFrame.Internal.Column
+import qualified DataFrame.Internal.Pretty as P
+import Type.Reflection (Typeable, typeOf, typeRep)
+
+{- | Operators are an open typeclass: built-ins get their own 'Typeable' type so the
+simplifier can match them by 'cast', and users can add instances. The generic
+'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.
+-}
+class (Typeable op) => UnaryOp op where
+ unaryFn :: op a b -> a -> b
+ unaryName :: op a b -> T.Text
+ unarySymbol :: op a b -> Maybe T.Text
+ unarySymbol _ = Nothing
+
+class (Typeable op) => BinaryOp op where
+ binaryFn :: op a b c -> a -> b -> c
+ binaryName :: op a b c -> T.Text
+ binarySymbol :: op a b c -> Maybe T.Text
+ binarySymbol _ = Nothing
+ binaryCommutative :: op a b c -> Bool
+ binaryCommutative _ = False
+ binaryPrecedence :: op a b c -> Int
+ binaryPrecedence _ = 9
+
+data UnUDF a b = MkUnaryOp
+ { unaryFn :: a -> b
+ , unaryName :: T.Text
+ , unarySymbol :: Maybe T.Text
+ }
+
+data BinUDF a b c = MkBinaryOp
+ { binaryFn :: a -> b -> c
+ , binaryName :: T.Text
+ , binarySymbol :: Maybe T.Text
+ , binaryCommutative :: Bool
+ , binaryPrecedence :: Int
+ }
+
+instance UnaryOp UnUDF where
+ unaryFn (MkUnaryOp{unaryFn = f}) = f
+ unaryName (MkUnaryOp{unaryName = n}) = n
+ unarySymbol (MkUnaryOp{unarySymbol = s}) = s
+
+instance BinaryOp BinUDF where
+ binaryFn (MkBinaryOp{binaryFn = f}) = f
+ binaryName (MkBinaryOp{binaryName = n}) = n
+ binarySymbol (MkBinaryOp{binarySymbol = s}) = s
+ binaryCommutative (MkBinaryOp{binaryCommutative = c}) = c
+ binaryPrecedence (MkBinaryOp{binaryPrecedence = p}) = p
+
+data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int
+ deriving (Show, Eq, Ord, Read)
+
+data AggStrategy a b where
+ CollectAgg ::
+ (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b
+ FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
+ MergeAgg ::
+ (Columnable acc) =>
+ T.Text ->
+ acc ->
+ (acc -> b -> acc) ->
+ (acc -> acc -> acc) ->
+ (acc -> a) ->
+ AggStrategy a b
+
+data Expr a where
+ Col :: (Columnable a) => T.Text -> Expr a
+ CastWith ::
+ (Columnable a, Columnable b, Read a) =>
+ T.Text ->
+ T.Text ->
+ (Either String a -> b) ->
+ Expr b
+ CastExprWith ::
+ (Columnable a, Columnable b, Columnable src, Read a) =>
+ T.Text ->
+ (Either String a -> b) ->
+ Expr src ->
+ Expr b
+ Lit :: (Columnable a) => a -> Expr a
+ Unary ::
+ (UnaryOp op, Columnable a, Columnable b) => op b a -> Expr b -> Expr a
+ Binary ::
+ (BinaryOp op, Columnable c, Columnable b, Columnable a) =>
+ op c b a -> Expr c -> Expr b -> Expr a
+ If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
+ Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a
+ Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a
+
+data UExpr where
+ UExpr :: (Columnable a) => Expr a -> UExpr
+
+instance Show UExpr where
+ show :: UExpr -> String
+ show (UExpr expr) = show expr
+
+type NamedExpr = (T.Text, UExpr)
+
+instance (Num a, Columnable a) => Num (Expr a) where
+ (+) :: Expr a -> Expr a -> Expr a
+ (+) =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = (+)
+ , binaryName = "add"
+ , binarySymbol = Just "+"
+ , binaryCommutative = True
+ , binaryPrecedence = 6
+ }
+ )
+
+ (-) :: Expr a -> Expr a -> Expr a
+ (-) =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = (-)
+ , binaryName = "sub"
+ , binarySymbol = Just "-"
+ , binaryCommutative = False
+ , binaryPrecedence = 6
+ }
+ )
+
+ (*) :: Expr a -> Expr a -> Expr a
+ (*) =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = (*)
+ , binaryName = "mult"
+ , binarySymbol = Just "*"
+ , binaryCommutative = True
+ , binaryPrecedence = 7
+ }
+ )
+
+ fromInteger :: Integer -> Expr a
+ fromInteger = Lit . fromInteger
+
+ negate :: Expr a -> Expr a
+ negate =
+ Unary
+ (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing})
+
+ abs :: (Num a) => Expr a -> Expr a
+ abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing})
+
+ signum :: (Num a) => Expr a -> Expr a
+ signum =
+ Unary
+ (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing})
+
+add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
+add = (+)
+
+sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
+sub = (-)
+
+mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
+mult = (*)
+
+instance (Fractional a, Columnable a) => Fractional (Expr a) where
+ fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
+ fromRational = Lit . fromRational
+
+ (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
+ (/) =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = (/)
+ , binaryName = "divide"
+ , binarySymbol = Just "/"
+ , binaryCommutative = False
+ , binaryPrecedence = 7
+ }
+ )
+
+divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
+divide = (/)
+
+instance (IsString a, Columnable a) => IsString (Expr a) where
+ fromString :: String -> Expr a
+ fromString s = Lit (fromString s)
+
+instance (Floating a, Columnable a) => Floating (Expr a) where
+ pi :: (Floating a, Columnable a) => Expr a
+ pi = Lit pi
+ exp :: (Floating a, Columnable a) => Expr a -> Expr a
+ exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing})
+ sqrt :: (Floating a, Columnable a) => Expr a -> Expr a
+ sqrt =
+ Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing})
+ (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
+ (**) =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = (**)
+ , binaryName = "exponentiate"
+ , binarySymbol = Just "**"
+ , binaryCommutative = False
+ , binaryPrecedence = 8
+ }
+ )
+ log :: (Floating a, Columnable a) => Expr a -> Expr a
+ log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing})
+ logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
+ logBase =
+ Binary
+ ( MkBinaryOp
+ { binaryFn = logBase
+ , binaryName = "logBase"
+ , binarySymbol = Nothing
+ , binaryCommutative = False
+ , binaryPrecedence = 1
+ }
+ )
+ sin :: (Floating a, Columnable a) => Expr a -> Expr a
+ sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing})
+ cos :: (Floating a, Columnable a) => Expr a -> Expr a
+ cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing})
+ tan :: (Floating a, Columnable a) => Expr a -> Expr a
+ tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing})
+ asin :: (Floating a, Columnable a) => Expr a -> Expr a
+ asin =
+ Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing})
+ acos :: (Floating a, Columnable a) => Expr a -> Expr a
+ acos =
+ Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing})
+ atan :: (Floating a, Columnable a) => Expr a -> Expr a
+ atan =
+ Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing})
+ sinh :: (Floating a, Columnable a) => Expr a -> Expr a
+ sinh =
+ Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing})
+ cosh :: (Floating a, Columnable a) => Expr a -> Expr a
+ cosh =
+ Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing})
+ asinh :: (Floating a, Columnable a) => Expr a -> Expr a
+ asinh =
+ Unary
+ (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing})
+ acosh :: (Floating a, Columnable a) => Expr a -> Expr a
+ acosh =
+ Unary
+ (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing})
+ atanh :: (Floating a, Columnable a) => Expr a -> Expr a
+ atanh =
+ Unary
+ (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing})
+
+instance (Show a) => Show (Expr a) where
+ show :: 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 ++ ")"
+ show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")"
+ show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
+ show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
+ show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
+ show (Over keys inner) = "(over " ++ show keys ++ " " ++ show inner ++ ")"
+
+normalize :: (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)
+ Binary op e1 e2
+ | binaryCommutative op ->
+ let n1 = normalize e1
+ n2 = normalize e2
+ in case testEquality (typeOf n1) (typeOf n2) of
+ Nothing -> expr
+ Just Refl ->
+ if compareExpr n1 n2 == GT
+ then Binary op n2 n1
+ else Binary op n1 n2
+ | otherwise -> Binary op (normalize e1) (normalize e2)
+ Agg strat e -> Agg strat (normalize e)
+ Over keys inner -> Over keys (normalize inner)
+
+-- Compare expressions for ordering (used in normalization)
+compareExpr :: Expr a -> Expr a -> Ordering
+compareExpr e1 e2 = compare (exprKey e1) (exprKey e2)
+ 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
+ exprKey (Binary op e1' e2') = "4:" ++ T.unpack (binaryName op) ++ exprKey e1' ++ exprKey e2'
+ exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e
+ exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e
+ exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e
+ exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e
+
+eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool
+eqExpr l r = eqNormalized (normalize l) (normalize r)
+ where
+ exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool
+ exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of
+ Just Refl -> eqExpr e1 e2
+ Nothing -> False
+ eqNormalized :: Expr a -> Expr a -> Bool
+ eqNormalized (Col n1) (Col n2) = n1 == n2
+ eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2
+ eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2
+ eqNormalized (Lit v1) (Lit v2) = v1 == v2
+ eqNormalized (If c1 t1 e1) (If c2 t2 e2) =
+ eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2
+ eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2
+ eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b
+ eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) =
+ n1 == n2 && e1 `exprEq` e2
+ eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) =
+ n1 == n2 && e1 `exprEq` e2
+ eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) =
+ n1 == n2 && e1 `exprEq` e2
+ eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2
+ eqNormalized _ _ = False
+
+replaceExpr ::
+ forall a b c.
+ (Columnable a, Columnable b, Columnable c) =>
+ Expr a -> Expr b -> Expr c -> Expr c
+replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of
+ Just Refl -> case testEquality (typeRep @a) (typeRep @c) of
+ Just Refl -> if eqExpr old expr then new else replace'
+ Nothing -> expr
+ Nothing -> replace'
+ 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)
+ (Unary op value) -> Unary op (replaceExpr new old value)
+ (Binary op l r) -> Binary op (replaceExpr new old l) (replaceExpr new old r)
+ (Agg op inner) -> Agg op (replaceExpr new old inner)
+ (Over keys inner) -> Over keys (replaceExpr new old inner)
+
+{- | Simultaneously substitute 'Col' references from a name→expression map in a
+single parallel pass, so a swap like @{a ↦ col b, b ↦ col a}@ works. Raw-text
+references (in 'CastWith', 'Over' keys) are left untouched; type mismatch raises.
+-}
+substituteColumns ::
+ forall a. (Columnable a) => M.Map T.Text UExpr -> Expr a -> Expr a
+substituteColumns subs = go
+ where
+ go :: forall b. (Columnable b) => Expr b -> Expr b
+ go e@(Col name) = case M.lookup name subs of
+ Nothing -> e
+ Just (UExpr (repl :: Expr c)) -> case testEquality (typeRep @b) (typeRep @c) of
+ Just Refl -> repl
+ Nothing ->
+ error $
+ "substituteColumns: type mismatch for column "
+ ++ show name
+ ++ "; column has type "
+ ++ show (typeRep @b)
+ ++ " but replacement has type "
+ ++ show (typeRep @c)
+ go e@(CastWith{}) = e
+ go (CastExprWith t f e) = CastExprWith t f (go e)
+ go e@(Lit _) = e
+ go (If cond l r) = If (go cond) (go l) (go r)
+ go (Unary op value) = Unary op (go value)
+ go (Binary op l r) = Binary op (go l) (go r)
+ go (Agg op inner) = Agg op (go inner)
+ go (Over keys inner) = Over keys (go inner)
+
+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
+eSize (Binary _ l r) = 1 + eSize l + eSize r
+eSize (Agg _strategy expr) = eSize expr + 1
+eSize (Over _ inner) = 1 + eSize inner
+
+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
+getColumns (Binary _op l r) = getColumns l <> getColumns r
+getColumns (Agg _strategy expr) = getColumns expr
+getColumns (Over keys inner) = keys <> getColumns inner
+
+{- | Render an expression as readable, width-aware pseudo-code at the default
+width ('P.defaultWidth'). See 'prettyPrintWidth' to control wrapping.
+-}
+prettyPrint :: Expr a -> String
+prettyPrint = prettyPrintWidth P.defaultWidth
+
+{- | Render an expression as readable, width-aware pseudo-code: long binary chains
+wrap onto aligned continuation lines, @if@/@then@/@else@ break onto their own lines
+(nested @else if@ form a flat ladder), and sub-exprs are parenthesized by precedence.
+-}
+prettyPrintWidth :: Int -> Expr a -> String
+prettyPrintWidth width = P.render width . toDoc 0
+ where
+ toDoc :: Int -> Expr x -> P.Doc
+ toDoc prec expr = case expr of
+ Col name -> P.text (T.unpack name)
+ CastWith name _ _ -> P.text (T.unpack name)
+ CastExprWith tag _ inner -> P.text (T.unpack tag) <> P.parens (toDoc 0 inner)
+ Lit value -> P.text (show value)
+ If{} -> renderIf prec expr
+ Unary op arg ->
+ let fn = fromMaybe (unaryName op) (unarySymbol op)
+ in P.text (T.unpack fn) <> P.parens (toDoc 0 arg)
+ Binary op l r -> case binarySymbol op of
+ Just sym -> renderBinary prec op (T.unpack sym) l r
+ Nothing ->
+ P.text (T.unpack (binaryName op))
+ <> P.parens (toDoc 0 l <> P.text ", " <> toDoc 0 r)
+ Agg (CollectAgg op _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)
+ Agg (FoldAgg op _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)
+ Agg (MergeAgg op _ _ _ _) arg -> P.text (T.unpack op) <> P.parens (toDoc 0 arg)
+ Over keys inner ->
+ toDoc 0 inner <> P.text (".over(" ++ show (map T.unpack keys) ++ ")")
+
+ renderBinary ::
+ (BinaryOp op) => Int -> op c b a -> String -> Expr c -> Expr b -> P.Doc
+ renderBinary prec op sym l r =
+ let p = binaryPrecedence op
+ body
+ | binaryCommutative op =
+ let operands =
+ flattenChain (binaryName op) p l
+ ++ flattenChain (binaryName op) p r
+ in case operands of
+ (o0 : os@(_ : _)) ->
+ P.group
+ ( o0
+ <> P.nest
+ 2
+ (P.hcat [P.line <> P.text sym P.<+> o | o <- os])
+ )
+ _ -> toDoc p l P.<+> P.text sym P.<+> toDoc p r
+ | otherwise =
+ P.group (toDoc p l <> P.nest 2 (P.line <> P.text sym P.<+> toDoc p r))
+ in if prec > p
+ then P.parens body
+ else if prec >= 1 && prec < p then P.parensWhenBroken body else body
+
+ flattenChain :: T.Text -> Int -> Expr x -> [P.Doc]
+ flattenChain name p e = case e of
+ Binary op' l' r'
+ | binaryName op' == name
+ , binaryPrecedence op' == p
+ , binaryCommutative op' ->
+ flattenChain name p l' ++ flattenChain name p r'
+ _ -> [toDoc p e]
+
+ renderIf :: Int -> Expr x -> P.Doc
+ renderIf prec (If c t e) =
+ let blk =
+ P.text "if" P.<+> P.nest 3 (P.group (toDoc 0 c))
+ <> P.hardline
+ <> P.text "then" P.<+> toDoc 0 t
+ <> P.hardline
+ <> renderElse e
+ in if prec > 0 then P.parens (P.nest 2 blk) else blk
+ renderIf _ _ = mempty
+
+ renderElse :: Expr x -> P.Doc
+ renderElse (If c t e) =
+ P.text "else if" P.<+> P.nest 8 (P.group (toDoc 0 c))
+ <> P.hardline
+ <> P.text "then" P.<+> toDoc 0 t
+ <> P.hardline
+ <> renderElse e
+ renderElse other = P.text "else" P.<+> toDoc 0 other
diff --git a/src-internal/DataFrame/Internal/Grouping.hs b/src-internal/DataFrame/Internal/Grouping.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Grouping.hs
@@ -0,0 +1,410 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE Strict #-}
+{-# LANGUAGE TypeApplications #-}
+
+module DataFrame.Internal.Grouping (
+ groupBy,
+ groupBySeq,
+ groupByPar,
+ buildRowToGroup,
+ changingPoints,
+) where
+
+import qualified Data.List as L
+import qualified Data.Map as M
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+import Control.Exception (throw)
+import Control.Monad
+import Control.Monad.ST (ST, runST)
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import DataFrame.Errors
+import DataFrame.Internal.Column (
+ Bitmap,
+ Column (..),
+ bitmapTestBit,
+ )
+import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))
+import DataFrame.Internal.DictEncode (dictEncodeColumnUpTo)
+import DataFrame.Internal.GroupingDirect (
+ DirectGrouping (..),
+ tryDirectGroupColumn,
+ )
+import DataFrame.Internal.GroupingPar (parallelAssignGroups, shouldParallelize)
+import DataFrame.Internal.Hash
+import DataFrame.Internal.HashTable (htInsert, newHashTable)
+import DataFrame.Internal.PackedText (
+ PackedTextData,
+ packedLength,
+ packedSlice,
+ sliceEqBytes,
+ )
+import DataFrame.Internal.RadixRank (rankByHash)
+import DataFrame.Internal.Types
+import System.IO.Unsafe (unsafePerformIO)
+import Type.Reflection (typeRep)
+
+{- | O(k * n) group the dataframe by the given key columns, bucketing rows with an
+open-addressing hash table that re-verifies keys on each hash hit. Groups are
+numbered in first-appearance order; 'valueIndices'/'offsets' follow by counting sort.
+-}
+groupBy ::
+ [T.Text] ->
+ DataFrame ->
+ GroupedDataFrame
+groupBy names df
+ | any (`notElem` columnNames df) names =
+ throw $
+ ColumnsNotFoundException
+ (names L.\\ columnNames df)
+ "groupBy"
+ (columnNames df)
+ | nRows df == 0 =
+ Grouped
+ df
+ names
+ VU.empty
+ (VU.fromList [0])
+ VU.empty
+ | Just dg <- tryDirectGroup names df = dg
+ | shouldParallelize n = groupByPar names df
+ | otherwise = groupBySeq names df
+ where
+ !n = nRows df
+
+{- | Low-cardinality direct-indexed grouping fast path
+('DataFrame.Internal.GroupingDirect'): fires only for a single clean small-range
+@Int@ key. Returns 'Nothing' on any other key shape, falling back to the hash path.
+-}
+tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame
+tryDirectGroup [name] df = do
+ col <- M.lookup name (columnIndices df) >>= \i -> columns df V.!? i
+ case tryDirectGroupColumn col of
+ Just dg ->
+ Just (Grouped df [name] (dgValueIndices dg) (dgOffsets dg) (dgRowToGroup dg))
+ Nothing -> tryDictGroup (nRows df) df [name] col
+tryDirectGroup _ _ = Nothing
+
+{- | Dictionary-encode a single text key to dense int codes, then derive
+@valueIndices@/@offsets@ by counting sort. Profiled slower than the fused hash
+group-by on every db-benchmark question, so it always falls back ('dictGroupEnabled').
+-}
+tryDictGroup ::
+ Int -> DataFrame -> [T.Text] -> Column -> Maybe GroupedDataFrame
+tryDictGroup n df names col
+ | dictGroupEnabled && not (shouldParallelize n) = do
+ (codes, card) <- dictEncodeColumnUpTo dictSingleThreshold col
+ let (vis, os) = indicesFromGroups codes card
+ Just (Grouped df names vis os codes)
+ | otherwise = Nothing
+
+{- | Master switch for the single-key dict-encode grouping path. 'False' because
+it profiled slower than the hash group-by on every db-benchmark group-by question
+(see 'tryDictGroup'); the path is kept compiled and tested but not taken.
+-}
+dictGroupEnabled :: Bool
+dictGroupEnabled = False
+
+{- | Cardinality ceiling for the single-key dict-encode probe: it bails to 'Nothing'
+once the distinct count passes this. Only consulted when 'dictGroupEnabled' is 'True'.
+-}
+dictSingleThreshold :: Int
+dictSingleThreshold = 4096
+
+{- | The sequential grouping path: a single open-addressing table over all rows,
+canonically remapped. Always available regardless of capabilities; the parallel
+path is verified equal to it by a property test.
+-}
+groupBySeq :: [T.Text] -> DataFrame -> GroupedDataFrame
+groupBySeq names df =
+ let !n = nRows df
+ indicesToGroup = keyColIndices names df
+ (rtg0, repHash, repRow) = assignGroups df indicesToGroup n
+ !nGroups = VU.length repHash
+ !remap = canonicalRemap repHash repRow
+ !rtg = VU.map (VU.unsafeIndex remap) rtg0
+ (vis, os) = indicesFromGroups rtg nGroups
+ in Grouped df names vis os rtg
+
+{- | The parallel partitioned grouping path (see 'DataFrame.Internal.GroupingPar'):
+forks one task per capability, producing output bit-for-bit identical to
+'groupBySeq'. Pure via 'unsafePerformIO' (deterministic thread fan-out only).
+-}
+groupByPar :: [T.Text] -> DataFrame -> GroupedDataFrame
+groupByPar names df =
+ let !n = nRows df
+ indicesToGroup = keyColIndices names df
+ !hashes = runST (computeHashes df indicesToGroup n)
+ !eqRow = eqKeyRow df indicesToGroup
+ (rtg, vis, os) = unsafePerformIO (parallelAssignGroups n hashes eqRow)
+ in Grouped df names vis os rtg
+{-# NOINLINE groupByPar #-}
+
+-- | Column indices of the requested key columns, in column order.
+keyColIndices :: [T.Text] -> DataFrame -> [Int]
+keyColIndices names df =
+ M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)
+
+{- | Assign every row to a dense group id in first-appearance order. Returns
+@(rowToGroup, repHash, repRow)@ — the hash and representative row of each group.
+'eqKeyRow' re-verifies the real key on each hash hit so colliding keys stay apart.
+-}
+assignGroups ::
+ DataFrame -> [Int] -> Int -> (VU.Vector Int, VU.Vector Int, VU.Vector Int)
+assignGroups df indicesToGroup n = runST $ do
+ hashes <- computeHashes df indicesToGroup n
+ let !eqRow = eqKeyRow df indicesToGroup
+ ht <- newHashTable n
+ rtg <- VUM.new n
+ repHashM <- VUM.new n
+ repRowM <- VUM.new n
+ let go !i !next
+ | i >= n = pure next
+ | otherwise = do
+ let !h = VU.unsafeIndex hashes i
+ (gid, isNew) <- htInsert ht eqRow next i h
+ VUM.unsafeWrite rtg i gid
+ when isNew $ do
+ VUM.unsafeWrite repHashM next h
+ VUM.unsafeWrite repRowM next i
+ go (i + 1) (if isNew then next + 1 else next)
+ !nGroups <- go 0 0
+ frozen <- VU.unsafeFreeze rtg
+ repHash <- VU.unsafeFreeze (VUM.slice 0 nGroups repHashM)
+ repRow <- VU.unsafeFreeze (VUM.slice 0 nGroups repRowM)
+ pure (frozen, repHash, repRow)
+
+{- | Map each first-appearance group id to its canonical id: groups ordered by
+ascending representative hash (tie-broken by representative row), making group
+order a deterministic function of the key set so set ops commute. O(g), no sort.
+-}
+canonicalRemap :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
+canonicalRemap repHash _repRow =
+ runST (rankByHash (pure . VU.unsafeIndex repHash) (VU.length repHash))
+
+{- | Compute the FNV row-hash of the key columns into a fresh unboxed vector,
+mixing 'nullSalt' for null slots so a missing value never collides with a
+present one of the same bits.
+-}
+computeHashes :: DataFrame -> [Int] -> Int -> ST s (VU.Vector Int)
+computeHashes df indicesToGroup n = do
+ mh <- VUM.replicate n fnvOffset
+ let selectedCols = map (columns df V.!) indicesToGroup
+ forM_ selectedCols $ \case
+ UnboxedColumn ubm (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> hashUnboxed mh ubm mixInt v
+ Nothing ->
+ case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> hashUnboxed mh ubm mixDouble v
+ Nothing ->
+ case sIntegral @a of
+ STrue ->
+ hashUnboxed mh ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v
+ SFalse ->
+ case sFloating @a of
+ STrue ->
+ hashUnboxed mh ubm (\h d -> mixDouble h (realToFrac d :: Double)) v
+ SFalse ->
+ hashUnboxed mh ubm mixShow v
+ BoxedColumn bm (v :: V.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @T.Text) of
+ Just Refl ->
+ V.imapM_
+ ( \i t -> do
+ !h <- VUM.unsafeRead mh i
+ let h' = case bm of
+ Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
+ _ -> mixText h t
+ VUM.unsafeWrite mh i h'
+ )
+ v
+ Nothing ->
+ V.imapM_
+ ( \i d -> do
+ !h <- VUM.unsafeRead mh i
+ let h' = case bm of
+ Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
+ _ -> mixShow h d
+ VUM.unsafeWrite mh i h'
+ )
+ v
+ PackedText bm p -> hashPacked mh bm p
+ VU.unsafeFreeze mh
+
+{- | Build the row-key equality predicate over the selected key columns.
+@eqKeyRow df idxs a b@ is 'True' iff rows @a@ and @b@ agree on all key columns
+(validity first, a null equals only a null). Used to reject hash collisions.
+-}
+eqKeyRow :: DataFrame -> [Int] -> Int -> Int -> Bool
+eqKeyRow df indicesToGroup =
+ let !preds = map (colEqRow . (columns df V.!)) indicesToGroup
+ go [] _ _ = True
+ go (p : ps) a b = p a b && go ps a b
+ in go preds
+
+{- | Per-column row equality respecting nulls. Two rows are equal at a column
+when both are null, or both are valid and their values compare equal.
+-}
+colEqRow :: Column -> (Int -> Int -> Bool)
+colEqRow (UnboxedColumn bm v) =
+ let eqV a b = VU.unsafeIndex v a == VU.unsafeIndex v b
+ in withNulls bm eqV
+colEqRow (BoxedColumn bm v) =
+ let eqV a b = V.unsafeIndex v a == V.unsafeIndex v b
+ in withNulls bm eqV
+colEqRow (PackedText bm p) =
+ let eqV a b =
+ let (arrA, oA, lA) = packedSlice p a
+ (arrB, oB, lB) = packedSlice p b
+ in sliceEqBytes arrA oA lA arrB oB lB
+ in withNulls bm eqV
+{-# INLINE colEqRow #-}
+
+{- | Wrap a value-equality with null handling: equal iff both valid and the
+values agree, or both null.
+-}
+withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> (Int -> Int -> Bool)
+withNulls Nothing eqV = eqV
+withNulls (Just bm) eqV = \a b ->
+ case (bitmapTestBit bm a, bitmapTestBit bm b) of
+ (True, True) -> eqV a b
+ (False, False) -> True
+ _ -> False
+{-# INLINE withNulls #-}
+
+{- | Derive @(valueIndices, offsets)@ from @rowToGroup@ via a stable counting
+sort on the group id: a per-group count, a prefix-sum into group offsets, then a
+single placement pass keeps rows in original order within each group.
+-}
+indicesFromGroups :: VU.Vector Int -> Int -> (VU.Vector Int, VU.Vector Int)
+indicesFromGroups rtg nGroups = runST $ do
+ let !n = VU.length rtg
+ counts <- VUM.replicate (nGroups + 1) 0
+ let countLoop !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !g = VU.unsafeIndex rtg i
+ c <- VUM.unsafeRead counts g
+ VUM.unsafeWrite counts g (c + 1)
+ countLoop (i + 1)
+ countLoop 0
+ offsM <- VUM.new (nGroups + 1)
+ let scan !k !acc
+ | k > nGroups = pure ()
+ | otherwise = do
+ VUM.unsafeWrite offsM k acc
+ c <- VUM.unsafeRead counts k
+ scan (k + 1) (acc + c)
+ scan 0 0
+ let seed !k
+ | k > nGroups = pure ()
+ | otherwise = do
+ s <- VUM.unsafeRead offsM k
+ VUM.unsafeWrite counts k s
+ seed (k + 1)
+ seed 0
+ vis <- VUM.new n
+ let place !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !g = VU.unsafeIndex rtg i
+ pos <- VUM.unsafeRead counts g
+ VUM.unsafeWrite vis pos i
+ VUM.unsafeWrite counts g (pos + 1)
+ place (i + 1)
+ place 0
+ offs <- VU.unsafeFreeze offsM
+ frozenVis <- VU.unsafeFreeze vis
+ pure (frozenVis, offs)
+
+{- | Fold a value-mix over an unboxed column into the running hash vector,
+respecting the null bitmap: a null slot mixes a fixed 'nullSalt' sentinel.
+-}
+hashUnboxed ::
+ (VU.Unbox a) =>
+ VUM.MVector s Int ->
+ Maybe Bitmap ->
+ (Int -> a -> Int) ->
+ VU.Vector a ->
+ ST s ()
+hashUnboxed mh ubm mix v = case ubm of
+ Nothing ->
+ VU.imapM_
+ ( \i x -> do
+ !h <- VUM.unsafeRead mh i
+ VUM.unsafeWrite mh i (mix h x)
+ )
+ v
+ Just bm ->
+ VU.imapM_
+ ( \i x -> do
+ !h <- VUM.unsafeRead mh i
+ VUM.unsafeWrite
+ mh
+ i
+ (if bitmapTestBit bm i then mix h x else mixInt h nullSalt)
+ )
+ v
+{-# INLINE hashUnboxed #-}
+
+{- | Hash a packed-text column over its raw UTF-8 byte slices (no per-row
+'Data.Text.Text'), mixing 'nullSalt' for null rows. Shares 'mixBytes' with
+'mixText' so packed and boxed Text columns hash identically.
+-}
+hashPacked ::
+ VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
+hashPacked mh bm p = go 0
+ where
+ !n = packedLength p
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ !h <- VUM.unsafeRead mh i
+ let h' = case bm of
+ Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
+ _ -> let (arr, o, l) = packedSlice p i in mixBytes h arr o l
+ VUM.unsafeWrite mh i h'
+ go (i + 1)
+{-# INLINE hashPacked #-}
+
+-- Inline accessors to avoid depending on Operations.Core
+
+columnNames :: DataFrame -> [T.Text]
+columnNames = M.keys . columnIndices
+
+nRows :: DataFrame -> Int
+nRows = fst . dataframeDimensions
+
+{- | Build the rowToGroup lookup vector from valueIndices and offsets.
+rowToGroup[i] = k means row i belongs to group k.
+-}
+buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int
+buildRowToGroup n vis os = runST $ do
+ rtg <- VUM.new n
+ let nGroups = VU.length os - 1
+ forM_ [0 .. nGroups - 1] $ \k ->
+ let s = VU.unsafeIndex os k
+ e = VU.unsafeIndex os (k + 1)
+ in forM_ [s .. e - 1] $ \i ->
+ VUM.unsafeWrite rtg (VU.unsafeIndex vis i) k
+ VU.unsafeFreeze rtg
+{-# NOINLINE buildRowToGroup #-}
+
+changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int
+changingPoints vs =
+ VU.reverse
+ (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs)))
+ where
+ initialState = ([0], snd (VU.head vs))
+ findChangePoints (!offs, !currentVal) index (_, !newVal)
+ | currentVal == newVal = (offs, currentVal)
+ | otherwise = (index : offs, newVal)
diff --git a/src-internal/DataFrame/Internal/GroupingDirect.hs b/src-internal/DataFrame/Internal/GroupingDirect.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/GroupingDirect.hs
@@ -0,0 +1,234 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | Low-cardinality direct-indexed grouping fast path: when the key is a single
+clean unboxed @Int@ column of small value range, the value itself indexes a dense
+accumulator (no hashing/probing). Emits groups in ascending value order.
+-}
+module DataFrame.Internal.GroupingDirect (
+ directGroupThreshold,
+ tryDirectGroupColumn,
+ DirectGrouping (..),
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import System.IO.Unsafe (unsafePerformIO)
+import Type.Reflection (typeRep)
+
+import DataFrame.Internal.Column (Column (..))
+
+{- | Largest key value RANGE (max - min + 1) the direct grouping path accepts. A
+@2^20@-slot histogram is 8MB; the low-cardinality questions sit far below it
+(id4 range 100, id6 range 1e5). Wider ranges fall back to the hash group-by.
+-}
+directGroupThreshold :: Int
+directGroupThreshold = 1048576
+
+{- | The grouping layout the hash path also produces: @rowToGroup@, the
+group-sorted @valueIndices@, the @offsets@ prefix array, and the group count.
+-}
+data DirectGrouping = DirectGrouping
+ { dgRowToGroup :: !(VU.Vector Int)
+ , dgValueIndices :: !(VU.Vector Int)
+ , dgOffsets :: !(VU.Vector Int)
+ , dgNGroups :: !Int
+ }
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+parThreshold :: Int
+parThreshold = 200000
+
+{- | Take the direct path if the (single) key column is a clean non-null unboxed
+@Int@ column with a small value range. Returns 'Nothing' to fall back to the
+hash group-by on anything else (boxed/text keys, nullable, wide ranges, empty).
+-}
+tryDirectGroupColumn :: Column -> Maybe DirectGrouping
+tryDirectGroupColumn (UnboxedColumn Nothing (v :: VU.Vector a))
+ | Just Refl <- testEquality (typeRep @a) (typeRep @Int)
+ , not (VU.null v) =
+ let (!mn, !mx) = rangeOf v
+ !range = mx - mn + 1
+ in if range >= 1 && range <= directGroupThreshold
+ then Just (directGroup v mn range)
+ else Nothing
+tryDirectGroupColumn _ = Nothing
+
+-- | Parallel min/max reduce (order-independent).
+rangeOf :: VU.Vector Int -> (Int, Int)
+rangeOf v
+ | not (shouldPar n) = rangeChunk v 0 n
+ | otherwise = unsafePerformIO $ do
+ let !caps = capabilities
+ !per = (n + caps - 1) `div` caps
+ spawn w = do
+ var <- newEmptyMVar
+ let !lo = min n (w * per)
+ !hi = min n (lo + per)
+ _ <- forkIO (try (pure $! rangeChunk v lo hi) >>= putMVar var)
+ pure var
+ vars <- mapM spawn [0 .. caps - 1]
+ rs <- mapM takeMVar vars
+ rs' <- mapM (either (throwIO @SomeException) pure) rs
+ pure (combineRanges (filter (\(a, _) -> a /= maxBound) rs'))
+ where
+ !n = VU.length v
+{-# NOINLINE rangeOf #-}
+
+rangeChunk :: VU.Vector Int -> Int -> Int -> (Int, Int)
+rangeChunk v lo hi = go lo maxBound minBound
+ where
+ go !i !mn !mx
+ | i >= hi = (mn, mx)
+ | otherwise =
+ let !x = VU.unsafeIndex v i
+ in go (i + 1) (min mn x) (max mx x)
+
+combineRanges :: [(Int, Int)] -> (Int, Int)
+combineRanges [] = (0, 0)
+combineRanges ((a0, b0) : rest) = foldr (\(a, b) (ma, mb) -> (min ma a, max mb b)) (a0, b0) rest
+
+shouldPar :: Int -> Bool
+shouldPar n = n >= parThreshold && capabilities > 1
+
+{- | Build the grouping by counting sort on @value - min@: a (parallel) per-value
+histogram, compaction of non-empty values into ascending dense ids, a scan into
+offsets, then a stable placement pass building @valueIndices@ and @rowToGroup@.
+-}
+directGroup :: VU.Vector Int -> Int -> Int -> DirectGrouping
+directGroup v mn range = unsafePerformIO $ do
+ let !n = VU.length v
+ hist <- buildHistogram v mn range n
+ valToGroup <- VUM.replicate range (-1 :: Int)
+ grpCount <- VUM.new range
+ nGroups <- compact hist range valToGroup grpCount
+ offsM <- VUM.new (nGroups + 1)
+ cursor <- VUM.new nGroups
+ scanOffsets grpCount nGroups offsM cursor
+ rtg <- VUM.new n
+ vis <- VUM.new n
+ place v mn n valToGroup cursor rtg vis
+ frozenRtg <- VU.unsafeFreeze rtg
+ frozenVis <- VU.unsafeFreeze vis
+ frozenOffs <- VU.unsafeFreeze offsM
+ pure (DirectGrouping frozenRtg frozenVis frozenOffs nGroups)
+{-# NOINLINE directGroup #-}
+
+{- | Parallel per-value histogram: each worker fills a private @range@-slot
+count over its row chunk, then the partials are summed (exact integers, so the
+merge order is irrelevant). Sequential single pass below 'parThreshold'.
+-}
+buildHistogram :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
+buildHistogram v mn range n
+ | not (shouldPar n) = histChunk v mn range 0 n
+ | otherwise = do
+ let !caps = capabilities
+ !per = (n + caps - 1) `div` caps
+ spawn w = do
+ var <- newEmptyMVar
+ let !lo = min n (w * per)
+ !hi = min n (lo + per)
+ _ <- forkIO (try (histChunk v mn range lo hi) >>= putMVar var)
+ pure var
+ vars <- mapM spawn [0 .. caps - 1]
+ rs <- mapM takeMVar vars
+ parts <- mapM (either (throwIO @SomeException) pure) rs
+ case parts of
+ [] -> VUM.replicate range 0
+ (p0 : rest) -> do
+ mapM_ (addInto p0 range) rest
+ pure p0
+
+histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
+histChunk v mn range lo hi = do
+ acc <- VUM.replicate range (0 :: Int)
+ let go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ let !k = VU.unsafeIndex v i - mn
+ c <- VUM.unsafeRead acc k
+ VUM.unsafeWrite acc k (c + 1)
+ go (i + 1)
+ go lo
+ pure acc
+
+addInto :: VUM.IOVector Int -> Int -> VUM.IOVector Int -> IO ()
+addInto dst range src = go 0
+ where
+ go !k
+ | k >= range = pure ()
+ | otherwise = do
+ a <- VUM.unsafeRead dst k
+ b <- VUM.unsafeRead src k
+ VUM.unsafeWrite dst k (a + b)
+ go (k + 1)
+
+{- | Walk the histogram in ascending value order, assigning a dense group id to
+each non-empty value and copying its count into @grpCount@ at that id. Returns
+the group count.
+-}
+compact ::
+ VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO Int
+compact hist range valToGroup grpCount = go 0 0
+ where
+ go !val !next
+ | val >= range = pure next
+ | otherwise = do
+ c <- VUM.unsafeRead hist val
+ if c == 0
+ then go (val + 1) next
+ else do
+ VUM.unsafeWrite valToGroup val next
+ VUM.unsafeWrite grpCount next c
+ go (val + 1) (next + 1)
+
+{- | Exclusive prefix scan of group counts into @offsM@ (length nGroups+1) and
+seed the per-group write @cursor@ at each group's start offset.
+-}
+scanOffsets ::
+ VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO ()
+scanOffsets grpCount nGroups offsM cursor = go 0 0
+ where
+ go !g !acc
+ | g >= nGroups = VUM.unsafeWrite offsM nGroups acc
+ | otherwise = do
+ VUM.unsafeWrite offsM g acc
+ VUM.unsafeWrite cursor g acc
+ c <- VUM.unsafeRead grpCount g
+ go (g + 1) (acc + c)
+
+{- | Stable placement pass: for each row in original order, look up its group id
+through the value map, write @rowToGroup@, and append the row to its group's run
+in @valueIndices@ via the advancing cursor (rows keep original order per group).
+-}
+place ::
+ VU.Vector Int ->
+ Int ->
+ Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ IO ()
+place v mn n valToGroup cursor rtg vis = go 0
+ where
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !val = VU.unsafeIndex v i - mn
+ g <- VUM.unsafeRead valToGroup val
+ VUM.unsafeWrite rtg i g
+ pos <- VUM.unsafeRead cursor g
+ VUM.unsafeWrite vis pos i
+ VUM.unsafeWrite cursor g (pos + 1)
+ go (i + 1)
diff --git a/src-internal/DataFrame/Internal/GroupingPar.hs b/src-internal/DataFrame/Internal/GroupingPar.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/GroupingPar.hs
@@ -0,0 +1,308 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE Strict #-}
+
+{- | Parallel partitioned group-by: rows are counting-sorted into partitions by the
+top hash bits, then one task per capability groups its partitions independently.
+Output is bit-for-bit identical to the sequential 'DataFrame.Internal.Grouping.groupBy'.
+-}
+module DataFrame.Internal.GroupingPar (
+ parallelAssignGroups,
+ shouldParallelize,
+ parThreshold,
+ numPartitionsFor,
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import Control.Monad (forM_, when)
+import Data.Bits (countLeadingZeros, unsafeShiftR)
+import Data.IORef (atomicModifyIORef', newIORef)
+import qualified Data.Vector as V
+import qualified Data.Vector.Mutable as VM
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import Data.Word (Word64)
+import DataFrame.Internal.HashTable (
+ htInsert,
+ newHashTable,
+ )
+import DataFrame.Internal.RadixRank (rankByHash)
+import System.IO.Unsafe (unsafePerformIO)
+
+{- | Below this many rows the partition/fork overhead is not worth it; 'groupBy'
+uses its sequential 'ST' path instead.
+-}
+parThreshold :: Int
+parThreshold = 200000
+
+{- | Whether 'groupBy' should take the parallel path: more than one capability
+and at least 'parThreshold' rows.
+-}
+shouldParallelize :: Int -> Bool
+shouldParallelize n = n >= parThreshold && capabilities > 1
+{-# NOINLINE shouldParallelize #-}
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+{- | Sign-preserving unsigned remap: ascending 'Word64' order of @key h@ equals
+ascending signed-'Int' order of @h@, so partitioning and sorting on it reproduce
+the sequential @compare \`on\` repHash@ ordering exactly.
+-}
+key :: Int -> Word64
+key h = fromIntegral h + 0x8000000000000000
+{-# INLINE key #-}
+
+-- | Partition index of a hash: the top @log2 p@ bits of its unsigned key.
+partIx :: Int -> Int -> Int
+partIx shift h = fromIntegral (key h `unsafeShiftR` shift)
+{-# INLINE partIx #-}
+
+{- | Number of partitions: a power of two, at least @4 * caps@ (P >> cores for
+skew tolerance), floored at 256.
+-}
+numPartitionsFor :: Int -> Int
+numPartitionsFor caps = go 1
+ where
+ target = max 256 (4 * caps)
+ go p
+ | p >= target = p
+ | otherwise = go (p * 2)
+
+-- | @floor (log2 x)@ for a power-of-two @x@.
+intLog2 :: Int -> Int
+intLog2 x = 63 - countLeadingZeros x
+{-# INLINE intLog2 #-}
+
+{- | Parallel group assignment. @parallelAssignGroups n hashes eqRow@ returns
+@(rowToGroup, valueIndices, offsets)@ in canonical group order. @eqRow a b@ must
+report whether rows @a@ and @b@ share all key columns (null-aware).
+-}
+parallelAssignGroups ::
+ Int ->
+ VU.Vector Int ->
+ (Int -> Int -> Bool) ->
+ IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)
+parallelAssignGroups n hashes eqRow = do
+ caps <- getNumCapabilities
+ let !p = numPartitionsFor caps
+ !shift = 64 - intLog2 p
+ (partStart, sortedRows) <- partitionRows n hashes p shift
+ localGid <- VUM.new (max 1 n)
+ canonBoxes <- VM.replicate p (VU.empty :: VU.Vector Int)
+ nLocalGroups <- VUM.replicate p (0 :: Int)
+ runPartitions
+ caps
+ p
+ partStart
+ sortedRows
+ hashes
+ eqRow
+ localGid
+ canonBoxes
+ nLocalGroups
+ (globalBase, canonOf, nGroups) <- canonicalize p canonBoxes nLocalGroups
+ assemble n p partStart sortedRows localGid globalBase canonOf nGroups
+
+-------------------------------------------------------------------------------
+-- Phase 1: counting sort by partition
+-------------------------------------------------------------------------------
+
+{- | Bucket every row index into its partition by a counting sort. Returns the
+exclusive prefix-sum @partStart@ (length @p+1@, @partStart[p] == n@) and the row
+indices laid out partition-by-partition in @sortedRows@.
+-}
+partitionRows ::
+ Int -> VU.Vector Int -> Int -> Int -> IO (VU.Vector Int, VU.Vector Int)
+partitionRows n hashes p shift = do
+ counts <- VUM.replicate (p + 1) (0 :: Int)
+ let countLoop !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !pp = partIx shift (VU.unsafeIndex hashes i)
+ c <- VUM.unsafeRead counts pp
+ VUM.unsafeWrite counts pp (c + 1)
+ countLoop (i + 1)
+ countLoop 0
+ partStartM <- VUM.new (p + 1)
+ let scan !k !acc
+ | k > p = pure ()
+ | otherwise = do
+ VUM.unsafeWrite partStartM k acc
+ c <- if k < p then VUM.unsafeRead counts k else pure 0
+ scan (k + 1) (acc + c)
+ scan 0 0
+ cursor <- VUM.new p
+ forM_ [0 .. p - 1] $ \k -> VUM.unsafeRead partStartM k >>= VUM.unsafeWrite cursor k
+ sortedM <- VUM.new (max 1 n)
+ let place !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !pp = partIx shift (VU.unsafeIndex hashes i)
+ pos <- VUM.unsafeRead cursor pp
+ VUM.unsafeWrite sortedM pos i
+ VUM.unsafeWrite cursor pp (pos + 1)
+ place (i + 1)
+ place 0
+ partStart <- VU.unsafeFreeze partStartM
+ sortedRows <- VU.unsafeFreeze sortedM
+ pure (partStart, sortedRows)
+
+-------------------------------------------------------------------------------
+-- Phase 2: per-partition grouping (parallel)
+-------------------------------------------------------------------------------
+
+{- | Group each partition with its own hash table, then rank its local groups into
+canonical order — all inside the parallel worker. Forks @caps@ workers pulling
+partition indices off a shared counter; disjoint keys mean no cross-partition merge.
+-}
+runPartitions ::
+ Int ->
+ Int ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ (Int -> Int -> Bool) ->
+ VUM.IOVector Int ->
+ VM.IOVector (VU.Vector Int) ->
+ VUM.IOVector Int ->
+ IO ()
+runPartitions caps p partStart sortedRows hashes eqRow localGid canonBoxes nLocalGroups = do
+ next <- newIORef 0
+ let groupPartition !pp = do
+ let !s = VU.unsafeIndex partStart pp
+ !e = VU.unsafeIndex partStart (pp + 1)
+ !sz = e - s
+ when (sz > 0) $ do
+ ht <- newHashTable sz
+ repHashM <- VUM.new sz
+ let loop !pos !nextGid
+ | pos >= e = pure nextGid
+ | otherwise = do
+ let !row = VU.unsafeIndex sortedRows pos
+ !h = VU.unsafeIndex hashes row
+ (gid, isNew) <- htInsert ht eqRow nextGid row h
+ VUM.unsafeWrite localGid pos gid
+ if isNew
+ then do
+ VUM.unsafeWrite repHashM nextGid h
+ loop (pos + 1) (nextGid + 1)
+ else loop (pos + 1) nextGid
+ ng <- loop s 0
+ VUM.unsafeWrite nLocalGroups pp ng
+ canon <- rankByHash (VUM.unsafeRead repHashM) ng
+ VM.unsafeWrite canonBoxes pp canon
+ worker = do
+ i <- atomicModifyIORef' next (\j -> (j + 1, j))
+ when (i < p) $ groupPartition i >> worker
+ forkJoin_ (replicate caps worker)
+
+-------------------------------------------------------------------------------
+-- Phase 3: global base ids + assembly
+-------------------------------------------------------------------------------
+
+{- | Exclusive prefix sum of the per-partition group counts into @globalBase@
+(@globalBase[pp]@ = first global id of partition @pp@). Ranks were computed in
+'runPartitions'; prepending the base to each yields the sequential order.
+-}
+canonicalize ::
+ Int ->
+ VM.IOVector (VU.Vector Int) ->
+ VUM.IOVector Int ->
+ IO (VU.Vector Int, V.Vector (VU.Vector Int), Int)
+canonicalize p canonBoxes nLocalGroups = do
+ globalBaseM <- VUM.new (p + 1)
+ let go !pp !base
+ | pp >= p = VUM.unsafeWrite globalBaseM p base >> pure base
+ | otherwise = do
+ VUM.unsafeWrite globalBaseM pp base
+ ng <- VUM.unsafeRead nLocalGroups pp
+ go (pp + 1) (base + ng)
+ total <- go 0 0
+ globalBase <- VU.unsafeFreeze globalBaseM
+ canonOf <- V.unsafeFreeze canonBoxes
+ pure (globalBase, canonOf, total)
+
+{- | Build the final @(rowToGroup, valueIndices, offsets)@: the global group id of a
+sorted position is @globalBase[pp] + canonOf[pp][localGid]@. @valueIndices@ orders
+rows by group, @offsets@ the boundaries, @rowToGroup@ the inverse per original row.
+-}
+assemble ::
+ Int ->
+ Int ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ VUM.IOVector Int ->
+ VU.Vector Int ->
+ V.Vector (VU.Vector Int) ->
+ Int ->
+ IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)
+assemble n p partStart sortedRows localGid globalBase canonOf nGroups = do
+ rtgM <- VUM.new (max 1 n)
+ counts <- VUM.replicate (nGroups + 1) (0 :: Int)
+ gidAt <- VUM.new (max 1 n)
+ let scanPos !pp
+ | pp >= p = pure ()
+ | otherwise = do
+ let !s = VU.unsafeIndex partStart pp
+ !e = VU.unsafeIndex partStart (pp + 1)
+ !base = VU.unsafeIndex globalBase pp
+ !canon = V.unsafeIndex canonOf pp
+ let inner !pos
+ | pos >= e = pure ()
+ | otherwise = do
+ lg <- VUM.unsafeRead localGid pos
+ let !g = base + VU.unsafeIndex canon lg
+ !row = VU.unsafeIndex sortedRows pos
+ VUM.unsafeWrite gidAt pos g
+ VUM.unsafeWrite rtgM row g
+ c <- VUM.unsafeRead counts g
+ VUM.unsafeWrite counts g (c + 1)
+ inner (pos + 1)
+ inner s
+ scanPos (pp + 1)
+ scanPos 0
+ offsM <- VUM.new (nGroups + 1)
+ let scan !k !acc
+ | k > nGroups = pure ()
+ | otherwise = do
+ VUM.unsafeWrite offsM k acc
+ c <- if k < nGroups then VUM.unsafeRead counts k else pure 0
+ scan (k + 1) (acc + c)
+ scan 0 0
+ cursor <- VUM.new (max 1 nGroups)
+ forM_ [0 .. nGroups - 1] $ \k -> VUM.unsafeRead offsM k >>= VUM.unsafeWrite cursor k
+ visM <- VUM.new (max 1 n)
+ let placeVis !pos
+ | pos >= n = pure ()
+ | otherwise = do
+ g <- VUM.unsafeRead gidAt pos
+ let !row = VU.unsafeIndex sortedRows pos
+ c <- VUM.unsafeRead cursor g
+ VUM.unsafeWrite visM c row
+ VUM.unsafeWrite cursor g (c + 1)
+ placeVis (pos + 1)
+ placeVis 0
+ rtg <- VU.unsafeFreeze rtgM
+ offs <- VU.unsafeFreeze offsM
+ vis <- VU.unsafeFreeze visM
+ pure (rtg, vis, offs)
+
+-------------------------------------------------------------------------------
+-- Thread fan-out (plain forkIO + MVar join, no sparks)
+-------------------------------------------------------------------------------
+
+-- | Run each action on its own thread; rethrow the first failure (in order).
+forkJoin_ :: [IO ()] -> IO ()
+forkJoin_ actions = do
+ vars <- mapM spawn actions
+ results <- mapM takeMVar vars
+ mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
+ where
+ spawn act = do
+ var <- newEmptyMVar
+ _ <- forkIO (try act >>= putMVar var)
+ pure var
diff --git a/src-internal/DataFrame/Internal/Hash.hs b/src-internal/DataFrame/Internal/Hash.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Hash.hs
@@ -0,0 +1,113 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE MagicHash #-}
+
+{- | A poor-man's hash used by 'DataFrame.Internal.Grouping' to bucket rows
+without depending on @hashable@. Each value is folded into an 'Int' with an
+FxHash-style step (rotate, xor, multiply); small and not cryptographic.
+-}
+module DataFrame.Internal.Hash (
+ fnvOffset,
+ nullSalt,
+ mixInt,
+ mixDouble,
+ mixBool,
+ mixChar,
+ mixText,
+ mixBytes,
+ mixShow,
+) where
+
+import Data.Bits (rotateL, unsafeShiftL, unsafeShiftR, xor)
+import Data.Char (ord)
+import qualified Data.Text as T
+import qualified Data.Text.Array as A
+#if MIN_VERSION_text(2,1,0)
+import Data.Array.Byte (ByteArray (ByteArray))
+#else
+import Data.Text.Array (Array (ByteArray))
+#endif
+import Data.Text.Internal (Text (Text))
+import GHC.Exts (Int (I#), indexWord8Array#, indexWord8ArrayAsWord64#)
+import GHC.Word (Word64 (W64#), Word8 (W8#))
+
+{- | FNV-1a 64-bit offset basis (used as the initial accumulator).
+The literal is unsigned and exceeds 'Int' range, so we round-trip through
+'Word64' to get the well-defined two's-complement bit pattern.
+-}
+fnvOffset :: Int
+fnvOffset = fromIntegral (0xcbf29ce484222325 :: Word64)
+
+-- | FNV-1a 64-bit prime.
+fnvPrime :: Int
+fnvPrime = 0x00000100000001b3
+
+{- | Sentinel mixed in for a /null/ slot, so @Nothing@ does not hash the same as
+a present value with equal bits (e.g. @Just 0@). A fixed distinctive constant
+keeps null hashing deterministic; a real value equal to it collides only rarely.
+-}
+nullSalt :: Int
+nullSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word64)
+
+{- | Mix an 'Int' into the accumulator with an FxHash-style step. The rotate
+diffuses each value's bits before the next is folded in, avoiding the structured
+collisions a plain xor-then-multiply produces on small/adjacent group keys.
+-}
+mixInt :: Int -> Int -> Int
+mixInt acc x = (rotateL acc 13 `xor` x) * fnvPrime
+{-# INLINE mixInt #-}
+
+{- | Mix a 'Double' into the accumulator. Loses sub-millisecond precision
+but matches the bucketing the old hashable-based code used.
+-}
+mixDouble :: Int -> Double -> Int
+mixDouble acc d = mixInt acc (floor (d * 1000))
+{-# INLINE mixDouble #-}
+
+mixBool :: Int -> Bool -> Int
+mixBool acc b = mixInt acc (if b then 1 else 0)
+{-# INLINE mixBool #-}
+
+mixChar :: Int -> Char -> Int
+mixChar acc = mixInt acc . ord
+{-# INLINE mixChar #-}
+
+{- | Mix a 'T.Text' value into the accumulator over its raw UTF-8 bytes, eight at
+a time. Reading a whole 'Word64' per step cuts the multiply count ~8x on long
+keys while staying collision-equivalent (UTF-8 is injective).
+-}
+mixText :: Int -> T.Text -> Int
+mixText !acc (Text arr off len) = mixBytes acc arr off len
+{-# INLINE mixText #-}
+
+{- | Mix a raw UTF-8 byte slice @[off, off+len)@ of a 'Data.Text.Array.Array'
+into the accumulator, eight bytes at a time. The shared kernel behind
+'mixText' and the packed-text hash path, so the two never drift.
+-}
+mixBytes :: Int -> A.Array -> Int -> Int -> Int
+mixBytes !acc arr off len = goBytes (goWords acc off) wordsEnd
+ where
+ !(ByteArray ba) = arr
+ !nWords = len `unsafeShiftR` 3
+ !wordsEnd = off + (nWords `unsafeShiftL` 3)
+ !end = off + len
+ goWords !h !i
+ | i >= wordsEnd = h
+ | otherwise =
+ let !(I# i#) = i
+ !w = fromIntegral (W64# (indexWord8ArrayAsWord64# ba i#)) :: Int
+ in goWords (mixInt h w) (i + 8)
+ goBytes !h !i
+ | i >= end = h
+ | otherwise =
+ let !(I# i#) = i
+ !b = fromIntegral (W8# (indexWord8Array# ba i#)) :: Int
+ in goBytes (mixInt h b) (i + 1)
+{-# INLINE mixBytes #-}
+
+{- | Fallback for arbitrary 'Show'-able values. Slower but covers types
+without a dedicated combinator (e.g. 'Day', 'UTCTime').
+-}
+mixShow :: (Show a) => Int -> a -> Int
+mixShow acc = mixText acc . T.pack . show
+{-# INLINE mixShow #-}
diff --git a/src-internal/DataFrame/Internal/HashTable.hs b/src-internal/DataFrame/Internal/HashTable.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/HashTable.hs
@@ -0,0 +1,94 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+{- | A flat, unboxed, open-addressing (linear-probe) hash table mapping a row's
+key-hash to a dense group id, re-verifying the real key on every hash hit to
+reject collisions. Runs in any 'PrimMonad' ('ST' for grouping, 'IO' per worker).
+-}
+module DataFrame.Internal.HashTable (
+ HashTable (..),
+ newHashTable,
+ htInsert,
+ nextPow2Above,
+) where
+
+import Control.Monad.Primitive (PrimMonad, PrimState)
+import Data.Bits ((.&.))
+import qualified Data.Vector.Unboxed.Mutable as VUM
+
+{- | An open-addressing linear-probe table. @htMask@ is @capacity - 1@ (capacity
+is a power of two) and maps a hash to its home slot.
+-}
+data HashTable s = HashTable
+ { htHash :: !(VUM.MVector s Int)
+ , htGroup :: !(VUM.MVector s Int)
+ , htRep :: !(VUM.MVector s Int)
+ , htMask :: !Int
+ }
+
+{- | Smallest power of two strictly greater than @n@, at least 2. Sizes the
+table so the load factor stays below ~0.5 even when every row is a distinct
+group.
+-}
+nextPow2Above :: Int -> Int
+nextPow2Above n = go 2
+ where
+ go !p
+ | p > n = p
+ | otherwise = go (p * 2)
+{-# INLINE nextPow2Above #-}
+
+{- | Allocate an empty table able to hold up to @n@ distinct groups while
+keeping the load factor under ~0.5 (capacity @= nextPow2Above (2*n)@). All
+group slots start empty (@-1@).
+-}
+newHashTable :: (PrimMonad m) => Int -> m (HashTable (PrimState m))
+newHashTable n = do
+ let !cap = nextPow2Above (2 * max 1 n)
+ h <- VUM.unsafeNew cap
+ g <- VUM.replicate cap (-1)
+ r <- VUM.unsafeNew cap
+ pure (HashTable h g r (cap - 1))
+{-# INLINE newHashTable #-}
+
+{- | Look up @row@ (with precomputed @hash@) and return its dense group id: an
+empty slot starts a new group via @nextGroup@, a stored-hash match is re-verified
+with @eqRow@ before reuse. The 'Bool' is 'True' when a new group was created.
+-}
+htInsert ::
+ (PrimMonad m) =>
+ HashTable (PrimState m) ->
+ -- | @eqRow a b@: do rows @a@ and @b@ have equal key columns?
+ (Int -> Int -> Bool) ->
+ -- | Next dense group id to assign if this row starts a new group.
+ Int ->
+ -- | Row index being inserted.
+ Int ->
+ -- | Precomputed hash of the row's key.
+ Int ->
+ m (Int, Bool)
+htInsert ht eqRow nextGroup row hash = go (hash .&. mask)
+ where
+ !mask = htMask ht
+ !hs = htHash ht
+ !gs = htGroup ht
+ !rs = htRep ht
+ go !slot = do
+ g <- VUM.unsafeRead gs slot
+ if g < 0
+ then do
+ VUM.unsafeWrite hs slot hash
+ VUM.unsafeWrite gs slot nextGroup
+ VUM.unsafeWrite rs slot row
+ pure (nextGroup, True)
+ else do
+ h <- VUM.unsafeRead hs slot
+ if h == hash
+ then do
+ rep <- VUM.unsafeRead rs slot
+ if eqRow rep row
+ then pure (g, False)
+ else go ((slot + 1) .&. mask)
+ else go ((slot + 1) .&. mask)
+{-# INLINE htInsert #-}
diff --git a/src-internal/DataFrame/Internal/Interpreter.hs b/src-internal/DataFrame/Internal/Interpreter.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Interpreter.hs
@@ -0,0 +1,1051 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module DataFrame.Internal.Interpreter (
+ -- * New core API
+ Value (..),
+ Ctx (..),
+ eval,
+ materialize,
+
+ -- * Backward-compatible API
+ interpret,
+ interpretAggregation,
+ AggregationResult (..),
+) where
+
+import Data.Bifunctor (first)
+import qualified Data.Map as M
+import qualified Data.Text as T
+import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
+import qualified Data.Vector as V
+import qualified Data.Vector.Generic as VG
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import DataFrame.Errors
+import DataFrame.Internal.Column
+import DataFrame.Internal.DataFrame
+import DataFrame.Internal.Expression
+import qualified DataFrame.Internal.Grouping as G
+import DataFrame.Internal.Types
+import Type.Reflection (
+ Typeable,
+ typeRep,
+ )
+
+import Data.Int (Int16, Int32, Int64, Int8)
+
+-- Specializations for common aggregation types to avoid dictionary overhead.
+-- foldLinearGroups: mean accumulator
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Double -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Float -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Int -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Int8 -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Int16 -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Int32 -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (MeanAcc -> Int64 -> MeanAcc) ->
+ MeanAcc ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+-- foldLinearGroups: count accumulator
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Double -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Float -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Int -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Int8 -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Int16 -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Int32 -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int -> Int64 -> Int) ->
+ Int ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+-- foldLinearGroups: sum/min/max (acc == elem)
+{-# SPECIALIZE foldLinearGroups ::
+ (Double -> Double -> Double) ->
+ Double ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Float -> Float -> Float) ->
+ Float ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int8 -> Int8 -> Int8) ->
+ Int8 ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int16 -> Int16 -> Int16) ->
+ Int16 ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int32 -> Int32 -> Int32) ->
+ Int32 ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE foldLinearGroups ::
+ (Int64 -> Int64 -> Int64) ->
+ Int64 ->
+ Column ->
+ VU.Vector Int ->
+ Int ->
+ Either DataFrameException Column
+ #-}
+
+-- mapColumn: finalize
+{-# SPECIALIZE mapColumn ::
+ (MeanAcc -> Double) -> Column -> Either DataFrameException Column
+ #-}
+{-# SPECIALIZE mapColumn ::
+ (Double -> Double) -> Column -> Either DataFrameException Column
+ #-}
+{-# SPECIALIZE mapColumn ::
+ (Float -> Float) -> Column -> Either DataFrameException Column
+ #-}
+{-# SPECIALIZE mapColumn ::
+ (Int -> Int) -> Column -> Either DataFrameException Column
+ #-}
+
+-- zipWithColumns: binary ops
+{-# SPECIALIZE zipWithColumns ::
+ (Double -> Double -> Double) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Float -> Float -> Float) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int16 -> Int16 -> Int16) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int32 -> Int32 -> Int32) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int64 -> Int64 -> Int64) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+-- Bool-returning binary comparators (hot path for Expr Bool used in
+-- DecisionTree splits)
+{-# SPECIALIZE zipWithColumns ::
+ (Double -> Double -> Bool) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Float -> Float -> Bool) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Int -> Int -> Bool) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+{-# SPECIALIZE zipWithColumns ::
+ (Bool -> Bool -> Bool) ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+ #-}
+
+-- Bool-mapping unary ops (e.g. 'not')
+{-# SPECIALIZE mapColumn ::
+ (Bool -> Bool) -> Column -> Either DataFrameException Column
+ #-}
+
+-------------------------------------------------------------------------------
+-- Value: the unified result type
+-------------------------------------------------------------------------------
+
+{- | The result of interpreting an expression. Keeps literals as scalars
+until the point where a concrete column is needed, avoiding premature
+broadcast allocations.
+-}
+data Value a where
+ -- | A single value, not yet broadcast to any length.
+ Scalar :: (Columnable a) => !a -> Value a
+ {- | A flat column (one element per row in the flat case, or one
+ element per group after aggregation).
+ -}
+ Flat :: (Columnable a) => !Column -> Value a
+ {- | A grouped column: one 'Column' slice per group. Only produced
+ when interpreting inside a 'GroupCtx'.
+ -}
+ Group :: (Columnable a) => !(V.Vector Column) -> Value a
+
+instance (Show a) => Show (Value a) where
+ show (Scalar v) = show v
+ show (Flat v) = show v
+ show (Group v) = show v
+
+-- | The interpretation context.
+data Ctx
+ = FlatCtx DataFrame
+ | GroupCtx GroupedDataFrame
+
+-------------------------------------------------------------------------------
+-- Materialisation
+-------------------------------------------------------------------------------
+
+{- | Force a 'Value' into a flat 'Column' of the given length. Scalars
+are broadcast; flat columns are returned as-is.
+-}
+materialize :: forall a. (Columnable a) => Int -> Value a -> Column
+materialize n (Scalar v) = broadcastScalar @a n v
+materialize _ (Flat c) = c
+materialize _ (Group _) =
+ error "materialize: cannot flatten a grouped value to a single column"
+
+{- | Replicate a scalar to a column of length @n@, choosing the most
+efficient representation.
+-}
+broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column
+broadcastScalar n v = case sUnbox @a of
+ STrue -> fromUnboxedVector (VU.replicate n v)
+ SFalse -> fromVector (V.replicate n v)
+
+-------------------------------------------------------------------------------
+-- Lifting: the core combinators
+-------------------------------------------------------------------------------
+
+-- | Apply a pure function to a 'Value'.
+liftValue ::
+ (Columnable b, Columnable a) =>
+ (b -> a) -> Value b -> Either DataFrameException (Value a)
+liftValue f (Scalar v) = Right (Scalar (f v))
+liftValue f (Flat col) = Flat <$> mapColumn f col
+liftValue f (Group gs) = Group <$> V.mapM (mapColumn f) gs
+{-# INLINEABLE liftValue #-}
+
+{- | Apply a binary function to two 'Value's. When one side is a 'Scalar' the
+operation degenerates to 'liftValue', recovering the old @Binary op (Lit l) right@
+special cases without explicit pattern matches.
+-}
+liftValue2 ::
+ (Columnable c, Columnable b, Columnable a) =>
+ (c -> b -> a) ->
+ Value c ->
+ Value b ->
+ Either DataFrameException (Value a)
+liftValue2 f (Scalar l) (Scalar r) = Right (Scalar (f l r))
+liftValue2 f (Scalar l) v = liftValue (f l) v
+liftValue2 f v (Scalar r) = liftValue (`f` r) v
+liftValue2 f (Flat l) (Flat r) = Flat <$> zipWithColumns f l r
+liftValue2 f (Group ls) (Group rs)
+ | V.length ls == V.length rs =
+ Group <$> V.zipWithM (zipWithColumns f) ls rs
+liftValue2 _ (Flat _) (Group _) =
+ Left $ AggregatedAndNonAggregatedException "aggregated" "non-aggregated"
+liftValue2 _ (Group _) (Flat _) =
+ Left $ AggregatedAndNonAggregatedException "non-aggregated" "aggregated"
+liftValue2 _ (Group _) (Group _) =
+ Left $ InternalException "Group count mismatch in binary operation"
+{-# INLINEABLE liftValue2 #-}
+
+-- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.
+branchValue ::
+ forall a.
+ (Columnable a) =>
+ Value Bool ->
+ Value a ->
+ Value a ->
+ Either DataFrameException (Value a)
+branchValue (Scalar True) l _ = Right l
+branchValue (Scalar False) _ r = Right r
+branchValue cond (Scalar l) (Scalar r) =
+ liftValue (\c -> if c then l else r) cond
+branchValue cond (Scalar l) r =
+ liftValue2 (\c rv -> if c then l else rv) cond r
+branchValue cond l (Scalar r) =
+ liftValue2 (\c lv -> if c then lv else r) cond l
+branchValue (Flat cc) (Flat lc) (Flat rc) =
+ Flat <$> branchColumn @a cc lc rc
+branchValue (Group cgs) (Group lgs) (Group rgs)
+ | V.length cgs == V.length lgs
+ && V.length lgs == V.length rgs =
+ Group
+ <$> V.generateM
+ (V.length cgs)
+ ( \i ->
+ branchColumn @a (cgs V.! i) (lgs V.! i) (rgs V.! i)
+ )
+branchValue _ _ _ =
+ Left $
+ AggregatedAndNonAggregatedException
+ "if-then-else branches"
+ "mismatched shapes"
+{-# INLINEABLE branchValue #-}
+
+{- | Low-level column branch: given a boolean column and two same-typed
+columns, produce the element-wise selection.
+-}
+branchColumn ::
+ forall a.
+ (Columnable a) =>
+ Column ->
+ Column ->
+ Column ->
+ Either DataFrameException Column
+branchColumn cc lc rc = do
+ cs <- toVector @Bool @V.Vector cc
+ ls <- toVector @a @V.Vector lc
+ rs <- toVector @a @V.Vector rc
+ pure $
+ fromVector @a $
+ V.zipWith3 (\c l r -> if c then l else r) cs ls rs
+
+-------------------------------------------------------------------------------
+-- Error enrichment
+-------------------------------------------------------------------------------
+
+{- | Wrap an interpretation step so that any 'TypeMismatchException' gets
+annotated with the expression that was being evaluated.
+-}
+addContext ::
+ (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b
+addContext expr = first (enrichError (show expr))
+
+enrichError :: String -> DataFrameException -> DataFrameException
+enrichError loc (TypeMismatchException ctx) =
+ TypeMismatchException
+ ctx
+ { callingFunctionName =
+ callingFunctionName ctx <|+> Just "eval"
+ , errorColumnName =
+ errorColumnName ctx <|+> Just loc
+ }
+ where
+ Nothing <|+> b = b
+ a <|+> _ = a
+enrichError _ e = e
+
+-------------------------------------------------------------------------------
+-- Group slicing
+-------------------------------------------------------------------------------
+
+{- | Given a flat column and grouping metadata, produce one 'Column' per
+group. Each result column is an O(1) slice into a sorted copy of the
+input — the sort happens once, not per-group.
+-}
+sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column
+sliceGroups col os indices = case col of
+ PackedText _ _ -> sliceGroups (materializePacked col) os indices
+ BoxedColumn bm vec ->
+ let !sorted =
+ V.generate
+ (VU.length indices)
+ ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`))
+ in V.generate nGroups $ \i ->
+ BoxedColumn
+ (fmap (bitmapSlice (start i) (len i)) bm)
+ (V.unsafeSlice (start i) (len i) sorted)
+ UnboxedColumn bm vec ->
+ let !sorted = VU.unsafeBackpermute vec indices
+ in V.generate nGroups $ \i ->
+ UnboxedColumn
+ (fmap (bitmapSlice (start i) (len i)) bm)
+ (VU.unsafeSlice (start i) (len i) sorted)
+ where
+ !nGroups = VU.length os - 1
+ start i = os `VU.unsafeIndex` i
+ len i = os `VU.unsafeIndex` (i + 1) - start i
+{-# INLINE sliceGroups #-}
+
+numGroups :: GroupedDataFrame -> Int
+numGroups gdf = VU.length (offsets gdf) - 1
+
+-- | Build the inverse of a permutation vector.
+invertPermutation :: VU.Vector Int -> VU.Vector Int
+invertPermutation perm = VU.create $ do
+ let !n = VU.length perm
+ inv <- VUM.new n
+ VU.imapM_ (flip (VUM.unsafeWrite inv)) perm
+ return inv
+{-# INLINE invertPermutation #-}
+
+-------------------------------------------------------------------------------
+-- promoteColumnWith: unified numeric / text coercion for CastWith
+-------------------------------------------------------------------------------
+
+{- | Coerce a column to type @a@, then apply @onResult@ to each element; the handler
+selects the mode (like @cast@, @castWithDefault@, or @castEither@). Handles Double/
+Float/Int coercion and 'reads'-parses Text; other mismatches return 'Left'.
+-}
+promoteColumnWith ::
+ forall a b.
+ (Columnable a, Columnable b, Read a) =>
+ (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 Nothing (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
+ UnboxedColumn (Just bm) (v :: VU.Vector c) ->
+ case sFloating @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Double))
+ else onResult (Left "null")
+ )
+ SFalse -> case sIntegral @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Double))
+ else onResult (Left "null")
+ )
+ SFalse -> castMismatch @c @b
+ BoxedColumn _ _ -> tryParseWith @Double onResult col
+ PackedText _ _ -> promoteToDoubleWith onResult (materializePacked col)
+
+promoteToFloatWith ::
+ forall b.
+ (Columnable b) =>
+ (Either String Float -> b) -> Column -> Either DataFrameException Column
+promoteToFloatWith onResult col = case col of
+ UnboxedColumn Nothing (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
+ UnboxedColumn (Just bm) (v :: VU.Vector c) ->
+ case sFloating @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Float))
+ else onResult (Left "null")
+ )
+ SFalse -> case sIntegral @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Float))
+ else onResult (Left "null")
+ )
+ SFalse -> castMismatch @c @b
+ BoxedColumn _ _ -> tryParseWith @Float onResult col
+ PackedText _ _ -> promoteToFloatWith onResult (materializePacked col)
+
+promoteToIntWith ::
+ forall b.
+ (Columnable b) =>
+ (Either String Int -> b) -> Column -> Either DataFrameException Column
+promoteToIntWith onResult col = case col of
+ UnboxedColumn Nothing (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
+ UnboxedColumn (Just bm) (v :: VU.Vector c) ->
+ case sFloating @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (round (realToFrac (VU.unsafeIndex v i) :: Double)))
+ else onResult (Left "null")
+ )
+ SFalse -> case sIntegral @c of
+ STrue ->
+ Right $
+ fromVector @b
+ ( V.generate (VU.length v) $ \i ->
+ if bitmapTestBit bm i
+ then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Int))
+ else onResult (Left "null")
+ )
+ SFalse -> castMismatch @c @b
+ BoxedColumn _ _ -> tryParseWith @Int onResult col
+ PackedText _ _ -> promoteToIntWith onResult (materializePacked 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)
+ _ -> case reads (show s) of
+ [(x, "")] -> f (Right x)
+ _ -> f (Left s)
+
+tryParseWith ::
+ forall a b.
+ (Columnable a, Columnable b, Read a) =>
+ (Either String a -> b) -> Column -> Either DataFrameException Column
+tryParseWith onResult col = case col of
+ PackedText _ _ -> tryParseWith onResult (materializePacked col)
+ BoxedColumn bm (v :: V.Vector c) ->
+ case testEquality (typeRep @c) (typeRep @String) of
+ Just Refl -> case bm of
+ Nothing -> Right $ fromVector @b $ V.map (parseWith onResult) v
+ Just bitmap ->
+ Right $
+ fromVector @b $
+ V.imap
+ ( \i x ->
+ if bitmapTestBit bitmap i then parseWith onResult x else onResult (Left "null")
+ )
+ v
+ Nothing ->
+ case testEquality (typeRep @c) (typeRep @T.Text) of
+ Just Refl -> case bm of
+ Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . T.unpack) v
+ Just bitmap ->
+ Right $
+ fromVector @b $
+ V.imap
+ ( \i x ->
+ if bitmapTestBit bitmap i
+ then parseWith onResult (T.unpack x)
+ else onResult (Left "null")
+ )
+ v
+ Nothing -> castMismatch @c @b
+ UnboxedColumn bm (v :: VU.Vector c) -> case bm of
+ Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . show) (V.convert v)
+ Just bitmap ->
+ Right $
+ fromVector @b $
+ V.imap
+ ( \i x ->
+ if bitmapTestBit bitmap i
+ then parseWith onResult (show x)
+ else onResult (Left "null")
+ )
+ (V.convert v)
+
+{- | When output type @b@ is @Maybe c@ (or @Maybe (Maybe c)@) and the column stores
+plain @c@, wrap each element in 'Just' (the double-Maybe case collapses to a single
+@Maybe c@). 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 Nothing (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 ->
+ case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of
+ Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped
+ Nothing -> Nothing
+ BoxedColumn Nothing (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
+ _ -> 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
+-------------------------------------------------------------------------------
+
+{- | Evaluate an expression in a given context, producing a 'Value'.
+This single function replaces both the old @interpret@ (flat) and
+@interpretAggregation@ (grouped) code paths.
+-}
+eval ::
+ forall a.
+ (Columnable a) =>
+ Ctx -> Expr a -> Either DataFrameException (Value a)
+eval _ (Lit v) = Right (Scalar v)
+eval (FlatCtx df) (Col name) =
+ case getColumn name df of
+ Nothing ->
+ Left $ ColumnsNotFoundException [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 $
+ ColumnsNotFoundException
+ [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 ()
+ )
+eval (FlatCtx df) (CastWith name _tag onResult) =
+ case getColumn name df of
+ Nothing ->
+ Left $
+ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
+ Just c -> Flat <$> promoteColumnWith onResult c
+eval (GroupCtx gdf) (CastWith name _tag onResult) =
+ case getColumn name (fullDataframe gdf) of
+ Nothing ->
+ Left $
+ ColumnsNotFoundException
+ [name]
+ ""
+ (M.keys $ columnIndices $ fullDataframe gdf)
+ Just c -> do
+ promoted <- promoteColumnWith onResult c
+ Right $ Group (sliceGroups promoted (offsets gdf) (valueIndices gdf))
+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
+eval ctx expr@(Unary op (inner :: Expr b)) = addContext expr $ do
+ v <- eval @b ctx inner
+ liftValue (unaryFn op) v
+eval ctx expr@(Binary op (left :: Expr c) (right :: Expr b)) =
+ addContext expr $ do
+ l <- eval @c ctx left
+ r <- eval @b ctx right
+ liftValue2 (binaryFn op) l r
+eval ctx expr@(If cond l r) = addContext expr $ do
+ c <- eval @Bool ctx cond
+ lv <- eval @a ctx l
+ rv <- eval @a ctx r
+ branchValue c lv rv
+eval (FlatCtx df) expr@(Over keys inner) = addContext expr $ do
+ let gdf = G.groupBy keys df
+ v <- eval (GroupCtx gdf) inner
+ case v of
+ Scalar s ->
+ Right (Scalar s)
+ Flat groupCol ->
+ Right (Flat (atIndicesStable (rowToGroup gdf) groupCol))
+ Group groupCols -> do
+ sorted <- V.fold1M' concatColumns groupCols
+ let inv = invertPermutation (valueIndices gdf)
+ Right (Flat (atIndicesStable inv sorted))
+eval (GroupCtx _) expr@(Over _ _) =
+ addContext expr $
+ Left
+ ( InternalException
+ "Over (window function) is not supported inside a grouped context"
+ )
+-- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.
+-- Avoids the O(n) backpermute in sliceGroups by folding directly over
+-- permuted indices. Only matches when inner is exactly (Col name).
+
+eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col name :: Expr b)) =
+ addContext expr $
+ case getColumn name (fullDataframe gdf) of
+ Nothing ->
+ Left $
+ ColumnsNotFoundException
+ [name]
+ ""
+ (M.keys $ columnIndices $ fullDataframe gdf)
+ Just col ->
+ Flat <$> foldLinearGroups @b @a f seed col (rowToGroup gdf) (numGroups gdf)
+-- Fast path: FoldAgg (seedless) on a bare Col in GroupCtx.
+
+eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name :: Expr b)) =
+ addContext expr $
+ case testEquality (typeRep @a) (typeRep @b) of
+ Nothing ->
+ Left $
+ InternalException
+ "Type mismatch in seedless fold: \
+ \accumulator and element types must match"
+ Just Refl ->
+ case getColumn name (fullDataframe gdf) of
+ Nothing ->
+ Left $
+ ColumnsNotFoundException
+ [name]
+ ""
+ (M.keys $ columnIndices $ fullDataframe gdf)
+ Just col ->
+ Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf)
+-- Fast path: MergeAgg on a bare Col in GroupCtx.
+
+eval
+ (GroupCtx gdf)
+ expr@( Agg
+ (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))
+ (Col name :: Expr b)
+ ) =
+ addContext expr $
+ case getColumn name (fullDataframe gdf) of
+ Nothing ->
+ Left $
+ ColumnsNotFoundException
+ [name]
+ ""
+ (M.keys $ columnIndices $ fullDataframe gdf)
+ Just col ->
+ Flat
+ <$> ( foldLinearGroups @b step seed col (rowToGroup gdf) (numGroups gdf)
+ >>= mapColumn finalize
+ )
+eval ctx expr@(Agg (CollectAgg _ (f :: v b -> a)) inner) =
+ addContext expr $ do
+ v <- eval @b ctx inner
+ case v of
+ Scalar _ ->
+ Left $
+ InternalException
+ "Cannot apply a collection aggregation to a scalar"
+ Flat col ->
+ Scalar <$> applyCollect @v @b @a f col
+ Group gs ->
+ Flat . fromVector
+ <$> V.mapM (applyCollect @v @b @a f) gs
+eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =
+ addContext expr $ do
+ v <- eval @b ctx inner
+ case v of
+ Scalar x -> Right (broadcastFold ctx seed f x)
+ Flat col ->
+ Scalar <$> foldlColumn @b @a f seed col
+ Group gs ->
+ Flat . fromVector
+ <$> V.mapM (foldlColumn @b @a f seed) gs
+eval
+ ctx
+ expr@( Agg
+ (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))
+ (inner :: Expr b)
+ ) =
+ addContext expr $ do
+ v <- eval @b ctx inner
+ case v of
+ Scalar x -> case broadcastFold ctx seed step x of
+ Scalar acc -> Right (Scalar (finalize acc))
+ Flat col -> Flat <$> mapColumn @acc @a finalize col
+ Group _ ->
+ Left
+ ( InternalException
+ "broadcastFold unexpectedly produced a Group value"
+ )
+ Flat col ->
+ Scalar . finalize <$> foldlColumn @b step seed col
+ Group gs ->
+ Flat . fromVector
+ <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs
+eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =
+ addContext expr $
+ case testEquality (typeRep @a) (typeRep @b) of
+ Nothing ->
+ Left $
+ InternalException
+ "Type mismatch in seedless fold: \
+ \accumulator and element types must match"
+ Just Refl -> do
+ v <- eval @b ctx inner
+ case v of
+ Scalar _ ->
+ Left $
+ InternalException
+ "fold1 requires at least one element"
+ Flat col ->
+ Scalar <$> foldl1Column @a f col
+ Group gs ->
+ Flat . fromVector
+ <$> V.mapM (foldl1Column @a f) gs
+
+broadcastFold ::
+ forall acc b.
+ (Columnable acc) =>
+ Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
+broadcastFold (FlatCtx df) seed step x =
+ let n = fst (dataframeDimensions df)
+ in Scalar (iterateStep n step seed x)
+broadcastFold (GroupCtx gdf) seed step x =
+ let offs = offsets gdf
+ ng = VU.length offs - 1
+ results =
+ V.generate ng $ \i ->
+ let sz = offs VU.! (i + 1) - offs VU.! i
+ in iterateStep sz step seed x
+ in Flat (fromVector results)
+
+iterateStep :: Int -> (acc -> b -> acc) -> acc -> b -> acc
+iterateStep n step = go n
+ where
+ go 0 !acc _ = acc
+ go k !acc x = go (k - 1) (step acc x) x
+
+{- | Apply a 'CollectAgg' function to a single column, extracting the
+appropriate vector type and applying the aggregation function.
+-}
+applyCollect ::
+ forall v b a.
+ (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>
+ (v b -> a) -> Column -> Either DataFrameException a
+applyCollect f col = f <$> toVector @b @v col
+
+{- | Result of interpreting an expression in a grouped context.
+Retained for backward compatibility with 'aggregate' and friends.
+-}
+data AggregationResult a
+ = UnAggregated Column
+ | Aggregated (TypedColumn a)
+
+{- | Interpret an expression against a flat 'DataFrame', producing a typed column.
+Calls 'eval' then 'materialize'; 'Lit' values are broadcast here at the boundary
+rather than eagerly.
+-}
+interpret ::
+ forall a.
+ (Columnable a) =>
+ DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
+interpret df expr = do
+ v <- eval (FlatCtx df) expr
+ pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v
+
+{- | Interpret an expression against a 'GroupedDataFrame',
+distinguishing aggregated results from bare column references.
+Internally calls 'eval'.
+-}
+interpretAggregation ::
+ forall a.
+ (Columnable a) =>
+ GroupedDataFrame ->
+ Expr a ->
+ Either DataFrameException (AggregationResult a)
+interpretAggregation gdf expr = do
+ v <- eval (GroupCtx gdf) expr
+ case v of
+ Scalar a ->
+ Right $
+ Aggregated $
+ TColumn $
+ broadcastScalar @a (numGroups gdf) a
+ Flat col ->
+ Right $ Aggregated $ TColumn col
+ Group _ ->
+ Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty
diff --git a/src-internal/DataFrame/Internal/Nullable.hs b/src-internal/DataFrame/Internal/Nullable.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Nullable.hs
@@ -0,0 +1,467 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE UndecidableSuperClasses #-}
+
+{- | Nullable-aware arithmetic and comparison operators ('.+', '.==', …) that work
+transparently across nullable (@Maybe a@) and non-nullable (@a@) operands.
+Functional dependencies infer the result type without annotations.
+
+@
+-- 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,
+ widenCmpOp,
+ 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
+
+{- | Arithmetic binary operations that work over nullable and non-nullable operand
+types. The functional dependency @a b -> c@ infers the result; the 'OVERLAPPABLE'
+non-nullable instance yields to the specific @(Maybe a, Maybe a)@ one.
+-}
+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
+
+{- | Comparison binary operations over nullable and non-nullable operands. No
+functional dependency on @e@; overlapping/overlappable instance pragmas pick 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' (applied
+directly when non-nullable, under 'Just' when @a = Maybe x@). 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
+result is @Maybe r@ if either operand is nullable, else @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 Int32 Float where widen1 = fromIntegral; widen2 = id
+instance NumericWidenOp Float Int32 where
+ widen1 = id
+ widen2 = fromIntegral
+instance NumericWidenOp Int32 Double where widen1 = fromIntegral; widen2 = id
+instance NumericWidenOp Double Int32 where
+ widen1 = id
+ widen2 = fromIntegral
+instance NumericWidenOp Int64 Float where widen1 = fromIntegral; widen2 = id
+instance NumericWidenOp Float Int64 where
+ widen1 = id
+ widen2 = fromIntegral
+instance NumericWidenOp Int64 Double where widen1 = fromIntegral; widen2 = id
+instance NumericWidenOp Double Int64 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)
+
+-- | Apply a comparison function after widening both operands to their common type.
+widenCmpOp ::
+ forall a b.
+ (NumericWidenOp a b) =>
+ (Promote a b -> Promote a b -> Bool) ->
+ a ->
+ b ->
+ Bool
+widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y)
+
+-- | Result type of a widening binary operator, accounting for nullable wrappers.
+type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))
+
+-- ---------------------------------------------------------------------------
+-- 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-internal/DataFrame/Internal/PackedText.hs b/src-internal/DataFrame/Internal/PackedText.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/PackedText.hs
@@ -0,0 +1,149 @@
+{-# LANGUAGE BangPatterns #-}
+
+{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares one
+UTF-8 byte buffer across all rows of a string column, with @n+1@ row offsets, so
+no per-row 'Data.Text.Text' header is materialized until decode is demanded.
+-}
+module DataFrame.Internal.PackedText (
+ PackedTextData (..),
+ mkPackedContiguous,
+ packedGather,
+ packedTake,
+ packedRowOffsetVec,
+ packedLength,
+ packedSlice,
+ packedIndexText,
+ sliceEqBytes,
+ sliceCmpBytes,
+) where
+
+import qualified Data.Text as T
+import qualified Data.Text.Array as A
+import qualified Data.Vector.Unboxed as VU
+
+import Data.Ord (comparing)
+import Data.Text.Internal (Text (Text))
+import DataFrame.Internal.Utf8 (isValidUtf8Slice, lenientDecodeSlice)
+
+{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes
+@[offsets!r, offsets!(r+1))@); validity lives in the column's bitmap. @ptSel@ is
+an optional selection layer letting a gather/join/sort result share the buffer.
+-}
+data PackedTextData = PackedTextData
+ { ptBytes :: {-# UNPACK #-} !A.Array
+ , ptOffsets :: {-# UNPACK #-} !(VU.Vector Int)
+ , ptSel :: !(Maybe (VU.Vector Int))
+ }
+
+-- | Build a contiguous packed payload (no selection): the freeze-path shape.
+mkPackedContiguous :: A.Array -> VU.Vector Int -> PackedTextData
+mkPackedContiguous arr offs = PackedTextData arr offs Nothing
+{-# INLINE mkPackedContiguous #-}
+
+{- | Reindex a packed payload by a selection vector, sharing the byte buffer;
+logical row @i@ becomes base row @indices!i@. A negative or out-of-range index
+decodes to the empty slice. Composes with an existing selection.
+-}
+packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData
+packedGather indices (PackedTextData arr offs msel) =
+ let !base = VU.length offs - 1
+ clamp r = if r >= 0 && r < base then r else -1
+ sel' = case msel of
+ Nothing -> VU.map clamp indices
+ Just s ->
+ VU.map
+ (\i -> if i >= 0 && i < VU.length s then clamp (VU.unsafeIndex s i) else -1)
+ indices
+ in PackedTextData arr offs (Just sel')
+{-# INLINE packedGather #-}
+
+{- | Take the first @k@ logical rows, sharing the byte buffer via a capped
+selection layer. O(k), no byte copy or decode — cheap @take@/display on a
+large packed column.
+-}
+packedTake :: Int -> PackedTextData -> PackedTextData
+packedTake k (PackedTextData arr offs msel) =
+ let !base = VU.length offs - 1
+ !k' = max 0 k
+ in case msel of
+ Just s -> PackedTextData arr offs (Just (VU.take k' s))
+ Nothing -> PackedTextData arr offs (Just (VU.enumFromN 0 (min k' base)))
+{-# INLINE packedTake #-}
+
+-- | Map a logical row index to its base row, honoring any selection layer.
+baseRow :: PackedTextData -> Int -> Int
+baseRow (PackedTextData _ _ Nothing) i = i
+baseRow (PackedTextData _ _ (Just sel)) i = VU.unsafeIndex sel i
+{-# INLINE baseRow #-}
+
+-- | Row count: @length sel@ when selected, else @length offsets - 1@.
+packedLength :: PackedTextData -> Int
+packedLength (PackedTextData _ offs Nothing) = VU.length offs - 1
+packedLength (PackedTextData _ _ (Just sel)) = VU.length sel
+{-# INLINE packedLength #-}
+
+-- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.
+packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int)
+packedSlice p@(PackedTextData arr offs _) i =
+ let !r = baseRow p i
+ in if r < 0
+ then (arr, 0, 0)
+ else
+ let o = VU.unsafeIndex offs r in (arr, o, VU.unsafeIndex offs (r + 1) - o)
+{-# INLINE packedSlice #-}
+
+{- | The shared buffer + contiguous @n+1@ offsets when the payload is the
+unselected base; a selected (gathered) payload returns 'Nothing' (its rows are
+non-contiguous). Lets the boxed-Text fallback take the fast contiguous path.
+-}
+packedRowOffsetVec :: PackedTextData -> Maybe (A.Array, VU.Vector Int)
+packedRowOffsetVec (PackedTextData arr offs Nothing) = Just (arr, offs)
+packedRowOffsetVec _ = Nothing
+{-# INLINE packedRowOffsetVec #-}
+
+{- | On-demand single 'Data.Text.Text' for row @i@, using the same
+validate-or-lenient decode as 'sliceTextVector' so output is bit-identical.
+-}
+packedIndexText :: PackedTextData -> Int -> T.Text
+packedIndexText p i =
+ let (arr, o, l) = packedSlice p i
+ in decodeField arr o l
+{-# INLINE packedIndexText #-}
+
+-- Decode one field exactly as 'sliceTextVector' does per row.
+decodeField :: A.Array -> Int -> Int -> T.Text
+decodeField arr o l
+ | l == 0 = T.empty
+ | isValidUtf8Slice arr o l = Text arr o l
+ | otherwise = lenientDecodeSlice arr o l
+{-# INLINE decodeField #-}
+
+{- | Byte-wise equality of two slices. UTF-8 is injective on valid scalar
+sequences and lenient decode is deterministic, so this agrees with
+@Text@'s '==' on the decoded values.
+-}
+sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool
+sliceEqBytes a ao al b bo bl
+ | al /= bl = False
+ | otherwise = go 0
+ where
+ go !k
+ | k >= al = True
+ | A.unsafeIndex a (ao + k) == A.unsafeIndex b (bo + k) = go (k + 1)
+ | otherwise = False
+{-# INLINE sliceEqBytes #-}
+
+{- | Unsigned byte-lexicographic comparison (memcmp semantics). For
+well-formed UTF-8 this matches 'Data.Text.compare' exactly, since UTF-8
+byte order equals codepoint order for all valid scalars.
+-}
+sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering
+sliceCmpBytes a ao al b bo bl = go 0
+ where
+ !m = min al bl
+ go !k
+ | k >= m = compare al bl
+ | otherwise = case comparing id (A.unsafeIndex a (ao + k)) (A.unsafeIndex b (bo + k)) of
+ EQ -> go (k + 1)
+ r -> r
+{-# INLINE sliceCmpBytes #-}
diff --git a/src-internal/DataFrame/Internal/ParRadixSort.hs b/src-internal/DataFrame/Internal/ParRadixSort.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/ParRadixSort.hs
@@ -0,0 +1,272 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+{- | Parallel stable sort of row indices by ascending unsigned order of a per-row
+'Int' hash, used by the join build side. A counting sort buckets rows into
+key-ordered partitions that workers LSD-radix-sort in parallel, with no merge step.
+-}
+module DataFrame.Internal.ParRadixSort (
+ parSortByHash,
+ parSortThreshold,
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import Control.Monad (forM_, when)
+import Data.Bits (countLeadingZeros, unsafeShiftR, (.&.))
+import Data.IORef (atomicModifyIORef', newIORef)
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import Data.Word (Word64)
+import DataFrame.Internal.RadixRank (sortKey)
+import System.IO.Unsafe (unsafePerformIO)
+
+{- | Below this many rows the partition/fork overhead is not worth it; the
+caller's sequential LSD radix path is used instead.
+-}
+parSortThreshold :: Int
+parSortThreshold = 500000
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+{- | Top-bits partition index of a hash: the high @64 - shift@ bits of its
+unsigned 'sortKey'. Ascending partition order equals ascending key order.
+-}
+partIx :: Int -> Int -> Int
+partIx shift h = fromIntegral ((fromIntegral (sortKey h) :: Word64) `unsafeShiftR` shift)
+{-# INLINE partIx #-}
+
+-- | Number of partitions: a power of two, at least @4 * caps@, floored at 256.
+numPartitionsFor :: Int -> Int
+numPartitionsFor caps = go 1
+ where
+ target = max 256 (4 * caps)
+ go p
+ | p >= target = p
+ | otherwise = go (p * 2)
+
+-- | @floor (log2 x)@ for a power-of-two @x@.
+intLog2 :: Int -> Int
+intLog2 x = 63 - countLeadingZeros x
+{-# INLINE intLog2 #-}
+
+{- | Parallel stable sort of @[0, n)@ by ascending unsigned hash order. See the
+module header for the ordering contract.
+-}
+parSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)
+parSortByHash n hashes
+ | n <= 1 =
+ (hashes, VU.enumFromN 0 n)
+ | n < parSortThreshold || capabilities <= 1 =
+ seqSortByHash n hashes
+ | otherwise = unsafePerformIO (parSortByHashIO n hashes)
+{-# NOINLINE parSortByHash #-}
+
+-------------------------------------------------------------------------------
+-- Sequential LSD radix sort (also the per-partition worker kernel)
+-------------------------------------------------------------------------------
+
+{- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash, 8
+bits per pass over the full 64-bit key. Returns @(sortedHashes, sortedIndices)@.
+-}
+seqSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)
+seqSortByHash n hashes = unsafePerformIO $ do
+ keysA <- VUM.new n
+ orderA <- VUM.new n
+ let seed !i
+ | i >= n = pure ()
+ | otherwise = do
+ VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes i))
+ VUM.unsafeWrite orderA i i
+ seed (i + 1)
+ seed 0
+ keysB <- VUM.new n
+ orderB <- VUM.new n
+ radixPasses n keysA orderA keysB orderB
+ order <- VU.unsafeFreeze orderA
+ pure (VU.unsafeBackpermute hashes order, order)
+
+{- | Run all eight stable 8-bit LSD passes, ping-ponging between the two
+key/order buffer pairs so the sorted order lands back in @(keysA, orderA)@.
+@keysA[i]@ must already hold @sortKey (hash of orderA[i])@ on entry.
+-}
+radixPasses ::
+ Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ IO ()
+radixPasses n keysA orderA keysB orderB = do
+ counts <- VUM.new 256
+ let pass ::
+ Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ IO ()
+ pass !shiftBits !srcK !srcO !dstK !dstO = do
+ VUM.set counts 0
+ let count !i
+ | i >= n = pure ()
+ | otherwise = do
+ k <- VUM.unsafeRead srcK i
+ let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
+ VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)
+ count (i + 1)
+ count 0
+ let scan !b !acc
+ | b >= 256 = pure ()
+ | otherwise = do
+ c <- VUM.unsafeRead counts b
+ VUM.unsafeWrite counts b acc
+ scan (b + 1) (acc + c)
+ scan 0 0
+ let place !i
+ | i >= n = pure ()
+ | otherwise = do
+ k <- VUM.unsafeRead srcK i
+ o <- VUM.unsafeRead srcO i
+ let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
+ pos <- VUM.unsafeRead counts b
+ VUM.unsafeWrite counts b (pos + 1)
+ VUM.unsafeWrite dstK pos k
+ VUM.unsafeWrite dstO pos o
+ place (i + 1)
+ place 0
+ pass 0 keysA orderA keysB orderB
+ pass 8 keysB orderB keysA orderA
+ pass 16 keysA orderA keysB orderB
+ pass 24 keysB orderB keysA orderA
+ pass 32 keysA orderA keysB orderB
+ pass 40 keysB orderB keysA orderA
+ pass 48 keysA orderA keysB orderB
+ pass 56 keysB orderB keysA orderA
+
+-------------------------------------------------------------------------------
+-- Parallel path: counting-sort partition, then per-partition sort in parallel
+-------------------------------------------------------------------------------
+
+parSortByHashIO :: Int -> VU.Vector Int -> IO (VU.Vector Int, VU.Vector Int)
+parSortByHashIO n hashes = do
+ caps <- getNumCapabilities
+ let !p = numPartitionsFor caps
+ !shift = 64 - intLog2 p
+ (partStart, partRows) <- partitionRows n hashes p shift
+ outOrder <- VUM.new n
+ outKeys <- VUM.new n
+ sortPartitions caps p partStart partRows hashes outOrder outKeys
+ order <- VU.unsafeFreeze outOrder
+ pure (VU.unsafeBackpermute hashes order, order)
+
+{- | Bucket every row index into its top-bits partition by a counting sort.
+Returns the exclusive prefix sum @partStart@ (length @p+1@, @partStart[p] == n@)
+and the row indices laid out partition-by-partition in ascending key order.
+-}
+partitionRows ::
+ Int -> VU.Vector Int -> Int -> Int -> IO (VU.Vector Int, VU.Vector Int)
+partitionRows n hashes p shift = do
+ counts <- VUM.replicate (p + 1) (0 :: Int)
+ let countLoop !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !pp = partIx shift (VU.unsafeIndex hashes i)
+ c <- VUM.unsafeRead counts pp
+ VUM.unsafeWrite counts pp (c + 1)
+ countLoop (i + 1)
+ countLoop 0
+ partStartM <- VUM.new (p + 1)
+ let scan !k !acc
+ | k > p = pure ()
+ | otherwise = do
+ VUM.unsafeWrite partStartM k acc
+ c <- if k < p then VUM.unsafeRead counts k else pure 0
+ scan (k + 1) (acc + c)
+ scan 0 0
+ cursor <- VUM.new p
+ forM_ [0 .. p - 1] $ \k -> VUM.unsafeRead partStartM k >>= VUM.unsafeWrite cursor k
+ rowsM <- VUM.new (max 1 n)
+ let place !i
+ | i >= n = pure ()
+ | otherwise = do
+ let !pp = partIx shift (VU.unsafeIndex hashes i)
+ pos <- VUM.unsafeRead cursor pp
+ VUM.unsafeWrite rowsM pos i
+ VUM.unsafeWrite cursor pp (pos + 1)
+ place (i + 1)
+ place 0
+ partStart <- VU.unsafeFreeze partStartM
+ partRows <- VU.unsafeFreeze rowsM
+ pure (partStart, partRows)
+
+{- | Stable-sort each partition by full key, writing sorted original indices
+into @outOrder@ and their hashes into @outKeys@ at the partition's slot range.
+Forks @caps@ workers that pull partition indices off a shared atomic counter.
+Within a partition the counting sort already left rows in ascending original
+order, so the LSD radix sort's stability reproduces the global @(key, row)@
+order. Partitions below two elements are already sorted (counting sort kept
+original order) and are copied directly.
+-}
+sortPartitions ::
+ Int ->
+ Int ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ VU.Vector Int ->
+ VUM.IOVector Int ->
+ VUM.IOVector Int ->
+ IO ()
+sortPartitions caps p partStart partRows hashes outOrder outKeys = do
+ next <- newIORef 0
+ let sortOne !pp = do
+ let !s = VU.unsafeIndex partStart pp
+ !e = VU.unsafeIndex partStart (pp + 1)
+ !sz = e - s
+ when (sz > 0) $
+ if sz == 1
+ then do
+ let !r = VU.unsafeIndex partRows s
+ VUM.unsafeWrite outOrder s r
+ VUM.unsafeWrite outKeys s (VU.unsafeIndex hashes r)
+ else do
+ keysA <- VUM.new sz
+ orderA <- VUM.new sz
+ let seed !i
+ | i >= sz = pure ()
+ | otherwise = do
+ let !r = VU.unsafeIndex partRows (s + i)
+ VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes r))
+ VUM.unsafeWrite orderA i r
+ seed (i + 1)
+ seed 0
+ keysB <- VUM.new sz
+ orderB <- VUM.new sz
+ radixPasses sz keysA orderA keysB orderB
+ let emit !i
+ | i >= sz = pure ()
+ | otherwise = do
+ o <- VUM.unsafeRead orderA i
+ VUM.unsafeWrite outOrder (s + i) o
+ VUM.unsafeWrite outKeys (s + i) (VU.unsafeIndex hashes o)
+ emit (i + 1)
+ emit 0
+ worker = do
+ i <- atomicModifyIORef' next (\j -> (j + 1, j))
+ when (i < p) $ sortOne i >> worker
+ forkJoin_ (replicate caps worker)
+
+-- | Run each action on its own thread; rethrow the first failure (in order).
+forkJoin_ :: [IO ()] -> IO ()
+forkJoin_ actions = do
+ vars <- mapM spawn actions
+ results <- mapM takeMVar vars
+ mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
+ where
+ spawn act = do
+ var <- newEmptyMVar
+ _ <- forkIO (try act >>= putMVar var)
+ pure var
diff --git a/src-internal/DataFrame/Internal/Pretty.hs b/src-internal/DataFrame/Internal/Pretty.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Pretty.hs
@@ -0,0 +1,129 @@
+{- | A minimal Wadler/Leijen-style document combinator and width-aware renderer.
+A 'Doc' describes a layout abstractly; 'render' chooses where soft breaks become
+newlines to fit a target width. 'Group' lays a region flat when it fits.
+-}
+module DataFrame.Internal.Pretty (
+ Doc,
+ text,
+ line,
+ hardline,
+ nest,
+ group,
+ (<+>),
+ hcat,
+ punctuate,
+ parens,
+ parensWhenBroken,
+ defaultWidth,
+ render,
+) where
+
+data Doc
+ = Empty
+ | Text String
+ | Line
+ | Cat Doc Doc
+ | Nest Int Doc
+ | Group Doc
+ | Hard
+ | Alt Doc Doc
+
+instance Semigroup Doc where
+ (<>) = Cat
+
+instance Monoid Doc where
+ mempty = Empty
+
+-- | A literal chunk of text. Must not contain newlines (use 'line'/'hardline').
+text :: String -> Doc
+text = Text
+
+{- | A soft break: a single space when its enclosing 'group' fits the width,
+otherwise a newline + current indentation.
+-}
+line :: Doc
+line = Line
+
+-- | A hard break that never flattens; any enclosing 'group' is forced to break.
+hardline :: Doc
+hardline = Hard
+
+-- | Add @k@ spaces to the indentation applied at line breaks inside @d@.
+nest :: Int -> Doc -> Doc
+nest = Nest
+
+-- | Lay the document out flat if it fits the remaining width, broken otherwise.
+group :: Doc -> Doc
+group = Group
+
+-- | Concatenate two documents separated by a single space.
+(<+>) :: Doc -> Doc -> Doc
+x <+> y = x <> Text " " <> y
+
+infixr 6 <+>
+
+hcat :: [Doc] -> Doc
+hcat = mconcat
+
+-- | Append @sep@ after every element but the last.
+punctuate :: Doc -> [Doc] -> [Doc]
+punctuate _ [] = []
+punctuate _ [d] = [d]
+punctuate sep (d : ds) = (d <> sep) : punctuate sep ds
+
+parens :: Doc -> Doc
+parens d = Text "(" <> d <> Text ")"
+
+{- | Render @d@ bare when it fits flat on the current line, wrapped in parens when
+it must break across lines. Keeps operator grouping unambiguous once a
+sub-expression wraps, without parenthesis noise on one-line expressions.
+-}
+parensWhenBroken :: Doc -> Doc
+parensWhenBroken d = Group (Alt d (parens d))
+
+defaultWidth :: Int
+defaultWidth = 80
+
+data Mode = Flat | Break
+
+-- | Render a document, breaking soft lines so output fits @width@ columns.
+render :: Int -> Doc -> String
+render width doc = layout 0 [(0, Break, doc)]
+ where
+ layout :: Int -> [(Int, Mode, Doc)] -> String
+ layout _ [] = ""
+ layout col ((i, m, d) : rest) = case d of
+ Empty -> layout col rest
+ Text s -> s ++ layout (col + length s) rest
+ Cat x y -> layout col ((i, m, x) : (i, m, y) : rest)
+ Nest j x -> layout col ((i + j, m, x) : rest)
+ Line -> case m of
+ Flat -> ' ' : layout (col + 1) rest
+ Break -> '\n' : replicate i ' ' ++ layout i rest
+ Hard -> '\n' : replicate i ' ' ++ layout i rest
+ Group x ->
+ if fits (width - col) ((i, Flat, x) : rest)
+ then layout col ((i, Flat, x) : rest)
+ else layout col ((i, Break, x) : rest)
+ Alt flat broken -> case m of
+ Flat -> layout col ((i, Flat, flat) : rest)
+ Break -> layout col ((i, Break, broken) : rest)
+
+ fits :: Int -> [(Int, Mode, Doc)] -> Bool
+ fits w _ | w < 0 = False
+ fits _ [] = True
+ fits w ((i, m, d) : rest) = case d of
+ Empty -> fits w rest
+ Text s -> fits (w - length s) rest
+ Cat x y -> fits w ((i, m, x) : (i, m, y) : rest)
+ Nest j x -> fits w ((i + j, m, x) : rest)
+ Line -> case m of
+ Flat -> fits (w - 1) rest
+ Break -> True
+ Hard -> case m of
+ Flat -> False
+ Break -> True
+ Group x -> fits w ((i, Flat, x) : rest)
+ Alt flat broken -> case m of
+ Flat -> fits w ((i, Flat, flat) : rest)
+ Break -> fits w ((i, Break, broken) : rest)
diff --git a/src-internal/DataFrame/Internal/RadixRank.hs b/src-internal/DataFrame/Internal/RadixRank.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/RadixRank.hs
@@ -0,0 +1,101 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+{- | Stable rank of a set of group representatives by ascending unsigned hash
+order. Shared by the sequential and parallel group-by canonical-ordering steps
+so they stay bit-for-bit identical. @O(ng)@ stable LSD radix sort.
+-}
+module DataFrame.Internal.RadixRank (
+ rankByHash,
+ sortKey,
+) where
+
+import Control.Monad (when)
+import Control.Monad.Primitive (PrimMonad)
+import Data.Bits (unsafeShiftR, (.&.))
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import Data.Word (Word64)
+
+{- | Unsigned sort key of a hash: ascending 'Word64' order of @sortKey h@ equals
+ascending signed-'Int' order of @h@. Reinterpreted to 'Int' for the byte-wise
+radix passes (the byte mask makes the sign extension irrelevant).
+-}
+sortKey :: Int -> Int
+sortKey h = fromIntegral (fromIntegral h + 0x8000000000000000 :: Word64)
+{-# INLINE sortKey #-}
+
+-- | See the module header. @readHash@ supplies the hash of local group @gid@.
+rankByHash ::
+ forall m. (PrimMonad m) => (Int -> m Int) -> Int -> m (VU.Vector Int)
+rankByHash readHash ng = do
+ rankM <- VUM.new (max 1 ng)
+ if ng <= 1
+ then when (ng == 1) (VUM.unsafeWrite rankM 0 0)
+ else do
+ keysA <- VUM.new ng
+ orderA <- VUM.new ng
+ let seed !i
+ | i >= ng = pure ()
+ | otherwise = do
+ h <- readHash i
+ VUM.unsafeWrite keysA i (sortKey h)
+ VUM.unsafeWrite orderA i i
+ seed (i + 1)
+ seed 0
+ keysB <- VUM.new ng
+ orderB <- VUM.new ng
+ counts <- VUM.new 256
+ let pass ::
+ Int ->
+ VUM.MVector (VUM.PrimState m) Int ->
+ VUM.MVector (VUM.PrimState m) Int ->
+ VUM.MVector (VUM.PrimState m) Int ->
+ VUM.MVector (VUM.PrimState m) Int ->
+ m ()
+ pass !shiftBits !srcK !srcO !dstK !dstO = do
+ VUM.set counts 0
+ let count !i
+ | i >= ng = pure ()
+ | otherwise = do
+ k <- VUM.unsafeRead srcK i
+ let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
+ VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)
+ count (i + 1)
+ count 0
+ let scan !b !acc
+ | b >= 256 = pure ()
+ | otherwise = do
+ c <- VUM.unsafeRead counts b
+ VUM.unsafeWrite counts b acc
+ scan (b + 1) (acc + c)
+ scan 0 0
+ let place !i
+ | i >= ng = pure ()
+ | otherwise = do
+ k <- VUM.unsafeRead srcK i
+ o <- VUM.unsafeRead srcO i
+ let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
+ pos <- VUM.unsafeRead counts b
+ VUM.unsafeWrite counts b (pos + 1)
+ VUM.unsafeWrite dstK pos k
+ VUM.unsafeWrite dstO pos o
+ place (i + 1)
+ place 0
+ pass 0 keysA orderA keysB orderB
+ pass 8 keysB orderB keysA orderA
+ pass 16 keysA orderA keysB orderB
+ pass 24 keysB orderB keysA orderA
+ pass 32 keysA orderA keysB orderB
+ pass 40 keysB orderB keysA orderA
+ pass 48 keysA orderA keysB orderB
+ pass 56 keysB orderB keysA orderA
+ let inv !r
+ | r >= ng = pure ()
+ | otherwise = do
+ g <- VUM.unsafeRead orderA r
+ VUM.unsafeWrite rankM g r
+ inv (r + 1)
+ inv 0
+ VU.unsafeFreeze rankM
+{-# INLINEABLE rankByHash #-}
diff --git a/src-internal/DataFrame/Internal/Row.hs b/src-internal/DataFrame/Internal/Row.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Row.hs
@@ -0,0 +1,181 @@
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+module DataFrame.Internal.Row where
+
+import qualified Data.List as L
+import qualified Data.Map as M
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+
+import Control.Exception (throw)
+import Data.Function (on)
+import Data.Maybe (catMaybes, fromMaybe, isNothing, mapMaybe)
+import Data.Type.Equality (TestEquality (..))
+import Data.Typeable (type (:~:) (..))
+import DataFrame.Errors (DataFrameException (..))
+import DataFrame.Internal.Column
+import DataFrame.Internal.DataFrame
+import DataFrame.Internal.Expression (Expr (..))
+import DataFrame.Internal.PackedText (packedIndexText, packedLength)
+import Type.Reflection (typeOf, typeRep)
+
+data Any where
+ Value :: (Columnable a) => a -> Any
+ -- Saves us the extra indirection we get from making Value (Maybe a)
+ -- and having to unpack it again to check for nulls.
+ -- Instead, we just have Null as a separate constructor.
+ Null :: Any
+
+instance Eq Any where
+ (==) :: Any -> Any -> Bool
+ (Value a) == (Value b) = fromMaybe False $ do
+ Refl <- testEquality (typeOf a) (typeOf b)
+ return $ a == b
+ Null == Null = True
+ _ == _ = False
+
+instance Show Any where
+ show :: Any -> String
+ show (Value a) = T.unpack (showValue a)
+ show Null = "null"
+
+showValue :: forall a. (Columnable a) => a -> T.Text
+showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of
+ Just Refl -> v
+ Nothing -> case testEquality (typeRep @a) (typeRep @String) of
+ Just Refl -> T.pack v
+ Nothing -> (T.pack . show) v
+
+-- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.
+toAny :: forall a. (Columnable a) => a -> Any
+toAny = Value
+
+-- | Unwraps a value from an \Any\ type. A 'Null' cell yields 'Nothing'.
+fromAny :: forall a. (Columnable a) => Any -> Maybe a
+fromAny Null = Nothing
+fromAny (Value (v :: b)) = do
+ Refl <- testEquality (typeRep @a) (typeRep @b)
+ pure v
+
+{- | Wrap a column cell into an 'Any', honouring the column's null bitmap: a slot
+marked invalid becomes 'Null', any other slot becomes a 'Value'. Only needs
+@Columnable a@ (the stored element type), not @Columnable (Maybe a)@.
+-}
+cellAny :: (Columnable a) => Maybe Bitmap -> Int -> a -> Any
+cellAny Nothing _ x = Value x
+cellAny (Just bm) i x = if bitmapTestBit bm i then Value x else Null
+
+type Row = V.Vector Any
+
+(!?) :: [a] -> Int -> Maybe a
+(!?) [] _ = Nothing
+(!?) (x : _) 0 = Just x
+(!?) (_x : xs) n = (!?) xs (n - 1)
+
+{- | Reconstruct column @i@ from a list of rows. The element type is taken from
+the first non-'Null' cell; a differently-typed cell is skipped. If any cell is
+'Null' the result is a nullable column, so a round-trip preserves nulls.
+-}
+mkColumnFromRow :: Int -> [[Any]] -> Column
+mkColumnFromRow i rows =
+ let cells = mapMaybe (!? i) rows
+ in case L.find isValue cells of
+ Nothing -> fromList ([] :: [T.Text])
+ Just (Value (_ :: a)) ->
+ let collect Null = Just (Nothing :: Maybe a)
+ collect (Value (v' :: b)) =
+ case testEquality (typeRep @a) (typeRep @b) of
+ Just Refl -> Just (Just v')
+ Nothing -> Nothing
+ maybes = mapMaybe collect cells
+ in if any isNothing maybes
+ then fromMaybeVec (V.fromList maybes)
+ else fromList (catMaybes maybes)
+ Just Null -> fromList ([] :: [T.Text])
+ where
+ isValue (Value _) = True
+ isValue Null = False
+
+{- | Convert the whole dataframe to a list of rows, one per row index in natural
+order; each row lists all columns ordered by column index. Materializes every
+row, so prefer 'toRowVector' for large frames.
+
+>>> toRowList df
+[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...]
+-}
+toRowList :: DataFrame -> [[(T.Text, Any)]]
+toRowList df =
+ let
+ names = map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df))
+ in
+ map
+ (zip names . V.toList . mkRowRep df names)
+ [0 .. (fst (dataframeDimensions df) - 1)]
+
+{- | Convert the dataframe to a vector of rows containing only the named columns,
+in the given order. An empty name list yields one empty row per dataframe row.
+
+>>> toRowVector ["name", "age"] df
+Vector of rows with only name and age fields
+-}
+toRowVector :: [T.Text] -> DataFrame -> V.Vector Row
+toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)
+
+{- | Given a row gets the value associated with a field.
+
+==== __Examples__
+
+>>> map (rowValue (F.col @Int "age")) (toRowList df)
+[25,30, ...]
+-}
+rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a
+rowValue (Col name) row = lookup name row >>= fromAny @a
+rowValue _ _ = error "Can only get rowValue of column reference"
+
+mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row
+mkRowFromArgs names df i = V.map get (V.fromList names)
+ where
+ get name = case getColumn name df of
+ Nothing ->
+ throw $
+ ColumnsNotFoundException
+ [name]
+ "[INTERNAL] mkRowFromArgs"
+ (M.keys $ columnIndices df)
+ Just (BoxedColumn bm column) -> cellAny bm i (column V.! i)
+ Just (UnboxedColumn bm column) -> cellAny bm i (column VU.! i)
+ Just (PackedText bm p) -> cellAny bm i (packedIndexText p i)
+
+-- Returns row values in the caller's requested column order, not the
+-- dataframe's storage order.
+mkRowRep :: DataFrame -> [T.Text] -> Int -> Row
+mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index))
+ where
+ names' = V.fromList names
+ throwError name =
+ error $
+ "Column "
+ ++ T.unpack name
+ ++ " has less items than "
+ ++ "the other columns at index "
+ ++ show i
+ get name = case getColumn name df of
+ Just (BoxedColumn bm c) -> case c V.!? i of
+ Just e -> cellAny bm i e
+ Nothing -> throwError name
+ Just (UnboxedColumn bm c) -> case c VU.!? i of
+ Just e -> cellAny bm i e
+ Nothing -> throwError name
+ Just (PackedText bm p)
+ | i < packedLength p -> cellAny bm i (packedIndexText p i)
+ | otherwise -> throwError name
+ Nothing ->
+ throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
diff --git a/src-internal/DataFrame/Internal/RowHash.hs b/src-internal/DataFrame/Internal/RowHash.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/RowHash.hs
@@ -0,0 +1,216 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | Row-hash kernels with a parallel driver, feeding grouping and the join
+build/probe. Each row's hash depends only on its own bytes, so hashing disjoint
+ranges in parallel is race-free and bit-identical to the sequential pass.
+-}
+module DataFrame.Internal.RowHash (
+ computeRowHashesIO,
+ hashRowRange,
+ parRowHashThreshold,
+) where
+
+import Control.Concurrent (forkIO, getNumCapabilities)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Exception (SomeException, throwIO, try)
+import qualified Data.Text as T
+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
+import qualified Data.Vector as V
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
+import System.IO.Unsafe (unsafePerformIO)
+import Type.Reflection (typeRep)
+
+import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)
+import DataFrame.Internal.Hash (
+ fnvOffset,
+ mixBytes,
+ mixDouble,
+ mixInt,
+ mixShow,
+ mixText,
+ nullSalt,
+ )
+import DataFrame.Internal.PackedText (
+ PackedTextData (..),
+ packedSlice,
+ )
+import DataFrame.Internal.Types (
+ SBool (..),
+ sFloating,
+ sIntegral,
+ )
+
+{- | At least this many rows make the fork/coordination overhead of the parallel
+hash worth it. Below it the sequential single range is used. Matches the
+grouping/join parallel thresholds so the whole pipeline switches together.
+-}
+parRowHashThreshold :: Int
+parRowHashThreshold = 200000
+
+capabilities :: Int
+capabilities = unsafePerformIO getNumCapabilities
+{-# NOINLINE capabilities #-}
+
+{- | Compute the per-row key hash over the selected key columns of an @n@-row
+frame. Forks one worker per capability over disjoint row ranges when the row
+count justifies it, else hashes the single full range; output is capability-independent.
+-}
+computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int)
+computeRowHashesIO n selected = do
+ mv <- VUM.unsafeNew (max 1 n)
+ let runRange lo hi = hashRowRange mv lo hi selected
+ if n >= parRowHashThreshold && capabilities > 1
+ then do
+ let !caps = capabilities
+ !per = (n + caps - 1) `div` caps
+ spawn w = do
+ var <- newEmptyMVar
+ let !lo = min n (w * per)
+ !hi = min n (lo + per)
+ _ <- forkIO (try (runRange lo hi) >>= putMVar var)
+ pure var
+ vars <- mapM spawn [0 .. caps - 1]
+ rs <- mapM takeMVar vars
+ mapM_ (either (throwIO @SomeException) pure) rs
+ else runRange 0 n
+ VU.unsafeFreeze (VUM.slice 0 n mv)
+
+{- | Mix every selected column over the row range @[lo, hi)@ into @mv@, seeding
+each slot with 'fnvOffset'. Must match the sequential grouping hash byte-for-byte
+so grouping and joins bucket identically.
+-}
+hashRowRange :: VUM.IOVector Int -> Int -> Int -> [Column] -> IO ()
+hashRowRange mv lo hi cols = do
+ seedRange mv lo hi
+ mapM_ (mixColumnRange mv lo hi) cols
+
+seedRange :: VUM.IOVector Int -> Int -> Int -> IO ()
+seedRange mv lo hi = go lo
+ where
+ go !i
+ | i >= hi = pure ()
+ | otherwise = VUM.unsafeWrite mv i fnvOffset >> go (i + 1)
+
+{- | Fold one column's values over @[lo, hi)@ into the running hashes. The branch
+structure mirrors the sequential grouping hash: typed unboxed fast paths, then a
+'mixShow' fallback, with the null bitmap mixing 'nullSalt'.
+-}
+mixColumnRange :: VUM.IOVector Int -> Int -> Int -> Column -> IO ()
+mixColumnRange mv lo hi = \case
+ UnboxedColumn ubm (v :: VU.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @Int) of
+ Just Refl -> unboxedRange mv lo hi ubm mixInt v
+ Nothing ->
+ case testEquality (typeRep @a) (typeRep @Double) of
+ Just Refl -> unboxedRange mv lo hi ubm mixDouble v
+ Nothing ->
+ case sIntegral @a of
+ STrue ->
+ unboxedRange mv lo hi ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v
+ SFalse ->
+ case sFloating @a of
+ STrue ->
+ unboxedRange mv lo hi ubm (\h d -> mixDouble h (realToFrac d :: Double)) v
+ SFalse ->
+ unboxedRange mv lo hi ubm mixShow v
+ BoxedColumn bm (v :: V.Vector a) ->
+ case testEquality (typeRep @a) (typeRep @T.Text) of
+ Just Refl -> boxedRange mv lo hi bm mixText v
+ Nothing -> boxedRange mv lo hi bm mixShow v
+ PackedText bm p -> packedRange mv lo hi bm p
+
+{- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. @INLINE@d to
+specialise on the element type and mixing function per call site.
+-}
+unboxedRange ::
+ (VU.Unbox a) =>
+ VUM.IOVector Int ->
+ Int ->
+ Int ->
+ Maybe Bitmap ->
+ (Int -> a -> Int) ->
+ VU.Vector a ->
+ IO ()
+unboxedRange mv lo hi ubm mix v = go lo
+ where
+ go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ h <- VUM.unsafeRead mv i
+ let !h' = case ubm of
+ Just bm | not (bitmapTestBit bm i) -> mixInt h nullSalt
+ _ -> mix h (VU.unsafeIndex v i)
+ VUM.unsafeWrite mv i h'
+ go (i + 1)
+{-# INLINE unboxedRange #-}
+
+boxedRange ::
+ VUM.IOVector Int ->
+ Int ->
+ Int ->
+ Maybe Bitmap ->
+ (Int -> a -> Int) ->
+ V.Vector a ->
+ IO ()
+boxedRange mv lo hi bm mix v = go lo
+ where
+ go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ h <- VUM.unsafeRead mv i
+ let !h' = case bm of
+ Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
+ _ -> mix h (V.unsafeIndex v i)
+ VUM.unsafeWrite mv i h'
+ go (i + 1)
+{-# INLINE boxedRange #-}
+
+{- | Mix a packed-text column's range over its raw UTF-8 byte slices. The
+unselected payload is the hot path (indexes the offset vector directly); a
+selected payload (a gather/join result) falls back to 'packedSlice'.
+-}
+packedRange ::
+ VUM.IOVector Int ->
+ Int ->
+ Int ->
+ Maybe Bitmap ->
+ PackedTextData ->
+ IO ()
+packedRange mv lo hi bm p =
+ case ptSel p of
+ Nothing -> contiguous (ptBytes p) (ptOffsets p)
+ Just _ -> selected
+ where
+ valid i = case bm of
+ Just bm' -> bitmapTestBit bm' i
+ Nothing -> True
+ contiguous !arr !offs = go lo
+ where
+ go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ h <- VUM.unsafeRead mv i
+ let !o = VU.unsafeIndex offs i
+ !l = VU.unsafeIndex offs (i + 1) - o
+ !h' = if valid i then mixBytes h arr o l else mixInt h nullSalt
+ VUM.unsafeWrite mv i h'
+ go (i + 1)
+ selected = go lo
+ where
+ go !i
+ | i >= hi = pure ()
+ | otherwise = do
+ h <- VUM.unsafeRead mv i
+ let !h' =
+ if valid i
+ then let (arr, o, l) = packedSlice p i in mixBytes h arr o l
+ else mixInt h nullSalt
+ VUM.unsafeWrite mv i h'
+ go (i + 1)
+{-# INLINE packedRange #-}
diff --git a/src-internal/DataFrame/Internal/Simplify.hs b/src-internal/DataFrame/Internal/Simplify.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Simplify.hs
@@ -0,0 +1,417 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+module DataFrame.Internal.Simplify (
+ simplify,
+ simplifyPredicatePair,
+
+ -- * Path-condition entailment (for fitted-tree pruning)
+ PredFact,
+ factTrue,
+ factFalse,
+ entails,
+) where
+
+import Control.Monad (guard)
+import Data.Maybe (fromMaybe)
+import Data.Type.Equality (testEquality, (:~:) (Refl))
+import Type.Reflection (eqTypeRep, typeRep, (:~~:) (HRefl), pattern App)
+
+import DataFrame.Internal.Column (Columnable)
+import DataFrame.Internal.Expression (
+ BinaryOp,
+ Expr (..),
+ UnaryOp (unaryName),
+ eqExpr,
+ normalize,
+ )
+import DataFrame.Operators (
+ NullAnd,
+ NullEq,
+ NullGeq,
+ NullGt,
+ NullLeq,
+ NullLt,
+ NullNeq,
+ NullOr,
+ (.==.),
+ )
+
+simplify :: forall a. (Columnable a) => Expr a -> Expr a
+simplify e
+ | isBoolish @a = fixpoint (10 :: Int) e
+ | otherwise = e
+ where
+ fixpoint 0 x = x
+ fixpoint n x = let x' = simplifyB x in if eqExpr x x' then x else fixpoint (n - 1) x'
+
+isBoolish :: forall a. (Columnable a) => Bool
+isBoolish =
+ case ( testEquality (typeRep @a) (typeRep @Bool)
+ , testEquality (typeRep @a) (typeRep @(Maybe Bool))
+ ) of
+ (Just Refl, _) -> True
+ (_, Just Refl) -> True
+ _ -> False
+
+data Conn = ConnAnd | ConnOr
+
+connOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Conn
+connOf _
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullAnd) = Just ConnAnd
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullOr) = Just ConnOr
+ | otherwise = Nothing
+
+simplifyB :: forall a. (Columnable a) => Expr a -> Expr a
+simplifyB expr = case expr of
+ Binary (op :: op c b a) l r
+ | Just conn <- connOf op
+ , Just Refl <- testEquality (typeRep @c) (typeRep @a)
+ , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->
+ let l' = simplifyB l; r' = simplifyB r
+ in fromMaybe (Binary op l' r') (combine conn l' r')
+ | otherwise -> expr
+ Unary (op :: op b a) inner
+ | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)
+ , Just Refl <- testEquality (typeRep @b) (typeRep @Bool)
+ , unaryName op == "not" ->
+ simplifyNot op (simplifyB inner)
+ | otherwise -> expr
+ If c t f ->
+ let c' = simplify c
+ t' = simplifyB t
+ f' = simplifyB f
+ in case asBoolLit c' of
+ Just True -> t'
+ Just False -> f'
+ Nothing
+ | eqExpr t' f' -> t'
+ | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)
+ , asBoolLit t' == Just True
+ , asBoolLit f' == Just False ->
+ c'
+ | otherwise -> If c' t' f'
+ _ -> expr
+
+simplifyNot :: (UnaryOp op) => op Bool Bool -> Expr Bool -> Expr Bool
+simplifyNot op inner = case asBoolLit inner of
+ Just b -> Lit (not b)
+ Nothing -> case inner of
+ Unary (op2 :: op2 b2 Bool) inner2
+ | unaryName op2 == "not"
+ , Just Refl <- testEquality (typeRep @b2) (typeRep @Bool) ->
+ inner2
+ _ -> Unary op inner
+
+combine :: (Columnable a) => Conn -> Expr a -> Expr a -> Maybe (Expr a)
+combine ConnAnd = combineAnd
+combine ConnOr = combineOr
+
+asBoolLit :: forall a. (Columnable a) => Expr a -> Maybe Bool
+asBoolLit (Lit v) =
+ case testEquality (typeRep @a) (typeRep @Bool) of
+ Just Refl -> Just v
+ Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of
+ Just Refl -> v
+ Nothing -> Nothing
+asBoolLit _ = Nothing
+
+{- | Polymorphic boolean literal: @Lit b@ for @Expr Bool@, @Lit (Just b)@ for
+@Expr (Maybe Bool)@.
+-}
+litBoolish :: forall a. (Columnable a) => Bool -> Maybe (Expr a)
+litBoolish v =
+ case testEquality (typeRep @a) (typeRep @Bool) of
+ Just Refl -> Just (Lit v)
+ Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of
+ Just Refl -> Just (Lit (Just v))
+ Nothing -> Nothing
+
+combineAnd :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)
+combineAnd l r
+ | eqExpr l r = Just l
+ | asBoolLit l == Just False = litBoolish False
+ | asBoolLit r == Just False = litBoolish False
+ | asBoolLit l == Just True = Just r
+ | asBoolLit r == Just True = Just l
+ | absorbs ConnOr l r = Just l
+ | absorbs ConnOr r l = Just r
+ | otherwise = simplifyPredicatePair True l r
+
+combineOr :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)
+combineOr l r
+ | eqExpr l r = Just l
+ | asBoolLit l == Just True = litBoolish True
+ | asBoolLit r == Just True = litBoolish True
+ | asBoolLit l == Just False = Just r
+ | asBoolLit r == Just False = Just l
+ | absorbs ConnAnd l r = Just l
+ | absorbs ConnAnd r l = Just r
+ | otherwise = simplifyPredicatePair False l r
+
+absorbs :: (Columnable a) => Conn -> Expr a -> Expr a -> Bool
+absorbs conn x (Binary (op :: op c b a) ya yb)
+ | Just c' <- connOf op
+ , sameConn conn c'
+ , Just Refl <- testEquality (typeRep @c) (typeRep @a)
+ , Just Refl <- testEquality (typeRep @b) (typeRep @a) =
+ eqExpr x ya || eqExpr x yb
+absorbs _ _ _ = False
+
+sameConn :: Conn -> Conn -> Bool
+sameConn ConnAnd ConnAnd = True
+sameConn ConnOr ConnOr = True
+sameConn _ _ = False
+
+data Cmp = CLt | CLeq | CGt | CGeq | CEq | CNeq deriving (Eq)
+
+data NullK = Total | FalseOnNull | UnknownOnNull deriving (Eq)
+
+data Atom = Atom
+ { aCmp :: Cmp
+ , aThr :: !Double
+ , aKey :: String
+ , aNull :: NullK
+ , aIntegral :: Bool
+ }
+
+cmpOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Cmp
+cmpOf _
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLt) = Just CLt
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLeq) = Just CLeq
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGt) = Just CGt
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGeq) = Just CGeq
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullEq) = Just CEq
+ | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullNeq) = Just CNeq
+ | otherwise = Nothing
+
+isLower, isUpper :: Cmp -> Bool
+isLower c = c == CGt || c == CGeq
+isUpper c = c == CLt || c == CLeq
+
+-- | True if @x@ is a @Maybe _@ type.
+isMaybeTy :: forall x. (Columnable x) => Bool
+isMaybeTy = case typeRep @x of
+ App con _ -> case eqTypeRep con (typeRep @Maybe) of Just HRefl -> True; _ -> False
+ _ -> False
+
+litDouble :: forall b. (Columnable b) => Expr b -> Maybe Double
+litDouble (Lit v) =
+ case testEquality (typeRep @b) (typeRep @Double) of
+ Just Refl -> Just v
+ Nothing -> case testEquality (typeRep @b) (typeRep @Int) of
+ Just Refl -> Just (fromIntegral v)
+ Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Double)) of
+ Just Refl -> v
+ Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Int)) of
+ Just Refl -> fromIntegral <$> v
+ Nothing -> Nothing
+litDouble _ = Nothing
+
+{- | True for a column lifted from an integral type (never NaN): @toDouble (col …)@
+or a column whose type is itself integral.
+-}
+integralColE :: forall c. (Columnable c) => Expr c -> Bool
+integralColE (Unary op _) = unaryName op == "toDouble"
+integralColE _ =
+ or
+ [ matches @Int
+ , matches @(Maybe Int)
+ ]
+ where
+ matches :: forall t. (Columnable t) => Bool
+ matches = case testEquality (typeRep @c) (typeRep @t) of Just Refl -> True; _ -> False
+
+atomOf :: forall a. (Columnable a) => Expr a -> Maybe Atom
+atomOf (Unary fm (Binary (op :: op c b r) (colE :: Expr c) litE))
+ | unaryName fm == "fromMaybe"
+ , Just cmp <- cmpOf op
+ , Just t <- litDouble litE =
+ Just (Atom cmp t (show (normalize colE)) FalseOnNull (integralColE colE))
+atomOf (Binary (op :: op c b a) (colE :: Expr c) litE)
+ | Just cmp <- cmpOf op
+ , Just t <- litDouble litE =
+ let nk = if isMaybeTy @c then UnknownOnNull else Total
+ in Just (Atom cmp t (show (normalize colE)) nk (integralColE colE))
+atomOf _ = Nothing
+
+simplifyPredicatePair ::
+ forall a. (Columnable a) => Bool -> Expr a -> Expr a -> Maybe (Expr a)
+simplifyPredicatePair isAnd a b = do
+ atomA <- atomOf a
+ atomB <- atomOf b
+ guard (aKey atomA == aKey atomB)
+ let nk = aNull atomA
+ integral = aIntegral atomA
+ if isAnd
+ then andAtoms a atomA b atomB nk integral
+ else orAtoms a atomA b atomB nk integral
+
+-- | Contradiction folds to a literal False unless null-rows make it unknown.
+litFalseGated :: (Columnable a) => NullK -> Maybe (Expr a)
+litFalseGated UnknownOnNull = Nothing
+litFalseGated _ = litBoolish False
+
+{- | Tautology to literal True is sound only for total (never-null) atoms; the
+exhaustive-cover form additionally needs a non-NaN (integral) column.
+-}
+litTrueTotal :: (Columnable a) => NullK -> Maybe (Expr a)
+litTrueTotal Total = litBoolish True
+litTrueTotal _ = Nothing
+
+andAtoms ::
+ (Columnable a) =>
+ Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)
+andAtoms a atomA b atomB nk _ =
+ let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB
+ in if
+ | isLower cA, isLower cB, cA == cB -> Just (if tA >= tB then a else b)
+ | isUpper cA, isUpper cB, cA == cB -> Just (if tA <= tB then a else b)
+ | isLower cA, isUpper cB -> lu cA tA cB tB
+ | isUpper cA, isLower cB -> lu cB tB cA tA
+ | cA == CEq, cB == CEq -> if tA == tB then Just a else litFalseGated nk
+ | cA == CEq, cB == CNeq -> if tA == tB then litFalseGated nk else Just a
+ | cA == CNeq, cB == CEq -> if tA == tB then litFalseGated nk else Just b
+ | cA == CEq -> if satisfies tA cB tB then Just a else litFalseGated nk
+ | cB == CEq -> if satisfies tB cA tA then Just b else litFalseGated nk
+ | cA == CNeq, cB == CNeq -> Nothing
+ | cA == CNeq -> if outside tA cB tB then Just b else Nothing
+ | cB == CNeq -> if outside tB cA tA then Just a else Nothing
+ | otherwise -> Nothing
+ where
+ lu lc lo uc hi
+ | lo > hi = litFalseGated nk
+ | lo == hi, lc == CGeq, uc == CLeq = pointEq a lo
+ | lo == hi = litFalseGated nk
+ | otherwise = Nothing
+
+orAtoms ::
+ (Columnable a) =>
+ Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)
+orAtoms a atomA b atomB nk integral =
+ let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB
+ in if
+ | isLower cA, isLower cB, cA == cB -> Just (if tA <= tB then a else b)
+ | isUpper cA, isUpper cB, cA == cB -> Just (if tA >= tB then a else b)
+ | isUpper cA
+ , isLower cB
+ , nk == Total
+ , integral
+ , covers cB tB cA tA ->
+ litTrueTotal nk
+ | isLower cA
+ , isUpper cB
+ , nk == Total
+ , integral
+ , covers cA tA cB tB ->
+ litTrueTotal nk
+ | cA == CNeq, cB == CNeq -> if tA == tB then Just a else litTrueTotal nk
+ | cA == CEq, cB == CNeq -> if tA == tB then litTrueTotal nk else Just b
+ | cA == CNeq, cB == CEq -> if tA == tB then litTrueTotal nk else Just a
+ | cA == CEq, cB == CEq -> if tA == tB then Just a else Nothing
+ | otherwise -> Nothing
+
+{- | Build @col == t@ for the point-collapse rule; only strict @Expr Bool@ over a
+@Double@ column (otherwise bail).
+-}
+pointEq :: forall a. (Columnable a) => Expr a -> Double -> Maybe (Expr a)
+pointEq atom lo = case testEquality (typeRep @a) (typeRep @Bool) of
+ Just Refl -> (\colE -> colE .==. Lit lo) <$> recoverColD atom
+ Nothing -> Nothing
+
+recoverColD :: Expr x -> Maybe (Expr Double)
+recoverColD (Binary _ (colE :: Expr c) _) =
+ case testEquality (typeRep @c) (typeRep @Double) of
+ Just Refl -> Just colE
+ _ -> Nothing
+recoverColD (Unary _ inner) = recoverColD inner
+recoverColD _ = Nothing
+
+covers :: Cmp -> Double -> Cmp -> Double -> Bool
+covers lowerCmp lo upperCmp hi =
+ lo < hi || (lo == hi && (lowerCmp == CGeq || upperCmp == CLeq))
+
+satisfies :: Double -> Cmp -> Double -> Bool
+satisfies t CGt tb = t > tb
+satisfies t CGeq tb = t >= tb
+satisfies t CLt tb = t < tb
+satisfies t CLeq tb = t <= tb
+satisfies _ _ _ = False
+
+outside :: Double -> Cmp -> Double -> Bool
+outside t CGt tb = t <= tb
+outside t CGeq tb = t < tb
+outside t CLt tb = t >= tb
+outside t CLeq tb = t > tb
+outside _ _ _ = False
+
+-- ---------------------------------------------------------------------------
+-- Path-condition entailment for fitted-tree pruning.
+-- ---------------------------------------------------------------------------
+
+-- | A known same-column threshold fact accumulated along a tree path.
+data PredFact = PredFact !String !Cmp !Double
+
+-- | The fact a branch's true edge establishes (the condition holds).
+factTrue :: Expr Bool -> Maybe PredFact
+factTrue e = (\a -> PredFact (aKey a) (aCmp a) (aThr a)) <$> atomOf e
+
+{- | The fact a branch's false edge establishes (the negated condition). Only
+sound for non-NaN (integral) columns — a NaN row takes the false edge too,
+so @¬(x>t)@ is not a clean @x<=t@ bound for floats.
+-}
+factFalse :: Expr Bool -> Maybe PredFact
+factFalse e = do
+ a <- atomOf e
+ guard (aIntegral a && aNull a == Total)
+ nc <- negCmp (aCmp a)
+ pure (PredFact (aKey a) nc (aThr a))
+
+negCmp :: Cmp -> Maybe Cmp
+negCmp CLt = Just CGeq
+negCmp CLeq = Just CGt
+negCmp CGt = Just CLeq
+negCmp CGeq = Just CLt
+negCmp _ = Nothing
+
+{- | @entails facts cond@: 'Just' 'True' when the path facts force @cond@ true,
+'Just' 'False' when they force it false, 'Nothing' when undecided.
+-}
+entails :: [PredFact] -> Expr Bool -> Maybe Bool
+entails facts cond = do
+ a <- atomOf cond
+ let decisions =
+ [ d
+ | PredFact fk fc ft <- facts
+ , fk == aKey a
+ , Just d <- [factImplies (fc, ft) (aCmp a, aThr a)]
+ ]
+ case decisions of
+ (d : _) -> Just d
+ [] -> Nothing
+
+{- | Does the fact's solution set sit inside @cond@ ('Just' 'True'), disjoint
+from it ('Just' 'False'), or neither ('Nothing')? Boundary strictness is
+honoured: e.g. @x<=t@ does NOT entail @x=t ∧ x<=t@ is not empty.
+-}
+factImplies :: (Cmp, Double) -> (Cmp, Double) -> Maybe Bool
+factImplies (fc, ft) (cc, tc)
+ | isLower fc, isLower cc, subset = Just True
+ | isUpper fc, isUpper cc, subset = Just True
+ | isLower fc, isUpper cc, disjointAtEq = Just False
+ | isUpper fc, isLower cc, disjointBelow = Just False
+ | otherwise = Nothing
+ where
+ fIncl = fc == CGeq || fc == CLeq
+ cIncl = cc == CGeq || cc == CLeq
+ subset =
+ (if isLower fc then ft > tc else ft < tc)
+ || (ft == tc && (not fIncl || cIncl))
+ disjointAtEq = ft > tc || (ft == tc && not (fIncl && cIncl))
+ disjointBelow = ft < tc || (ft == tc && not (fIncl && cIncl))
diff --git a/src-internal/DataFrame/Internal/Types.hs b/src-internal/DataFrame/Internal/Types.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Types.hs
@@ -0,0 +1,161 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+module DataFrame.Internal.Types where
+
+import Data.Int (Int16, Int32, Int64, Int8)
+import Data.Kind (Constraint, Type)
+import Data.Typeable (Typeable)
+import qualified Data.Vector.Unboxed as VU
+import Data.Word (Word16, Word32, Word64, Word8)
+
+type Columnable' a = (Typeable a, Show a, Eq a)
+
+{- | Inline replacement for @Data.These.These@ to keep @dataframe-core@ free
+of the @these@ package dependency. Only the three constructors and the
+derived classes are used internally.
+-}
+data These a b = This a | That b | These a b
+ deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)
+
+{- | A type with column representations used to select the
+"right" representation when specializing the `toColumn` function.
+-}
+data Rep
+ = RBoxed
+ | RUnboxed
+ | RNullableBoxed
+
+-- | Type-level if statement.
+type family If (cond :: Bool) (yes :: k) (no :: k) :: k where
+ If 'True yes _ = yes
+ If 'False _ no = no
+
+-- | All unboxable types (according to the `vector` package).
+type family Unboxable (a :: Type) :: Bool where
+ Unboxable Int = 'True
+ Unboxable Int8 = 'True
+ Unboxable Int16 = 'True
+ Unboxable Int32 = 'True
+ Unboxable Int64 = 'True
+ Unboxable Word = 'True
+ Unboxable Word8 = 'True
+ Unboxable Word16 = 'True
+ Unboxable Word32 = 'True
+ Unboxable Word64 = 'True
+ Unboxable Char = 'True
+ Unboxable Bool = 'True
+ Unboxable Double = 'True
+ Unboxable Float = 'True
+ Unboxable _ = 'False
+
+type family Numeric (a :: Type) :: Bool where
+ Numeric Integer = 'True
+ Numeric Int = 'True
+ Numeric Int8 = 'True
+ Numeric Int16 = 'True
+ Numeric Int32 = 'True
+ Numeric Int64 = 'True
+ Numeric Word = 'True
+ Numeric Word8 = 'True
+ Numeric Word16 = 'True
+ Numeric Word32 = 'True
+ Numeric Word64 = 'True
+ Numeric Double = 'True
+ Numeric Float = 'True
+ Numeric _ = 'False
+
+-- | Compute the column representation tag for any 'a'.
+type family KindOf a :: Rep where
+ KindOf (Maybe a) = 'RNullableBoxed
+ KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed
+
+-- | Type-level boolean for constraint/type comparison.
+data SBool (b :: Bool) where
+ STrue :: SBool 'True
+ SFalse :: SBool 'False
+
+-- | The runtime witness for our type-level branching.
+class SBoolI (b :: Bool) where
+ sbool :: SBool b
+
+instance SBoolI 'True where sbool = STrue
+instance SBoolI 'False where sbool = SFalse
+
+-- | Runtime witness for whether @a@ is unboxable.
+sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)
+sUnbox = sbool @(Unboxable a)
+
+sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)
+sNumeric = sbool @(Numeric a)
+
+type family When (flag :: Bool) (c :: Constraint) :: Constraint where
+ When 'True c = c
+ When 'False c = ()
+
+type UnboxIf a = When (Unboxable a) (VU.Unbox a)
+
+type family IntegralTypes (a :: Type) :: Bool where
+ IntegralTypes Integer = 'True
+ IntegralTypes Int = 'True
+ IntegralTypes Int8 = 'True
+ IntegralTypes Int16 = 'True
+ IntegralTypes Int32 = 'True
+ IntegralTypes Int64 = 'True
+ IntegralTypes Word = 'True
+ IntegralTypes Word8 = 'True
+ IntegralTypes Word16 = 'True
+ IntegralTypes Word32 = 'True
+ IntegralTypes Word64 = 'True
+ IntegralTypes _ = 'False
+
+sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)
+sIntegral = sbool @(IntegralTypes a)
+
+type IntegralIf a = When (IntegralTypes a) (Integral a)
+
+type family FloatingTypes (a :: Type) :: Bool where
+ FloatingTypes Float = 'True
+ FloatingTypes Double = 'True
+ FloatingTypes _ = 'False
+
+sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)
+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
diff --git a/src-internal/DataFrame/Internal/Utf8.hs b/src-internal/DataFrame/Internal/Utf8.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Internal/Utf8.hs
@@ -0,0 +1,95 @@
+{-# LANGUAGE BangPatterns #-}
+
+{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by
+'DataFrame.Internal.ColumnBuilder' to turn shared byte buffers into 'Text'.
+-}
+module DataFrame.Internal.Utf8 (
+ isValidUtf8Slice,
+ isUtf8Boundary,
+ lenientDecodeSlice,
+ sliceTextVector,
+) where
+
+import qualified Data.Text as T
+import qualified Data.Text.Array as A
+import qualified Data.Vector as VB
+import qualified Data.Vector.Mutable as VBM
+import qualified Data.Vector.Unboxed as VU
+
+import Data.Text.Internal (Text (..))
+import Data.Text.Internal.Encoding.Utf8 (
+ DecoderResult (..),
+ utf8DecodeContinue,
+ utf8DecodeStart,
+ )
+import Data.Text.Internal.Validate (isValidUtf8ByteArray)
+import Data.Word (Word8)
+
+-- | Whether @len@ bytes starting at @off@ are well-formed UTF-8.
+isValidUtf8Slice :: A.Array -> Int -> Int -> Bool
+isValidUtf8Slice = isValidUtf8ByteArray
+{-# INLINE isValidUtf8Slice #-}
+
+{- | Whether a byte may start a code point (i.e. is not a continuation
+byte). Field slices of a valid buffer are themselves valid iff every
+field starts on a boundary.
+-}
+isUtf8Boundary :: Word8 -> Bool
+isUtf8Boundary w = w < 0x80 || w >= 0xC0
+{-# INLINE isUtf8Boundary #-}
+
+{- | Decode a byte slice exactly like @decodeUtf8Lenient@: greedy decode at
+each position; any byte that cannot begin a complete, valid sequence within
+the slice becomes one U+FFFD and decoding resumes at the next byte.
+-}
+lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text
+lenientDecodeSlice arr off len = T.pack (go off)
+ where
+ !end = off + len
+ go !i
+ | i >= end = []
+ | otherwise = case tryDecode i of
+ Just (c, i') -> c : go i'
+ Nothing -> '\xFFFD' : go (i + 1)
+ tryDecode !i = loop (utf8DecodeStart (A.unsafeIndex arr i)) (i + 1)
+ where
+ loop (Accept c) !j = Just (c, j)
+ loop Reject _ = Nothing
+ loop (Incomplete st cp) !j
+ | j >= end = Nothing
+ | otherwise = loop (utf8DecodeContinue (A.unsafeIndex arr j) st cp) (j + 1)
+
+{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes
+@[offs!i, offs!(i+1))@. Fast path validates the whole span once when every field
+starts on a code-point boundary; else per-field validation with lenient decode.
+-}
+sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text
+sliceTextVector arr offs = VB.create $ do
+ mv <- VBM.unsafeNew n
+ let fill dec = go 0
+ where
+ go !i
+ | i >= n = pure ()
+ | otherwise = do
+ let o = VU.unsafeIndex offs i
+ !t = dec o (VU.unsafeIndex offs (i + 1) - o)
+ VBM.unsafeWrite mv i t
+ go (i + 1)
+ if fast then fill mkSlice else fill decodeField
+ pure mv
+ where
+ n = VU.length offs - 1
+ base = VU.unsafeIndex offs 0
+ used = VU.unsafeIndex offs n
+ boundariesOk !i
+ | i >= n = True
+ | otherwise =
+ let o = VU.unsafeIndex offs i
+ in (o >= used || isUtf8Boundary (A.unsafeIndex arr o))
+ && boundariesOk (i + 1)
+ fast = isValidUtf8Slice arr base (used - base) && boundariesOk 0
+ mkSlice o l = if l == 0 then T.empty else Text arr o l
+ decodeField o l
+ | l == 0 = T.empty
+ | isValidUtf8Slice arr o l = Text arr o l
+ | otherwise = lenientDecodeSlice arr o l
diff --git a/src-internal/DataFrame/Operators.hs b/src-internal/DataFrame/Operators.hs
new file mode 100644
--- /dev/null
+++ b/src-internal/DataFrame/Operators.hs
@@ -0,0 +1,425 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+
+module DataFrame.Operators where
+
+import Data.Function ((&))
+import qualified Data.Text as T
+import DataFrame.Internal.Column (Columnable)
+import DataFrame.Internal.Expression (
+ BinUDF (MkBinaryOp),
+ BinaryOp (
+ binaryCommutative,
+ binaryFn,
+ binaryName,
+ binaryPrecedence,
+ binarySymbol
+ ),
+ Expr (Binary, Col, If, Lit, Unary),
+ NamedExpr,
+ UExpr (UExpr),
+ UnUDF (MkUnaryOp),
+ )
+import DataFrame.Internal.Nullable (
+ BaseType,
+ DivWidenOp,
+ NullCmpResult,
+ NullLift2Op (applyNull2),
+ NullableCmpOp (nullCmpOp),
+ NumericWidenOp,
+ WidenResult,
+ WidenResultDiv,
+ divArithOp,
+ widenArithOp,
+ widenCmpOp,
+ )
+import DataFrame.Internal.Types (Promote, PromoteDiv)
+
+infixr 8 .^^, .^^., .^, .^.
+infixl 7 .*, ./, .*., ./.
+infixl 6 .+, .-, .+., .-.
+infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=.
+infixr 3 .&&, .&&.
+infixr 2 .||, .||.
+infixr 0 .=
+
+(|>) :: a -> (a -> b) -> b
+(|>) = (&)
+
+as :: (Columnable a) => Expr a -> T.Text -> NamedExpr
+as expr colName = (colName, UExpr expr)
+
+name :: (Show a) => Expr a -> T.Text
+name (Col n) = n
+name other =
+ error $
+ "You must call `name` on a column reference. Not the expression: " ++ show other
+
+col :: (Columnable a) => T.Text -> Expr a
+col = Col
+
+ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
+ifThenElse = If
+
+lit :: (Columnable a) => a -> Expr a
+lit = Lit
+
+(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr
+(.=) = flip as
+
+liftDecorated ::
+ (Columnable a, Columnable b) =>
+ (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b
+liftDecorated f opName rep = Unary (MkUnaryOp f opName rep)
+
+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 opName rep comm prec =
+ Binary (MkBinaryOp f opName rep comm prec)
+
+data NullEq a b c where
+ NullEq ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Eq (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullEq a b (NullCmpResult a b)
+
+data NullNeq a b c where
+ NullNeq ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Eq (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullNeq a b (NullCmpResult a b)
+
+data NullLt a b c where
+ NullLt ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullLt a b (NullCmpResult a b)
+
+data NullGt a b c where
+ NullGt ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullGt a b (NullCmpResult a b)
+
+data NullLeq a b c where
+ NullLeq ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullLeq a b (NullCmpResult a b)
+
+data NullGeq a b c where
+ NullGeq ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ NullGeq a b (NullCmpResult a b)
+
+data NullAnd a b c where
+ NullAnd ::
+ (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
+ NullAnd a b (NullCmpResult a b)
+
+data NullOr a b c where
+ NullOr ::
+ (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
+ NullOr a b (NullCmpResult a b)
+
+instance BinaryOp NullEq where
+ binaryFn NullEq = applyNull2 (widenCmpOp (==))
+ binaryName NullEq = "eq"
+ binarySymbol NullEq = Just ".=="
+ binaryCommutative NullEq = True
+ binaryPrecedence NullEq = 4
+instance BinaryOp NullNeq where
+ binaryFn NullNeq = applyNull2 (widenCmpOp (/=))
+ binaryName NullNeq = "neq"
+ binarySymbol NullNeq = Just "./="
+ binaryCommutative NullNeq = True
+ binaryPrecedence NullNeq = 4
+instance BinaryOp NullLt where
+ binaryFn NullLt = applyNull2 (widenCmpOp (<))
+ binaryName NullLt = "lt"
+ binarySymbol NullLt = Just ".<"
+ binaryPrecedence NullLt = 4
+instance BinaryOp NullGt where
+ binaryFn NullGt = applyNull2 (widenCmpOp (>))
+ binaryName NullGt = "gt"
+ binarySymbol NullGt = Just ".>"
+ binaryPrecedence NullGt = 4
+instance BinaryOp NullLeq where
+ binaryFn NullLeq = applyNull2 (widenCmpOp (<=))
+ binaryName NullLeq = "leq"
+ binarySymbol NullLeq = Just ".<="
+ binaryPrecedence NullLeq = 4
+instance BinaryOp NullGeq where
+ binaryFn NullGeq = applyNull2 (widenCmpOp (>=))
+ binaryName NullGeq = "geq"
+ binarySymbol NullGeq = Just ".>="
+ binaryPrecedence NullGeq = 4
+instance BinaryOp NullAnd where
+ binaryFn NullAnd = nullCmpOp (&&)
+ binaryName NullAnd = "nulland"
+ binarySymbol NullAnd = Just ".&&"
+ binaryCommutative NullAnd = True
+ binaryPrecedence NullAnd = 3
+instance BinaryOp NullOr where
+ binaryFn NullOr = nullCmpOp (||)
+ binaryName NullOr = "nullor"
+ binarySymbol NullOr = Just ".||"
+ binaryCommutative NullOr = True
+ binaryPrecedence NullOr = 2
+
+(.==.) ::
+ (Columnable a, Eq a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4
+
+(./=.) ::
+ (Columnable a, Eq a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4
+
+(.<.) ::
+ (Columnable a, Ord a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4
+
+(.>.) ::
+ (Columnable a, Ord a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4
+
+(.<=.) ::
+ (Columnable a, Ord a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4
+
+(.>=.) ::
+ (Columnable a, Ord a) =>
+ Expr a ->
+ Expr a ->
+ Expr Bool
+(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4
+
+(.+.) :: (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. Widens numeric operands to their common type,
+so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either
+operand is nullable.
+-}
+(.==) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Eq (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(.==) = Binary NullEq
+
+-- | Nullable-aware inequality. Widens numeric operands to their common type.
+(./=) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Eq (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(./=) = Binary NullNeq
+
+-- | Nullable-aware less-than. Widens numeric operands to their common type.
+(.<) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(.<) = Binary NullLt
+
+-- | Nullable-aware greater-than. Widens numeric operands to their common type.
+(.>) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(.>) = Binary NullGt
+
+{- | Nullable-aware less-than-or-equal. Widens numeric operands to their
+common type, so @Expr Double .<= Expr Int@ typechecks.
+-}
+(.<=) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(.<=) = Binary NullLeq
+
+-- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.
+(.>=) ::
+ ( NumericWidenOp (BaseType a) (BaseType b)
+ , NullLift2Op a b Bool (NullCmpResult a b)
+ , Ord (Promote (BaseType a) (BaseType b))
+ ) =>
+ Expr a ->
+ Expr b ->
+ Expr (NullCmpResult a b)
+(.>=) = Binary NullGeq
+
+(.&&.) :: 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)
+(.&&) = Binary NullAnd
+
+-- | 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)
+(.||) = Binary NullOr
+
+(.^^) ::
+ ( 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/Core.hs b/src/DataFrame/Core.hs
new file mode 100644
--- /dev/null
+++ b/src/DataFrame/Core.hs
@@ -0,0 +1,104 @@
+{- | The curated public surface of @dataframe-core@: the interchange types
+('DataFrame', 'Column', 'Row', 'Expr'), element constraints, and the
+rendering/serialization verbs. Internal plumbing stays in @dataframe-core:internal@.
+-}
+module DataFrame.Core (
+ -- * The DataFrame
+ DataFrame,
+ GroupedDataFrame,
+ empty,
+ fromNamedColumns,
+ insertColumn,
+ columnNames,
+ null,
+
+ -- * Columns
+ Column,
+ fromList,
+ fromVector,
+ fromUnboxedVector,
+ mkRandom,
+ toList,
+ toVector,
+ hasElemType,
+ hasMissing,
+ isNumeric,
+
+ -- * Element constraints
+ Columnable,
+ Columnable',
+
+ -- * Rows
+ Row,
+ Any,
+ toAny,
+ fromAny,
+ rowValue,
+ toRowList,
+ toRowVector,
+
+ -- * Expressions
+ Expr,
+ NamedExpr,
+ eSize,
+ prettyPrint,
+ prettyPrintWidth,
+
+ -- * Rendering & serialization
+ TruncateConfig (..),
+ defaultTruncateConfig,
+ toCsv,
+ toCsv',
+ toSeparated,
+ toMarkdown,
+ toMarkdown',
+) where
+
+import Prelude hiding (null)
+
+import DataFrame.Internal.Column (
+ Column,
+ Columnable,
+ fromList,
+ fromUnboxedVector,
+ fromVector,
+ hasElemType,
+ hasMissing,
+ isNumeric,
+ mkRandom,
+ toList,
+ toVector,
+ )
+import DataFrame.Internal.DataFrame (
+ DataFrame,
+ GroupedDataFrame,
+ TruncateConfig (..),
+ columnNames,
+ defaultTruncateConfig,
+ empty,
+ fromNamedColumns,
+ insertColumn,
+ null,
+ toCsv,
+ toCsv',
+ toMarkdown,
+ toMarkdown',
+ toSeparated,
+ )
+import DataFrame.Internal.Expression (
+ Expr,
+ NamedExpr,
+ eSize,
+ prettyPrint,
+ prettyPrintWidth,
+ )
+import DataFrame.Internal.Row (
+ Any,
+ Row,
+ fromAny,
+ rowValue,
+ toAny,
+ toRowList,
+ toRowVector,
+ )
+import DataFrame.Internal.Types (Columnable')
diff --git a/src/DataFrame/Display/Terminal/Colours.hs b/src/DataFrame/Display/Terminal/Colours.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Terminal/Colours.hs
+++ /dev/null
@@ -1,7 +0,0 @@
-module DataFrame.Display.Terminal.Colours where
-
-red, green, brightGreen, brightBlue :: String -> String
-red = id
-green = id
-brightGreen = id
-brightBlue = id
diff --git a/src/DataFrame/Display/Terminal/PrettyPrint.hs b/src/DataFrame/Display/Terminal/PrettyPrint.hs
deleted file mode 100644
--- a/src/DataFrame/Display/Terminal/PrettyPrint.hs
+++ /dev/null
@@ -1,118 +0,0 @@
-{-# LANGUAGE OverloadedStrings #-}
-
-module DataFrame.Display.Terminal.PrettyPrint where
-
-import qualified Data.Text as T
-import qualified Data.Vector as V
-
-{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with
-ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for
-notebooks.
--}
-data RenderFormat = Plain | Markdown
- deriving (Show, Eq)
-
--- Utility functions to show a DataFrame as a Markdown-ish table.
-
--- Adapted from: https://stackoverflow.com/questions/5929377/format-list-output-in-haskell
--- a type for fill functions
-type Filler = Int -> T.Text -> T.Text
-
--- a type for describing table columns
-data ColDesc t = ColDesc
- { colTitleFill :: Filler
- , colTitle :: T.Text
- , colValueFill :: Filler
- }
-
--- functions that fill a string (s) to a given width (n) by adding pad
--- character (c) to align left, right, or center
-fillLeft :: Char -> Int -> T.Text -> T.Text
-fillLeft c n s = s <> T.replicate (n - T.length s) (T.singleton c)
-
-fillRight :: Char -> Int -> T.Text -> T.Text
-fillRight c n s = T.replicate (n - T.length s) (T.singleton c) <> s
-
-fillCenter :: Char -> Int -> T.Text -> T.Text
-fillCenter c n s =
- T.replicate l (T.singleton c) <> s <> T.replicate r (T.singleton c)
- where
- x = n - T.length s
- l = x `div` 2
- r = x - l
-
--- functions that fill with spaces
-left :: Int -> T.Text -> T.Text
-left = fillLeft ' '
-
-right :: Int -> T.Text -> T.Text
-right = fillRight ' '
-
-center :: Int -> T.Text -> T.Text
-center = fillCenter ' '
-
-{- | Render a table from column-major data. @columns@ has one 'V.Vector' per
-column; widths are computed in one pass per column (no row-major transpose),
-and row lines are built by indexing each column at row @i@.
--}
-showTable ::
- RenderFormat ->
- [T.Text] ->
- [T.Text] ->
- [V.Vector T.Text] ->
- T.Text
-showTable fmt header types columns =
- let isMarkdown = fmt == Markdown
- -- In a GitHub pipe table a literal '|' ends the cell and a newline ends
- -- the row, so a value like the operator name <|> would split the row into
- -- the wrong columns. Escape both for the Markdown format only.
- esc = if isMarkdown then escapeMarkdownCell else id
- hdr = map esc header
- tys = map esc types
- cols = map (V.map esc) columns
- consolidatedHeader =
- if isMarkdown
- then zipWith (\h t -> h <> "
" <> t) hdr tys
- else hdr
- cs = map (\h -> ColDesc center h left) consolidatedHeader
- nRows = case cols of
- (c : _) -> V.length c
- [] -> 0
- columnMaxWidth col
- | V.null col = 0
- | otherwise = V.foldl' (\acc x -> max acc (T.length x)) 0 col
- widths =
- zipWith3
- (\h t col -> T.length h `max` T.length t `max` columnMaxWidth col)
- consolidatedHeader
- tys
- cols
- dashesOf w = T.replicate w "-"
- border = T.intercalate "---" (map dashesOf widths)
- separator = T.intercalate "-|-" (map dashesOf widths)
- fillCells fill cells =
- T.intercalate " | " (zipWith3 fill cs widths cells)
- rowCells i = map (V.! i) cols
- rowLines = [fillCells colValueFill (rowCells i) | i <- [0 .. nRows - 1]]
- wrapMd t = T.concat ["| ", t, " |"]
- outputLines =
- if isMarkdown
- then
- wrapMd (fillCells colTitleFill consolidatedHeader)
- : wrapMd separator
- : map wrapMd rowLines
- else
- border
- : fillCells colTitleFill consolidatedHeader
- : separator
- : fillCells colTitleFill tys
- : separator
- : rowLines
- in T.unlines outputLines
-
-{- | Escape a value for a GitHub-flavoured Markdown table cell: a bare @|@ would
-end the cell and a newline would end the row, so both are neutralised (the pipe
-backslash-escaped, the newline turned into a @\@).
--}
-escapeMarkdownCell :: T.Text -> T.Text
-escapeMarkdownCell = T.replace "\n" "
" . T.replace "|" "\\|"
diff --git a/src/DataFrame/Errors.hs b/src/DataFrame/Errors.hs
deleted file mode 100644
--- a/src/DataFrame/Errors.hs
+++ /dev/null
@@ -1,187 +0,0 @@
-{-# LANGUAGE DeriveAnyClass #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE InstanceSigs #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RankNTypes #-}
-
-module DataFrame.Errors where
-
-import qualified Data.Map.Lazy as ML
-import qualified Data.Text as T
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-
-import Control.Exception
-import qualified Data.List as L
-import Data.Typeable (Typeable)
-import DataFrame.Display.Terminal.Colours
-import Type.Reflection (TypeRep)
-
-data TypeErrorContext a b = MkTypeErrorContext
- { userType :: Either String (TypeRep a)
- , expectedType :: Either String (TypeRep b)
- , errorColumnName :: Maybe String
- , callingFunctionName :: Maybe String
- }
-
-data DataFrameException where
- TypeMismatchException ::
- forall a b.
- (Typeable a, Typeable b) =>
- TypeErrorContext a b ->
- DataFrameException
- AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException
- ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException
- EmptyDataSetException :: T.Text -> DataFrameException
- InternalException :: T.Text -> DataFrameException
- NonColumnReferenceException :: T.Text -> DataFrameException
- UnaggregatedException :: T.Text -> DataFrameException
- WrongQuantileNumberException :: Int -> DataFrameException
- WrongQuantileIndexException :: VU.Vector Int -> Int -> DataFrameException
- deriving (Exception)
-
-instance Show DataFrameException where
- show :: DataFrameException -> String
- show (TypeMismatchException context) =
- let
- errorString =
- typeMismatchError
- (either id show (userType context))
- (either id show (expectedType context))
- in
- addCallPointInfo
- (errorColumnName context)
- (callingFunctionName context)
- errorString
- show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns
- show (EmptyDataSetException callPoint) = emptyDataSetError callPoint
- show (WrongQuantileNumberException q) = wrongQuantileNumberError q
- show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q
- show (InternalException msg) = "Internal error: " ++ T.unpack msg
- show (NonColumnReferenceException msg) = "Expression must be a column reference in: " ++ T.unpack msg
- show (UnaggregatedException expr) = "Expression is not fully aggregated: " ++ T.unpack expr
- show (AggregatedAndNonAggregatedException expr1 expr2) =
- "Cannot combine aggregated and non-aggregated expressions: \n"
- ++ T.unpack expr1
- ++ "\n"
- ++ T.unpack expr2
-
-columnNotFound :: T.Text -> T.Text -> [T.Text] -> String
-columnNotFound missingColumn = columnsNotFound [missingColumn]
-
-columnsNotFound :: [T.Text] -> T.Text -> [T.Text] -> String
-columnsNotFound missingColumns callPoint availableColumns =
- red "\n\n[ERROR] "
- ++ missingColumnsLabel missingColumns
- ++ ": "
- ++ T.unpack (T.intercalate ", " missingColumns)
- ++ " for operation "
- ++ T.unpack callPoint
- ++ formatSuggestions missingColumns availableColumns
- ++ "\n\n"
- where
- missingColumnsLabel [_] = "Column not found"
- missingColumnsLabel _ = "Columns not found"
-
- formatSuggestions [missingColumn] columns =
- case guessColumnName missingColumn columns of
- "" -> ""
- guessed ->
- "\n\tDid you mean "
- ++ T.unpack guessed
- ++ "?"
- formatSuggestions names columns =
- case traverse (`suggestColumnName` columns) names of
- Just guessedColumns
- | not (null guessedColumns) ->
- "\n\tDid you mean "
- ++ formatColumnSuggestions guessedColumns
- ++ "?"
- _ -> ""
-
- suggestColumnName missingColumn columns = case guessColumnName missingColumn columns of
- "" -> Nothing
- guessed -> Just guessed
-
- formatColumnSuggestions guessedColumns =
- "["
- ++ L.intercalate ", " (map (show . T.unpack) guessedColumns)
- ++ "]"
-
-typeMismatchError :: String -> String -> String
-typeMismatchError givenType expType =
- red $
- red "\n\n[Error]: Type Mismatch"
- ++ "\n\tWhile running your code I tried to "
- ++ "get a column of type: "
- ++ red (show givenType)
- ++ " but the column in the dataframe was actually of type: "
- ++ green (show expType)
-
-emptyDataSetError :: T.Text -> String
-emptyDataSetError callPoint =
- red "\n\n[ERROR] "
- ++ T.unpack callPoint
- ++ " cannot be called on empty data sets"
-
-wrongQuantileNumberError :: Int -> String
-wrongQuantileNumberError q =
- red "\n\n[ERROR] "
- ++ "Quantile number q should satisfy "
- ++ "q >= 2, but here q is "
- ++ show q
-
-wrongQuantileIndexError :: VU.Vector Int -> Int -> String
-wrongQuantileIndexError qs q =
- red "\n\n[ERROR] "
- ++ "For quantile number q, "
- ++ "each quantile index i "
- ++ "should satisfy 0 <= i <= q, "
- ++ "but here q is "
- ++ show q
- ++ " and indexes are "
- ++ show qs
-
-addCallPointInfo :: Maybe String -> Maybe String -> String -> String
-addCallPointInfo (Just name) (Just cp) err =
- err
- ++ ( "\n\tThis happened when calling function "
- ++ brightGreen cp
- ++ " on "
- ++ brightGreen name
- )
-addCallPointInfo Nothing (Just cp) err =
- err
- ++ ( "\n\tThis happened when calling function "
- ++ brightGreen cp
- )
-addCallPointInfo (Just name) Nothing err =
- err
- ++ ( "\n\tOn "
- ++ name
- ++ "\n\n"
- )
-addCallPointInfo Nothing Nothing err = err
-
-guessColumnName :: T.Text -> [T.Text] -> T.Text
-guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of
- [] -> ""
- res -> (snd . minimum) res
-
-editDistance :: T.Text -> T.Text -> Int
-editDistance xs ys = table ML.! (m, n)
- where
- (m, n) = (T.length xs, T.length ys)
- xv = V.fromList (T.unpack xs)
- yv = V.fromList (T.unpack ys)
- table :: ML.Map (Int, Int) Int
- table = ML.fromList [((i, j), dist i j) | i <- [0 .. m], j <- [0 .. n]]
- dist 0 j = j
- dist i 0 = i
- dist i j =
- minimum
- [ table ML.! (i - 1, j) + 1
- , table ML.! (i, j - 1) + 1
- , (if xv V.! (i - 1) == yv V.! (j - 1) then 0 else 1)
- + table ML.! (i - 1, j - 1)
- ]
diff --git a/src/DataFrame/Internal/AggKernel.hs b/src/DataFrame/Internal/AggKernel.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/AggKernel.hs
+++ /dev/null
@@ -1,308 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Vectorized scatter-accumulate aggregation kernel.
-
-The grouping layer ('DataFrame.Internal.Grouping') hands us a dense
-@rowToGroup@ vector (group id per row, canonical order) plus the number of
-groups. For the common reductions this kernel replaces the per-group boxed
-expression-interpreter fold with a single unboxed linear pass that scatters
-each row's value into primitive per-group accumulator arrays indexed by the
-group id. No boxed accumulator record, no per-element dictionary closure: the
-element type is resolved once per column by a 'typeRep' switch and the inner
-loop is a monomorphic primop on a 'VU.Vector'.
-
-Result columns are length @nGroups@ in canonical group order, so they line up
-with the key columns 'aggregate' gathers with @selectIndices@.
-
-The kernel is strictly a FAST PATH: the matcher 'DataFrame.Internal.AggPlan.planAgg'
-recognises a small set of expression shapes; anything it does not recognise
-keeps the existing interpreter, so the general @aggregate@ API stays correct for
-arbitrary expressions.
--}
-module DataFrame.Internal.AggKernel (
- Reduction (..),
- scatterReduce,
- scatterColumnToDouble,
-) where
-
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Monad (when)
-import Control.Monad.ST (ST, runST)
-import DataFrame.Internal.Column (
- Column (..),
- Columnable,
- fromUnboxedVector,
- materializePacked,
- )
-import Type.Reflection (typeRep)
-
-{- | A recognised fast-path reduction over a single value column. The element
-type (Int vs Double) is resolved at scatter time; sum/min/max preserve the
-column's element type, everything else produces a Double column.
--}
-data Reduction
- = RSum
- | RCount
- | RMin
- | RMax
- | RMean
- | RStd
- | RVar
- | RTop2Sum
- deriving (Eq, Show)
-
--------------------------------------------------------------------------------
--- Column extraction
--------------------------------------------------------------------------------
-
-{- | Coerce an unboxed Int or Double column to an unboxed Double vector for the
-moment/mean/sd/median family. Returns 'Nothing' for boxed, nullable, or other
-element types (the caller then falls back to the interpreter).
--}
-scatterColumnToDouble :: Column -> Maybe (VU.Vector Double)
-scatterColumnToDouble = \case
- UnboxedColumn Nothing (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> Just v
- Nothing -> case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> Just (VU.map fromIntegral v)
- Nothing -> Nothing
- p@(PackedText _ _) -> scatterColumnToDouble (materializePacked p)
- _ -> Nothing
-
--------------------------------------------------------------------------------
--- Scatter reductions
--------------------------------------------------------------------------------
-
-{- | Run one fast-path reduction. Returns 'Nothing' when the value column is
-not a non-null unboxed Int/Double column (then the caller falls back).
--}
-scatterReduce ::
- Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column
-scatterReduce red g nGroups col = case col of
- UnboxedColumn Nothing (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> Just (reduceTyped red g nGroups v intIdent)
- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> Just (reduceTyped red g nGroups v dblIdent)
- Nothing -> Nothing
- p@(PackedText _ _) -> scatterReduce red g nGroups (materializePacked p)
- _ -> Nothing
-{-# INLINEABLE scatterReduce #-}
-
--- | Per-type seed identities for the order-preserving reductions.
-data Idents a = Idents {minSeed :: !a, maxSeed :: !a}
-
-intIdent :: Idents Int
-intIdent = Idents maxBound minBound
-
-dblIdent :: Idents Double
-dblIdent = Idents (1 / 0) (negate (1 / 0))
-
-{- | The monomorphic reduction body. @count@ always yields an Int column;
-@sum@/@min@/@max@ preserve the element type; @mean@/@std@/@var@ produce Double;
-@top2Sum@ produces Double.
--}
-reduceTyped ::
- forall a.
- (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>
- Reduction -> VU.Vector Int -> Int -> VU.Vector a -> Idents a -> Column
-reduceTyped red g nGroups v idents = case red of
- RCount -> fromUnboxedVector (countScatter g nGroups)
- RSum -> fromUnboxedVector (sumScatter g nGroups v)
- RMin -> fromUnboxedVector (extremaScatter min (minSeed idents) g nGroups v)
- RMax -> fromUnboxedVector (extremaScatter max (maxSeed idents) g nGroups v)
- RMean -> fromUnboxedVector (meanScatter g nGroups v)
- RVar -> fromUnboxedVector (varScatter False g nGroups v)
- RStd -> fromUnboxedVector (varScatter True g nGroups v)
- RTop2Sum -> fromUnboxedVector (top2Scatter g nGroups v)
-{-# INLINE reduceTyped #-}
-
-countScatter :: VU.Vector Int -> Int -> VU.Vector Int
-countScatter g nGroups = runST $ do
- cnt <- VUM.replicate nGroups (0 :: Int)
- let n = VU.length g
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- c <- VUM.unsafeRead cnt k
- VUM.unsafeWrite cnt k (c + 1)
- go (i + 1)
- go 0
- VU.unsafeFreeze cnt
-
-sumScatter ::
- (VU.Unbox a, Num a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a
-sumScatter g nGroups v = runST $ do
- s <- VUM.replicate nGroups 0
- let n = VU.length v
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- cur <- VUM.unsafeRead s k
- VUM.unsafeWrite s k (cur + VU.unsafeIndex v i)
- go (i + 1)
- go 0
- VU.unsafeFreeze s
-{-# INLINE sumScatter #-}
-
-extremaScatter ::
- (VU.Unbox a) =>
- (a -> a -> a) -> a -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector a
-extremaScatter combine seed g nGroups v = runST $ do
- m <- VUM.replicate nGroups seed
- let n = VU.length v
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- cur <- VUM.unsafeRead m k
- VUM.unsafeWrite m k (combine cur (VU.unsafeIndex v i))
- go (i + 1)
- go 0
- VU.unsafeFreeze m
-{-# INLINE extremaScatter #-}
-
-meanScatter ::
- (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
-meanScatter g nGroups v = runST $ do
- s <- VUM.replicate nGroups (0 :: Double)
- cnt <- VUM.replicate nGroups (0 :: Int)
- scatterSumCount g v s cnt
- finalizeMean nGroups s cnt
-{-# INLINE meanScatter #-}
-
-{- | One pass filling running sum and count from value column @v@ over groups
-@g@ into the supplied accumulator arrays.
--}
-scatterSumCount ::
- (VU.Unbox a, Real a) =>
- VU.Vector Int ->
- VU.Vector a ->
- VUM.MVector s Double ->
- VUM.MVector s Int ->
- ST s ()
-scatterSumCount g v s cnt = go 0
- where
- n = VU.length v
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- !x = realToFrac (VU.unsafeIndex v i)
- curS <- VUM.unsafeRead s k
- VUM.unsafeWrite s k (curS + x)
- curC <- VUM.unsafeRead cnt k
- VUM.unsafeWrite cnt k (curC + 1)
- go (i + 1)
-{-# INLINE scatterSumCount #-}
-
-finalizeMean ::
- Int -> VUM.MVector s Double -> VUM.MVector s Int -> ST s (VU.Vector Double)
-finalizeMean nGroups s cnt = do
- out <- VUM.new nGroups
- let go !k
- | k >= nGroups = pure ()
- | otherwise = do
- sv <- VUM.unsafeRead s k
- c <- VUM.unsafeRead cnt k
- VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)
- go (k + 1)
- go 0
- VU.unsafeFreeze out
-
-{- | Sample variance (or its square root for sd) via a per-group Welford
-recurrence, scattered into three unboxed arrays @(count, mean, m2)@. This is the
-same numerically-stable update as the interpreter's @varianceStep@, so the
-result is byte-identical to the existing CollectAgg path (and the db-benchmark
-checksum is unchanged); the win is the unboxed scatter replacing the per-group
-boxed @VarAcc@ fold over a materialized slice. Degenerate groups (n < 2) yield
-0, matching @computeVariance@.
--}
-varScatter ::
- (VU.Unbox a, Real a) =>
- Bool -> VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
-varScatter takeSqrt g nGroups v = runST $ do
- cnt <- VUM.replicate nGroups (0 :: Int)
- meanV <- VUM.replicate nGroups (0 :: Double)
- m2 <- VUM.replicate nGroups (0 :: Double)
- let n = VU.length v
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- !x = realToFrac (VU.unsafeIndex v i)
- c <- VUM.unsafeRead cnt k
- mu <- VUM.unsafeRead meanV k
- mm <- VUM.unsafeRead m2 k
- let !c' = c + 1
- !delta = x - mu
- !mu' = mu + delta / fromIntegral c'
- !mm' = mm + delta * (x - mu')
- VUM.unsafeWrite cnt k c'
- VUM.unsafeWrite meanV k mu'
- VUM.unsafeWrite m2 k mm'
- go (i + 1)
- go 0
- out <- VUM.new nGroups
- let fin !k
- | k >= nGroups = pure ()
- | otherwise = do
- c <- VUM.unsafeRead cnt k
- mm <- VUM.unsafeRead m2 k
- let var = if c < 2 then 0 else mm / fromIntegral (c - 1)
- VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)
- fin (k + 1)
- fin 0
- VU.unsafeFreeze out
-{-# INLINE varScatter #-}
-
-{- | Per-group sum of the two largest values: a 2-slot scatter holding the
-running first and second maximum, then @m1 + m2@. Matches the @take 2 . sortBy
-(flip compare)@ definition used by the benchmark's @top2Sum@.
--}
-top2Scatter ::
- (VU.Unbox a, Real a) => VU.Vector Int -> Int -> VU.Vector a -> VU.Vector Double
-top2Scatter g nGroups v = runST $ do
- let ninf = negate (1 / 0) :: Double
- m1 <- VUM.replicate nGroups ninf
- m2 <- VUM.replicate nGroups ninf
- let n = VU.length v
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- !x = realToFrac (VU.unsafeIndex v i)
- a1 <- VUM.unsafeRead m1 k
- if x > a1
- then do
- VUM.unsafeWrite m1 k x
- VUM.unsafeWrite m2 k a1
- else do
- a2 <- VUM.unsafeRead m2 k
- when (x > a2) (VUM.unsafeWrite m2 k x)
- go (i + 1)
- go 0
- out <- VUM.new nGroups
- let fin !k
- | k >= nGroups = pure ()
- | otherwise = do
- a1 <- VUM.unsafeRead m1 k
- a2 <- VUM.unsafeRead m2 k
- let s = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)
- VUM.unsafeWrite out k s
- fin (k + 1)
- fin 0
- VU.unsafeFreeze out
-{-# INLINE top2Scatter #-}
diff --git a/src/DataFrame/Internal/AggKernelDirect.hs b/src/DataFrame/Internal/AggKernelDirect.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/AggKernelDirect.hs
+++ /dev/null
@@ -1,373 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Low-cardinality DIRECT-INDEXED accumulator fast path (research #4).
-
-When the group-by key resolves to a small dense domain (@nGroups@ below
-'directThreshold'), the dense @rowToGroup@ vector handed down by the grouping
-layer already IS the group id per row, so we can bypass the @valueIndices@
-gather entirely: scan @rowToGroup@ and the value column /in lockstep, in
-original row order/, scattering each value straight into a per-group accumulator
-array indexed by the dense id. Both arrays are read sequentially (no random
-gather through @valueIndices@), and for a small domain the accumulator stays
-cache-resident.
-
-For parallelism we use the two-phase morsel pattern (#2): split the ROW range
-into one contiguous chunk per capability, give each worker a private tiny
-accumulator array, scan its chunk sequentially, then merge the @caps@ partials in
-a cheap O(caps * nGroups) pass. This makes few-group questions scale (the work
-per worker is a tight sequential pass) instead of being dominated by parallel
-fan-out, which is exactly the regression the group-range gather suffered on Q4.
-
-CRITICAL — byte-identical at any @-N@: the two-phase merge only changes the
-fold/combine order, so it is admitted ONLY for reductions whose result is
-independent of that order: integer sum (exact, associative), count, min, max,
-and integer mean (an exact integer sum and count divided once at finalize). The
-order-sensitive float reductions (Double sum/mean, variance, sd, top2) are NOT
-routed here — the caller keeps the order-preserving group-range kernel for them,
-so the db-benchmark checksums stay byte-identical between -N1 and -N8.
--}
-module DataFrame.Internal.AggKernelDirect (
- directThreshold,
- directReduce,
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import System.IO.Unsafe (unsafePerformIO)
-import Type.Reflection (typeRep)
-
-import DataFrame.Internal.AggKernel (Reduction (..))
-import DataFrame.Internal.Column (
- Column (..),
- fromUnboxedVector,
- materializePacked,
- )
-
-{- | Group-domain size at or below which the direct-indexed accumulator path is
-taken. The db-benchmark low-cardinality questions (id1=100, id4=100, id6=1e5,
-id2:id4=1e4) sit at or below it; wider domains keep the group-range kernel. The
-admitted reductions are order-independent, so the per-worker accumulator merge is
-exact regardless of size.
--}
-directThreshold :: Int
-directThreshold = 262144
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
-{- | Below this many rows the parallel fan-out is not worth it; a single
-sequential direct pass runs instead (tiny accumulator, one tight loop). Matches
-the grouping/scatter parallel threshold.
--}
-parThreshold :: Int
-parThreshold = 200000
-
-{- | Run a recognised reduction through the direct-indexed path. Returns
-'Nothing' (so the caller falls back to the order-preserving kernel) unless BOTH:
-(a) the reduction's result is order-independent at this element type — see the
-module note — and (b) the value column is a clean unboxed Int/Double column.
-
-@g@ is the dense @rowToGroup@ vector (group id per row, original row order);
-@nGroups@ is the dense domain size, already verified @<= directThreshold@ by the
-caller.
--}
-directReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column
-directReduce red g nGroups col = case col of
- UnboxedColumn Nothing (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> directInt red g nGroups v
- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> directDouble red g nGroups v
- Nothing -> Nothing
- p@(PackedText _ _) -> directReduce red g nGroups (materializePacked p)
- _ -> Nothing
-{-# INLINEABLE directReduce #-}
-
--- | The order-independent reductions over an Int column.
-directInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column
-directInt red g nGroups v = case red of
- RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))
- RSum -> Just (fromUnboxedVector (sumIntDirect g nGroups v))
- RMin -> Just (fromUnboxedVector (extremaIntDirect True g nGroups v))
- RMax -> Just (fromUnboxedVector (extremaIntDirect False g nGroups v))
- RMean -> Just (fromUnboxedVector (meanIntDirect g nGroups v))
- _ -> Nothing
-
-{- | Over a Double column only @count@ is order-independent; the float
-sum/mean/variance reductions must keep the order-preserving kernel.
--}
-directDouble ::
- Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column
-directDouble red g nGroups v = case red of
- RCount -> Just (fromUnboxedVector (countDirect g nGroups (VU.length v)))
- _ -> Nothing
-
--- | Whether to fan out at this row count.
-shouldPar :: Int -> Bool
-shouldPar n = n >= parThreshold && capabilities > 1
-
-{- | Fork @caps@ workers over disjoint contiguous ROW ranges @[lo, hi)@ of
-@[0, n)@, balanced evenly by row count; each worker @w@ runs @fill lo hi@
-producing its OWN private accumulator (thread-local, no shared array, no sync).
-Returns the partials in worker order for the caller's merge. Rethrows the first
-worker failure.
--}
-runPartialsOver ::
- Int -> Int -> (Int -> Int -> IO (VUM.IOVector Int)) -> IO [VUM.IOVector Int]
-runPartialsOver n caps fill = do
- let !per = (n + caps - 1) `div` caps
- spawn w = do
- var <- newEmptyMVar
- let !lo = min n (w * per)
- !hi = min n (lo + per)
- _ <- forkIO (try (fill lo hi) >>= putMVar var)
- pure var
- vars <- mapM spawn [0 .. caps - 1]
- results <- mapM takeMVar vars
- mapM (either (throwIO @SomeException) pure) results
-
-{- | As 'runPartialsOver' but each worker produces a PAIR of accumulators (e.g.
-sum and count for the fused integer mean).
--}
-runPartialsPairOver ::
- Int ->
- Int ->
- (Int -> Int -> IO (VUM.IOVector Int, VUM.IOVector Int)) ->
- IO [(VUM.IOVector Int, VUM.IOVector Int)]
-runPartialsPairOver n caps fill = do
- let !per = (n + caps - 1) `div` caps
- spawn w = do
- var <- newEmptyMVar
- let !lo = min n (w * per)
- !hi = min n (lo + per)
- _ <- forkIO (try (fill lo hi) >>= putMVar var)
- pure var
- vars <- mapM spawn [0 .. caps - 1]
- results <- mapM takeMVar vars
- mapM (either (throwIO @SomeException) pure) results
-
--------------------------------------------------------------------------------
--- Count (order-independent: per-group row count)
--------------------------------------------------------------------------------
-
-countDirect :: VU.Vector Int -> Int -> Int -> VU.Vector Int
-countDirect g nGroups n
- | not (shouldPar n) =
- unsafePerformIO (countChunk g nGroups 0 n >>= VU.unsafeFreeze)
- | otherwise = unsafePerformIO $ do
- parts <- runPartialsOver n capabilities (countChunk g nGroups)
- mergeIntSum nGroups parts
-{-# NOINLINE countDirect #-}
-
-countChunk :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
-countChunk g nGroups lo hi = do
- acc <- VUM.replicate nGroups (0 :: Int)
- let go !i
- | i >= hi = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- c <- VUM.unsafeRead acc k
- VUM.unsafeWrite acc k (c + 1)
- go (i + 1)
- go lo
- pure acc
-
--------------------------------------------------------------------------------
--- Integer sum (exact: merge order irrelevant)
--------------------------------------------------------------------------------
-
-sumIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
-sumIntDirect g nGroups v
- | not (shouldPar n) =
- unsafePerformIO (sumIntChunk g v nGroups 0 n >>= VU.unsafeFreeze)
- | otherwise = unsafePerformIO $ do
- parts <- runPartialsOver n capabilities (sumIntChunk g v nGroups)
- mergeIntSum nGroups parts
- where
- !n = VU.length v
-{-# NOINLINE sumIntDirect #-}
-
-sumIntChunk ::
- VU.Vector Int -> VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
-sumIntChunk g v nGroups lo hi = do
- acc <- VUM.replicate nGroups (0 :: Int)
- let go !i
- | i >= hi = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- c <- VUM.unsafeRead acc k
- VUM.unsafeWrite acc k (c + VU.unsafeIndex v i)
- go (i + 1)
- go lo
- pure acc
-
--------------------------------------------------------------------------------
--- Integer min / max (order-independent)
--------------------------------------------------------------------------------
-
-extremaIntDirect ::
- Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
-extremaIntDirect isMin g nGroups v
- | not (shouldPar n) =
- unsafePerformIO (extremaIntChunk isMin g v nGroups 0 n >>= VU.unsafeFreeze)
- | otherwise = unsafePerformIO $ do
- parts <- runPartialsOver n capabilities (extremaIntChunk isMin g v nGroups)
- mergeExtremaInt isMin nGroups parts
- where
- !n = VU.length v
-{-# NOINLINE extremaIntDirect #-}
-
-extremaIntChunk ::
- Bool ->
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- Int ->
- Int ->
- IO (VUM.IOVector Int)
-extremaIntChunk isMin g v nGroups lo hi = do
- let !seed = if isMin then maxBound else minBound
- combine a b = if isMin then min a b else max a b
- acc <- VUM.replicate nGroups seed
- let go !i
- | i >= hi = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- c <- VUM.unsafeRead acc k
- VUM.unsafeWrite acc k (combine c (VU.unsafeIndex v i))
- go (i + 1)
- go lo
- pure acc
-
--------------------------------------------------------------------------------
--- Integer mean (exact integer sum + count, divided once -> order-independent)
--------------------------------------------------------------------------------
-
-{- | Integer mean in ONE fused pass: a running integer sum and count per group,
-divided once at finalize. The integer sum is exact, so the parallel partial
-merge is byte-identical to the sequential single pass at any @-N@.
--}
-meanIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double
-meanIntDirect g nGroups v
- | not (shouldPar n) = unsafePerformIO $ do
- (s, c) <- meanIntChunk g v nGroups 0 n
- finalizeMeanInt nGroups s c
- | otherwise = unsafePerformIO $ do
- parts <- runPartialsPairOver n capabilities (meanIntChunk g v nGroups)
- (s, c) <- mergePair nGroups parts
- finalizeMeanInt nGroups s c
- where
- !n = VU.length v
-{-# NOINLINE meanIntDirect #-}
-
-meanIntChunk ::
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- Int ->
- Int ->
- IO (VUM.IOVector Int, VUM.IOVector Int)
-meanIntChunk g v nGroups lo hi = do
- s <- VUM.replicate nGroups (0 :: Int)
- c <- VUM.replicate nGroups (0 :: Int)
- let go !i
- | i >= hi = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- sv <- VUM.unsafeRead s k
- VUM.unsafeWrite s k (sv + VU.unsafeIndex v i)
- cv <- VUM.unsafeRead c k
- VUM.unsafeWrite c k (cv + 1)
- go (i + 1)
- go lo
- pure (s, c)
-
-finalizeMeanInt ::
- Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO (VU.Vector Double)
-finalizeMeanInt nGroups s c = do
- out <- VUM.new nGroups
- let go !k
- | k >= nGroups = pure ()
- | otherwise = do
- sv <- VUM.unsafeRead s k
- cv <- VUM.unsafeRead c k
- VUM.unsafeWrite
- out
- k
- (if cv == 0 then 0 / 0 else fromIntegral sv / fromIntegral cv)
- go (k + 1)
- go 0
- VU.unsafeFreeze out
-
--------------------------------------------------------------------------------
--- Partial accumulation + merge
--------------------------------------------------------------------------------
-
-mergeIntSum :: Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
-mergeIntSum nGroups parts = case parts of
- [] -> VU.unsafeFreeze =<< VUM.replicate nGroups 0
- (p0 : rest) -> do
- let add !p = do
- let go !k
- | k >= nGroups = pure ()
- | otherwise = do
- a <- VUM.unsafeRead p0 k
- b <- VUM.unsafeRead p k
- VUM.unsafeWrite p0 k (a + b)
- go (k + 1)
- go 0
- mapM_ add rest
- VU.unsafeFreeze p0
-
-{- | Merge per-worker (sum, count) partials into the first worker's pair by
-exact integer addition; returns the accumulated pair for finalize.
--}
-mergePair ::
- Int ->
- [(VUM.IOVector Int, VUM.IOVector Int)] ->
- IO (VUM.IOVector Int, VUM.IOVector Int)
-mergePair nGroups parts = case parts of
- [] -> (,) <$> VUM.replicate nGroups 0 <*> VUM.replicate nGroups 0
- ((s0, c0) : rest) -> do
- let add (s, c) = do
- let go !k
- | k >= nGroups = pure ()
- | otherwise = do
- sa <- VUM.unsafeRead s0 k
- sb <- VUM.unsafeRead s k
- VUM.unsafeWrite s0 k (sa + sb)
- ca <- VUM.unsafeRead c0 k
- cb <- VUM.unsafeRead c k
- VUM.unsafeWrite c0 k (ca + cb)
- go (k + 1)
- go 0
- mapM_ add rest
- pure (s0, c0)
-
-mergeExtremaInt :: Bool -> Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
-mergeExtremaInt isMin nGroups parts = case parts of
- [] ->
- VU.unsafeFreeze =<< VUM.replicate nGroups (if isMin then maxBound else minBound)
- (p0 : rest) -> do
- let combine a b = if isMin then min a b else max a b
- add !p = do
- let go !k
- | k >= nGroups = pure ()
- | otherwise = do
- a <- VUM.unsafeRead p0 k
- b <- VUM.unsafeRead p k
- VUM.unsafeWrite p0 k (combine a b)
- go (k + 1)
- go 0
- mapM_ add rest
- VU.unsafeFreeze p0
diff --git a/src/DataFrame/Internal/AggKernelPar.hs b/src/DataFrame/Internal/AggKernelPar.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/AggKernelPar.hs
+++ /dev/null
@@ -1,454 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Parallel scatter-accumulate aggregation kernel.
-
-This builds on the sequential kernel ('DataFrame.Internal.AggKernel') and the
-Round-5 grouping layout: 'groupBy' hands us @valueIndices@ (rows ordered by
-group) and @offsets@ (per-group boundaries), so each group's rows are a
-contiguous run @valueIndices[offsets[g] .. offsets[g+1])@ in their original row
-order.
-
-The parallel driver splits the dense group-id range @[0, nGroups)@ into one
-contiguous chunk per capability, balanced by row count. Because group ranges are
-disjoint and their @valueIndices@ runs are disjoint, every worker reads and
-writes its own slice of the shared output array(s) — there is NO cross-worker
-overlap and NO merge. And because each group's rows are visited in the same
-original-row order as the sequential @rowToGroup@ scan, the per-group fold order
-is unchanged, so the result is /byte-identical/ to the sequential kernel at any
-@-N@ (no fold-order drift, even for the float sums).
-
-Forks plain 'forkIO' workers (no sparks), one per capability; falls back to the
-sequential 'DataFrame.Internal.AggKernel' path when @caps == 1@ or the row count
-is below 'parThreshold'.
--}
-module DataFrame.Internal.AggKernelPar (
- scatterReducePar,
- momentScatterPar,
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import Control.Monad (when)
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import System.IO.Unsafe (unsafePerformIO)
-import Type.Reflection (typeRep)
-
-import DataFrame.Internal.AggKernel (
- Reduction (..),
- scatterColumnToDouble,
- scatterReduce,
- )
-import DataFrame.Internal.AggPlan (Moments (..), momentScatter)
-import DataFrame.Internal.Column (
- Column (..),
- Columnable,
- fromUnboxedVector,
- materializePacked,
- )
-
--------------------------------------------------------------------------------
--- Parallelisation policy
--------------------------------------------------------------------------------
-
-{- | Below this many rows the fork overhead is not worth it; the sequential
-kernel runs instead. Mirrors 'DataFrame.Internal.GroupingPar.parThreshold'.
--}
-parThreshold :: Int
-parThreshold = 200000
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
--- | Whether to take the parallel path at this row count.
-shouldPar :: Int -> Bool
-shouldPar n = n >= parThreshold && capabilities > 1
-
--------------------------------------------------------------------------------
--- Group-range partitioning by row count
--------------------------------------------------------------------------------
-
-{- | Split @[0, nGroups)@ into @caps@ contiguous group ranges, each holding a
-near-equal share of rows. Returns @caps + 1@ boundaries @b@ with @b[0] == 0@ and
-@b[caps] == nGroups@; worker @w@ owns groups @[b[w], b[w+1])@. A row-balanced
-split keeps skew low even when group sizes vary wildly.
--}
-groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int
-groupRangeBounds offs nGroups caps = VU.create $ do
- b <- VUM.new (caps + 1)
- let !nRows = VU.unsafeIndex offs nGroups
- !per = max 1 ((nRows + caps - 1) `div` caps)
- -- First group whose start offset reaches @target@ (>= prev), or nGroups.
- adv !target !gg
- | gg >= nGroups = nGroups
- | VU.unsafeIndex offs gg >= target = gg
- | otherwise = adv target (gg + 1)
- go !w !prev
- | w >= caps = VUM.unsafeWrite b caps nGroups
- | otherwise = do
- let !target = min nRows (w * per)
- !g = adv target prev
- VUM.unsafeWrite b w g
- go (w + 1) g
- VUM.unsafeWrite b 0 0
- go 1 0
- pure b
-
-{- | Fork @caps@ workers, worker @w@ running @act (b[w]) (b[w+1])@ over its
-disjoint group range; join and rethrow the first failure. Sequential when
-@caps == 1@.
--}
-forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
-forEachRange bounds caps act
- | caps <= 1 = act (VU.unsafeIndex bounds 0) (VU.unsafeIndex bounds caps)
- | otherwise = do
- vars <- mapM spawn [0 .. caps - 1]
- results <- mapM takeMVar vars
- mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
- where
- spawn w = do
- var <- newEmptyMVar
- let !s = VU.unsafeIndex bounds w
- !e = VU.unsafeIndex bounds (w + 1)
- _ <- forkIO (try (act s e) >>= putMVar var)
- pure var
-
--------------------------------------------------------------------------------
--- Parallel single-column reductions
--------------------------------------------------------------------------------
-
-{- | Parallel counterpart of 'scatterReduce'. Returns 'Nothing' on the same
-columns the sequential kernel rejects (boxed/nullable/non-Int-Double); on the
-sequential path or tiny inputs it delegates to 'scatterReduce'. The result is
-byte-identical to 'scatterReduce' (same per-group fold order).
--}
-scatterReducePar ::
- Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> Maybe Column
-scatterReducePar red vis offs nGroups col
- | not (shouldPar (VU.length vis)) || nGroups <= 1 =
- scatterReduce red (rtgFromVis vis offs nGroups) nGroups col
- | otherwise = case col of
- UnboxedColumn Nothing (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> Just (reduceParTyped red vis offs nGroups v intIdent)
- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> Just (reduceParTyped red vis offs nGroups v dblIdent)
- Nothing -> Nothing
- p@(PackedText _ _) -> scatterReducePar red vis offs nGroups (materializePacked p)
- _ -> Nothing
-{-# NOINLINE scatterReducePar #-}
-
-{- | Reconstruct @rowToGroup@ from the group layout for the sequential delegate.
-Only used on the small/-N1 path, so the extra pass is negligible.
--}
-rtgFromVis :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int
-rtgFromVis vis offs nGroups = VU.create $ do
- let n = VU.length vis
- rtg <- VUM.new (max 1 n)
- let go !g
- | g >= nGroups = pure ()
- | otherwise = do
- let !e = VU.unsafeIndex offs (g + 1)
- inner !pos
- | pos >= e = pure ()
- | otherwise = do
- VUM.unsafeWrite rtg (VU.unsafeIndex vis pos) g
- inner (pos + 1)
- inner (VU.unsafeIndex offs g)
- go (g + 1)
- go 0
- pure rtg
-
-data Idents a = Idents {minSeed :: !a, maxSeed :: !a}
-
-intIdent :: Idents Int
-intIdent = Idents maxBound minBound
-
-dblIdent :: Idents Double
-dblIdent = Idents (1 / 0) (negate (1 / 0))
-
-{- | The monomorphic parallel reduction body. Each scatter allocates its full
-@nGroups@ output array(s) once, then workers fill disjoint group ranges in
-parallel. The result type follows the sequential kernel exactly.
--}
-reduceParTyped ::
- forall a.
- (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>
- Reduction ->
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Idents a ->
- Column
-reduceParTyped red vis offs nGroups v idents =
- let !caps = capabilities
- !bounds = groupRangeBounds offs nGroups caps
- in case red of
- RCount -> fromUnboxedVector (unsafePerformIO (countPar vis offs nGroups caps bounds))
- RSum -> fromUnboxedVector (unsafePerformIO (sumPar vis offs nGroups v caps bounds))
- RMin ->
- fromUnboxedVector
- (unsafePerformIO (extremaPar min (minSeed idents) vis offs nGroups v caps bounds))
- RMax ->
- fromUnboxedVector
- (unsafePerformIO (extremaPar max (maxSeed idents) vis offs nGroups v caps bounds))
- RMean -> fromUnboxedVector (unsafePerformIO (meanPar vis offs nGroups v caps bounds))
- RVar ->
- fromUnboxedVector
- (unsafePerformIO (varPar False vis offs nGroups v caps bounds))
- RStd ->
- fromUnboxedVector (unsafePerformIO (varPar True vis offs nGroups v caps bounds))
- RTop2Sum -> fromUnboxedVector (unsafePerformIO (top2Par vis offs nGroups v caps bounds))
-{-# INLINE reduceParTyped #-}
-
--- | Iterate the rows of groups @[gs, ge)@ in @valueIndices@/group order.
-overGroups ::
- VU.Vector Int -> VU.Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
-overGroups vis offs gs ge step = grp gs
- where
- grp !g
- | g >= ge = pure ()
- | otherwise = do
- let !e = VU.unsafeIndex offs (g + 1)
- inner !pos
- | pos >= e = pure ()
- | otherwise = step g (VU.unsafeIndex vis pos) >> inner (pos + 1)
- inner (VU.unsafeIndex offs g)
- grp (g + 1)
-{-# INLINE overGroups #-}
-
-countPar ::
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector Int)
-countPar _vis offs nGroups caps bounds = do
- out <- VUM.replicate nGroups (0 :: Int)
- forEachRange bounds caps $ \gs ge ->
- let grp !g
- | g >= ge = pure ()
- | otherwise = do
- let !c = VU.unsafeIndex offs (g + 1) - VU.unsafeIndex offs g
- VUM.unsafeWrite out g c
- grp (g + 1)
- in grp gs
- VU.unsafeFreeze out
-
-sumPar ::
- (VU.Unbox a, Num a) =>
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector a)
-sumPar vis offs nGroups v caps bounds = do
- out <- VUM.replicate nGroups 0
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- cur <- VUM.unsafeRead out g
- VUM.unsafeWrite out g (cur + VU.unsafeIndex v row)
- VU.unsafeFreeze out
-{-# INLINE sumPar #-}
-
-extremaPar ::
- (VU.Unbox a) =>
- (a -> a -> a) ->
- a ->
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector a)
-extremaPar combine seed vis offs nGroups v caps bounds = do
- out <- VUM.replicate nGroups seed
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- cur <- VUM.unsafeRead out g
- VUM.unsafeWrite out g (combine cur (VU.unsafeIndex v row))
- VU.unsafeFreeze out
-{-# INLINE extremaPar #-}
-
-meanPar ::
- (VU.Unbox a, Real a) =>
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector Double)
-meanPar vis offs nGroups v caps bounds = do
- s <- VUM.replicate nGroups (0 :: Double)
- cnt <- VUM.replicate nGroups (0 :: Int)
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- let !x = realToFrac (VU.unsafeIndex v row)
- cs <- VUM.unsafeRead s g
- VUM.unsafeWrite s g (cs + x)
- cc <- VUM.unsafeRead cnt g
- VUM.unsafeWrite cnt g (cc + 1)
- out <- VUM.new nGroups
- let fin !k
- | k >= nGroups = pure ()
- | otherwise = do
- sv <- VUM.unsafeRead s k
- c <- VUM.unsafeRead cnt k
- VUM.unsafeWrite out k (if c == 0 then 0 / 0 else sv / fromIntegral c)
- fin (k + 1)
- fin 0
- VU.unsafeFreeze out
-{-# INLINE meanPar #-}
-
-{- | Per-group Welford variance/sd, parallel by group range. The recurrence is
-applied in original-row order within each group, identical to the sequential
-@varScatter@, so the result is byte-identical (no parallel-combine needed: each
-group lives wholly inside one worker's range).
--}
-varPar ::
- (VU.Unbox a, Real a) =>
- Bool ->
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector Double)
-varPar takeSqrt vis offs nGroups v caps bounds = do
- cnt <- VUM.replicate nGroups (0 :: Int)
- meanV <- VUM.replicate nGroups (0 :: Double)
- m2 <- VUM.replicate nGroups (0 :: Double)
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- let !x = realToFrac (VU.unsafeIndex v row)
- c <- VUM.unsafeRead cnt g
- mu <- VUM.unsafeRead meanV g
- mm <- VUM.unsafeRead m2 g
- let !c' = c + 1
- !delta = x - mu
- !mu' = mu + delta / fromIntegral c'
- !mm' = mm + delta * (x - mu')
- VUM.unsafeWrite cnt g c'
- VUM.unsafeWrite meanV g mu'
- VUM.unsafeWrite m2 g mm'
- out <- VUM.new nGroups
- let fin !k
- | k >= nGroups = pure ()
- | otherwise = do
- c <- VUM.unsafeRead cnt k
- mm <- VUM.unsafeRead m2 k
- let var = if c < 2 then 0 else mm / fromIntegral (c - 1)
- VUM.unsafeWrite out k (if takeSqrt then sqrt var else var)
- fin (k + 1)
- fin 0
- VU.unsafeFreeze out
-{-# INLINE varPar #-}
-
-top2Par ::
- (VU.Unbox a, Real a) =>
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector a ->
- Int ->
- VU.Vector Int ->
- IO (VU.Vector Double)
-top2Par vis offs nGroups v caps bounds = do
- let ninf = negate (1 / 0) :: Double
- m1 <- VUM.replicate nGroups ninf
- m2 <- VUM.replicate nGroups ninf
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- let !x = realToFrac (VU.unsafeIndex v row)
- a1 <- VUM.unsafeRead m1 g
- if x > a1
- then do
- VUM.unsafeWrite m1 g x
- VUM.unsafeWrite m2 g a1
- else do
- a2 <- VUM.unsafeRead m2 g
- when (x > a2) (VUM.unsafeWrite m2 g x)
- out <- VUM.new nGroups
- let fin !k
- | k >= nGroups = pure ()
- | otherwise = do
- a1 <- VUM.unsafeRead m1 k
- a2 <- VUM.unsafeRead m2 k
- let sm = (if isInfinite a1 then 0 else a1) + (if isInfinite a2 then 0 else a2)
- VUM.unsafeWrite out k sm
- fin (k + 1)
- fin 0
- VU.unsafeFreeze out
-{-# INLINE top2Par #-}
-
--------------------------------------------------------------------------------
--- Parallel fused two-column moments (Q9)
--------------------------------------------------------------------------------
-
-{- | Parallel counterpart of 'momentScatter': one fused pass over both columns,
-each group's six sums accumulated in original-row order inside one worker's
-range. Byte-identical to 'momentScatter'; delegates to it on the sequential
-path. Returns 'Nothing' unless both columns are non-null unboxed Int/Double.
--}
-momentScatterPar ::
- VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Moments
-momentScatterPar vis offs nGroups colX colY
- | not (shouldPar (VU.length vis)) || nGroups <= 1 =
- momentScatter (rtgFromVis vis offs nGroups) nGroups colX colY
- | otherwise = do
- xs <- scatterColumnToDouble colX
- ys <- scatterColumnToDouble colY
- let !caps = capabilities
- !bounds = groupRangeBounds offs nGroups caps
- pure (unsafePerformIO (momentPar vis offs nGroups xs ys caps bounds))
-{-# NOINLINE momentScatterPar #-}
-
-momentPar ::
- VU.Vector Int ->
- VU.Vector Int ->
- Int ->
- VU.Vector Double ->
- VU.Vector Double ->
- Int ->
- VU.Vector Int ->
- IO Moments
-momentPar vis offs nGroups xs ys caps bounds = do
- cnt <- VUM.replicate nGroups (0 :: Int)
- sx <- VUM.replicate nGroups (0 :: Double)
- sy <- VUM.replicate nGroups (0 :: Double)
- sxx <- VUM.replicate nGroups (0 :: Double)
- syy <- VUM.replicate nGroups (0 :: Double)
- sxy <- VUM.replicate nGroups (0 :: Double)
- let bump arr g d = VUM.unsafeRead arr g >>= \c -> VUM.unsafeWrite arr g (c + d)
- forEachRange bounds caps $ \gs ge ->
- overGroups vis offs gs ge $ \g row -> do
- let !x = VU.unsafeIndex xs row
- !y = VU.unsafeIndex ys row
- VUM.unsafeRead cnt g >>= \c -> VUM.unsafeWrite cnt g (c + 1)
- bump sx g x
- bump sy g y
- bump sxx g (x * x)
- bump syy g (y * y)
- bump sxy g (x * y)
- Moments . fromUnboxedVector
- <$> VU.unsafeFreeze cnt
- <*> (fromUnboxedVector <$> VU.unsafeFreeze sx)
- <*> (fromUnboxedVector <$> VU.unsafeFreeze sy)
- <*> (fromUnboxedVector <$> VU.unsafeFreeze sxx)
- <*> (fromUnboxedVector <$> VU.unsafeFreeze syy)
- <*> (fromUnboxedVector <$> VU.unsafeFreeze sxy)
diff --git a/src/DataFrame/Internal/AggPlan.hs b/src/DataFrame/Internal/AggPlan.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/AggPlan.hs
+++ /dev/null
@@ -1,316 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | The aggregation fast-path planner and the two-column moment scatter.
-
-'planAgg' inspects a named output expression and, on a recognised shape over a
-clean (non-null, unboxed Int/Double) column, returns the 'AggPlan' the caller
-runs through the scatter kernel ('DataFrame.Internal.AggKernel'); anything it
-does not recognise returns 'Nothing' so the caller keeps the existing
-interpreter. Recognition is by the 'AggStrategy' name tag plus the shape of the
-inner 'Expr' — no new constructor is needed, and the general @aggregate@ API is
-unchanged.
-
-'momentScatter' fuses the additive moment sums of two columns (count, Sx, Sy,
-Sxx, Syy, Sxy) into one pass — the sufficient statistics for the Q9 regression
-family. It is exposed for callers that want to collapse the six separate folds
-into a single pass.
--}
-module DataFrame.Internal.AggPlan (
- AggPlan (..),
- planAgg,
- Moments (..),
- momentScatter,
- MomentPlan (..),
- planMoments,
-) where
-
-import qualified Data.Map.Strict as M
-import qualified Data.Text as T
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Monad.ST (runST)
-import DataFrame.Internal.AggKernel (Reduction (..), scatterColumnToDouble)
-import DataFrame.Internal.Column (Column (..), fromUnboxedVector)
-import DataFrame.Internal.DataFrame (
- DataFrame (derivingExpressions),
- GroupedDataFrame (..),
- getColumn,
- )
-import DataFrame.Internal.Expression (
- AggStrategy (..),
- BinaryOp (binaryCommutative, binaryName),
- Expr (..),
- UExpr (..),
- )
-import Type.Reflection (Typeable, typeRep)
-
-{- | The plan 'planAgg' produces for a recognised output expression. The median
-plan carries only the column name (the holistic grouped sort lives in the
-operations layer, where @vector-algorithms@ is available).
--}
-data AggPlan
- = -- | A single scatter reduction over one named column.
- PlanScatter Reduction T.Text
- | -- | @max a - min b@ (Q7): two scatters then a vectorized combine.
- PlanMaxMinusMin T.Text T.Text
- | -- | Holistic median over one named column.
- PlanMedian T.Text
-
-{- | Inspect a named output expression. On a recognised shape over a present
-clean column return @Just plan@; otherwise 'Nothing'. Nullable (bitmap) or
-non-Int/Double value columns are rejected here so the scatter only ever sees a
-clean unboxed vector.
--}
-planAgg :: GroupedDataFrame -> UExpr -> Maybe AggPlan
-planAgg gdf (UExpr (expr :: Expr a)) = case expr of
- Agg (FoldAgg tag _ _) (Col name) -> foldPlan tag name
- Agg (MergeAgg tag _ _ _ _) (Col name) -> mergePlan tag name
- Agg (CollectAgg tag _) (Col name) -> collectPlan tag name
- Binary
- op
- (Agg (FoldAgg lt Nothing _) (Col a))
- (Agg (FoldAgg rt Nothing _) (Col b)) ->
- if binaryName op == "sub" && lt == "maximum" && rt == "minimum"
- then requireBoth a b (PlanMaxMinusMin a b)
- else Nothing
- _ -> Nothing
- where
- foldPlan tag name = case tag of
- "sum" -> require name (PlanScatter RSum name)
- "minimum" -> require name (PlanScatter RMin name)
- "maximum" -> require name (PlanScatter RMax name)
- _ -> Nothing
- mergePlan tag name = case tag of
- "mean" -> outputType @Double >> require name (PlanScatter RMean name)
- "count" -> outputType @Int >> require name (PlanScatter RCount name)
- _ -> Nothing
- outputType :: forall t. (Typeable t) => Maybe ()
- outputType = case testEquality (typeRep @a) (typeRep @t) of
- Just Refl -> Just ()
- Nothing -> Nothing
- collectPlan tag name = case tag of
- "stddev" -> require name (PlanScatter RStd name)
- "variance" -> require name (PlanScatter RVar name)
- "top2Sum" -> require name (PlanScatter RTop2Sum name)
- "median" -> require name (PlanMedian name)
- _ -> Nothing
- require name plan = colUnboxedNumeric name >> Just plan
- requireBoth a b plan = colUnboxedNumeric a >> colUnboxedNumeric b >> Just plan
- colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of
- Just c | isUnboxedNumeric c -> Just ()
- _ -> Nothing
-
--- | The matcher only fires on non-null unboxed Int/Double columns.
-isUnboxedNumeric :: Column -> Bool
-isUnboxedNumeric = \case
- UnboxedColumn Nothing (_ :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> True
- Nothing -> case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> True
- Nothing -> False
- _ -> False
-
-{- | A recognised moment (Q9 regression) aggregate group: six output columns
-that together form the sufficient statistics of two base columns @x@ and @y@.
-The caller runs the fused 'momentScatter'/'momentScatterPar' once over
-@(colX, colY)@ and binds each output name to the named field of the result.
--}
-data MomentPlan = MomentPlan
- { mpColX :: T.Text
- , mpColY :: T.Text
- , mpNName :: T.Text
- , mpSxName :: T.Text
- , mpSyName :: T.Text
- , mpSxxName :: T.Text
- , mpSyyName :: T.Text
- , mpSxyName :: T.Text
- }
-
-{- | The shape of a sum's argument once unary coercions are peeled and derived
-columns are resolved through @derivingExpressions@: either linear in one base
-column or the product of two base columns (sorted).
--}
-data Term
- = Lin T.Text
- | Prod T.Text T.Text
- deriving (Eq, Ord, Show)
-
-{- | Recognise the multi-aggregate moment shape across a whole @aggregate@ list:
-exactly @count(_)@, @sum(x)@, @sum(y)@, @sum(x*x)@, @sum(y*y)@, @sum(x*y)@ over
-two distinct base columns @x@ and @y@ (after resolving derived product columns
-through @derivingExpressions@). Returns 'Nothing' on any other set so the caller
-falls back to the per-expression planner. Both base columns must be clean
-unboxed Int/Double, the same gate the single-column scatter uses.
--}
-planMoments :: GroupedDataFrame -> [(T.Text, UExpr)] -> Maybe MomentPlan
-planMoments gdf aggs
- | length aggs /= 6 = Nothing
- | otherwise = do
- let exprs = derivingExpressions (fullDataframe gdf)
- roles <- traverse (classify exprs) aggs
- let names = M.fromList [(r, nm) | (nm, r) <- roles]
- nName <- M.lookup RoleN names
- (x, y) <- pickBaseColumns roles
- sxName <- M.lookup (RoleLin x) names
- syName <- M.lookup (RoleLin y) names
- sxxName <- M.lookup (RoleProd x x) names
- syyName <- M.lookup (RoleProd y y) names
- sxyName <- M.lookup (RoleProd x y) names
- _ <- if x /= y then Just () else Nothing
- _ <- colUnboxedNumeric x
- _ <- colUnboxedNumeric y
- pure
- MomentPlan
- { mpColX = x
- , mpColY = y
- , mpNName = nName
- , mpSxName = sxName
- , mpSyName = syName
- , mpSxxName = sxxName
- , mpSyyName = syyName
- , mpSxyName = sxyName
- }
- where
- colUnboxedNumeric name = case getColumn name (fullDataframe gdf) of
- Just c | isUnboxedNumeric c -> Just ()
- _ -> Nothing
-
--- | The output role each named aggregation plays in the moment shape.
-data Role
- = RoleN
- | RoleLin T.Text
- | RoleProd T.Text T.Text
- deriving (Eq, Ord, Show)
-
--- | Tag a single named aggregation with its moment role, or reject the group.
-classify :: M.Map T.Text UExpr -> (T.Text, UExpr) -> Maybe (T.Text, Role)
-classify exprs (name, UExpr expr) = case expr of
- Agg (MergeAgg "count" _ _ _ _) _ -> Just (name, RoleN)
- Agg (FoldAgg "sum" _ _) arg -> (\t -> (name, termRole t)) <$> resolveTerm exprs (UExpr arg)
- _ -> Nothing
-
-termRole :: Term -> Role
-termRole (Lin a) = RoleLin a
-termRole (Prod a b) = RoleProd a b
-
-{- | Resolve a (sum-argument) expression to its 'Term'. Peels @toDouble@-style
-unary coercions, follows a derived column to its stored expression, and
-recognises a commutative product of two linear terms.
--}
-resolveTerm :: M.Map T.Text UExpr -> UExpr -> Maybe Term
-resolveTerm exprs = go (8 :: Int)
- where
- go 0 _ = Nothing
- go fuel (UExpr e) = case e of
- Col nm -> case M.lookup nm exprs of
- Just ue -> go (fuel - 1) ue
- Nothing -> Just (Lin nm)
- Unary _ inner -> go (fuel - 1) (UExpr inner)
- Binary op l r
- | binaryName op == "mult" && binaryCommutative op -> do
- Lin a <- go (fuel - 1) (UExpr l)
- Lin b <- go (fuel - 1) (UExpr r)
- Just (sortProd a b)
- _ -> Nothing
-
--- | Products are unordered: store the pair sorted so @x*y@ and @y*x@ unify.
-sortProd :: T.Text -> T.Text -> Term
-sortProd a b
- | a <= b = Prod a b
- | otherwise = Prod b a
-
-{- | From the classified roles, find the unordered pair of base columns that the
-linear sums name. There must be exactly two distinct linear-sum columns.
--}
-pickBaseColumns :: [(T.Text, Role)] -> Maybe (T.Text, T.Text)
-pickBaseColumns roles =
- case lins of
- [a, b] | a /= b -> Just (a, b)
- _ -> Nothing
- where
- lins = M.keys (M.fromList [(c, ()) | (_, RoleLin c) <- roles])
-
-{- | The additive moment sums of two columns, each an @nGroups@-length column:
-@(n, Sx, Sy, Sxx, Syy, Sxy)@.
--}
-data Moments = Moments
- { mN :: Column
- , mSx :: Column
- , mSy :: Column
- , mSxx :: Column
- , mSyy :: Column
- , mSxy :: Column
- }
-
-{- | One pass over two Double-coercible columns @x@ and @y@ filling the count
-and the five sums. Collapses the Q9 regression family's six independent folds
-(and three derive passes) into a single fused pass. Returns 'Nothing' unless
-both columns are non-null unboxed Int/Double.
--}
-momentScatter :: VU.Vector Int -> Int -> Column -> Column -> Maybe Moments
-momentScatter g nGroups colX colY = do
- xs <- scatterColumnToDouble colX
- ys <- scatterColumnToDouble colY
- let (cnt, sx, sy, sxx, syy, sxy) = momentPass g nGroups xs ys
- pure
- Moments
- { mN = fromUnboxedVector cnt
- , mSx = fromUnboxedVector sx
- , mSy = fromUnboxedVector sy
- , mSxx = fromUnboxedVector sxx
- , mSyy = fromUnboxedVector syy
- , mSxy = fromUnboxedVector sxy
- }
-
-momentPass ::
- VU.Vector Int ->
- Int ->
- VU.Vector Double ->
- VU.Vector Double ->
- ( VU.Vector Int
- , VU.Vector Double
- , VU.Vector Double
- , VU.Vector Double
- , VU.Vector Double
- , VU.Vector Double
- )
-momentPass g nGroups xs ys = runST $ do
- cnt <- VUM.replicate nGroups (0 :: Int)
- sx <- VUM.replicate nGroups (0 :: Double)
- sy <- VUM.replicate nGroups (0 :: Double)
- sxx <- VUM.replicate nGroups (0 :: Double)
- syy <- VUM.replicate nGroups (0 :: Double)
- sxy <- VUM.replicate nGroups (0 :: Double)
- let n = VU.length xs
- bump arr k d = VUM.unsafeRead arr k >>= \c -> VUM.unsafeWrite arr k (c + d)
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex g i
- !x = VU.unsafeIndex xs i
- !y = VU.unsafeIndex ys i
- VUM.unsafeRead cnt k >>= \c -> VUM.unsafeWrite cnt k (c + 1)
- bump sx k x
- bump sy k y
- bump sxx k (x * x)
- bump syy k (y * y)
- bump sxy k (x * y)
- go (i + 1)
- go 0
- (,,,,,)
- <$> VU.unsafeFreeze cnt
- <*> VU.unsafeFreeze sx
- <*> VU.unsafeFreeze sy
- <*> VU.unsafeFreeze sxx
- <*> VU.unsafeFreeze syy
- <*> VU.unsafeFreeze sxy
diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Column.hs
+++ /dev/null
@@ -1,1852 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ConstraintKinds #-}
-{-# LANGUAGE DataKinds #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE InstanceSigs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE PolyKinds #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE UndecidableInstances #-}
-
-module DataFrame.Internal.Column where
-
-import qualified Data.Text as T
-import qualified Data.Vector as VB
-import qualified Data.Vector.Generic as VG
-import qualified Data.Vector.Mutable as VBM
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Exception (throw)
-import Control.Monad (forM_, when)
-import Control.Monad.ST (ST, runST)
-import Data.Bits (
- complement,
- popCount,
- setBit,
- shiftL,
- shiftR,
- testBit,
- (.&.),
- )
-import Data.Kind (Type)
-import Data.Maybe
-import Data.Type.Equality (TestEquality (..))
-import Data.Word (Word8)
-import DataFrame.Errors
-import DataFrame.Internal.PackedText (
- PackedTextData (..),
- packedGather,
- packedIndexText,
- packedLength,
- packedRowOffsetVec,
- packedSlice,
- packedTake,
- sliceEqBytes,
- )
-import DataFrame.Internal.Types
-import DataFrame.Internal.Utf8 (sliceTextVector)
-import System.IO.Unsafe (unsafePerformIO)
-import System.Random
-import Type.Reflection
-
--- | A bit-packed validity bitmap. Bit @i@ = 1 means row @i@ is valid (not null).
-type Bitmap = VU.Vector Word8
-
-{- | Our representation of a column is a GADT that can store data based on the underlying data.
-
-This allows us to pattern match on data kinds and limit some operations to only some
-kinds of vectors. Nullability is represented via an optional bit-packed 'Bitmap':
-@Nothing@ = no nulls; @Just bm@ = bit @i@ of @bm@ is 1 iff row @i@ is valid.
--}
-data Column where
- BoxedColumn :: (Columnable a) => Maybe Bitmap -> VB.Vector a -> Column
- UnboxedColumn ::
- (Columnable a, VU.Unbox a) => Maybe Bitmap -> VU.Vector a -> Column
- -- Bit-packed Text: a shared UTF-8 byte buffer + (n+1) row offsets + an
- -- optional validity bitmap. No per-row Text header; Text is produced only
- -- on demand. Behaves as a column of Text (or Maybe Text when a bitmap is
- -- present). Only the CSV ingest path emits this; user-built Text columns
- -- stay 'BoxedColumn'.
- PackedText :: Maybe Bitmap -> {-# UNPACK #-} !PackedTextData -> 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
-
--- ---------------------------------------------------------------------------
--- Bitmap helpers
--- ---------------------------------------------------------------------------
-
--- | Test whether row @i@ is valid (not null) in a bitmap.
-bitmapTestBit :: Bitmap -> Int -> Bool
-bitmapTestBit bm i = testBit (VU.unsafeIndex bm (i `shiftR` 3)) (i .&. 7)
-{-# INLINE bitmapTestBit #-}
-
--- | Build a fully-valid bitmap for @n@ rows (all bits set).
-allValidBitmap :: Int -> Bitmap
-allValidBitmap n =
- let bytes = (n + 7) `shiftR` 3
- lastBits = n .&. 7
- full = VU.replicate (bytes - 1) 0xFF
- lastByte = if lastBits == 0 then 0xFF else (1 `shiftL` lastBits) - 1
- in if bytes == 0 then VU.empty else VU.snoc full lastByte
-{-# INLINE allValidBitmap #-}
-
-{- | Build a bitmap from a @VU.Vector Word8@ validity vector
-(1 = valid, 0 = null), as produced by Arrow / Parquet decoders.
--}
-buildBitmapFromValid :: VU.Vector Word8 -> Bitmap
-buildBitmapFromValid valid =
- let n = VU.length valid
- bytes = (n + 7) `shiftR` 3
- in VU.generate bytes $ \b ->
- let base = b `shiftL` 3
- setBitIf acc bit =
- let idx = base + bit
- in if idx < n && VU.unsafeIndex valid idx /= 0
- then setBit acc bit
- else acc
- in foldl setBitIf (0 :: Word8) [0 .. 7]
-
-{- | Build a bitmap from a list of null-row indices.
-@nullIdxs@ are the positions that are NULL.
--}
-buildBitmapFromNulls :: Int -> [Int] -> Bitmap
-buildBitmapFromNulls n nullIdxs =
- let base = allValidBitmap n
- in VU.modify
- ( \mv ->
- forM_ nullIdxs $ \i -> do
- let byteIdx = i `shiftR` 3
- bitIdx = i .&. 7
- v <- VUM.unsafeRead mv byteIdx
- VUM.unsafeWrite mv byteIdx (clearBit8 v bitIdx)
- )
- base
- where
- clearBit8 :: Word8 -> Int -> Word8
- clearBit8 b bit = b .&. complement (1 `shiftL` bit)
-
--- | Slice a bitmap for rows @[start .. start+len-1]@.
-bitmapSlice :: Int -> Int -> Bitmap -> Bitmap
-bitmapSlice start len bm
- | start .&. 7 == 0 =
- -- byte-aligned: simple slice; clamp so we never ask for more bytes than exist
- let startByte = start `shiftR` 3
- bytes = min ((len + 7) `shiftR` 3) (VU.length bm - startByte)
- in VU.slice startByte bytes bm
- | otherwise =
- -- non-aligned: unpack bit-by-bit and repack
- let n = min len (VU.length bm `shiftL` 3 - start)
- in buildBitmapFromValid $
- VU.generate n $
- \i -> if bitmapTestBit bm (start + i) then 1 else 0
-
--- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.
-bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap
-bitmapConcat n1 bm1 n2 bm2 =
- buildBitmapFromValid $
- VU.generate (n1 + n2) $ \i ->
- if i < n1
- then if bitmapTestBit bm1 i then 1 else 0
- else if bitmapTestBit bm2 (i - n1) then 1 else 0
-
--- | Combine two bitmaps with AND (both must be valid for result to be valid).
-mergeBitmaps :: Bitmap -> Bitmap -> Bitmap
-mergeBitmaps = VU.zipWith (.&.)
-
-{- | Materialize a nullable column from @VB.Vector (Maybe a)@.
-When @a@ is unboxable, creates an 'UnboxedColumn' (more compact).
-Otherwise creates a 'BoxedColumn'.
-Always attaches a bitmap so the column is recognized as nullable even when
-no 'Nothing' values are present (preserves the Maybe type marker).
--}
-fromMaybeVec :: forall a. (Columnable a) => VB.Vector (Maybe a) -> Column
-fromMaybeVec v = case sUnbox @a of
- STrue -> fromMaybeVecUnboxed v
- SFalse ->
- let n = VB.length v
- nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]
- bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs
- dat = VB.map (fromMaybe (errorWithoutStackTrace "fromMaybeVec: Nothing slot")) v
- in BoxedColumn (Just bm) dat
-
-{- | Materialize a nullable 'UnboxedColumn' to @VB.Vector (Maybe a)@ using runST.
-Always attaches a bitmap so the column is recognized as nullable even when
-no 'Nothing' values are present (preserves the Maybe type marker).
--}
-fromMaybeVecUnboxed ::
- forall a. (Columnable a, VU.Unbox a) => VB.Vector (Maybe a) -> Column
-fromMaybeVecUnboxed v =
- let n = VB.length v
- nullIdxs = [i | i <- [0 .. n - 1], isNothing (VB.unsafeIndex v i)]
- bm = if null nullIdxs then allValidBitmap n else buildBitmapFromNulls n nullIdxs
- dat = runST $ do
- mv <- VUM.new n
- VG.iforM_ v $ \i mx -> forM_ mx (VUM.unsafeWrite mv i)
- VU.unsafeFreeze mv
- in UnboxedColumn (Just bm) dat
-
--- | Materialize an element from a column at index @i@, respecting the bitmap.
-columnElemIsNull :: Column -> Int -> Bool
-columnElemIsNull (BoxedColumn (Just bm) _) i = not (bitmapTestBit bm i)
-columnElemIsNull (UnboxedColumn (Just bm) _) i = not (bitmapTestBit bm i)
-columnElemIsNull (PackedText (Just bm) _) i = not (bitmapTestBit bm i)
-columnElemIsNull _ _ = False
-
--- | Return the 'Maybe Bitmap' from a column.
-columnBitmap :: Column -> Maybe Bitmap
-columnBitmap (BoxedColumn bm _) = bm
-columnBitmap (UnboxedColumn bm _) = bm
-columnBitmap (PackedText bm _) = bm
-
-{- | The universal cold fallback: decode a 'PackedText' into a
-@BoxedColumn Text@ using the exact 'sliceTextVector' path the boxed-Text
-builder used, so the result is bit-identical to materializing at freeze.
-Identity on every other column.
--}
-materializePacked :: Column -> Column
-materializePacked (PackedText bm p) = case packedRowOffsetVec p of
- Just (arr, offs) -> BoxedColumn bm (sliceTextVector arr offs)
- -- Gathered/selected payload: rows are non-contiguous, decode per row.
- Nothing -> BoxedColumn bm (VB.generate (packedLength p) (packedIndexText p))
-materializePacked c = c
-{-# INLINE materializePacked #-}
-
--- | Whether a column is a 'PackedText'.
-isPackedText :: Column -> Bool
-isPackedText (PackedText _ _) = True
-isPackedText _ = False
-{-# INLINE isPackedText #-}
-
--- ---------------------------------------------------------------------------
--- End bitmap helpers
--- ---------------------------------------------------------------------------
-
-{- | 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
-
-instance (Eq a) => Eq (TypedColumn a) where
- (==) :: (Eq a) => TypedColumn a -> TypedColumn a -> Bool
- (==) (TColumn a) (TColumn b) = a == b
-
--- | Gets the underlying value from a TypedColumn.
-unwrapTypedColumn :: TypedColumn a -> Column
-unwrapTypedColumn (TColumn value) = value
-
--- | Gets the underlying vector from a TypedColumn.
-vectorFromTypedColumn :: TypedColumn a -> VB.Vector a
-vectorFromTypedColumn (TColumn value) = either throw id (toVector value)
-
--- | Checks if a column contains missing values (has a bitmap).
-hasMissing :: Column -> Bool
-hasMissing (BoxedColumn (Just _) _) = True
-hasMissing (UnboxedColumn (Just _) _) = True
-hasMissing (PackedText (Just _) _) = True
-hasMissing _ = False
-
--- | Checks if a column contains only missing values.
-allMissing :: Column -> Bool
-allMissing (BoxedColumn (Just bm) col) = VU.all (== 0) bm && not (VB.null col)
-allMissing (UnboxedColumn (Just bm) col) = VU.all (== 0) bm && not (VU.null col)
-allMissing (PackedText (Just bm) p) = VU.all (== 0) bm && packedLength p > 0
-allMissing _ = False
-
--- | Checks if a column contains numeric values.
-isNumeric :: Column -> Bool
-isNumeric (UnboxedColumn _ (_vec :: VU.Vector a)) = case sNumeric @a of
- STrue -> True
- _ -> False
-isNumeric (BoxedColumn _ (_vec :: VB.Vector a)) = case testEquality (typeRep @a) (typeRep @Integer) of
- Nothing -> False
- Just Refl -> True
-isNumeric (PackedText _ _) = False
-
-{- | Checks if a column is of a given type values.
-For nullable columns (@BoxedColumn (Just _)@ or @UnboxedColumn (Just _)@),
-also returns @True@ when @a = Maybe b@ and the column stores @b@ internally.
--}
-hasElemType :: forall a. (Columnable a) => Column -> Bool
-hasElemType = \case
- BoxedColumn bm (_column :: VB.Vector b) -> checkBoxed bm (typeRep @b)
- UnboxedColumn bm (_column :: VU.Vector b) -> checkUnboxed bm (typeRep @b)
- PackedText bm _ -> checkBoxed bm (typeRep @T.Text)
- where
- -- Direct type match
- directMatch :: forall (b :: Type). TypeRep b -> Bool
- directMatch = isJust . testEquality (typeRep @a)
- -- For a nullable column (has bitmap), also accept a = Maybe b
- checkMaybe :: forall (b :: Type). TypeRep b -> Bool
- checkMaybe tb = case typeRep @a of
- App tMaybe tInner -> case eqTypeRep tMaybe (typeRep @Maybe) of
- Just HRefl -> isJust (testEquality tInner tb)
- Nothing -> False
- _ -> False
- checkBoxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
- checkBoxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)
- checkUnboxed :: forall (b :: Type). Maybe Bitmap -> TypeRep b -> Bool
- checkUnboxed bm tb = directMatch tb || (isJust bm && checkMaybe tb)
-
--- | An internal/debugging function to get the column type of a column.
-columnVersionString :: Column -> String
-columnVersionString column = case column of
- BoxedColumn Nothing _ -> "Boxed"
- BoxedColumn (Just _) _ -> "NullableBoxed"
- UnboxedColumn Nothing _ -> "Unboxed"
- UnboxedColumn (Just _) _ -> "NullableUnboxed"
- PackedText Nothing _ -> "Boxed"
- PackedText (Just _) _ -> "NullableBoxed"
-
-{- | An internal/debugging function to get the type stored in the outermost vector
-of a column.
--}
-columnTypeString :: Column -> String
-columnTypeString column = case column of
- BoxedColumn Nothing (_ :: VB.Vector a) -> show (typeRep @a)
- BoxedColumn (Just _) (_ :: VB.Vector a) -> showMaybeType @a
- UnboxedColumn Nothing (_ :: VU.Vector a) -> show (typeRep @a)
- UnboxedColumn (Just _) (_ :: VU.Vector a) -> showMaybeType @a
- PackedText Nothing _ -> show (typeRep @T.Text)
- PackedText (Just _) _ -> showMaybeType @T.Text
- where
- showMaybeType :: forall a. (Typeable a) => String
- showMaybeType =
- let s = show (typeRep @a)
- in "Maybe " ++ if ' ' `elem` s then "(" ++ s ++ ")" else s
-
-instance (Show a) => Show (TypedColumn a) where
- show :: (Show a) => TypedColumn a -> String
- show (TColumn col) = show col
-
-{- | Force evaluation of all elements in a column. Replacement for the removed
-@instance NFData Column@; used by the IO and lazy-executor strict paths.
--}
-forceColumn :: Column -> ()
-forceColumn (BoxedColumn Nothing (v :: VB.Vector a)) = VB.foldl' (const (`seq` ())) () v
-forceColumn (BoxedColumn (Just bm) (v :: VB.Vector a)) =
- let n = VB.length v
- go !i
- | i >= n = ()
- | bitmapTestBit bm i = VB.unsafeIndex v i `seq` go (i + 1)
- | otherwise = go (i + 1)
- in go 0
-forceColumn (UnboxedColumn _ v) = v `seq` ()
-forceColumn (PackedText _ (PackedTextData arr offs sel)) = arr `seq` offs `seq` sel `seq` ()
-
-instance Show Column where
- show :: Column -> String
- show (BoxedColumn Nothing column) = show column
- show (BoxedColumn (Just bm) column) =
- let n = VB.length column
- elems =
- [ if bitmapTestBit bm i then show (VB.unsafeIndex column i) else "null"
- | i <- [0 .. n - 1]
- ]
- in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"
- show (UnboxedColumn Nothing column) = show column
- show (UnboxedColumn (Just bm) column) =
- let n = VU.length column
- elems =
- [ if bitmapTestBit bm i then show (VU.unsafeIndex column i) else "null"
- | i <- [0 .. n - 1]
- ]
- in "[" ++ foldl (\acc e -> if null acc then e else acc ++ "," ++ e) "" elems ++ "]"
- show c@(PackedText _ _) = show (materializePacked c)
-
-{- | Compare two nullable boxed columns element by element, skipping null slots.
-Uses a manual loop to avoid stream fusion forcing null-slot error thunks.
--}
-eqBoxedCols ::
- (Eq a) => Maybe Bitmap -> VB.Vector a -> Maybe Bitmap -> VB.Vector a -> Bool
-eqBoxedCols bm1 a bm2 b
- | VB.length a /= VB.length b = False
- | otherwise = go 0
- where
- !n = VB.length a
- go !i
- | i >= n = True
- | nullA || nullB = (nullA == nullB) && go (i + 1)
- | VB.unsafeIndex a i == VB.unsafeIndex b i = go (i + 1)
- | otherwise = False
- where
- nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
- nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
-{-# INLINE eqBoxedCols #-}
-
-instance Eq Column where
- (==) :: Column -> Column -> Bool
- (==) (BoxedColumn bm1 (a :: VB.Vector t1)) (BoxedColumn bm2 (b :: VB.Vector t2)) =
- case testEquality (typeRep @t1) (typeRep @t2) of
- Nothing -> False
- Just Refl -> eqBoxedCols bm1 a bm2 b
- (==) (UnboxedColumn bm1 (a :: VU.Vector t1)) (UnboxedColumn bm2 (b :: VU.Vector t2)) =
- case testEquality (typeRep @t1) (typeRep @t2) of
- Nothing -> False
- Just Refl ->
- VU.length a == VU.length b
- && VU.and
- ( VU.imap
- ( \i x ->
- let nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
- nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
- in if nullA || nullB then nullA == nullB else x == VU.unsafeIndex b i
- )
- a
- )
- (==) (PackedText bm1 p1) (PackedText bm2 p2) = eqPackedCols bm1 p1 bm2 p2
- (==) lhs@(PackedText _ _) rhs = materializePacked lhs == rhs
- (==) lhs rhs@(PackedText _ _) = lhs == materializePacked rhs
- (==) _ _ = False
-
-{- | Byte-slice equality of two packed-text columns, skipping null slots
-(a null compares equal only to a null), mirroring 'eqBoxedCols'.
--}
-eqPackedCols ::
- Maybe Bitmap -> PackedTextData -> Maybe Bitmap -> PackedTextData -> Bool
-eqPackedCols bm1 p1 bm2 p2
- | packedLength p1 /= packedLength p2 = False
- | otherwise = go 0
- where
- !n = packedLength p1
- go !i
- | i >= n = True
- | nullA || nullB = (nullA == nullB) && go (i + 1)
- | otherwise =
- let (a1, o1, l1) = packedSlice p1 i
- (a2, o2, l2) = packedSlice p2 i
- in sliceEqBytes a1 o1 l1 a2 o2 l2 && go (i + 1)
- where
- nullA = maybe False (\bm -> not (bitmapTestBit bm i)) bm1
- nullB = maybe False (\bm -> not (bitmapTestBit bm i)) bm2
-{-# INLINE eqPackedCols #-}
-
-{- | A class for converting a vector to a column of the appropriate type.
-Given each Rep we tell the `toColumnRep` function which Column type to pick.
--}
-class ColumnifyRep (r :: Rep) a where
- toColumnRep :: VB.Vector a -> Column
-
--- | Constraint synonym for what we can put into columns.
-type Columnable a =
- ( Columnable' a
- , ColumnifyRep (KindOf a) a
- , UnboxIf a
- , IntegralIf a
- , FloatingIf a
- , SBoolI (Unboxable a)
- , SBoolI (Numeric a)
- , SBoolI (IntegralTypes a)
- , SBoolI (FloatingTypes a)
- )
-
-instance
- (Columnable a, VU.Unbox a) =>
- ColumnifyRep 'RUnboxed a
- where
- toColumnRep :: (Columnable a, VUM.Unbox a) => VB.Vector a -> Column
- toColumnRep v = UnboxedColumn Nothing (VU.convert v)
-
-instance
- (Columnable a) =>
- ColumnifyRep 'RBoxed a
- where
- toColumnRep :: (Columnable a) => VB.Vector a -> Column
- toColumnRep = BoxedColumn Nothing
-
-instance
- (Columnable a) =>
- ColumnifyRep 'RNullableBoxed (Maybe a)
- where
- toColumnRep :: (Columnable a) => VB.Vector (Maybe a) -> Column
- toColumnRep = fromMaybeVec
-
-{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.
-
-__Examples:__
-
-@
-> import qualified Data.Vector as V
-> fromVector (VB.fromList [(1 :: Int), 2, 3, 4])
-[1,2,3,4]
-@
--}
-fromVector ::
- forall a.
- (Columnable a, ColumnifyRep (KindOf a) a) =>
- VB.Vector a -> Column
-fromVector = toColumnRep @(KindOf a)
-
-{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.
-
-__Examples:__
-
-@
-> import qualified Data.Vector.Unboxed as V
-> fromUnboxedVector (VB.fromList [(1 :: Int), 2, 3, 4])
-[1,2,3,4]
-@
--}
-fromUnboxedVector ::
- forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column
-fromUnboxedVector = UnboxedColumn Nothing
-
-{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.
-
-__Examples:__
-
-@
-> fromList [(1 :: Int), 2, 3, 4]
-[1,2,3,4]
-@
--}
-fromList ::
- forall a.
- (Columnable a, ColumnifyRep (KindOf a) a) =>
- [a] -> Column
-fromList = toColumnRep @(KindOf a) . VB.fromList
-
-{- | O(n) Create a column of random elements within a range.
-
-Takes a random number generator, a length, and a lower and upper bound for the random values.
-
-__Examples:__
-
-@
-> import System.Random (mkStdGen)
-> mkRandom (mkStdGen 42) 4 0 10
-[4,2,6,5]
-@
--}
-mkRandom ::
- (RandomGen g, Columnable a, ColumnifyRep (KindOf a) a, UniformRange a) =>
- g -> Int -> a -> a -> Column
-mkRandom pureGen k lo hi = fromList $ go pureGen k
- where
- go _g 0 = []
- go g n =
- let
- (!v, !g') = uniformR (lo, hi) g
- in
- v : go g' (n - 1)
-
--- An internal helper for type errors
-throwTypeMismatch ::
- forall (a :: Type) (b :: Type).
- (Typeable a, Typeable b) => Either DataFrameException Column
-throwTypeMismatch =
- Left $
- TypeMismatchException
- MkTypeErrorContext
- { userType = Right (typeRep @b)
- , expectedType = Right (typeRep @a)
- , callingFunctionName = Nothing
- , errorColumnName = Nothing
- }
-
--- | An internal function to map a function over the values of a column.
-mapColumn ::
- forall b c.
- (Columnable b, Columnable c) =>
- (b -> c) -> Column -> Either DataFrameException Column
-mapColumn f = \case
- BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col
- UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col
- c@(PackedText _ _) -> mapColumn f (materializePacked c)
- where
- runBoxed ::
- forall a.
- (Columnable a) =>
- Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
- runBoxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of
- -- user maps over Maybe a (nullable column as Maybe)
- Just Refl ->
- let !n = VB.length col
- in -- Build result directly without intermediate Maybe vector to avoid
- -- fusion forcing null slots via VU.convert.
- Right $ case sUnbox @c of
- STrue -> UnboxedColumn Nothing $
- VU.generate n $ \i ->
- f
- ( if maybe True (`bitmapTestBit` i) bm
- then Just (VB.unsafeIndex col i)
- else Nothing
- )
- SFalse -> fromVector @c $
- VB.generate n $ \i ->
- f
- ( if maybe True (`bitmapTestBit` i) bm
- then Just (VB.unsafeIndex col i)
- else Nothing
- )
- Nothing -> case testEquality (typeRep @a) (typeRep @b) of
- Just Refl ->
- -- user maps over inner type a; preserve bitmap
- Right $ case sUnbox @c of
- STrue -> UnboxedColumn bm (VU.generate (VB.length col) (f . VB.unsafeIndex col))
- SFalse -> BoxedColumn bm (VB.map f col)
- Nothing -> throwTypeMismatch @a @b
-
- runUnboxed ::
- forall a.
- (Columnable a, VU.Unbox a) =>
- Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
- runUnboxed bm col = case testEquality (typeRep @b) (typeRep @(Maybe a)) of
- Just Refl ->
- let !n = VU.length col
- in Right $ case sUnbox @c of
- STrue -> UnboxedColumn Nothing $
- VU.generate n $ \i ->
- f
- ( if maybe True (`bitmapTestBit` i) bm
- then Just (VU.unsafeIndex col i)
- else Nothing
- )
- SFalse -> fromVector @c $
- VB.generate n $ \i ->
- f
- ( if maybe True (`bitmapTestBit` i) bm
- then Just (VU.unsafeIndex col i)
- else Nothing
- )
- Nothing -> case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Right $ case sUnbox @c of
- STrue -> UnboxedColumn bm (VU.map f col)
- SFalse -> BoxedColumn bm (VB.generate (VU.length col) (f . VU.unsafeIndex col))
- Nothing -> throwTypeMismatch @a @b
-{-# INLINEABLE mapColumn #-}
-
--- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.
-imapColumn ::
- forall b c.
- (Columnable b, Columnable c) =>
- (Int -> b -> c) -> Column -> Either DataFrameException Column
-imapColumn f = \case
- BoxedColumn bm (col :: VB.Vector a) -> runBoxed bm col
- UnboxedColumn bm (col :: VU.Vector a) -> runUnboxed bm col
- c@(PackedText _ _) -> imapColumn f (materializePacked c)
- where
- runBoxed ::
- forall a.
- (Columnable a) =>
- Maybe Bitmap -> VB.Vector a -> Either DataFrameException Column
- runBoxed bm col = case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Right $ case sUnbox @c of
- STrue ->
- UnboxedColumn
- bm
- (VU.generate (VB.length col) (\i -> f i (VB.unsafeIndex col i)))
- SFalse -> BoxedColumn bm (VB.imap f col)
- Nothing -> throwTypeMismatch @a @b
-
- runUnboxed ::
- forall a.
- (Columnable a, VU.Unbox a) =>
- Maybe Bitmap -> VU.Vector a -> Either DataFrameException Column
- runUnboxed bm col = case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Right $ case sUnbox @c of
- STrue -> UnboxedColumn bm (VU.imap f col)
- SFalse -> BoxedColumn bm (VB.imap f (VG.convert col))
- Nothing -> throwTypeMismatch @a @b
-
--- | O(1) Gets the number of elements in the column.
-columnLength :: Column -> Int
-columnLength (BoxedColumn _ xs) = VB.length xs
-columnLength (UnboxedColumn _ xs) = VU.length xs
-columnLength (PackedText _ p) = packedLength p
-{-# INLINE columnLength #-}
-
--- | O(n) Gets the number of non-null elements in the column.
-numElements :: Column -> Int
-numElements (BoxedColumn Nothing xs) = VB.length xs
-numElements (BoxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
-numElements (UnboxedColumn Nothing xs) = VU.length xs
-numElements (UnboxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
-numElements (PackedText Nothing p) = packedLength p
-numElements (PackedText (Just bm) _p) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
-{-# INLINE numElements #-}
-
--- | O(n) Takes the first n values of a column.
-takeColumn :: Int -> Column -> Column
-takeColumn n (BoxedColumn bm xs) =
- BoxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)
-takeColumn n (UnboxedColumn bm xs) =
- UnboxedColumn (fmap (bitmapSlice 0 n) bm) (VG.take n xs)
-takeColumn n (PackedText bm p) =
- PackedText (fmap (bitmapSlice 0 n) bm) (packedTake n p)
-{-# INLINE takeColumn #-}
-
--- | O(n) Takes the last n values of a column.
-takeLastColumn :: Int -> Column -> Column
-takeLastColumn n column = sliceColumn (columnLength column - n) n column
-{-# INLINE takeLastColumn #-}
-
--- | O(n) Takes n values after a given column index.
-sliceColumn :: Int -> Int -> Column -> Column
-sliceColumn start n (BoxedColumn bm xs) =
- BoxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)
-sliceColumn start n (UnboxedColumn bm xs) =
- UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs)
-sliceColumn start n c@(PackedText _ _) = sliceColumn start n (materializePacked c)
-{-# INLINE sliceColumn #-}
-
--- | 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 bm column) =
- BoxedColumn
- ( fmap
- ( \bm0 ->
- buildBitmapFromValid $
- VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
- )
- bm
- )
- ( VB.generate
- (VU.length indexes)
- ((column `VB.unsafeIndex`) . (indexes `VU.unsafeIndex`))
- )
-atIndicesStable indexes (UnboxedColumn bm column) =
- UnboxedColumn
- ( fmap
- ( \bm0 ->
- buildBitmapFromValid $
- VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
- )
- bm
- )
- (VU.unsafeBackpermute column indexes)
-atIndicesStable indexes (PackedText bm p) =
- -- Slice-preserving: share the byte buffer, permute via a selection vector
- -- instead of materializing boxed Text. Bitmap follows the same gather.
- PackedText
- ( fmap
- ( \bm0 ->
- buildBitmapFromValid $
- VU.map (\i -> if bitmapTestBit bm0 i then 1 else 0) indexes
- )
- bm
- )
- (packedGather indexes p)
-{-# INLINE atIndicesStable #-}
-
-{- | Like 'atIndicesStable' but treats negative indices as null.
-Keeps the index vector fully unboxed (no @VB.Vector (Maybe Int)@).
--}
-gatherWithSentinel :: VU.Vector Int -> Column -> Column
-gatherWithSentinel indices col =
- let !n = VU.length indices
- newBm = buildBitmapFromValid $ VU.generate n $ \i ->
- if VU.unsafeIndex indices i < 0 then 0 else 1
- in case col of
- PackedText srcBm p ->
- -- Slice-preserving sentinel gather: share the buffer, permute
- -- offsets (negative sentinel -> empty slice), build the null
- -- bitmap so -1 rows read as null in left/outer joins.
- let bm = case srcBm of
- Nothing -> Just newBm
- Just sb ->
- Just
- ( mergeBitmaps
- newBm
- ( buildBitmapFromValid $ VU.generate n $ \i ->
- let idx = VU.unsafeIndex indices i
- in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
- )
- )
- in PackedText bm (packedGather indices p)
- BoxedColumn srcBm v ->
- let dat = VB.generate n $ \i ->
- let !idx = VU.unsafeIndex indices i
- in if idx < 0 then VB.unsafeIndex v 0 else VB.unsafeIndex v idx
- bm = case srcBm of
- Nothing -> Just newBm
- Just sb ->
- Just
- ( mergeBitmaps
- newBm
- ( buildBitmapFromValid $ VU.generate n $ \i ->
- let idx = VU.unsafeIndex indices i
- in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
- )
- )
- in BoxedColumn bm dat
- UnboxedColumn srcBm v ->
- let dat = runST $ do
- mv <- VUM.new n
- VG.iforM_ indices $ \i idx ->
- when (idx >= 0) $ VUM.unsafeWrite mv i (VU.unsafeIndex v idx)
- VU.unsafeFreeze mv
- bm = case srcBm of
- Nothing -> Just newBm
- Just sb ->
- Just
- ( mergeBitmaps
- newBm
- ( buildBitmapFromValid $ VU.generate n $ \i ->
- let idx = VU.unsafeIndex indices i
- in if idx >= 0 && bitmapTestBit sb idx then 1 else 0
- )
- )
- in UnboxedColumn bm dat
-{-# INLINE gatherWithSentinel #-}
-
--- | 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))
-{-# INLINE getIndices #-}
-
--- | Internal helper to get indices in an unboxed vector.
-getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a
-getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i))
-{-# INLINE getIndicesUnboxed #-}
-
-findIndices ::
- forall a.
- (Columnable a) =>
- (a -> Bool) ->
- Column ->
- Either DataFrameException (VU.Vector Int)
-findIndices predicate = \case
- BoxedColumn _ (v :: VB.Vector b) -> run v VG.convert
- UnboxedColumn _ (v :: VU.Vector b) -> run v id
- c@(PackedText _ _) -> findIndices predicate (materializePacked c)
- where
- run ::
- forall b v.
- (Typeable b, VG.Vector v b, VG.Vector v Int) =>
- v b ->
- (v Int -> VU.Vector Int) ->
- Either DataFrameException (VU.Vector Int)
- run column finalize = case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Right . finalize $ VG.findIndices predicate column
- Nothing ->
- Left $
- TypeMismatchException
- MkTypeErrorContext
- { userType = Right (typeRep @a)
- , expectedType = Right (typeRep @b)
- , callingFunctionName = Just "findIndices"
- , errorColumnName = Nothing
- }
-
--- | Fold (right) column with index.
-ifoldrColumn ::
- forall a b.
- (Columnable a, Columnable b) =>
- (Int -> a -> b -> b) -> b -> Column -> Either DataFrameException b
-ifoldrColumn f acc = \case
- BoxedColumn _ column -> foldrWorker column
- UnboxedColumn _ column -> foldrWorker column
- c@(PackedText _ _) -> ifoldrColumn f acc (materializePacked c)
- 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 @c)
- , callingFunctionName = Just "ifoldrColumn"
- , errorColumnName = Nothing
- }
- )
-
-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
- c@(PackedText _ _) -> foldlColumn f acc (materializePacked c)
- 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 @c)
- , callingFunctionName = Just "ifoldrColumn"
- , errorColumnName = Nothing
- }
- )
-
-foldl1Column ::
- forall a.
- (Columnable a) =>
- (a -> a -> a) -> Column -> Either DataFrameException a
-foldl1Column f = \case
- BoxedColumn _ column -> foldl1Worker column
- UnboxedColumn _ column -> foldl1Worker column
- c@(PackedText _ _) -> foldl1Column f (materializePacked c)
- 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 @c)
- , callingFunctionName = Just "foldl1Column"
- , errorColumnName = Nothing
- }
- )
-
-{- | 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.
--}
-foldl1DirectGroups ::
- forall a.
- (Columnable a) =>
- (a -> a -> a) ->
- Column ->
- VU.Vector Int ->
- VU.Vector Int ->
- Either DataFrameException Column
-foldl1DirectGroups f col valueIndices offsets
- | VU.length offsets <= 1 = pure $ fromVector @a VB.empty
- | otherwise = case col of
- UnboxedColumn _ (vec :: VU.Vector d) -> UnboxedColumn Nothing <$> foldl1Worker vec
- BoxedColumn _ (vec :: VB.Vector d) -> BoxedColumn Nothing <$> foldl1Worker vec
- PackedText _ _ -> foldl1DirectGroups f (materializePacked col) valueIndices offsets
- 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.
-rowToGroup[i] = group index for row i.
-Avoids random column reads; random writes go to the accumulator array which is
-small (nGroups entries) and typically cache-resident.
-When @acc@ is unboxable, uses an unboxed mutable vector for the accumulator
-array, eliminating pointer indirection on every read/write.
--}
-foldLinearGroups ::
- forall b acc.
- (Columnable b, Columnable acc) =>
- (acc -> b -> acc) ->
- acc ->
- Column ->
- VU.Vector Int -> -- rowToGroup (length n)
- Int -> -- nGroups
- Either DataFrameException Column
-foldLinearGroups f seed col rowToGroup nGroups
- | nGroups == 0 = Right (fromVector @acc VB.empty)
- | otherwise = case col of
- UnboxedColumn _ (vec :: VU.Vector d) -> foldLinearWorker vec
- BoxedColumn _ (vec :: VB.Vector d) -> foldLinearWorker vec
- PackedText _ _ ->
- foldLinearGroups f seed (materializePacked col) rowToGroup nGroups
- 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 —
- -- avoiding a round-trip through VB.Vector.
- runWith :: ((Int -> IO acc) -> (Int -> acc -> IO ()) -> IO ()) -> IO Column
- runWith body = case sUnbox @acc of
- STrue -> do
- accs <- VUM.replicate nGroups seed
- body (VUM.unsafeRead accs) (VUM.unsafeWrite accs)
- UnboxedColumn Nothing <$> VU.unsafeFreeze accs
- SFalse -> do
- accs <- VBM.replicate nGroups seed
- body (VBM.unsafeRead accs) (VBM.unsafeWrite accs)
- fromVector @acc <$> VB.unsafeFreeze accs
- {-# INLINE runWith #-}
-{-# INLINEABLE foldLinearGroups #-}
-
-headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a
-headColumn = \case
- BoxedColumn _ col -> headWorker col
- UnboxedColumn _ col -> headWorker col
- c@(PackedText _ _) -> headColumn (materializePacked c)
- 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
-zipColumns l@(PackedText _ _) r = zipColumns (materializePacked l) r
-zipColumns l r@(PackedText _ _) = zipColumns l (materializePacked r)
-zipColumns (BoxedColumn _ column) (BoxedColumn _ other) = BoxedColumn Nothing (VG.zip column other)
-zipColumns (BoxedColumn _ column) (UnboxedColumn _ other) =
- BoxedColumn
- Nothing
- ( VB.generate
- (min (VG.length column) (VG.length other))
- (\i -> (column VG.! i, other VG.! i))
- )
-zipColumns (UnboxedColumn _ column) (BoxedColumn _ other) =
- BoxedColumn
- Nothing
- ( VB.generate
- (min (VG.length column) (VG.length other))
- (\i -> (column VG.! i, other VG.! i))
- )
-zipColumns (UnboxedColumn _ column) (UnboxedColumn _ other) = UnboxedColumn Nothing (VG.zip column other)
-{-# INLINE zipColumns #-}
-
--- | Merge two columns using `These`.
-mergeColumns :: Column -> Column -> Column
-mergeColumns colA colB = case (colA, colB) of
- (PackedText _ _, _) -> mergeColumns (materializePacked colA) colB
- (_, PackedText _ _) -> mergeColumns colA (materializePacked colB)
- (BoxedColumn bmA c1, BoxedColumn bmB c2) -> case (bmA, bmB) of
- (Just ba, Just bb) ->
- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
- let nullA = not (bitmapTestBit ba i)
- nullB = not (bitmapTestBit bb i)
- in case (nullA, nullB) of
- (True, True) -> error "mergeColumns: both null"
- (False, True) -> This v1
- (True, False) -> That v2
- (False, False) -> These v1 v2
- (Just ba, Nothing) ->
- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
- if not (bitmapTestBit ba i) then That v2 else These v1 v2
- (Nothing, Just bb) ->
- BoxedColumn Nothing $ mkVec c1 c2 $ \i v1 v2 ->
- if not (bitmapTestBit bb i) then This v1 else These v1 v2
- (Nothing, Nothing) ->
- BoxedColumn Nothing $ mkVecSimple c1 c2 These
- (BoxedColumn _ c1, UnboxedColumn _ c2) ->
- BoxedColumn Nothing $ mkVecSimple c1 c2 These
- (UnboxedColumn _ c1, BoxedColumn _ c2) ->
- BoxedColumn Nothing $ mkVecSimple c1 c2 These
- (UnboxedColumn _ c1, UnboxedColumn _ c2) ->
- BoxedColumn Nothing $ mkVecSimple c1 c2 These
- where
- mkVec c1 c2 combineElements =
- VB.generate
- (min (VG.length c1) (VG.length c2))
- (\i -> combineElements i (c1 VG.! i) (c2 VG.! i))
- {-# INLINE mkVec #-}
-
- mkVecSimple c1 c2 f =
- VB.generate
- (min (VG.length c1) (VG.length c2))
- (\i -> f (c1 VG.! i) (c2 VG.! i))
- {-# INLINE mkVecSimple #-}
-{-# INLINE mergeColumns #-}
-
--- | An internal, column version of zipWith.
-zipWithColumns ::
- forall a b c.
- (Columnable a, Columnable b, Columnable c) =>
- (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
-zipWithColumns f (UnboxedColumn bmL (column :: VU.Vector d)) (UnboxedColumn bmR (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of
- Just Refl -> case testEquality (typeRep @b) (typeRep @e) of
- Just Refl
- -- Fast path: both plain unboxed, no bitmaps involved in the output type
- | isNothing bmL
- , isNothing bmR ->
- pure $ case sUnbox @c of
- STrue -> UnboxedColumn Nothing (VU.zipWith f column other)
- SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)
- -- Type mismatch or bitmap involvement: fall through to general toVector path
- _ -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)
- Nothing -> zipWithColumnsGeneral f (UnboxedColumn bmL column) (UnboxedColumn bmR other)
--- TODO: mchavinda - reuse pattern from interpret where we augment the
--- error at the end.
-zipWithColumns f left right = zipWithColumnsGeneral f left right
-
-zipWithColumnsGeneral ::
- forall a b c.
- (Columnable a, Columnable b, Columnable c) =>
- (a -> b -> c) -> Column -> Column -> Either DataFrameException Column
-zipWithColumnsGeneral f left right = case toVector @a left of
- Left (TypeMismatchException context) ->
- Left $
- TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})
- Left e -> Left e
- Right left' -> case toVector @b right of
- Left (TypeMismatchException context) ->
- Left $
- TypeMismatchException (context{callingFunctionName = Just "zipWithColumns"})
- Left e -> Left e
- Right right' -> pure $ fromVector $ VB.zipWith f left' right'
-{-# INLINE zipWithColumnsGeneral #-}
-{-# INLINE zipWithColumns #-}
-
--- writeColumn and freezeColumn' (CSV-ingest helpers) moved to
--- DataFrame.IO.Internal.MutableColumn so the core column module does not
--- need to depend on DataFrame.Internal.Parsing.
-
-{- | Freeze a mutable column into an @Either Text a@ column: every recorded
-null position becomes @Left rawText@ (preserving the original input), every
-other position becomes @Right v@. Used by CSV readers under 'EitherRead' mode.
--}
-freezeColumnEither :: [(Int, T.Text)] -> MutableColumn -> IO Column
-freezeColumnEither nulls (MBoxedColumn col) = do
- frozen <- VB.unsafeFreeze col
- let nullMap = nulls
- pure $
- BoxedColumn Nothing $
- VB.imap
- ( \i v -> case lookup i nullMap of
- Just t -> Left t
- Nothing -> Right v
- )
- frozen
-freezeColumnEither nulls (MUnboxedColumn col) = do
- c <- VU.unsafeFreeze col
- let nullMap = nulls
- pure $
- BoxedColumn Nothing $
- VB.generate (VU.length c) $ \i ->
- case lookup i nullMap of
- Just t -> Left t
- Nothing -> Right (c VU.! i)
-{-# INLINE freezeColumnEither #-}
-
-{- | Promote a non-nullable column to a nullable one (add an all-valid bitmap).
-No-op when already nullable.
--}
-ensureOptional :: Column -> Column
-ensureOptional c@(BoxedColumn (Just _) _) = c
-ensureOptional (BoxedColumn Nothing col) =
- BoxedColumn (Just (allValidBitmap (VB.length col))) col
-ensureOptional c@(UnboxedColumn (Just _) _) = c
-ensureOptional (UnboxedColumn Nothing col) =
- UnboxedColumn (Just (allValidBitmap (VU.length col))) col
-ensureOptional c@(PackedText (Just _) _) = c
-ensureOptional (PackedText Nothing p) =
- PackedText (Just (allValidBitmap (packedLength p))) p
-
--- | Fills the end of a column, up to n, with null rows. Does nothing if column has length >= n.
-expandColumn :: Int -> Column -> Column
-expandColumn n c@(PackedText _ p)
- | n <= packedLength p = c
- | otherwise = expandColumn n (materializePacked c)
-expandColumn n column@(BoxedColumn bm col)
- | n <= VG.length col = column
- | otherwise =
- let extra = n - VG.length col
- newBm = case bm of
- Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])
- Just b ->
- Just
- (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))
- -- pad data with default (undefined slot, protected by bitmap)
- newCol = col <> VB.replicate extra (errorWithoutStackTrace "expandColumn: null slot")
- in BoxedColumn newBm newCol
-expandColumn n column@(UnboxedColumn bm col)
- | n <= VG.length col = column
- | otherwise =
- let extra = n - VG.length col
- newBm = case bm of
- Nothing -> Just (buildBitmapFromNulls n [VG.length col .. n - 1])
- Just b ->
- Just
- (bitmapConcat (VG.length col) b extra (VU.replicate ((extra + 7) `shiftR` 3) 0))
- newCol = runST $ do
- mv <- VUM.new n
- VU.imapM_ (VUM.unsafeWrite mv) col
- VU.unsafeFreeze mv
- in UnboxedColumn newBm newCol
-
--- | Fills the beginning of a column, up to n, with null rows. Does nothing if column has length >= n.
-leftExpandColumn :: Int -> Column -> Column
-leftExpandColumn n c@(PackedText _ p)
- | n <= packedLength p = c
- | otherwise = leftExpandColumn n (materializePacked c)
-leftExpandColumn n column@(BoxedColumn bm col)
- | n <= VG.length col = column
- | otherwise =
- let extra = n - VG.length col
- origLen = VG.length col
- newBm = case bm of
- Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])
- Just b ->
- let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0
- in Just (bitmapConcat extra nullPart origLen b)
- newCol =
- VB.replicate extra (errorWithoutStackTrace "leftExpandColumn: null slot") <> col
- in BoxedColumn newBm newCol
-leftExpandColumn n column@(UnboxedColumn bm col)
- | n <= VG.length col = column
- | otherwise =
- let extra = n - VG.length col
- origLen = VG.length col
- newBm = case bm of
- Nothing -> Just (buildBitmapFromNulls n [0 .. extra - 1])
- Just b ->
- let nullPart = VU.replicate ((extra + 7) `shiftR` 3) 0
- in Just (bitmapConcat extra nullPart origLen b)
- newCol = runST $ do
- mv <- VUM.new n
- VU.imapM_ (\i x -> VUM.unsafeWrite mv (extra + i) x) col
- VU.unsafeFreeze mv
- in UnboxedColumn newBm newCol
-
-{- | Concatenates two columns.
-Returns Nothing if the columns are of different types.
--}
-concatColumns :: Column -> Column -> Either DataFrameException Column
-concatColumns left right = case (left, right) of
- (PackedText _ _, _) -> concatColumns (materializePacked left) right
- (_, PackedText _ _) -> concatColumns left (materializePacked right)
- (BoxedColumn bmL l, BoxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of
- Just Refl ->
- let newBm = case (bmL, bmR) of
- (Nothing, Nothing) -> Nothing
- (Just bl, Nothing) ->
- Just
- (bitmapConcat (VB.length l) bl (VB.length r) (allValidBitmap (VB.length r)))
- (Nothing, Just br) ->
- Just
- (bitmapConcat (VB.length l) (allValidBitmap (VB.length l)) (VB.length r) br)
- (Just bl, Just br) -> Just (bitmapConcat (VB.length l) bl (VB.length r) br)
- in pure (BoxedColumn newBm (l <> r))
- Nothing -> Left (mismatchErr (typeOf r) (typeOf l))
- (UnboxedColumn bmL l, UnboxedColumn bmR r) -> case testEquality (typeOf l) (typeOf r) of
- Just Refl ->
- let newBm = case (bmL, bmR) of
- (Nothing, Nothing) -> Nothing
- (Just bl, Nothing) ->
- Just
- (bitmapConcat (VU.length l) bl (VU.length r) (allValidBitmap (VU.length r)))
- (Nothing, Just br) ->
- Just
- (bitmapConcat (VU.length l) (allValidBitmap (VU.length l)) (VU.length r) br)
- (Just bl, Just br) -> Just (bitmapConcat (VU.length l) bl (VU.length r) br)
- in pure (UnboxedColumn newBm (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.
-
-Works similar to 'concatColumns', but unlike that function, it will also combine columns of different types
-by wrapping the values in an Either.
-
-E.g. combining Column containing [1,2] with Column containing ["a","b"]
-will result in a Column containing [Left 1, Left 2, Right "a", Right "b"].
--}
-
-{- | O(n) Concatenate a list of same-type columns in a single allocation.
-All columns must have the same constructor and element type (as they will
-within a single Parquet column). Calls 'error' on mismatch.
--}
-concatManyColumns :: [Column] -> Column
-concatManyColumns [] = fromList ([] :: [Maybe Int])
-concatManyColumns [c] = c
-concatManyColumns all'
- | any isPackedText all' =
- concatManyColumns (map materializePacked all')
-concatManyColumns (c0 : cs) = case c0 of
- BoxedColumn bm0 v0 ->
- let getCol (BoxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of
- Just Refl -> (bm, v)
- Nothing -> error "concatManyColumns: BoxedColumn type mismatch"
- getCol _ = error "concatManyColumns: column constructor mismatch"
- rest = map getCol cs
- allVecs = v0 : map snd rest
- allBms = bm0 : map fst rest
- newBm
- | all isNothing allBms = Nothing
- | otherwise =
- let pairs = zip allVecs allBms
- expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VB.length v)) mb) pairs
- go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
- concatBms [] = VU.empty
- concatBms [(b, _v)] = b
- concatBms ((b1, v1) : (b2, v2) : rest') =
- let merged = go b1 (VB.length v1) b2 (VB.length v2)
- in concatBms ((merged, v1 <> v2) : rest')
- in Just $ concatBms (zip expandedBms allVecs)
- in BoxedColumn newBm (VB.concat allVecs)
- UnboxedColumn bm0 v0 ->
- let getCol (UnboxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of
- Just Refl -> (bm, v)
- Nothing -> error "concatManyColumns: UnboxedColumn type mismatch"
- getCol _ = error "concatManyColumns: column constructor mismatch"
- rest = map getCol cs
- allVecs = v0 : map snd rest
- allBms = bm0 : map fst rest
- newBm
- | all isNothing allBms = Nothing
- | otherwise =
- let pairs = zip allVecs allBms
- expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VU.length v)) mb) pairs
- go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
- concatBms [] = VU.empty
- concatBms [(b, _)] = b
- concatBms ((b1, v1) : (b2, v2) : rest') =
- let merged = go b1 (VU.length v1) b2 (VU.length v2)
- in concatBms ((merged, v1 <> v2) : rest')
- in Just $ concatBms (zip expandedBms allVecs)
- in UnboxedColumn newBm (VU.concat allVecs)
- PackedText _ _ -> concatManyColumns (map materializePacked (c0 : cs))
-
-concatColumnsEither :: Column -> Column -> Column
-concatColumnsEither l@(PackedText _ _) r = concatColumnsEither (materializePacked l) r
-concatColumnsEither l r@(PackedText _ _) = concatColumnsEither l (materializePacked r)
-concatColumnsEither (BoxedColumn bmL left) (BoxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of
- Nothing ->
- BoxedColumn Nothing $ fmap Left left <> fmap Right right
- Just Refl ->
- let newBm = case (bmL, bmR) of
- (Nothing, Nothing) -> Nothing
- (Just bl, Nothing) ->
- Just
- ( bitmapConcat
- (VB.length left)
- bl
- (VB.length right)
- (allValidBitmap (VB.length right))
- )
- (Nothing, Just br) ->
- Just
- ( bitmapConcat
- (VB.length left)
- (allValidBitmap (VB.length left))
- (VB.length right)
- br
- )
- (Just bl, Just br) -> Just (bitmapConcat (VB.length left) bl (VB.length right) br)
- in BoxedColumn newBm $ left <> right
-concatColumnsEither (UnboxedColumn bmL left) (UnboxedColumn bmR right) = case testEquality (typeOf left) (typeOf right) of
- Nothing ->
- BoxedColumn Nothing $
- fmap Left (VG.convert left) <> fmap Right (VG.convert right)
- Just Refl ->
- let newBm = case (bmL, bmR) of
- (Nothing, Nothing) -> Nothing
- (Just bl, Nothing) ->
- Just
- ( bitmapConcat
- (VU.length left)
- bl
- (VU.length right)
- (allValidBitmap (VU.length right))
- )
- (Nothing, Just br) ->
- Just
- ( bitmapConcat
- (VU.length left)
- (allValidBitmap (VU.length left))
- (VU.length right)
- br
- )
- (Just bl, Just br) -> Just (bitmapConcat (VU.length left) bl (VU.length right) br)
- in UnboxedColumn newBm $ left <> right
-concatColumnsEither (BoxedColumn _ left) (UnboxedColumn _ right) =
- BoxedColumn Nothing $ fmap Left left <> fmap Right (VG.convert right)
-concatColumnsEither (UnboxedColumn _ left) (BoxedColumn _ right) =
- BoxedColumn Nothing $ fmap Left (VG.convert left) <> fmap Right right
-
--- | Allocate a mutable column of size @n@ matching the constructor/type of the given column.
-newMutableColumn :: Int -> Column -> IO MutableColumn
-newMutableColumn n (BoxedColumn _ (_ :: VB.Vector a)) =
- MBoxedColumn <$> (VBM.new n :: IO (VBM.IOVector a))
-newMutableColumn n (UnboxedColumn _ (_ :: VU.Vector a)) =
- MUnboxedColumn <$> (VUM.new n :: IO (VUM.IOVector a))
-newMutableColumn n c@(PackedText _ _) = newMutableColumn n (materializePacked c)
-
--- | Copy a column chunk into a mutable column starting at offset @off@.
-copyIntoMutableColumn :: MutableColumn -> Int -> Column -> IO ()
-copyIntoMutableColumn (MBoxedColumn (mv :: VBM.IOVector b)) off (BoxedColumn _ (v :: VB.Vector a)) =
- case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> VG.imapM_ (\i x -> VBM.unsafeWrite mv (off + i) x) v
- Nothing -> error "copyIntoMutableColumn: Boxed type mismatch"
-copyIntoMutableColumn (MUnboxedColumn (mv :: VUM.IOVector b)) off (UnboxedColumn _ (v :: VU.Vector a)) =
- case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> VG.imapM_ (\i x -> VUM.unsafeWrite mv (off + i) x) v
- Nothing -> error "copyIntoMutableColumn: Unboxed type mismatch"
-copyIntoMutableColumn mc off c@(PackedText _ _) =
- copyIntoMutableColumn mc off (materializePacked c)
-copyIntoMutableColumn _ _ _ =
- error "copyIntoMutableColumn: constructor mismatch"
-
--- | Freeze a mutable column into an immutable column.
-freezeMutableColumn :: MutableColumn -> IO Column
-freezeMutableColumn (MBoxedColumn mv) = BoxedColumn Nothing <$> VB.unsafeFreeze mv
-freezeMutableColumn (MUnboxedColumn mv) = UnboxedColumn Nothing <$> VU.unsafeFreeze mv
-
-{- | O(n) Converts a column to a list. Throws an exception if the wrong type is specified.
-
-__Examples:__
-
-@
-> column = fromList [(1 :: Int), 2, 3, 4]
-> toList @Int column
-[1,2,3,4]
-> toList @Double column
-exception: ...
-@
--}
-toList :: forall a. (Columnable a) => Column -> [a]
-toList xs = case toVector @a xs of
- Left err -> throw err
- Right val -> VB.toList val
-
-{- | Converts a column to a vector of a specific type.
-
-This is a type-safe conversion that requires the column's element type
-to exactly match the requested type. You must specify the desired type
-via type applications.
-
-==== __Type Parameters__
-
-[@a@] The element type to convert to
-[@v@] The vector type (e.g., 'VU.Vector', 'VB.Vector')
-
-==== __Examples__
-
->>> toVector @Int @VU.Vector column
-Right (unboxed vector of Ints)
-
->>> toVector @Text @VB.Vector column
-Right (boxed vector of Text)
-
-==== __Returns__
-
-* 'Right' - The converted vector if types match
-* 'Left' 'TypeMismatchException' - If the column's type doesn't match the requested type
-
-==== __See also__
-
-For numeric conversions with automatic type coercion, see 'toDoubleVector',
-'toFloatVector', and 'toIntVector'.
--}
-toVector ::
- forall a v.
- (VG.Vector v a, Columnable a) => Column -> Either DataFrameException (v a)
-toVector col = case col of
- PackedText _ _ -> toVector (materializePacked col)
- BoxedColumn bm (inner :: VB.Vector c) ->
- -- Check if user wants Maybe c (nullable) or c directly
- case testEquality (typeRep @a) (typeRep @c) of
- Just Refl -> Right $ VG.convert inner
- Nothing ->
- -- Try: a = Maybe c
- case testEquality (typeRep @a) (typeRep @(Maybe c)) of
- Just Refl ->
- -- Use VB.generate to avoid fusion forcing null slots
- let !n = VB.length inner
- maybeVec = case bm of
- Nothing -> VB.generate n (Just . VB.unsafeIndex inner)
- Just bitmap -> VB.generate n $ \i ->
- if bitmapTestBit bitmap i then Just (VB.unsafeIndex inner i) else Nothing
- in Right $ VG.convert maybeVec
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @a)
- , expectedType = Right (typeRep @c)
- , callingFunctionName = Just "toVector"
- , errorColumnName = Nothing
- }
- )
- UnboxedColumn bm (inner :: VU.Vector c) ->
- case testEquality (typeRep @a) (typeRep @c) of
- Just Refl -> Right $ VG.convert inner
- Nothing ->
- case testEquality (typeRep @a) (typeRep @(Maybe c)) of
- Just Refl ->
- let maybeVec = case bm of
- Nothing -> VB.generate (VU.length inner) (Just . VU.unsafeIndex inner)
- Just bitmap -> VB.generate (VU.length inner) $ \i ->
- if bitmapTestBit bitmap i then Just (VU.unsafeIndex inner i) else Nothing
- in Right $ VG.convert maybeVec
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @a)
- , expectedType = Right (typeRep @c)
- , callingFunctionName = Just "toVector"
- , errorColumnName = Nothing
- }
- )
-
--- Some common types we will use for numerical computing.
-
-{- | Converts a column to an unboxed vector of 'Double' values.
-
-This function performs intelligent type coercion for numeric types:
-
-* If the column is already 'Double', returns it directly
-* If the column contains other floating-point types, converts via 'realToFrac'
-* If the column contains integral types, converts via 'fromIntegral' (beware of overflow if the type is `Integer`).
-
-==== __Optional column handling__
-
-For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).
-This allows optional numeric data to be represented in the resulting vector.
-
-==== __Returns__
-
-* 'Right' - The converted 'Double' vector
-* 'Left' 'TypeMismatchException' - If the column is not numeric
--}
-toDoubleVector :: Column -> Either DataFrameException (VU.Vector Double)
-toDoubleVector column =
- case column of
- PackedText _ _ -> toDoubleVector (materializePacked column)
- UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> case bm of
- Nothing -> Right f
- Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f
- Nothing -> case sFloating @a of
- STrue ->
- Right
- ( VU.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> realToFrac x
- )
- f
- )
- SFalse -> case sIntegral @a of
- STrue ->
- Right
- ( VU.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> fromIntegral x
- )
- f
- )
- SFalse ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Double)
- , expectedType = Right (typeRep @a)
- , callingFunctionName = Just "toDoubleVector"
- , errorColumnName = Nothing
- }
- )
- BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
- Just Refl ->
- Right
- ( VB.convert $
- VB.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> fromIntegral x
- )
- f
- )
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Double)
- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
- , callingFunctionName = Just "toDoubleVector"
- , errorColumnName = Nothing
- }
- )
-
-{- | Converts a column to an unboxed vector of 'Float' values.
-
-This function performs intelligent type coercion for numeric types:
-
-* If the column is already 'Float', returns it directly
-* If the column contains other floating-point types, converts via 'realToFrac'
-* If the column contains integral types, converts via 'fromIntegral'
-* If the column is boxed 'Integer', converts via 'fromIntegral' (beware of overflow for 64-bit integers and `Integer`)
-
-==== __Optional column handling__
-
-For 'OptionalColumn' types, 'Nothing' values are converted to @NaN@ (Not a Number).
-This allows optional numeric data to be represented in the resulting vector.
-
-==== __Returns__
-
-* 'Right' - The converted 'Float' vector
-* 'Left' 'TypeMismatchException' - If the column is not numeric
-
-==== __Precision warning__
-
-Converting from 'Double' to 'Float' may result in loss of precision.
--}
-toFloatVector :: Column -> Either DataFrameException (VU.Vector Float)
-toFloatVector column =
- case column of
- PackedText _ _ -> toFloatVector (materializePacked column)
- UnboxedColumn bm (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Float) of
- Just Refl -> case bm of
- Nothing -> Right f
- Just bitmap -> Right $ VU.imap (\i x -> if bitmapTestBit bitmap i then x else read "NaN") f
- Nothing -> case sFloating @a of
- STrue ->
- Right
- ( VU.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> realToFrac x
- )
- f
- )
- SFalse -> case sIntegral @a of
- STrue ->
- Right
- ( VU.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> fromIntegral x
- )
- f
- )
- SFalse ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Float)
- , expectedType = Right (typeRep @a)
- , callingFunctionName = Just "toFloatVector"
- , errorColumnName = Nothing
- }
- )
- BoxedColumn bm (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
- Just Refl ->
- Right
- ( VB.convert $
- VB.imap
- ( \i x -> case bm of
- Just bitmap | not (bitmapTestBit bitmap i) -> read "NaN"
- _ -> fromIntegral x
- )
- f
- )
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Float)
- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
- , callingFunctionName = Just "toFloatVector"
- , errorColumnName = Nothing
- }
- )
-
-{- | Converts a column to an unboxed vector of 'Int' values.
-
-This function performs intelligent type coercion for numeric types:
-
-* If the column is already 'Int', returns it directly
-* If the column contains floating-point types, rounds via 'round' and converts
-* If the column contains other integral types, converts via 'fromIntegral'
-* If the column is boxed 'Integer', converts via 'fromIntegral'
-
-==== __Returns__
-
-* 'Right' - The converted 'Int' vector
-* 'Left' 'TypeMismatchException' - If the column is not numeric
-
-==== __Note__
-
-Unlike 'toDoubleVector' and 'toFloatVector', this function does NOT support
-'OptionalColumn'. Optional columns must be handled separately.
-
-==== __Rounding behavior__
-
-Floating-point values are rounded to the nearest integer using 'round'.
-For example: 2.5 rounds to 2, 3.5 rounds to 4 (banker's rounding).
--}
-toIntVector :: Column -> Either DataFrameException (VU.Vector Int)
-toIntVector column =
- case column of
- PackedText _ _ -> toIntVector (materializePacked column)
- UnboxedColumn _ (f :: VU.Vector a) -> case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> Right f
- Nothing -> case sFloating @a of
- STrue -> Right (VU.map (round . (realToFrac :: a -> Double)) f)
- SFalse -> case sIntegral @a of
- STrue -> Right (VU.map fromIntegral f)
- SFalse ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Int)
- , expectedType = Right (typeRep @a)
- , callingFunctionName = Just "toIntVector"
- , errorColumnName = Nothing
- }
- )
- BoxedColumn _ (f :: VB.Vector a) -> case testEquality (typeRep @a) (typeRep @Integer) of
- Just Refl -> Right (VB.convert $ VB.map fromIntegral f)
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Int)
- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
- , callingFunctionName = Just "toIntVector"
- , errorColumnName = Nothing
- }
- )
-
-toUnboxedVector ::
- forall a.
- (Columnable a, VU.Unbox a) => Column -> Either DataFrameException (VU.Vector a)
-toUnboxedVector column =
- case column of
- UnboxedColumn _ (f :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Right f
- Nothing ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @Int)
- , expectedType = Right (typeRep @a)
- , callingFunctionName = Just "toUnboxedVector"
- , errorColumnName = Nothing
- }
- )
- _ ->
- Left $
- TypeMismatchException
- ( MkTypeErrorContext
- { userType = Right (typeRep @a)
- , expectedType = Left (columnTypeString column) :: Either String (TypeRep ())
- , callingFunctionName = Just "toUnboxedVector"
- , errorColumnName = Nothing
- }
- )
-{-# INLINE toUnboxedVector #-}
-
--- Shared finaliser for the two parseUnboxedColumn* helpers. Freezes
--- the mutable data vector, and only materialises the bitmap when the
--- column actually had nulls.
-{-# INLINE finalizeParseResult #-}
-finalizeParseResult ::
- (VU.Unbox a) =>
- VUM.STVector s a ->
- VUM.STVector s Word8 ->
- Bool ->
- ST s (Maybe (Maybe Bitmap, VU.Vector a))
-finalizeParseResult values vmask anyNull
- | anyNull = do
- vs <- VU.unsafeFreeze values
- vm <- VU.unsafeFreeze vmask
- return (Just (Just (buildBitmapFromValid vm), vs))
- | otherwise = do
- vs <- VU.unsafeFreeze values
- return (Just (Nothing, vs))
diff --git a/src/DataFrame/Internal/ColumnBuilder.hs b/src/DataFrame/Internal/ColumnBuilder.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/ColumnBuilder.hs
+++ /dev/null
@@ -1,299 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE UndecidableInstances #-}
-
-{- | Mutable, growable column builders for high-throughput ingest. No
-per-append @IORef@ traffic: hot counters live in an unboxed vector, payloads
-double on demand, and validity is only materialized once a null is seen.
--}
-module DataFrame.Internal.ColumnBuilder (
- ColumnBuilder (..),
- NumBuilder,
- IntBuilder,
- DoubleBuilder,
- TextBuilder,
- TextChunk (..),
- newIntBuilder,
- newDoubleBuilder,
- newNumBuilder,
- newTextBuilder,
- appendInt,
- appendDouble,
- appendNum,
- appendText,
- appendTextSlice,
- appendTextSliceFromPtr,
- freezeTextChunk,
- mergeColumns,
- mergeTextChunks,
-) where
-
-import qualified Data.Text as T
-import qualified Data.Text.Array as A
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Monad (when)
-import Control.Monad.ST (ST)
-import Data.Bits (shiftR)
-import Data.STRef
-import Data.Text.Internal (Text (..))
-import Data.Word (Word8)
-import DataFrame.Internal.Column hiding (mergeColumns)
-import DataFrame.Internal.ColumnMerge (
- TextChunk (..),
- mergeColumns,
- mergeTextChunks,
- packValidity,
- )
-import Foreign.Ptr (Ptr)
-
-{- | Operations shared by all column builders. A builder must not be used
-again after 'freezeBuilder' (its storage is frozen in place, not copied).
--}
-class ColumnBuilder b where
- -- | Append a null row (sentinel payload + invalid bit).
- appendNull :: b s -> ST s ()
-
- -- | Rows appended so far.
- builderLength :: b s -> ST s Int
-
- -- | Freeze into a fully-forced 'Column'; bitmap only when a null was seen.
- freezeBuilder :: b s -> ST s Column
-
--- Counter slots shared by the builders: rows, any-null flag, text bytes used.
-cRows, cAnyNull, cBytes :: Int
-cRows = 0
-cAnyNull = 1
-cBytes = 2
-
-{- | Builder for unboxed numeric payloads ('Int', 'Double', ...). 'nbNull'
-is the sentinel written into null slots (protected by the bitmap).
--}
-data NumBuilder a s = NumBuilder
- { nbNull :: !a
- , nbCounters :: !(VUM.MVector s Int)
- , nbArrays :: !(STRef s (NumArrays a s))
- }
-
-data NumArrays a s = NumArrays
- { naData :: !(VUM.MVector s a)
- , naValid :: !(VUM.MVector s Word8)
- }
-
-type IntBuilder = NumBuilder Int
-
-type DoubleBuilder = NumBuilder Double
-
--- | New numeric builder with a row-capacity hint and a null sentinel.
-newNumBuilder :: (VU.Unbox a) => a -> Int -> ST s (NumBuilder a s)
-newNumBuilder nullValue hint = do
- let cap = max 16 hint
- counters <- VUM.replicate 2 0
- dat <- VUM.unsafeNew cap
- val <- VUM.unsafeNew cap
- NumBuilder nullValue counters <$> newSTRef (NumArrays dat val)
-
-newIntBuilder :: Int -> ST s (IntBuilder s)
-newIntBuilder = newNumBuilder 0
-
-newDoubleBuilder :: Int -> ST s (DoubleBuilder s)
-newDoubleBuilder = newNumBuilder 0
-
-appendNum :: (VU.Unbox a) => NumBuilder a s -> a -> ST s ()
-appendNum b !x = do
- n <- VUM.unsafeRead (nbCounters b) cRows
- anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
- NumArrays dat val <- reserveNum b n
- VUM.unsafeWrite dat n x
- when (anyNull /= 0) $ VUM.unsafeWrite val n 1
- VUM.unsafeWrite (nbCounters b) cRows (n + 1)
-{-# INLINE appendNum #-}
-
-appendInt :: IntBuilder s -> Int -> ST s ()
-appendInt = appendNum
-{-# INLINE appendInt #-}
-
-appendDouble :: DoubleBuilder s -> Double -> ST s ()
-appendDouble = appendNum
-{-# INLINE appendDouble #-}
-
--- Fetch the arrays, growing (doubling) first if row @n@ would not fit.
-reserveNum :: (VU.Unbox a) => NumBuilder a s -> Int -> ST s (NumArrays a s)
-reserveNum b n = do
- arrs <- readSTRef (nbArrays b)
- if n < VUM.length (naData arrs) then pure arrs else growNum b arrs
-{-# INLINE reserveNum #-}
-
-growNum ::
- (VU.Unbox a) => NumBuilder a s -> NumArrays a s -> ST s (NumArrays a s)
-growNum b (NumArrays dat val) = do
- let cap = VUM.length dat
- dat' <- VUM.unsafeGrow dat cap
- val' <- VUM.unsafeGrow val cap
- let arrs = NumArrays dat' val'
- writeSTRef (nbArrays b) arrs
- pure arrs
-
-instance (Columnable a, VU.Unbox a) => ColumnBuilder (NumBuilder a) where
- appendNull b = do
- n <- VUM.unsafeRead (nbCounters b) cRows
- anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
- NumArrays dat val <- reserveNum b n
- VUM.unsafeWrite dat n (nbNull b)
- when (anyNull == 0) $ do
- VUM.set (VUM.slice 0 n val) 1
- VUM.unsafeWrite (nbCounters b) cAnyNull 1
- VUM.unsafeWrite val n 0
- VUM.unsafeWrite (nbCounters b) cRows (n + 1)
- {-# INLINE appendNull #-}
-
- builderLength b = VUM.unsafeRead (nbCounters b) cRows
-
- freezeBuilder b = do
- n <- VUM.unsafeRead (nbCounters b) cRows
- anyNull <- VUM.unsafeRead (nbCounters b) cAnyNull
- NumArrays dat val <- readSTRef (nbArrays b)
- !vs <- freezeTrimmed n dat
- if anyNull /= 0
- then do
- !bm <- packValidity n val
- pure $! UnboxedColumn (Just bm) vs
- else pure $! UnboxedColumn Nothing vs
-
--- Zero-copy freeze; copies to exact size when slack exceeds a quarter of n.
-freezeTrimmed :: (VU.Unbox a) => Int -> VUM.MVector s a -> ST s (VU.Vector a)
-freezeTrimmed n mv
- | VUM.length mv - n <= n `shiftR` 2 = VU.unsafeFreeze (VUM.slice 0 n mv)
- | otherwise = VU.freeze (VUM.slice 0 n mv)
-
-{- | Builder for 'Text' columns. All field bytes go into one exponentially
-grown byte array; rows are recorded as offsets, so an append is a memcpy
-and freezing slices 'Text' values off the shared array without copying.
--}
-data TextBuilder s = TextBuilder
- { tbCounters :: !(VUM.MVector s Int)
- , tbArrays :: !(STRef s (TextArrays s))
- }
-
-data TextArrays s = TextArrays
- { taBytes :: !(A.MArray s)
- , taByteCap :: !Int
- , taOffsets :: !(VUM.MVector s Int)
- -- ^ Row @i@ spans bytes @[offsets!i, offsets!(i+1))@.
- , taValid :: !(VUM.MVector s Word8)
- }
-
--- | New text builder with row-count and total-byte capacity hints.
-newTextBuilder :: Int -> Int -> ST s (TextBuilder s)
-newTextBuilder rowHint byteHint = do
- let rcap = max 16 rowHint
- bcap = max 64 byteHint
- counters <- VUM.replicate 3 0
- bytes <- A.new bcap
- offsets <- VUM.unsafeNew (rcap + 1)
- VUM.unsafeWrite offsets 0 0
- val <- VUM.unsafeNew rcap
- TextBuilder counters <$> newSTRef (TextArrays bytes bcap offsets val)
-
--- | Append @len@ raw bytes at @off@ in @src@ as one field (one memcpy).
-appendTextSlice :: TextBuilder s -> A.Array -> Int -> Int -> ST s ()
-appendTextSlice b src off len = do
- (n, pos, arrs) <- reserveText b len
- A.copyI len (taBytes arrs) pos src off
- finishTextAppend b arrs n (pos + len)
-{-# INLINE appendTextSlice #-}
-
--- | 'appendTextSlice' from foreign memory (e.g. an mmapped file buffer).
-appendTextSliceFromPtr :: TextBuilder s -> Ptr Word8 -> Int -> ST s ()
-appendTextSliceFromPtr b ptr len = do
- (n, pos, arrs) <- reserveText b len
- A.copyFromPointer (taBytes arrs) pos ptr len
- finishTextAppend b arrs n (pos + len)
-{-# INLINE appendTextSliceFromPtr #-}
-
--- | Append an already-decoded 'Text' (its bytes are UTF-8 already).
-appendText :: TextBuilder s -> T.Text -> ST s ()
-appendText b (Text src off len) = appendTextSlice b src off len
-{-# INLINE appendText #-}
-
-finishTextAppend :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s ()
-finishTextAppend b arrs n endPos = do
- anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
- when (anyNull /= 0) $ VUM.unsafeWrite (taValid arrs) n 1
- VUM.unsafeWrite (taOffsets arrs) (n + 1) endPos
- VUM.unsafeWrite (tbCounters b) cRows (n + 1)
- VUM.unsafeWrite (tbCounters b) cBytes endPos
-{-# INLINE finishTextAppend #-}
-
-reserveText :: TextBuilder s -> Int -> ST s (Int, Int, TextArrays s)
-reserveText b extra = do
- n <- VUM.unsafeRead (tbCounters b) cRows
- pos <- VUM.unsafeRead (tbCounters b) cBytes
- arrs <- readSTRef (tbArrays b)
- arrs' <-
- if n < VUM.length (taValid arrs) && pos + extra <= taByteCap arrs
- then pure arrs
- else growText b arrs (n + 1) (pos + extra)
- pure (n, pos, arrs')
-{-# INLINE reserveText #-}
-
-growText :: TextBuilder s -> TextArrays s -> Int -> Int -> ST s (TextArrays s)
-growText b (TextArrays bytes bcap offsets val) needRows needBytes = do
- let rcap = VUM.length val
- (offsets', val') <-
- if needRows > rcap
- then do
- let rcap' = max (2 * rcap) needRows
- o <- VUM.unsafeGrow offsets (rcap' - rcap)
- v <- VUM.unsafeGrow val (rcap' - rcap)
- pure (o, v)
- else pure (offsets, val)
- (bytes', bcap') <-
- if needBytes > bcap
- then do
- let cap' = max (2 * bcap) needBytes
- bs <- A.resizeM bytes cap'
- pure (bs, cap')
- else pure (bytes, bcap)
- let arrs = TextArrays bytes' bcap' offsets' val'
- writeSTRef (tbArrays b) arrs
- pure arrs
-
-{- | Freeze a 'TextBuilder' into a raw 'TextChunk' for byte-level merging
-('mergeTextChunks'): no 'T.Text' values are created until chunks merge.
--}
-freezeTextChunk :: TextBuilder s -> ST s TextChunk
-freezeTextChunk b = do
- n <- VUM.unsafeRead (tbCounters b) cRows
- anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
- used <- VUM.unsafeRead (tbCounters b) cBytes
- TextArrays bytes bcap offsets val <- readSTRef (tbArrays b)
- when (used < bcap) (A.shrinkM bytes used)
- arr <- A.unsafeFreeze bytes
- offs <- VU.unsafeFreeze (VUM.slice 0 (n + 1) offsets)
- bm <-
- if anyNull /= 0
- then Just <$> packValidity n val
- else pure Nothing
- pure (TextChunk arr used offs bm)
-
-instance ColumnBuilder TextBuilder where
- appendNull b = do
- (n, pos, arrs) <- reserveText b 0
- anyNull <- VUM.unsafeRead (tbCounters b) cAnyNull
- when (anyNull == 0) $ do
- VUM.set (VUM.slice 0 n (taValid arrs)) 1
- VUM.unsafeWrite (tbCounters b) cAnyNull 1
- VUM.unsafeWrite (taValid arrs) n 0
- VUM.unsafeWrite (taOffsets arrs) (n + 1) pos
- VUM.unsafeWrite (tbCounters b) cRows (n + 1)
- VUM.unsafeWrite (tbCounters b) cBytes pos
- {-# INLINE appendNull #-}
-
- builderLength b = VUM.unsafeRead (tbCounters b) cRows
-
- freezeBuilder b = do
- chunk <- freezeTextChunk b
- pure $! mergeTextChunks [chunk]
diff --git a/src/DataFrame/Internal/ColumnMerge.hs b/src/DataFrame/Internal/ColumnMerge.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/ColumnMerge.hs
+++ /dev/null
@@ -1,179 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Concatenation of per-chunk 'Column's (e.g. from parallel CSV chunks),
-re-exported through 'DataFrame.Internal.ColumnBuilder'. Text columns can
-merge at the byte level via 'TextChunk' \/ 'mergeTextChunks', so no
-per-chunk 'Data.Text.Text' values are ever materialized.
--}
-module DataFrame.Internal.ColumnMerge (
- TextChunk (..),
- mergeColumns,
- mergeTextChunks,
- packedFromTextChunk,
- packValidity,
- spliceBitmaps,
- tcRows,
-) where
-
-import qualified Data.Text.Array as A
-import qualified Data.Vector as VB
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Monad (foldM_, forM_, when)
-import Control.Monad.ST (ST, runST)
-import Data.Bits (shiftL, shiftR, (.&.), (.|.))
-import Data.Maybe (fromMaybe, isNothing)
-import Data.Type.Equality (testEquality, (:~:) (Refl))
-import Data.Word (Word8)
-import DataFrame.Internal.Column (
- Bitmap,
- Column (..),
- Columnable,
- allValidBitmap,
- materializePacked,
- )
-import DataFrame.Internal.PackedText (mkPackedContiguous)
-import Type.Reflection (typeRep)
-
-{- | A frozen text-builder chunk: raw UTF-8 bytes plus row offsets (row @i@
-spans bytes @[offsets!i, offsets!(i+1))@) and an optional validity bitmap.
-'Data.Text.Text' values are only created when chunks merge into a 'Column'.
--}
-data TextChunk = TextChunk
- { tcBytes :: !A.Array
- , tcUsed :: !Int
- , tcOffsets :: !(VU.Vector Int)
- , tcBitmap :: !(Maybe Bitmap)
- }
-
-tcRows :: TextChunk -> Int
-tcRows c = VU.length (tcOffsets c) - 1
-
-{- | Freeze a builder chunk directly into a packed-text column: NO
-'Data.Text.Text' materialization, NO UTF-8 validation pass (deferred to
-decode time). Not yet called by any reader (a later ingest stage flips
-'mergeTextChunks' to use it).
--}
-packedFromTextChunk :: TextChunk -> Column
-packedFromTextChunk (TextChunk arr _used offs bm) =
- PackedText bm (mkPackedContiguous arr offs)
-
-{- | Merge text chunks into one packed-text 'Column': one byte-array copy
-per chunk, one offset rebase, then wrap the merged shared buffer + offsets
-as 'PackedText' (no per-row 'Data.Text.Text' header, no eager UTF-8
-validation pass — decode is deferred to the last mile).
--}
-mergeTextChunks :: [TextChunk] -> Column
-mergeTextChunks [] = error "DataFrame.Internal.ColumnMerge.mergeTextChunks: empty list"
-mergeTextChunks [c] = packedFromTextChunk c
-mergeTextChunks cs = runST $ do
- let totalBytes = sum (map tcUsed cs)
- totalRows = sum (map tcRows cs)
- arr <- A.new (max 1 totalBytes)
- offs <- VUM.unsafeNew (totalRows + 1)
- VUM.unsafeWrite offs 0 0
- let splice !byteBase !rowBase c = do
- let n = tcRows c
- co = tcOffsets c
- A.copyI (tcUsed c) arr byteBase (tcBytes c) 0
- forM_ [1 .. n] $ \i ->
- VUM.unsafeWrite offs (rowBase + i) (byteBase + VU.unsafeIndex co i)
- pure (byteBase + tcUsed c, rowBase + n)
- foldM_ (\(b, r) c -> splice b r c) (0, 0) cs
- farr <- A.unsafeFreeze arr
- foffs <- VU.unsafeFreeze offs
- let !bm = spliceBitmaps [(tcBitmap c, tcRows c) | c <- cs]
- pure (PackedText bm (mkPackedContiguous farr foffs))
-
-{- | Merge per-chunk columns into one column: one allocation + memcpy per
-payload, with bitmaps spliced across non-byte-aligned chunk boundaries.
-All chunks must have the same element type.
--}
-mergeColumns :: [Column] -> Column
-mergeColumns [] = error "DataFrame.Internal.ColumnBuilder.mergeColumns: empty list"
-mergeColumns [c] = c
-mergeColumns cols@(c0 : _) = case c0 of
- PackedText _ _ -> mergeColumns (map materializePacked cols)
- UnboxedColumn _ (_ :: VU.Vector a) ->
- let parts = map (unboxedPart @a) cols
- !merged = VU.concat (map snd parts)
- !bm = spliceBitmaps [(mb, VU.length v) | (mb, v) <- parts]
- in UnboxedColumn bm merged
- BoxedColumn _ (_ :: VB.Vector a) ->
- let parts = map (boxedPart @a) cols
- !merged = VB.concat (map snd parts)
- !bm = spliceBitmaps [(mb, VB.length v) | (mb, v) <- parts]
- in BoxedColumn bm merged
-
-unboxedPart ::
- forall a. (Columnable a, VU.Unbox a) => Column -> (Maybe Bitmap, VU.Vector a)
-unboxedPart (UnboxedColumn mb (v :: VU.Vector b)) =
- case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> (mb, v)
- Nothing -> mergeMismatch
-unboxedPart _ = mergeMismatch
-
-boxedPart ::
- forall a. (Columnable a) => Column -> (Maybe Bitmap, VB.Vector a)
-boxedPart (BoxedColumn mb (v :: VB.Vector b)) =
- case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> (mb, v)
- Nothing -> mergeMismatch
-boxedPart _ = mergeMismatch
-
-mergeMismatch :: a
-mergeMismatch =
- error "DataFrame.Internal.ColumnBuilder.mergeColumns: chunk column types differ"
-
-{- | Splice chunk bitmaps end to end at the bit level. 'Nothing' if no chunk
-carries a bitmap; chunks without one count as all-valid otherwise.
--}
-spliceBitmaps :: [(Maybe Bitmap, Int)] -> Maybe Bitmap
-spliceBitmaps parts
- | all (isNothing . fst) parts = Nothing
- | otherwise = Just $ VU.create $ do
- let total = sum (map snd parts)
- outBytes = (total + 7) `shiftR` 3
- mv <- VUM.replicate outBytes 0
- let orInto i w =
- when (i < outBytes && w /= 0) $ do
- old <- VUM.unsafeRead mv i
- VUM.unsafeWrite mv i (old .|. w)
- splice !bitPos (mb, len) = do
- let bm = fromMaybe (allValidBitmap len) mb
- sh = bitPos .&. 7
- byte0 = bitPos `shiftR` 3
- lastIdx = ((len + 7) `shiftR` 3) - 1
- tailBits = len .&. 7
- lastMask =
- if tailBits == 0 then 0xFF else (1 `shiftL` tailBits) - 1
- forM_ [0 .. lastIdx] $ \k -> do
- let raw = VU.unsafeIndex bm k
- masked = if k == lastIdx then raw .&. lastMask else raw
- w = fromIntegral masked :: Word
- orInto (byte0 + k) (fromIntegral (w `shiftL` sh))
- when (sh /= 0) $
- orInto (byte0 + k + 1) (fromIntegral (w `shiftR` (8 - sh)))
- pure (bitPos + len)
- foldM_ splice 0 parts
- pure mv
-
--- | Pack a 0\/1 byte-per-row validity prefix into a bit-packed 'Bitmap'.
-packValidity :: Int -> VUM.MVector s Word8 -> ST s Bitmap
-packValidity n val = do
- bytes <- VU.unsafeFreeze (VUM.slice 0 n val)
- let assemble b =
- let base = b `shiftL` 3
- m = min 8 (n - base)
- go !acc !k
- | k >= m = acc
- | VU.unsafeIndex bytes (base + k) /= 0 =
- go (acc .|. (1 `shiftL` k)) (k + 1)
- | otherwise = go acc (k + 1)
- in go (0 :: Word8) 0
- pure $! VU.generate ((n + 7) `shiftR` 3) assemble
diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/DataFrame.hs
+++ /dev/null
@@ -1,378 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE InstanceSigs #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE PatternSynonyms #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-module DataFrame.Internal.DataFrame where
-
-import qualified Data.Map as M
-import qualified Data.Text as T
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-
-import Control.Exception (throw)
-import Data.Function (on)
-import Data.List (sortBy, (\\))
-import Data.Maybe (fromMaybe)
-import Data.Type.Equality (
- TestEquality (testEquality),
- type (:~:) (Refl),
- type (:~~:) (HRefl),
- )
-import DataFrame.Display.Terminal.PrettyPrint
-import DataFrame.Errors
-import DataFrame.Internal.Column
-import DataFrame.Internal.Expression
-import DataFrame.Internal.PackedText (packedIndexText)
-import Text.Printf
-import Type.Reflection (Typeable, eqTypeRep, typeRep, pattern App)
-import Prelude hiding (null)
-
-data DataFrame = DataFrame
- { columns :: V.Vector Column
- {- ^ Our main data structure stores a dataframe as
- a vector of columns. This improv
- -}
- , columnIndices :: M.Map T.Text Int
- -- ^ Keeps the column names in the order they were inserted in.
- , dataframeDimensions :: (Int, Int)
- -- ^ (rows, columns)
- , derivingExpressions :: M.Map T.Text UExpr
- }
-
-{- | Force evaluation of all columns in a DataFrame. Replacement for the removed
-@instance NFData DataFrame@; used by the IO and lazy-executor strict paths.
--}
-forceDataFrame :: DataFrame -> DataFrame
-forceDataFrame df@(DataFrame cols idx dims _exprs) =
- V.foldl' (\() c -> forceColumn c) () cols `seq` idx `seq` dims `seq` df
-
-{- | A record that contains information about how and what
-rows are grouped in the dataframe. This can only be used with
-`aggregate`.
--}
-data GroupedDataFrame = Grouped
- { fullDataframe :: DataFrame
- , groupedColumns :: [T.Text]
- , valueIndices :: VU.Vector Int
- , offsets :: VU.Vector Int
- , rowToGroup :: VU.Vector Int
- {- ^ rowToGroup[i] = group index for row i. Length n (one per row).
- Built once in 'groupBy'; reused by every aggregation.
- -}
- }
-
-instance Show GroupedDataFrame where
- show (Grouped df cols _indices _os _rtg) =
- printf
- "{ keyColumns: %s groupedColumns: %s }"
- (show cols)
- (show (M.keys (columnIndices df) \\ cols))
-
-instance Eq GroupedDataFrame where
- (==) (Grouped df cols _indices _os _rtg) (Grouped df' cols' _indices' _os' _rtg') = (df == df') && (cols == cols')
-
-instance Eq DataFrame where
- (==) :: DataFrame -> DataFrame -> Bool
- a == b =
- M.keys (columnIndices a) == M.keys (columnIndices b)
- && foldr
- ( \(name, index) acc -> acc && (columns a V.!? index == (columns b V.!? (columnIndices b M.! name)))
- )
- True
- (M.toList $ columnIndices a)
-
-instance Show DataFrame where
- show :: DataFrame -> String
- show d =
- let (r, _) = dataframeDimensions d
- cfg = defaultTruncateConfig
- shown = if maxRows cfg > 0 then min (maxRows cfg) r else r
- body = asTextWith Plain (Just cfg) d
- footer
- | shown < r =
- "\nShowing "
- <> T.pack (show shown)
- <> " rows out of "
- <> T.pack (show r)
- | otherwise = T.empty
- in T.unpack (body <> footer)
-
-{- | Configures how a 'DataFrame' is rendered as text. A non-positive value on
-any field means \"no limit\" on that axis.
-
-* 'maxRows' — render at most this many rows from the top of the frame.
-* 'maxColumns' — when the frame has more columns than this, the middle columns
- are collapsed into a single ellipsis column.
-* 'maxCellWidth' — text in any individual cell (including headers and type
- rows) longer than this is truncated with a trailing ellipsis.
--}
-data TruncateConfig = TruncateConfig
- { maxRows :: Int
- , maxColumns :: Int
- , maxCellWidth :: Int
- }
- deriving (Show, Eq)
-
--- | Sensible defaults for GHCi: 20 rows, 10 columns, 30 characters per cell.
-defaultTruncateConfig :: TruncateConfig
-defaultTruncateConfig =
- TruncateConfig{maxRows = 20, maxColumns = 10, maxCellWidth = 30}
-
--- | Ellipsis character used to mark elided columns and clipped cells.
-ellipsisText :: T.Text
-ellipsisText = "\x2026"
-
--- | For showing the dataframe as markdown in notebooks.
-toMarkdown :: DataFrame -> T.Text
-toMarkdown = asText Markdown
-
--- | For showing the dataframe as a string markdown in notebooks.
-toMarkdown' :: DataFrame -> String
-toMarkdown' = T.unpack . toMarkdown
-
-asText :: RenderFormat -> DataFrame -> T.Text
-asText fmt = asTextWith fmt Nothing
-
-asTextWith :: RenderFormat -> Maybe TruncateConfig -> DataFrame -> T.Text
-asTextWith fmt mTrunc d =
- let allHeaders =
- map fst (sortBy (compare `on` snd) (M.toList (columnIndices d)))
- nCols = length allHeaders
- (totalRows, _) = dataframeDimensions d
-
- rowCap = case mTrunc of
- Just cfg | maxRows cfg > 0 -> min totalRows (maxRows cfg)
- _ -> totalRows
-
- (visibleHeaders, ellipsisAt) = pickColumns mTrunc nCols allHeaders
-
- lookupCol name =
- fmap
- (takeColumn rowCap)
- ((V.!?) (columns d) ((M.!) (columnIndices d) name))
- survivingCols = map lookupCol visibleHeaders
- survivingTypes = map (maybe "" getType) survivingCols
- survivingData = map get survivingCols
-
- clipCell = case mTrunc of
- Just cfg | maxCellWidth cfg > 0 -> truncateCell (maxCellWidth cfg)
- _ -> id
-
- (finalHeaders, finalTypes, finalCols) = case ellipsisAt of
- Nothing -> (visibleHeaders, survivingTypes, survivingData)
- Just i ->
- let ellipsisCol = V.replicate rowCap ellipsisText
- in ( insertAt i ellipsisText visibleHeaders
- , insertAt i ellipsisText survivingTypes
- , insertAt i ellipsisCol survivingData
- )
-
- getType :: Column -> T.Text
- showMaybeType :: forall a. (Typeable a) => String
- showMaybeType =
- let s = show (typeRep @a)
- in "Maybe " <> if ' ' `elem` s then "(" <> s <> ")" else s
- getType (BoxedColumn Nothing (_ :: V.Vector a)) = T.pack $ show (typeRep @a)
- getType (BoxedColumn (Just _) (_ :: V.Vector a)) = T.pack $ showMaybeType @a
- getType (UnboxedColumn Nothing (_ :: VU.Vector a)) = T.pack $ show (typeRep @a)
- getType (UnboxedColumn (Just _) (_ :: VU.Vector a)) = T.pack $ showMaybeType @a
- getType (PackedText Nothing _) = T.pack $ show (typeRep @T.Text)
- getType (PackedText (Just _) _) = T.pack $ showMaybeType @T.Text
-
- -- Separate out cases dynamically so we don't end up making round trip
- -- string copies.
- get :: Maybe Column -> V.Vector T.Text
- get (Just (BoxedColumn (Just bm) (column :: V.Vector a))) =
- V.generate (V.length column) $ \i ->
- if bitmapTestBit bm i
- then T.pack (show (Just (V.unsafeIndex column i)))
- else "Nothing"
- get (Just (BoxedColumn Nothing (column :: V.Vector a))) =
- case testEquality (typeRep @a) (typeRep @T.Text) of
- Just Refl -> column
- Nothing -> case testEquality (typeRep @a) (typeRep @String) of
- Just Refl -> V.map T.pack column
- Nothing -> V.map (T.pack . show) column
- get (Just (UnboxedColumn (Just bm) column)) =
- V.generate (VU.length column) $ \i ->
- if bitmapTestBit bm i
- then T.pack (show (Just (VU.unsafeIndex column i)))
- else "Nothing"
- get (Just (UnboxedColumn Nothing column)) =
- V.generate (VU.length column) (T.pack . show . VU.unsafeIndex column)
- get (Just c@(PackedText _ _)) = get (Just (materializePacked c))
- get Nothing = V.empty
- in showTable
- fmt
- (map clipCell finalHeaders)
- (map clipCell finalTypes)
- (map (V.map clipCell) finalCols)
-
-{- | Decide which columns survive horizontal truncation and where (if anywhere)
-to splice in the ellipsis column. The split puts the extra column on the
-left for odd 'maxColumns'; the ellipsis is only inserted when it actually
-saves space (i.e. the frame has more than 'maxColumns' + 1 columns).
--}
-pickColumns ::
- Maybe TruncateConfig ->
- Int ->
- [a] ->
- ([a], Maybe Int)
-pickColumns mTrunc nCols xs = case mTrunc of
- Just cfg
- | let c = maxColumns cfg
- , c > 0
- , nCols > c + 1 ->
- let leftN = (c + 1) `div` 2
- rightN = c - leftN
- in ( Prelude.take leftN xs ++ Prelude.drop (nCols - rightN) xs
- , Just leftN
- )
- _ -> (xs, Nothing)
-
--- | Splice @x@ into @xs@ at index @i@ (0-based), shifting later elements right.
-insertAt :: Int -> a -> [a] -> [a]
-insertAt i x xs = let (l, r) = splitAt i xs in l ++ x : r
-
--- | Cap a single cell's rendered length, appending an ellipsis when shortened.
-truncateCell :: Int -> T.Text -> T.Text
-truncateCell n t
- | n <= 0 = t
- | T.compareLength t n /= GT = t
- | n == 1 = ellipsisText
- | otherwise = T.take (n - 1) t <> ellipsisText
-
--- | O(1) Creates an empty dataframe
-empty :: DataFrame
-empty =
- DataFrame
- { columns = V.empty
- , columnIndices = M.empty
- , dataframeDimensions = (0, 0)
- , derivingExpressions = M.empty
- }
-
--- | O(k) Get column names of the DataFrame in order of insertion.
-columnNames :: DataFrame -> [T.Text]
-columnNames = map fst . sortBy (compare `on` snd) . M.toList . columnIndices
-{-# INLINE columnNames #-}
-
-{- | Insert a column into a DataFrame. If a column with the same name already
-exists it is replaced in-place; otherwise the column is appended at the end.
-Other columns are expanded (padded with nulls) to match the new row count.
--}
-insertColumn :: T.Text -> Column -> DataFrame -> DataFrame
-insertColumn name column d =
- 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 ->
- DataFrame
- (V.map (expandColumn n) (columns d V.// [(i, column)]))
- (columnIndices d)
- (n, c)
- exprs
- Nothing ->
- DataFrame
- (V.map (expandColumn n) (columns d `V.snoc` column))
- (M.insert name c (columnIndices d))
- (n, c + 1)
- exprs
-
--- | Build a DataFrame from a list of @(name, column)@ pairs using 'insertColumn'.
-fromNamedColumns :: [(T.Text, Column)] -> DataFrame
-fromNamedColumns = foldl (\df (name, column) -> insertColumn name column df) empty
-
-{- | Safely retrieves a column by name from the dataframe.
-
-Returns 'Nothing' if the column does not exist.
-
-==== __Examples__
-
->>> getColumn "age" df
-Just (UnboxedColumn ...)
-
->>> getColumn "nonexistent" df
-Nothing
--}
-getColumn :: T.Text -> DataFrame -> Maybe Column
-getColumn name df
- | null df = Nothing
- | otherwise = do
- i <- columnIndices df M.!? name
- columns df V.!? i
-
-{- | Retrieves a column by name from the dataframe, throwing an exception if not found.
-
-This is an unsafe version of 'getColumn' that throws 'ColumnsNotFoundException'
-if the column does not exist. Use this when you are certain the column exists.
-
-==== __Throws__
-
-* 'ColumnsNotFoundException' - if the column with the given name does not exist
--}
-unsafeGetColumn :: T.Text -> DataFrame -> Column
-unsafeGetColumn name df = case getColumn name df of
- Nothing -> throw $ ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
- Just col -> col
-
-{- | Checks if the dataframe is empty (has no columns).
-
-Returns 'True' if the dataframe has no columns, 'False' otherwise.
-Note that a dataframe with columns but no rows is not considered null.
--}
-null :: DataFrame -> Bool
-null df = V.null (columns df)
-
--- | Convert a DataFrame to a CSV (comma-separated) text.
-toCsv :: DataFrame -> T.Text
-toCsv = toSeparated ','
-
--- | Convert a DataFrame to a CSV (comma-separated) string.
-toCsv' :: DataFrame -> String
-toCsv' = T.unpack . toSeparated ','
-
--- | Convert a DataFrame to a text representation with a custom separator.
-toSeparated :: Char -> DataFrame -> T.Text
-toSeparated sep df
- | null df = T.empty
- | otherwise =
- let (rows, _) = dataframeDimensions df
- headers = map fst (sortBy (compare `on` snd) (M.toList (columnIndices df)))
- sepText = T.singleton sep
- headerLine = T.intercalate sepText headers
- dataLines = map (T.intercalate sepText . getRowAsText df) [0 .. rows - 1]
- in T.unlines (headerLine : dataLines)
-
-getRowAsText :: DataFrame -> Int -> [T.Text]
-getRowAsText df i = map (`showElement` i) (V.toList (columns df))
-
-showElement :: Column -> Int -> T.Text
-showElement (BoxedColumn _ (c :: V.Vector a)) i = case c V.!? i of
- Nothing -> error $ "Column index out of bounds at row " ++ show i
- Just e
- | Just Refl <- testEquality (typeRep @a) (typeRep @T.Text) -> e
- | App t1 t2 <- typeRep @a
- , Just HRefl <- eqTypeRep t1 (typeRep @Maybe) ->
- case testEquality t2 (typeRep @T.Text) of
- Just Refl -> fromMaybe "null" e
- Nothing -> stripJust (T.pack (show e))
- | otherwise -> T.pack (show e)
-showElement (UnboxedColumn _ c) i = case c VU.!? i of
- Nothing -> error $ "Column index out of bounds at row " ++ show i
- Just e -> T.pack (show e)
-showElement (PackedText bm p) i = case bm of
- Just b | not (bitmapTestBit b i) -> "null"
- _ -> packedIndexText p i
-
-stripJust :: T.Text -> T.Text
-stripJust = fromMaybe "null" . T.stripPrefix "Just "
diff --git a/src/DataFrame/Internal/DictEncode.hs b/src/DataFrame/Internal/DictEncode.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/DictEncode.hs
+++ /dev/null
@@ -1,159 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Dictionary-encode a text (or factor) group key to dense @Int@ codes
-(research #4 / #5).
-
-A text group key is hashed and bucketed exactly as the grouping hash table does,
-but instead of carrying the full grouping output it only assigns each row a dense
-first-appearance code @0 .. card-1@ (a NULL row gets its own reserved code) and
-reports the cardinality. The intent was to feed those codes to the int fast path
-(direct-indexed low-card, or a packed composite key) so a text group key reaches
-it.
-
-VERDICT (this round): routing through the codes profiled SLOWER than the existing
-hash group-by on EVERY db-benchmark group-by question, so
-'DataFrame.Internal.Grouping' does not take the dict path (see 'tryDictGroup'
-there). The dict-build is its own hash pass over every row, and the hash
-group-by already fuses hashing and grouping into one pass; substituting int codes
-adds the encode pass without removing the dominant grouping work. Single low-card
-(Q1 id1), single high-card (Q3/Q7 id3) and the composites (Q2 id1:id2, Q10 six
-keys) were each measured and all lost.
-
-This module remains a correct, unit-tested building block: it produces the codes
-and the cardinality only; the routing decision lives in
-'DataFrame.Internal.Grouping'.
--}
-module DataFrame.Internal.DictEncode (
- dictEncodeColumn,
- dictEncodeColumnUpTo,
- dictMaxCardinality,
-) where
-
-import Control.Monad.ST (runST)
-import qualified Data.Text as T
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import Type.Reflection (typeRep)
-
-import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)
-import DataFrame.Internal.Hash (fnvOffset, mixBytes, mixText, nullSalt)
-import DataFrame.Internal.HashTable (htInsert, newHashTable)
-import DataFrame.Internal.PackedText (
- PackedTextData,
- packedLength,
- packedSlice,
- sliceEqBytes,
- )
-
-{- | Largest distinct-value count we will dictionary-encode. Above this the codes
-no longer index a reasonable direct accumulator and the dict-build pass is pure
-overhead, so the caller keeps the plain hash group-by. Matches the direct-group
-histogram budget.
--}
-dictMaxCardinality :: Int
-dictMaxCardinality = 1048576
-
-{- | Dictionary-encode a text-like column to dense first-appearance @Int@ codes.
-
-Returns @Just (codes, cardinality)@ where @codes!i@ is the dense id of row @i@'s
-value (a NULL row, when the column is nullable, is assigned its own reserved code
-distinct from every present value) and @cardinality@ is the number of distinct
-codes used. Returns 'Nothing' for any non-text column or when the cardinality
-exceeds 'dictMaxCardinality' (so the caller falls back to the hash group-by).
-
-Only 'PackedText' and boxed 'Data.Text.Text' columns are encoded; everything else
-is 'Nothing'.
--}
-dictEncodeColumn :: Column -> Maybe (VU.Vector Int, Int)
-dictEncodeColumn = dictEncodeColumnUpTo dictMaxCardinality
-
-{- | Dictionary-encode like 'dictEncodeColumn' but bail to 'Nothing' as soon as
-the distinct count would exceed @maxCard@. The early bail lets a low-cardinality
-PROBE (the single-key direct path) avoid a full high-cardinality pass when the
-column turns out to be high-card.
--}
-dictEncodeColumnUpTo :: Int -> Column -> Maybe (VU.Vector Int, Int)
-dictEncodeColumnUpTo maxCard (PackedText bm p) = encodePacked maxCard bm p
-dictEncodeColumnUpTo maxCard (BoxedColumn bm (v :: V.Vector a)) =
- case testEquality (typeRep @a) (typeRep @T.Text) of
- Just Refl -> encodeBoxedText maxCard bm v
- Nothing -> Nothing
-dictEncodeColumnUpTo _ _ = Nothing
-
-{- | Encode a packed-text column. Hashes each row's raw UTF-8 bytes (the same
-'mixBytes' the grouping hash uses) and re-verifies byte equality on collisions,
-assigning dense codes in first-appearance order. A null row hashes 'nullSalt' and
-re-verifies as equal-to-null only.
--}
-encodePacked ::
- Int -> Maybe Bitmap -> PackedTextData -> Maybe (VU.Vector Int, Int)
-encodePacked maxCard bm p =
- let !n = packedLength p
- valid i = case bm of
- Just b -> bitmapTestBit b i
- Nothing -> True
- hashAt i =
- if valid i
- then let (arr, o, l) = packedSlice p i in mixBytes fnvOffset arr o l
- else nullSalt
- eqAt a b =
- case (valid a, valid b) of
- (True, True) ->
- let (arrA, oA, lA) = packedSlice p a
- (arrB, oB, lB) = packedSlice p b
- in sliceEqBytes arrA oA lA arrB oB lB
- (False, False) -> True
- _ -> False
- in buildCodes maxCard n hashAt eqAt
-
-{- | Encode a boxed 'Data.Text.Text' column, mirroring 'encodePacked' but over
-boxed values (used when a user-built Text column is grouped).
--}
-encodeBoxedText ::
- Int -> Maybe Bitmap -> V.Vector T.Text -> Maybe (VU.Vector Int, Int)
-encodeBoxedText maxCard bm v =
- let !n = V.length v
- valid i = case bm of
- Just b -> bitmapTestBit b i
- Nothing -> True
- hashAt i =
- if valid i then mixText fnvOffset (V.unsafeIndex v i) else nullSalt
- eqAt a b =
- case (valid a, valid b) of
- (True, True) -> V.unsafeIndex v a == V.unsafeIndex v b
- (False, False) -> True
- _ -> False
- in buildCodes maxCard n hashAt eqAt
-
-{- | The shared code-assignment loop. Buckets every row through an
-open-addressing table on its precomputed hash, re-verifying the real value with
-@eqAt@ on a hash hit, assigning dense first-appearance codes. Bails to 'Nothing'
-the moment the distinct count would exceed 'dictMaxCardinality'.
--}
-buildCodes ::
- Int -> Int -> (Int -> Int) -> (Int -> Int -> Bool) -> Maybe (VU.Vector Int, Int)
-buildCodes maxCard n hashAt eqAt
- | n == 0 = Just (VU.empty, 0)
- | otherwise = runST $ do
- ht <- newHashTable (min n (maxCard + 1))
- codes <- VUM.new n
- let go !i !next
- | i >= n = pure (Just next)
- | next > maxCard = pure Nothing
- | otherwise = do
- let !h = hashAt i
- (code, isNew) <- htInsert ht eqAt next i h
- VUM.unsafeWrite codes i code
- go (i + 1) (if isNew then next + 1 else next)
- mres <- go 0 0
- case mres of
- Nothing -> pure Nothing
- Just card -> do
- frozen <- VU.unsafeFreeze codes
- pure (Just (frozen, card))
diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Expression.hs
+++ /dev/null
@@ -1,472 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE DisambiguateRecordFields #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE InstanceSigs #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE UndecidableInstances #-}
-{-# LANGUAGE NoFieldSelectors #-}
-
-module DataFrame.Internal.Expression where
-
-import qualified Data.Map.Strict as M
-import Data.String
-import qualified Data.Text as T
-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
-import qualified Data.Vector.Generic as VG
-import DataFrame.Internal.Column
-import Type.Reflection (Typeable, typeOf, typeRep)
-
-{- | Operators are an open typeclass: built-ins get their own 'Typeable' type so
-the simplifier can match them by 'cast', and users can add @instance@s. The
-generic 'UnUDF'/'BinUDF' carriers cover UDFs, dynamic-named, and arithmetic ops.
-Method names match the carrier record fields ('NoFieldSelectors' frees them), so
-existing construction and read sites are unchanged.
--}
-class (Typeable op) => UnaryOp op where
- unaryFn :: op a b -> a -> b
- unaryName :: op a b -> T.Text
- unarySymbol :: op a b -> Maybe T.Text
- unarySymbol _ = Nothing
-
-class (Typeable op) => BinaryOp op where
- binaryFn :: op a b c -> a -> b -> c
- binaryName :: op a b c -> T.Text
- binarySymbol :: op a b c -> Maybe T.Text
- binarySymbol _ = Nothing
- binaryCommutative :: op a b c -> Bool
- binaryCommutative _ = False
- binaryPrecedence :: op a b c -> Int
- binaryPrecedence _ = 9
-
-data UnUDF a b = MkUnaryOp
- { unaryFn :: a -> b
- , unaryName :: T.Text
- , unarySymbol :: Maybe T.Text
- }
-
-data BinUDF a b c = MkBinaryOp
- { binaryFn :: a -> b -> c
- , binaryName :: T.Text
- , binarySymbol :: Maybe T.Text
- , binaryCommutative :: Bool
- , binaryPrecedence :: Int
- }
-
-instance UnaryOp UnUDF where
- unaryFn (MkUnaryOp{unaryFn = f}) = f
- unaryName (MkUnaryOp{unaryName = n}) = n
- unarySymbol (MkUnaryOp{unarySymbol = s}) = s
-
-instance BinaryOp BinUDF where
- binaryFn (MkBinaryOp{binaryFn = f}) = f
- binaryName (MkBinaryOp{binaryName = n}) = n
- binarySymbol (MkBinaryOp{binarySymbol = s}) = s
- binaryCommutative (MkBinaryOp{binaryCommutative = c}) = c
- binaryPrecedence (MkBinaryOp{binaryPrecedence = p}) = p
-
-data MeanAcc = MeanAcc {-# UNPACK #-} !Double {-# UNPACK #-} !Int
- deriving (Show, Eq, Ord, Read)
-
-data AggStrategy a b where
- CollectAgg ::
- (VG.Vector v b, Typeable v) => T.Text -> (v b -> a) -> AggStrategy a b
- FoldAgg :: T.Text -> Maybe a -> (a -> b -> a) -> AggStrategy a b
- MergeAgg ::
- (Columnable acc) =>
- T.Text ->
- acc ->
- (acc -> b -> acc) ->
- (acc -> acc -> acc) ->
- (acc -> a) ->
- AggStrategy a b
-
-data Expr a where
- Col :: (Columnable a) => T.Text -> Expr a
- CastWith ::
- (Columnable a, Columnable b, Read a) =>
- T.Text ->
- T.Text ->
- (Either String a -> b) ->
- Expr b
- CastExprWith ::
- (Columnable a, Columnable b, Columnable src, Read a) =>
- T.Text ->
- (Either String a -> b) ->
- Expr src ->
- Expr b
- Lit :: (Columnable a) => a -> Expr a
- Unary ::
- (UnaryOp op, Columnable a, Columnable b) => op b a -> Expr b -> Expr a
- Binary ::
- (BinaryOp op, Columnable c, Columnable b, Columnable a) =>
- op c b a -> Expr c -> Expr b -> Expr a
- If :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
- Agg :: (Columnable a, Columnable b) => AggStrategy a b -> Expr b -> Expr a
- Over :: (Columnable a) => [T.Text] -> Expr a -> Expr a
-
-data UExpr where
- UExpr :: (Columnable a) => Expr a -> UExpr
-
-instance Show UExpr where
- show :: UExpr -> String
- show (UExpr expr) = show expr
-
-type NamedExpr = (T.Text, UExpr)
-
-instance (Num a, Columnable a) => Num (Expr a) where
- (+) :: Expr a -> Expr a -> Expr a
- (+) =
- Binary
- ( MkBinaryOp
- { binaryFn = (+)
- , binaryName = "add"
- , binarySymbol = Just "+"
- , binaryCommutative = True
- , binaryPrecedence = 6
- }
- )
-
- (-) :: Expr a -> Expr a -> Expr a
- (-) =
- Binary
- ( MkBinaryOp
- { binaryFn = (-)
- , binaryName = "sub"
- , binarySymbol = Just "-"
- , binaryCommutative = False
- , binaryPrecedence = 6
- }
- )
-
- (*) :: Expr a -> Expr a -> Expr a
- (*) =
- Binary
- ( MkBinaryOp
- { binaryFn = (*)
- , binaryName = "mult"
- , binarySymbol = Just "*"
- , binaryCommutative = True
- , binaryPrecedence = 7
- }
- )
-
- fromInteger :: Integer -> Expr a
- fromInteger = Lit . fromInteger
-
- negate :: Expr a -> Expr a
- negate =
- Unary
- (MkUnaryOp{unaryFn = negate, unaryName = "negate", unarySymbol = Nothing})
-
- abs :: (Num a) => Expr a -> Expr a
- abs = Unary (MkUnaryOp{unaryFn = abs, unaryName = "abs", unarySymbol = Nothing})
-
- signum :: (Num a) => Expr a -> Expr a
- signum =
- Unary
- (MkUnaryOp{unaryFn = signum, unaryName = "signum", unarySymbol = Nothing})
-
-add :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
-add = (+)
-
-sub :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
-sub = (-)
-
-mult :: (Num a, Columnable a) => Expr a -> Expr a -> Expr a
-mult = (*)
-
-instance (Fractional a, Columnable a) => Fractional (Expr a) where
- fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
- fromRational = Lit . fromRational
-
- (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
- (/) =
- Binary
- ( MkBinaryOp
- { binaryFn = (/)
- , binaryName = "divide"
- , binarySymbol = Just "/"
- , binaryCommutative = False
- , binaryPrecedence = 7
- }
- )
-
-divide :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
-divide = (/)
-
-instance (IsString a, Columnable a) => IsString (Expr a) where
- fromString :: String -> Expr a
- fromString s = Lit (fromString s)
-
-instance (Floating a, Columnable a) => Floating (Expr a) where
- pi :: (Floating a, Columnable a) => Expr a
- pi = Lit pi
- exp :: (Floating a, Columnable a) => Expr a -> Expr a
- exp = Unary (MkUnaryOp{unaryFn = exp, unaryName = "exp", unarySymbol = Nothing})
- sqrt :: (Floating a, Columnable a) => Expr a -> Expr a
- sqrt =
- Unary (MkUnaryOp{unaryFn = sqrt, unaryName = "sqrt", unarySymbol = Nothing})
- (**) :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
- (**) =
- Binary
- ( MkBinaryOp
- { binaryFn = (**)
- , binaryName = "exponentiate"
- , binarySymbol = Just "**"
- , binaryCommutative = False
- , binaryPrecedence = 8
- }
- )
- log :: (Floating a, Columnable a) => Expr a -> Expr a
- log = Unary (MkUnaryOp{unaryFn = log, unaryName = "log", unarySymbol = Nothing})
- logBase :: (Floating a, Columnable a) => Expr a -> Expr a -> Expr a
- logBase =
- Binary
- ( MkBinaryOp
- { binaryFn = logBase
- , binaryName = "logBase"
- , binarySymbol = Nothing
- , binaryCommutative = False
- , binaryPrecedence = 1
- }
- )
- sin :: (Floating a, Columnable a) => Expr a -> Expr a
- sin = Unary (MkUnaryOp{unaryFn = sin, unaryName = "sin", unarySymbol = Nothing})
- cos :: (Floating a, Columnable a) => Expr a -> Expr a
- cos = Unary (MkUnaryOp{unaryFn = cos, unaryName = "cos", unarySymbol = Nothing})
- tan :: (Floating a, Columnable a) => Expr a -> Expr a
- tan = Unary (MkUnaryOp{unaryFn = tan, unaryName = "tan", unarySymbol = Nothing})
- asin :: (Floating a, Columnable a) => Expr a -> Expr a
- asin =
- Unary (MkUnaryOp{unaryFn = asin, unaryName = "asin", unarySymbol = Nothing})
- acos :: (Floating a, Columnable a) => Expr a -> Expr a
- acos =
- Unary (MkUnaryOp{unaryFn = acos, unaryName = "acos", unarySymbol = Nothing})
- atan :: (Floating a, Columnable a) => Expr a -> Expr a
- atan =
- Unary (MkUnaryOp{unaryFn = atan, unaryName = "atan", unarySymbol = Nothing})
- sinh :: (Floating a, Columnable a) => Expr a -> Expr a
- sinh =
- Unary (MkUnaryOp{unaryFn = sinh, unaryName = "sinh", unarySymbol = Nothing})
- cosh :: (Floating a, Columnable a) => Expr a -> Expr a
- cosh =
- Unary (MkUnaryOp{unaryFn = cosh, unaryName = "cosh", unarySymbol = Nothing})
- asinh :: (Floating a, Columnable a) => Expr a -> Expr a
- asinh =
- Unary
- (MkUnaryOp{unaryFn = asinh, unaryName = "asinh", unarySymbol = Nothing})
- acosh :: (Floating a, Columnable a) => Expr a -> Expr a
- acosh =
- Unary
- (MkUnaryOp{unaryFn = acosh, unaryName = "acosh", unarySymbol = Nothing})
- atanh :: (Floating a, Columnable a) => Expr a -> Expr a
- atanh =
- Unary
- (MkUnaryOp{unaryFn = atanh, unaryName = "atanh", unarySymbol = Nothing})
-
-instance (Show a) => Show (Expr a) where
- show :: 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 ++ ")"
- show (Binary op a b) = "(" ++ T.unpack (binaryName op) ++ " " ++ show a ++ " " ++ show b ++ ")"
- show (Agg (CollectAgg op _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
- show (Agg (FoldAgg op _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
- show (Agg (MergeAgg op _ _ _ _) expr) = "(" ++ T.unpack op ++ " " ++ show expr ++ ")"
- show (Over keys inner) = "(over " ++ show keys ++ " " ++ show inner ++ ")"
-
-normalize :: (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)
- Binary op e1 e2
- | binaryCommutative op ->
- let n1 = normalize e1
- n2 = normalize e2
- in case testEquality (typeOf n1) (typeOf n2) of
- Nothing -> expr
- Just Refl ->
- if compareExpr n1 n2 == GT
- then Binary op n2 n1 -- Swap to canonical order
- else Binary op n1 n2
- | otherwise -> Binary op (normalize e1) (normalize e2)
- Agg strat e -> Agg strat (normalize e)
- Over keys inner -> Over keys (normalize inner)
-
--- Compare expressions for ordering (used in normalization)
-compareExpr :: Expr a -> Expr a -> Ordering
-compareExpr e1 e2 = compare (exprKey e1) (exprKey e2)
- 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
- exprKey (Binary op e1' e2') = "4:" ++ T.unpack (binaryName op) ++ exprKey e1' ++ exprKey e2'
- exprKey (Agg (CollectAgg name _) e) = "5:" ++ T.unpack name ++ exprKey e
- exprKey (Agg (FoldAgg name _ _) e) = "5:" ++ T.unpack name ++ exprKey e
- exprKey (Agg (MergeAgg name _ _ _ _) e) = "5:" ++ T.unpack name ++ exprKey e
- exprKey (Over keys e) = "6:over:" ++ show keys ++ exprKey e
-
-eqExpr :: forall a. (Columnable a) => Expr a -> Expr a -> Bool
-eqExpr l r = eqNormalized (normalize l) (normalize r)
- where
- exprEq :: (Columnable b, Columnable c) => Expr b -> Expr c -> Bool
- exprEq e1 e2 = case testEquality (typeOf e1) (typeOf e2) of
- Just Refl -> eqExpr e1 e2
- Nothing -> False
- eqNormalized :: Expr a -> Expr a -> Bool
- eqNormalized (Col n1) (Col n2) = n1 == n2
- eqNormalized (CastWith n1 t1 _) (CastWith n2 t2 _) = n1 == n2 && t1 == t2
- eqNormalized (CastExprWith t1 _ e1) (CastExprWith t2 _ e2) = t1 == t2 && e1 `exprEq` e2
- eqNormalized (Lit v1) (Lit v2) = v1 == v2
- eqNormalized (If c1 t1 e1) (If c2 t2 e2) =
- eqExpr c1 c2 && t1 `exprEq` t2 && e1 `exprEq` e2
- eqNormalized (Unary op1 e1) (Unary op2 e2) = unaryName op1 == unaryName op2 && e1 `exprEq` e2
- eqNormalized (Binary op1 e1a e1b) (Binary op2 e2a e2b) = binaryName op1 == binaryName op2 && e1a `exprEq` e2a && e1b `exprEq` e2b
- eqNormalized (Agg (CollectAgg n1 _) e1) (Agg (CollectAgg n2 _) e2) =
- n1 == n2 && e1 `exprEq` e2
- eqNormalized (Agg (FoldAgg n1 _ _) e1) (Agg (FoldAgg n2 _ _) e2) =
- n1 == n2 && e1 `exprEq` e2
- eqNormalized (Agg (MergeAgg n1 _ _ _ _) e1) (Agg (MergeAgg n2 _ _ _ _) e2) =
- n1 == n2 && e1 `exprEq` e2
- eqNormalized (Over k1 e1) (Over k2 e2) = k1 == k2 && e1 `exprEq` e2
- eqNormalized _ _ = False
-
-replaceExpr ::
- forall a b c.
- (Columnable a, Columnable b, Columnable c) =>
- Expr a -> Expr b -> Expr c -> Expr c
-replaceExpr new old expr = case testEquality (typeRep @b) (typeRep @c) of
- Just Refl -> case testEquality (typeRep @a) (typeRep @c) of
- Just Refl -> if eqExpr old expr then new else replace'
- Nothing -> expr
- Nothing -> replace'
- 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)
- (Unary op value) -> Unary op (replaceExpr new old value)
- (Binary op l r) -> Binary op (replaceExpr new old l) (replaceExpr new old r)
- (Agg op inner) -> Agg op (replaceExpr new old inner)
- (Over keys inner) -> Over keys (replaceExpr new old inner)
-
-{- | Simultaneously substitute column references using a map from column name to
-replacement expression. Unlike folding 'replaceExpr' over the bindings, this is a
-single parallel pass: every 'Col' reference is resolved against the original map,
-so a swap such as @{a ↦ col b, b ↦ col a}@ is handled correctly (sequential
-replacement would collapse both columns onto one).
-
-Only 'Col' references are substituted; raw-text references inside 'CastWith' and
-'Over' partition keys are left untouched (documented limitation — the fitted ML
-transforms that use this only ever emit @Col@/@Lit@/@Unary@/@Binary@/@If@). A
-binding whose replacement type does not match the referenced column's type is a
-programmer error and raises an exception.
--}
-substituteColumns ::
- forall a. (Columnable a) => M.Map T.Text UExpr -> Expr a -> Expr a
-substituteColumns subs = go
- where
- go :: forall b. (Columnable b) => Expr b -> Expr b
- go e@(Col name) = case M.lookup name subs of
- Nothing -> e
- Just (UExpr (repl :: Expr c)) -> case testEquality (typeRep @b) (typeRep @c) of
- Just Refl -> repl
- Nothing ->
- error $
- "substituteColumns: type mismatch for column "
- ++ show name
- ++ "; column has type "
- ++ show (typeRep @b)
- ++ " but replacement has type "
- ++ show (typeRep @c)
- go e@(CastWith{}) = e
- go (CastExprWith t f e) = CastExprWith t f (go e)
- go e@(Lit _) = e
- go (If cond l r) = If (go cond) (go l) (go r)
- go (Unary op value) = Unary op (go value)
- go (Binary op l r) = Binary op (go l) (go r)
- go (Agg op inner) = Agg op (go inner)
- go (Over keys inner) = Over keys (go inner)
-
-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
-eSize (Binary _ l r) = 1 + eSize l + eSize r
-eSize (Agg _strategy expr) = eSize expr + 1
-eSize (Over _ inner) = 1 + eSize inner
-
-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
-getColumns (Binary _op l r) = getColumns l <> getColumns r
-getColumns (Agg _strategy expr) = getColumns expr
-getColumns (Over keys inner) = keys <> getColumns inner
-
-prettyPrint :: Expr a -> String
-prettyPrint = go 0 0
- where
- indent :: Int -> String
- indent n = replicate (n * 2) ' '
-
- 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 =
- "if "
- ++ go (depth + 1) 0 cond
- ++ "\n"
- ++ indent (depth + 1)
- ++ "then "
- ++ go (depth + 1) 0 t
- ++ "\n"
- ++ indent (depth + 1)
- ++ "else "
- ++ go (depth + 1) 0 e
- in if prec > 0 then "(" ++ inner ++ ")" else inner
- Unary op arg -> case unarySymbol op of
- Nothing -> T.unpack (unaryName op) ++ "(" ++ go depth 0 arg ++ ")"
- Just sym -> T.unpack sym ++ "(" ++ go depth 0 arg ++ ")"
- Binary op l r ->
- let p = binaryPrecedence op
- inner = case binarySymbol op of
- Just name -> go depth p l ++ " " ++ T.unpack name ++ " " ++ go depth p r
- Nothing ->
- T.unpack (binaryName op) ++ "(" ++ go depth p l ++ ", " ++ go depth p r ++ ")"
- in if prec > p then "(" ++ inner ++ ")" else inner
- Agg (CollectAgg op _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"
- Agg (FoldAgg op _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"
- Agg (MergeAgg op _ _ _ _) arg -> T.unpack op ++ "(" ++ go depth 0 arg ++ ")"
- Over keys inner -> go depth 0 inner ++ ".over(" ++ show (map T.unpack keys) ++ ")"
diff --git a/src/DataFrame/Internal/Grouping.hs b/src/DataFrame/Internal/Grouping.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Grouping.hs
+++ /dev/null
@@ -1,454 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE Strict #-}
-{-# LANGUAGE TypeApplications #-}
-
-module DataFrame.Internal.Grouping (
- groupBy,
- groupBySeq,
- groupByPar,
- buildRowToGroup,
- changingPoints,
-) where
-
-import qualified Data.List as L
-import qualified Data.Map as M
-import qualified Data.Text as T
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-import Control.Exception (throw)
-import Control.Monad
-import Control.Monad.ST (ST, runST)
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import DataFrame.Errors
-import DataFrame.Internal.Column (
- Bitmap,
- Column (..),
- bitmapTestBit,
- )
-import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))
-import DataFrame.Internal.DictEncode (dictEncodeColumnUpTo)
-import DataFrame.Internal.GroupingDirect (
- DirectGrouping (..),
- tryDirectGroupColumn,
- )
-import DataFrame.Internal.GroupingPar (parallelAssignGroups, shouldParallelize)
-import DataFrame.Internal.Hash
-import DataFrame.Internal.HashTable (htInsert, newHashTable)
-import DataFrame.Internal.PackedText (
- PackedTextData,
- packedLength,
- packedSlice,
- sliceEqBytes,
- )
-import DataFrame.Internal.RadixRank (rankByHash)
-import DataFrame.Internal.Types
-import System.IO.Unsafe (unsafePerformIO)
-import Type.Reflection (typeRep)
-
-{- | O(k * n) groups the dataframe by the given rows aggregating the remaining rows
-into vector that should be reduced later.
-
-Rows are bucketed with an unboxed open-addressing hash table
-('DataFrame.Internal.HashTable') that maps each row's key-hash to a dense group
-id, re-verifying the real key columns on every hash hit. This both cuts the
-per-row boxed allocation of the previous 'Data.IntMap' bucketing (less GC) and
-fixes a latent collision bug where two distinct keys sharing a hash were merged.
-Groups are numbered in first-appearance order; 'valueIndices' / 'offsets' are
-then derived by a stable counting sort on the group id.
--}
-groupBy ::
- [T.Text] ->
- DataFrame ->
- GroupedDataFrame
-groupBy names df
- | any (`notElem` columnNames df) names =
- throw $
- ColumnsNotFoundException
- (names L.\\ columnNames df)
- "groupBy"
- (columnNames df)
- | nRows df == 0 =
- Grouped
- df
- names
- VU.empty
- (VU.fromList [0])
- VU.empty
- | Just dg <- tryDirectGroup names df = dg
- | shouldParallelize n = groupByPar names df
- | otherwise = groupBySeq names df
- where
- !n = nRows df
-
-{- | The low-cardinality direct-indexed grouping fast path
-('DataFrame.Internal.GroupingDirect'). Fires only for a SINGLE clean small-range
-@Int@ key column; one shared function feeds both -N1 and -N8 so the result is
-identical at any capability count (parallel==sequential by construction). Returns
-'Nothing' on any other key shape, falling through to the hash group-by.
--}
-tryDirectGroup :: [T.Text] -> DataFrame -> Maybe GroupedDataFrame
-tryDirectGroup [name] df = do
- col <- M.lookup name (columnIndices df) >>= \i -> columns df V.!? i
- case tryDirectGroupColumn col of
- Just dg ->
- Just (Grouped df [name] (dgValueIndices dg) (dgOffsets dg) (dgRowToGroup dg))
- Nothing -> tryDictGroup (nRows df) df [name] col
-tryDirectGroup _ _ = Nothing
-
-{- | Dictionary-encode a single text key to dense int codes (the codes ARE the
-first-appearance group ids), then derive @valueIndices@/@offsets@ from those
-codes with one counting sort.
-
-PROFILED AS A LOSS, so this currently always falls back ('Nothing'). The
-dict-build is its own hash pass over every row, and the hash group-by already
-hashes and groups in one fused pass: at -N1 the dict path measured ~0.44s vs
-~0.33s for the hash path on Q1 (id1, 100 groups) at 1e7 rows, and at -N8 the
-parallel partitioned grouping is far faster than any sequential dict-build. The
-high-card single keys (Q3/Q7 id3 ~1e5) and the multi-key composites (Q2 id1:id2,
-Q10 six keys) were each tried and also lost — substituting int codes does not
-remove the dominant grouping passes, it adds the encode passes on top. The
-'DataFrame.Internal.DictEncode.dictEncodeColumnUpTo' step is kept and unit-tested
-as a correct building block; the routing here is deliberately disabled. The
-@_n@/@_df@/@_names@/@_col@ wiring is retained so re-enabling is a one-line change
-should a parallel dict-encode ever change the verdict.
--}
-tryDictGroup ::
- Int -> DataFrame -> [T.Text] -> Column -> Maybe GroupedDataFrame
-tryDictGroup n df names col
- | dictGroupEnabled && not (shouldParallelize n) = do
- (codes, card) <- dictEncodeColumnUpTo dictSingleThreshold col
- let (vis, os) = indicesFromGroups codes card
- Just (Grouped df names vis os codes)
- | otherwise = Nothing
-
-{- | Master switch for the single-key dict-encode grouping path. 'False' because
-it profiled slower than the hash group-by on every db-benchmark group-by question
-(see 'tryDictGroup'); the path is kept compiled and tested but not taken.
--}
-dictGroupEnabled :: Bool
-dictGroupEnabled = False
-
-{- | Cardinality ceiling the single-key dict-encode probe would use: it bails to
-'Nothing' once the distinct count passes this, so a high-card key never pays for
-a full encode pass. Only consulted when 'dictGroupEnabled' is 'True'.
--}
-dictSingleThreshold :: Int
-dictSingleThreshold = 4096
-
-{- | The sequential grouping path: a single open-addressing table over all rows,
-canonically remapped. Always available regardless of capabilities; the parallel
-path is verified equal to it by a property test.
--}
-groupBySeq :: [T.Text] -> DataFrame -> GroupedDataFrame
-groupBySeq names df =
- let !n = nRows df
- indicesToGroup = keyColIndices names df
- (rtg0, repHash, repRow) = assignGroups df indicesToGroup n
- !nGroups = VU.length repHash
- !remap = canonicalRemap repHash repRow
- !rtg = VU.map (VU.unsafeIndex remap) rtg0
- (vis, os) = indicesFromGroups rtg nGroups
- in Grouped df names vis os rtg
-
-{- | The parallel partitioned grouping path (see
-'DataFrame.Internal.GroupingPar'). Forks one task per capability; produces an
-output bit-for-bit identical to 'groupBySeq'. Pure via 'unsafePerformIO' (the IO
-is deterministic thread fan-out only).
--}
-groupByPar :: [T.Text] -> DataFrame -> GroupedDataFrame
-groupByPar names df =
- let !n = nRows df
- indicesToGroup = keyColIndices names df
- !hashes = runST (computeHashes df indicesToGroup n)
- !eqRow = eqKeyRow df indicesToGroup
- (rtg, vis, os) = unsafePerformIO (parallelAssignGroups n hashes eqRow)
- in Grouped df names vis os rtg
-{-# NOINLINE groupByPar #-}
-
--- | Column indices of the requested key columns, in column order.
-keyColIndices :: [T.Text] -> DataFrame -> [Int]
-keyColIndices names df =
- M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)
-
-{- | Assign every row to a dense group id via the open-addressing table, in
-first-appearance order. Returns @(rowToGroup, repHash, repRow)@ where @repHash@ /
-@repRow@ are the hash and representative row index of each group (indexed by the
-first-appearance id). The table re-verifies the real key with 'eqKeyRow' on each
-hash hit so colliding keys are kept apart.
--}
-assignGroups ::
- DataFrame -> [Int] -> Int -> (VU.Vector Int, VU.Vector Int, VU.Vector Int)
-assignGroups df indicesToGroup n = runST $ do
- hashes <- computeHashes df indicesToGroup n
- let !eqRow = eqKeyRow df indicesToGroup
- ht <- newHashTable n
- rtg <- VUM.new n
- -- At most n groups; trimmed to the actual count on freeze.
- repHashM <- VUM.new n
- repRowM <- VUM.new n
- let go !i !next
- | i >= n = pure next
- | otherwise = do
- let !h = VU.unsafeIndex hashes i
- (gid, isNew) <- htInsert ht eqRow next i h
- VUM.unsafeWrite rtg i gid
- when isNew $ do
- VUM.unsafeWrite repHashM next h
- VUM.unsafeWrite repRowM next i
- go (i + 1) (if isNew then next + 1 else next)
- !nGroups <- go 0 0
- frozen <- VU.unsafeFreeze rtg
- repHash <- VU.unsafeFreeze (VUM.slice 0 nGroups repHashM)
- repRow <- VU.unsafeFreeze (VUM.slice 0 nGroups repRowM)
- pure (frozen, repHash, repRow)
-
-{- | Map each first-appearance group id to its canonical id: groups are ordered
-by ascending representative hash, tie-broken by representative row index. This
-makes the emitted group order a deterministic function of the key set (not of
-input row order), so set operations like @union a b@ and @union b a@ agree, and
-reproduces the ascending-hash order of the previous 'Data.IntMap' grouping.
-Returns @remap@ with @remap[firstAppearanceId] = canonicalId@.
-
-The ordering is the stable hash-rank of 'DataFrame.Internal.RadixRank': @repRow@
-is strictly ascending in first-appearance id order (a new group's representative
-is the first row that reaches it, scanned in increasing row index), so the stable
-sort's equal-hash tie-break reproduces the old @(hash, repRow)@ comparison. O(g)
-with no boxed-tuple comparison sort — this keeps the @1e7@-distinct-group case
-(Q10) off an @n log n@ list sort.
--}
-canonicalRemap :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
-canonicalRemap repHash _repRow =
- runST (rankByHash (pure . VU.unsafeIndex repHash) (VU.length repHash))
-
-{- | Compute the FNV row-hash of the key columns into a fresh unboxed vector,
-mixing 'nullSalt' for null slots so a missing value never collides with a
-present one of the same bits.
--}
-computeHashes :: DataFrame -> [Int] -> Int -> ST s (VU.Vector Int)
-computeHashes df indicesToGroup n = do
- mh <- VUM.replicate n fnvOffset
- let selectedCols = map (columns df V.!) indicesToGroup
- forM_ selectedCols $ \case
- UnboxedColumn ubm (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> hashUnboxed mh ubm mixInt v
- Nothing ->
- case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> hashUnboxed mh ubm mixDouble v
- Nothing ->
- case sIntegral @a of
- STrue ->
- hashUnboxed mh ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v
- SFalse ->
- case sFloating @a of
- STrue ->
- hashUnboxed mh ubm (\h d -> mixDouble h (realToFrac d :: Double)) v
- SFalse ->
- hashUnboxed mh ubm mixShow v
- BoxedColumn bm (v :: V.Vector a) ->
- case testEquality (typeRep @a) (typeRep @T.Text) of
- Just Refl ->
- V.imapM_
- ( \i t -> do
- !h <- VUM.unsafeRead mh i
- let h' = case bm of
- Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
- _ -> mixText h t
- VUM.unsafeWrite mh i h'
- )
- v
- Nothing ->
- V.imapM_
- ( \i d -> do
- !h <- VUM.unsafeRead mh i
- let h' = case bm of
- Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
- _ -> mixShow h d
- VUM.unsafeWrite mh i h'
- )
- v
- PackedText bm p -> hashPacked mh bm p
- VU.unsafeFreeze mh
-
-{- | Build the row-key equality predicate over the selected key columns.
-@eqKeyRow df idxs a b@ is 'True' iff rows @a@ and @b@ are equal across all key
-columns, comparing validity bits first (a null equals only another null) then
-the underlying value. Used by the hash table to reject hash collisions.
--}
-eqKeyRow :: DataFrame -> [Int] -> Int -> Int -> Bool
-eqKeyRow df indicesToGroup =
- let !preds = map (colEqRow . (columns df V.!)) indicesToGroup
- go [] _ _ = True
- go (p : ps) a b = p a b && go ps a b
- in go preds
-
-{- | Per-column row equality respecting nulls. Two rows are equal at a column
-when both are null, or both are valid and their values compare equal.
--}
-colEqRow :: Column -> (Int -> Int -> Bool)
-colEqRow (UnboxedColumn bm v) =
- let eqV a b = VU.unsafeIndex v a == VU.unsafeIndex v b
- in withNulls bm eqV
-colEqRow (BoxedColumn bm v) =
- let eqV a b = V.unsafeIndex v a == V.unsafeIndex v b
- in withNulls bm eqV
-colEqRow (PackedText bm p) =
- let eqV a b =
- let (arrA, oA, lA) = packedSlice p a
- (arrB, oB, lB) = packedSlice p b
- in sliceEqBytes arrA oA lA arrB oB lB
- in withNulls bm eqV
-{-# INLINE colEqRow #-}
-
-{- | Wrap a value-equality with null handling: equal iff both valid and the
-values agree, or both null.
--}
-withNulls :: Maybe Bitmap -> (Int -> Int -> Bool) -> (Int -> Int -> Bool)
-withNulls Nothing eqV = eqV
-withNulls (Just bm) eqV = \a b ->
- case (bitmapTestBit bm a, bitmapTestBit bm b) of
- (True, True) -> eqV a b
- (False, False) -> True
- _ -> False
-{-# INLINE withNulls #-}
-
-{- | Derive @(valueIndices, offsets)@ from @rowToGroup@ via a stable counting
-sort on the group id: a per-group count, a prefix-sum into group offsets, then a
-single placement pass keeps rows in original order within each group.
--}
-indicesFromGroups :: VU.Vector Int -> Int -> (VU.Vector Int, VU.Vector Int)
-indicesFromGroups rtg nGroups = runST $ do
- let !n = VU.length rtg
- -- counts[g] = size of group g (g in [0, nGroups)). Slot nGroups stays 0 so
- -- the exclusive scan below lands n in offsets[nGroups].
- counts <- VUM.replicate (nGroups + 1) 0
- let countLoop !i
- | i >= n = pure ()
- | otherwise = do
- let !g = VU.unsafeIndex rtg i
- c <- VUM.unsafeRead counts g
- VUM.unsafeWrite counts g (c + 1)
- countLoop (i + 1)
- countLoop 0
- -- Exclusive prefix scan of counts into offsets: offsets[k] is the start of
- -- group k and offsets[nGroups] == n.
- offsM <- VUM.new (nGroups + 1)
- let scan !k !acc
- | k > nGroups = pure ()
- | otherwise = do
- VUM.unsafeWrite offsM k acc
- c <- VUM.unsafeRead counts k
- scan (k + 1) (acc + c)
- scan 0 0
- -- 'counts' is repurposed as a per-group write cursor seeded at each group's
- -- start offset, giving a stable placement (rows keep original order).
- let seed !k
- | k > nGroups = pure ()
- | otherwise = do
- s <- VUM.unsafeRead offsM k
- VUM.unsafeWrite counts k s
- seed (k + 1)
- seed 0
- vis <- VUM.new n
- let place !i
- | i >= n = pure ()
- | otherwise = do
- let !g = VU.unsafeIndex rtg i
- pos <- VUM.unsafeRead counts g
- VUM.unsafeWrite vis pos i
- VUM.unsafeWrite counts g (pos + 1)
- place (i + 1)
- place 0
- offs <- VU.unsafeFreeze offsM
- frozenVis <- VU.unsafeFreeze vis
- pure (frozenVis, offs)
-
-{- | Fold a value-mix over an unboxed column into the running hash vector,
-respecting the null bitmap: a null slot mixes a fixed 'nullSalt' sentinel.
--}
-hashUnboxed ::
- (VU.Unbox a) =>
- VUM.MVector s Int ->
- Maybe Bitmap ->
- (Int -> a -> Int) ->
- VU.Vector a ->
- ST s ()
-hashUnboxed mh ubm mix v = case ubm of
- Nothing ->
- VU.imapM_
- ( \i x -> do
- !h <- VUM.unsafeRead mh i
- VUM.unsafeWrite mh i (mix h x)
- )
- v
- Just bm ->
- VU.imapM_
- ( \i x -> do
- !h <- VUM.unsafeRead mh i
- VUM.unsafeWrite
- mh
- i
- (if bitmapTestBit bm i then mix h x else mixInt h nullSalt)
- )
- v
-{-# INLINE hashUnboxed #-}
-
-{- | Hash a packed-text column over its raw UTF-8 byte slices (no per-row
-'Data.Text.Text'), mixing 'nullSalt' for null rows. Shares 'mixBytes' with
-'mixText' so packed and boxed Text columns hash identically.
--}
-hashPacked ::
- VUM.MVector s Int -> Maybe Bitmap -> PackedTextData -> ST s ()
-hashPacked mh bm p = go 0
- where
- !n = packedLength p
- go !i
- | i >= n = pure ()
- | otherwise = do
- !h <- VUM.unsafeRead mh i
- let h' = case bm of
- Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
- _ -> let (arr, o, l) = packedSlice p i in mixBytes h arr o l
- VUM.unsafeWrite mh i h'
- go (i + 1)
-{-# INLINE hashPacked #-}
-
--- Inline accessors to avoid depending on Operations.Core
-
-columnNames :: DataFrame -> [T.Text]
-columnNames = M.keys . columnIndices
-
-nRows :: DataFrame -> Int
-nRows = fst . dataframeDimensions
-
-{- | Build the rowToGroup lookup vector from valueIndices and offsets.
-rowToGroup[i] = k means row i belongs to group k.
--}
-buildRowToGroup :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int
-buildRowToGroup n vis os = runST $ do
- rtg <- VUM.new n
- let nGroups = VU.length os - 1
- forM_ [0 .. nGroups - 1] $ \k ->
- let s = VU.unsafeIndex os k
- e = VU.unsafeIndex os (k + 1)
- in forM_ [s .. e - 1] $ \i ->
- VUM.unsafeWrite rtg (VU.unsafeIndex vis i) k
- VU.unsafeFreeze rtg
-{-# NOINLINE buildRowToGroup #-}
-
-changingPoints :: VU.Vector (Int, Int) -> VU.Vector Int
-changingPoints vs =
- VU.reverse
- (VU.fromList (VU.length vs : fst (VU.ifoldl' findChangePoints initialState vs)))
- where
- initialState = ([0], snd (VU.head vs))
- findChangePoints (!offs, !currentVal) index (_, !newVal)
- | currentVal == newVal = (offs, currentVal)
- | otherwise = (index : offs, newVal)
diff --git a/src/DataFrame/Internal/GroupingDirect.hs b/src/DataFrame/Internal/GroupingDirect.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/GroupingDirect.hs
+++ /dev/null
@@ -1,260 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Low-cardinality DIRECT-INDEXED grouping fast path (research #4).
-
-When the group-by key is a single clean (non-null) unboxed @Int@ column whose
-value /range/ is small (@max - min + 1 <= directGroupThreshold@), the value
-itself indexes a dense accumulator: there is no hashing, no key re-verification,
-and no open-addressing probe. This is the X100 / ClickHouse FixedHashMap idea
-applied to the grouping step — the dominant cost of the low-cardinality
-db-benchmark questions (id4=100, id6=1e5) is the hash group-by, not the
-aggregate scatter, so bypassing the hash here is the real lever.
-
-The pipeline is three linear passes plus an O(range) compaction:
-
- 1. min/max of the key (parallel range-reduce, order-independent).
- 2. a per-value histogram (parallel per-thread histograms, exact-integer merge).
- 3. compact non-empty values into dense ids in ASCENDING value order, then a
- stable placement pass building @valueIndices@.
-
-The emitted group order is ascending key value — a deterministic function of the
-key set (like the hash path's canonical order), so set operations stay
-commutative and the db-benchmark checksums (order-independent sums) are
-unchanged. Crucially this single function is shared by both the sequential and
-parallel 'groupBy' entry points, so the parallel==sequential parity is automatic
-(identical output by construction) at any @-N@.
--}
-module DataFrame.Internal.GroupingDirect (
- directGroupThreshold,
- tryDirectGroupColumn,
- DirectGrouping (..),
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import System.IO.Unsafe (unsafePerformIO)
-import Type.Reflection (typeRep)
-
-import DataFrame.Internal.Column (Column (..))
-
-{- | Largest key value RANGE (max - min + 1) the direct grouping path accepts. A
-@2^20@-slot histogram is 8MB; the low-cardinality questions sit far below it
-(id4 range 100, id6 range 1e5). Wider ranges fall back to the hash group-by.
--}
-directGroupThreshold :: Int
-directGroupThreshold = 1048576
-
-{- | The grouping layout the hash path also produces: @rowToGroup@, the
-group-sorted @valueIndices@, the @offsets@ prefix array, and the group count.
--}
-data DirectGrouping = DirectGrouping
- { dgRowToGroup :: !(VU.Vector Int)
- , dgValueIndices :: !(VU.Vector Int)
- , dgOffsets :: !(VU.Vector Int)
- , dgNGroups :: !Int
- }
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
-parThreshold :: Int
-parThreshold = 200000
-
-{- | Take the direct path if the (single) key column is a clean non-null unboxed
-@Int@ column with a small value range. Returns 'Nothing' to fall back to the
-hash group-by on anything else (boxed/text keys, nullable, wide ranges, empty).
--}
-tryDirectGroupColumn :: Column -> Maybe DirectGrouping
-tryDirectGroupColumn (UnboxedColumn Nothing (v :: VU.Vector a))
- | Just Refl <- testEquality (typeRep @a) (typeRep @Int)
- , not (VU.null v) =
- let (!mn, !mx) = rangeOf v
- !range = mx - mn + 1
- in if range >= 1 && range <= directGroupThreshold
- then Just (directGroup v mn range)
- else Nothing
-tryDirectGroupColumn _ = Nothing
-
--- | Parallel min/max reduce (order-independent).
-rangeOf :: VU.Vector Int -> (Int, Int)
-rangeOf v
- | not (shouldPar n) = rangeChunk v 0 n
- | otherwise = unsafePerformIO $ do
- let !caps = capabilities
- !per = (n + caps - 1) `div` caps
- spawn w = do
- var <- newEmptyMVar
- let !lo = min n (w * per)
- !hi = min n (lo + per)
- _ <- forkIO (try (pure $! rangeChunk v lo hi) >>= putMVar var)
- pure var
- vars <- mapM spawn [0 .. caps - 1]
- rs <- mapM takeMVar vars
- rs' <- mapM (either (throwIO @SomeException) pure) rs
- pure (combineRanges (filter (\(a, _) -> a /= maxBound) rs'))
- where
- !n = VU.length v
-{-# NOINLINE rangeOf #-}
-
-rangeChunk :: VU.Vector Int -> Int -> Int -> (Int, Int)
-rangeChunk v lo hi = go lo maxBound minBound
- where
- go !i !mn !mx
- | i >= hi = (mn, mx)
- | otherwise =
- let !x = VU.unsafeIndex v i
- in go (i + 1) (min mn x) (max mx x)
-
-combineRanges :: [(Int, Int)] -> (Int, Int)
-combineRanges [] = (0, 0)
-combineRanges ((a0, b0) : rest) = foldr (\(a, b) (ma, mb) -> (min ma a, max mb b)) (a0, b0) rest
-
-shouldPar :: Int -> Bool
-shouldPar n = n >= parThreshold && capabilities > 1
-
-{- | Build the grouping by counting sort on @value - min@: a (parallel) per-value
-histogram, compaction of non-empty values into ascending dense ids, an exclusive
-scan into offsets, then a stable placement pass building @valueIndices@ and
-@rowToGroup@.
--}
-directGroup :: VU.Vector Int -> Int -> Int -> DirectGrouping
-directGroup v mn range = unsafePerformIO $ do
- let !n = VU.length v
- -- 1. Per-value histogram over the dense value index (parallel, exact).
- hist <- buildHistogram v mn range n
- -- 2. Compact non-empty values -> dense group ids (ascending value order),
- -- recording each value's group id and the group's row count.
- valToGroup <- VUM.replicate range (-1 :: Int)
- grpCount <- VUM.new range
- nGroups <- compact hist range valToGroup grpCount
- -- 3. Exclusive prefix scan of group counts -> offsets (length nGroups + 1).
- offsM <- VUM.new (nGroups + 1)
- cursor <- VUM.new nGroups
- scanOffsets grpCount nGroups offsM cursor
- -- 4. Stable placement: rowToGroup[i] and valueIndices in group order.
- rtg <- VUM.new n
- vis <- VUM.new n
- place v mn n valToGroup cursor rtg vis
- frozenRtg <- VU.unsafeFreeze rtg
- frozenVis <- VU.unsafeFreeze vis
- frozenOffs <- VU.unsafeFreeze offsM
- pure (DirectGrouping frozenRtg frozenVis frozenOffs nGroups)
-{-# NOINLINE directGroup #-}
-
-{- | Parallel per-value histogram: each worker fills a private @range@-slot
-count over its row chunk, then the partials are summed (exact integers, so the
-merge order is irrelevant). Sequential single pass below 'parThreshold'.
--}
-buildHistogram :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
-buildHistogram v mn range n
- | not (shouldPar n) = histChunk v mn range 0 n
- | otherwise = do
- let !caps = capabilities
- !per = (n + caps - 1) `div` caps
- spawn w = do
- var <- newEmptyMVar
- let !lo = min n (w * per)
- !hi = min n (lo + per)
- _ <- forkIO (try (histChunk v mn range lo hi) >>= putMVar var)
- pure var
- vars <- mapM spawn [0 .. caps - 1]
- rs <- mapM takeMVar vars
- parts <- mapM (either (throwIO @SomeException) pure) rs
- case parts of
- [] -> VUM.replicate range 0
- (p0 : rest) -> do
- mapM_ (addInto p0 range) rest
- pure p0
-
-histChunk :: VU.Vector Int -> Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
-histChunk v mn range lo hi = do
- acc <- VUM.replicate range (0 :: Int)
- let go !i
- | i >= hi = pure ()
- | otherwise = do
- let !k = VU.unsafeIndex v i - mn
- c <- VUM.unsafeRead acc k
- VUM.unsafeWrite acc k (c + 1)
- go (i + 1)
- go lo
- pure acc
-
-addInto :: VUM.IOVector Int -> Int -> VUM.IOVector Int -> IO ()
-addInto dst range src = go 0
- where
- go !k
- | k >= range = pure ()
- | otherwise = do
- a <- VUM.unsafeRead dst k
- b <- VUM.unsafeRead src k
- VUM.unsafeWrite dst k (a + b)
- go (k + 1)
-
-{- | Walk the histogram in ascending value order, assigning a dense group id to
-each non-empty value and copying its count into @grpCount@ at that id. Returns
-the group count.
--}
-compact ::
- VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO Int
-compact hist range valToGroup grpCount = go 0 0
- where
- go !val !next
- | val >= range = pure next
- | otherwise = do
- c <- VUM.unsafeRead hist val
- if c == 0
- then go (val + 1) next
- else do
- VUM.unsafeWrite valToGroup val next
- VUM.unsafeWrite grpCount next c
- go (val + 1) (next + 1)
-
-{- | Exclusive prefix scan of group counts into @offsM@ (length nGroups+1) and
-seed the per-group write @cursor@ at each group's start offset.
--}
-scanOffsets ::
- VUM.IOVector Int -> Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO ()
-scanOffsets grpCount nGroups offsM cursor = go 0 0
- where
- go !g !acc
- | g >= nGroups = VUM.unsafeWrite offsM nGroups acc
- | otherwise = do
- VUM.unsafeWrite offsM g acc
- VUM.unsafeWrite cursor g acc
- c <- VUM.unsafeRead grpCount g
- go (g + 1) (acc + c)
-
-{- | Stable placement pass: for each row in original order, look up its group id
-through the value map, write @rowToGroup@, and append the row to its group's run
-in @valueIndices@ via the advancing cursor (rows keep original order per group).
--}
-place ::
- VU.Vector Int ->
- Int ->
- Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- IO ()
-place v mn n valToGroup cursor rtg vis = go 0
- where
- go !i
- | i >= n = pure ()
- | otherwise = do
- let !val = VU.unsafeIndex v i - mn
- g <- VUM.unsafeRead valToGroup val
- VUM.unsafeWrite rtg i g
- pos <- VUM.unsafeRead cursor g
- VUM.unsafeWrite vis pos i
- VUM.unsafeWrite cursor g (pos + 1)
- go (i + 1)
diff --git a/src/DataFrame/Internal/GroupingPar.hs b/src/DataFrame/Internal/GroupingPar.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/GroupingPar.hs
+++ /dev/null
@@ -1,355 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE Strict #-}
-
-{- |
-Parallel partitioned group-by assignment. Row indices are partitioned by the
-high bits of their key-hash (a counting sort: per-range histogram, prefix-sum,
-scatter into one index buffer laid out partition-by-partition). One task per
-capability then groups its partitions with its OWN open-addressing hash table
-('DataFrame.Internal.HashTable') — keys are disjoint across partitions, so the
-per-partition group sets concatenate with NO merge.
-
-The output @(rowToGroup, valueIndices, offsets)@ is /bit-for-bit identical/ to
-the sequential 'DataFrame.Internal.Grouping.groupBy': groups are emitted in
-ascending @(repHash, repRow)@ order (signed-Int order on the hash, tie-broken by
-the representative row). Because the partition key is the top bits of a
-sign-preserving unsigned remap of the same hash, partition order already agrees
-with that global order; within a partition we sort the local groups by the same
-key. This is the parallel==sequential correctness gate.
-
-The driver forks plain 'forkIO' workers (no sparks) over a shared atomic-counter
-work queue, so partition skew is balanced. A sequential fallback is used when
-there is a single capability or the row count is below 'parThreshold' (decided in
-'shouldParallelize', which 'groupBy' consults before calling here).
--}
-module DataFrame.Internal.GroupingPar (
- parallelAssignGroups,
- shouldParallelize,
- parThreshold,
- numPartitionsFor,
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import Control.Monad (forM_, when)
-import Data.Bits (countLeadingZeros, unsafeShiftR)
-import Data.IORef (atomicModifyIORef', newIORef)
-import qualified Data.Vector as V
-import qualified Data.Vector.Mutable as VM
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import Data.Word (Word64)
-import DataFrame.Internal.HashTable (
- htInsert,
- newHashTable,
- )
-import DataFrame.Internal.RadixRank (rankByHash)
-import System.IO.Unsafe (unsafePerformIO)
-
-{- | Below this many rows the partition/fork overhead is not worth it; 'groupBy'
-uses its sequential 'ST' path instead.
--}
-parThreshold :: Int
-parThreshold = 200000
-
-{- | Whether 'groupBy' should take the parallel path: more than one capability
-and at least 'parThreshold' rows.
--}
-shouldParallelize :: Int -> Bool
-shouldParallelize n = n >= parThreshold && capabilities > 1
-{-# NOINLINE shouldParallelize #-}
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
-{- | Sign-preserving unsigned remap: ascending 'Word64' order of @key h@ equals
-ascending signed-'Int' order of @h@, so partitioning and sorting on it reproduce
-the sequential @compare \`on\` repHash@ ordering exactly.
--}
-key :: Int -> Word64
-key h = fromIntegral h + 0x8000000000000000
-{-# INLINE key #-}
-
--- | Partition index of a hash: the top @log2 p@ bits of its unsigned key.
-partIx :: Int -> Int -> Int
-partIx shift h = fromIntegral (key h `unsafeShiftR` shift)
-{-# INLINE partIx #-}
-
-{- | Number of partitions: a power of two, at least @4 * caps@ (P >> cores for
-skew tolerance), floored at 256.
--}
-numPartitionsFor :: Int -> Int
-numPartitionsFor caps = go 1
- where
- target = max 256 (4 * caps)
- go p
- | p >= target = p
- | otherwise = go (p * 2)
-
--- | @floor (log2 x)@ for a power-of-two @x@.
-intLog2 :: Int -> Int
-intLog2 x = 63 - countLeadingZeros x
-{-# INLINE intLog2 #-}
-
-{- | Parallel group assignment. @parallelAssignGroups n hashes eqRow@ returns
-@(rowToGroup, valueIndices, offsets)@ in canonical group order. @eqRow a b@ must
-report whether rows @a@ and @b@ share all key columns (null-aware).
--}
-parallelAssignGroups ::
- Int ->
- VU.Vector Int ->
- (Int -> Int -> Bool) ->
- IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)
-parallelAssignGroups n hashes eqRow = do
- caps <- getNumCapabilities
- let !p = numPartitionsFor caps
- !shift = 64 - intLog2 p
- -- Phase 1: counting sort of row indices by partition.
- (partStart, sortedRows) <- partitionRows n hashes p shift
- -- Phase 2: per-partition grouping + per-partition canonical ranking (both
- -- inside the parallel worker). localGid[pos] = local group id of the row at
- -- sorted position 'pos'; canonBoxes[part] = its rank vector.
- localGid <- VUM.new (max 1 n)
- canonBoxes <- VM.replicate p (VU.empty :: VU.Vector Int)
- nLocalGroups <- VUM.replicate p (0 :: Int)
- runPartitions
- caps
- p
- partStart
- sortedRows
- hashes
- eqRow
- localGid
- canonBoxes
- nLocalGroups
- -- Phase 3: global base ids (serial prefix sum; the ranking is already done).
- (globalBase, canonOf, nGroups) <- canonicalize p canonBoxes nLocalGroups
- assemble n p partStart sortedRows localGid globalBase canonOf nGroups
-
--------------------------------------------------------------------------------
--- Phase 1: counting sort by partition
--------------------------------------------------------------------------------
-
-{- | Bucket every row index into its partition by a counting sort. Returns the
-exclusive prefix-sum @partStart@ (length @p+1@, @partStart[p] == n@) and the row
-indices laid out partition-by-partition in @sortedRows@.
--}
-partitionRows ::
- Int -> VU.Vector Int -> Int -> Int -> IO (VU.Vector Int, VU.Vector Int)
-partitionRows n hashes p shift = do
- counts <- VUM.replicate (p + 1) (0 :: Int)
- let countLoop !i
- | i >= n = pure ()
- | otherwise = do
- let !pp = partIx shift (VU.unsafeIndex hashes i)
- c <- VUM.unsafeRead counts pp
- VUM.unsafeWrite counts pp (c + 1)
- countLoop (i + 1)
- countLoop 0
- partStartM <- VUM.new (p + 1)
- let scan !k !acc
- | k > p = pure ()
- | otherwise = do
- VUM.unsafeWrite partStartM k acc
- c <- if k < p then VUM.unsafeRead counts k else pure 0
- scan (k + 1) (acc + c)
- scan 0 0
- cursor <- VUM.new p
- forM_ [0 .. p - 1] $ \k -> VUM.unsafeRead partStartM k >>= VUM.unsafeWrite cursor k
- sortedM <- VUM.new (max 1 n)
- let place !i
- | i >= n = pure ()
- | otherwise = do
- let !pp = partIx shift (VU.unsafeIndex hashes i)
- pos <- VUM.unsafeRead cursor pp
- VUM.unsafeWrite sortedM pos i
- VUM.unsafeWrite cursor pp (pos + 1)
- place (i + 1)
- place 0
- partStart <- VU.unsafeFreeze partStartM
- sortedRows <- VU.unsafeFreeze sortedM
- pure (partStart, sortedRows)
-
--------------------------------------------------------------------------------
--- Phase 2: per-partition grouping (parallel)
--------------------------------------------------------------------------------
-
-{- | Group each partition with its own hash table, then rank its local groups
-into canonical order — all inside the parallel worker. Forks @caps@ workers that
-pull partition indices off a shared counter. For partition @pp@ spanning
-@[partStart[pp], partStart[pp+1])@ of @sortedRows@ a worker assigns dense local
-group ids (first-appearance order) into @localGid@ at the same sorted positions,
-recording each new group's representative hash and row into unboxed
-first-appearance vectors. It then computes @canonBoxes[pp]@ — @canon[localGid] =
-within-partition canonical rank — via a stable radix sort on the unsigned
-representative hash (ties keep first-appearance order, which is ascending repRow,
-so the @(key hash, repRow)@ order of the old comparison sort is reproduced with
-no boxed tuples). @nLocalGroups[pp]@ holds the group count.
--}
-runPartitions ::
- Int ->
- Int ->
- VU.Vector Int ->
- VU.Vector Int ->
- VU.Vector Int ->
- (Int -> Int -> Bool) ->
- VUM.IOVector Int ->
- VM.IOVector (VU.Vector Int) ->
- VUM.IOVector Int ->
- IO ()
-runPartitions caps p partStart sortedRows hashes eqRow localGid canonBoxes nLocalGroups = do
- next <- newIORef 0
- let groupPartition !pp = do
- let !s = VU.unsafeIndex partStart pp
- !e = VU.unsafeIndex partStart (pp + 1)
- !sz = e - s
- when (sz > 0) $ do
- ht <- newHashTable sz
- -- repHash indexed by local gid (first-appearance order). The
- -- representative row is implicitly ascending in gid order (sorted
- -- positions are in original-row order, a stable counting sort),
- -- so the stable hash-rank below needs no explicit repRow.
- repHashM <- VUM.new sz
- let loop !pos !nextGid
- | pos >= e = pure nextGid
- | otherwise = do
- let !row = VU.unsafeIndex sortedRows pos
- !h = VU.unsafeIndex hashes row
- (gid, isNew) <- htInsert ht eqRow nextGid row h
- VUM.unsafeWrite localGid pos gid
- if isNew
- then do
- VUM.unsafeWrite repHashM nextGid h
- loop (pos + 1) (nextGid + 1)
- else loop (pos + 1) nextGid
- ng <- loop s 0
- VUM.unsafeWrite nLocalGroups pp ng
- -- Rank this partition's groups into canonical order (shared with
- -- the sequential path). repRow is ascending in gid order (stable
- -- counting sort keeps sorted positions in original-row order), so
- -- the stable hash-rank's tie-break reproduces (hash, repRow).
- canon <- rankByHash (VUM.unsafeRead repHashM) ng
- VM.unsafeWrite canonBoxes pp canon
- worker = do
- i <- atomicModifyIORef' next (\j -> (j + 1, j))
- when (i < p) $ groupPartition i >> worker
- forkJoin_ (replicate caps worker)
-
--------------------------------------------------------------------------------
--- Phase 3: global base ids + assembly
--------------------------------------------------------------------------------
-
-{- | Exclusive prefix sum of the per-partition group counts into @globalBase@
-(length @p+1@, @globalBase[pp]@ is the first global id of partition @pp@,
-@globalBase[p]@ the total). The per-partition canonical ranks were already
-computed in 'runPartitions'; partitions are in ascending key order so prepending
-@globalBase[pp]@ to each rank yields the sequential @canonicalRemap@ order.
--}
-canonicalize ::
- Int ->
- VM.IOVector (VU.Vector Int) ->
- VUM.IOVector Int ->
- IO (VU.Vector Int, V.Vector (VU.Vector Int), Int)
-canonicalize p canonBoxes nLocalGroups = do
- globalBaseM <- VUM.new (p + 1)
- let go !pp !base
- | pp >= p = VUM.unsafeWrite globalBaseM p base >> pure base
- | otherwise = do
- VUM.unsafeWrite globalBaseM pp base
- ng <- VUM.unsafeRead nLocalGroups pp
- go (pp + 1) (base + ng)
- total <- go 0 0
- globalBase <- VU.unsafeFreeze globalBaseM
- canonOf <- V.unsafeFreeze canonBoxes
- pure (globalBase, canonOf, total)
-
-{- | Build the final @(rowToGroup, valueIndices, offsets)@. For each sorted
-position we know its partition, its local group id and the canonical maps, so the
-global group id is @globalBase[pp] + canonOf[pp][localGid]@. @valueIndices@ is the
-rows ordered by global group; @offsets@ the per-group boundaries; @rowToGroup@ the
-inverse mapping per original row.
--}
-assemble ::
- Int ->
- Int ->
- VU.Vector Int ->
- VU.Vector Int ->
- VUM.IOVector Int ->
- VU.Vector Int ->
- V.Vector (VU.Vector Int) ->
- Int ->
- IO (VU.Vector Int, VU.Vector Int, VU.Vector Int)
-assemble n p partStart sortedRows localGid globalBase canonOf nGroups = do
- rtgM <- VUM.new (max 1 n)
- -- Global group id of each sorted position, plus per-group counts.
- counts <- VUM.replicate (nGroups + 1) (0 :: Int)
- gidAt <- VUM.new (max 1 n)
- let scanPos !pp
- | pp >= p = pure ()
- | otherwise = do
- let !s = VU.unsafeIndex partStart pp
- !e = VU.unsafeIndex partStart (pp + 1)
- !base = VU.unsafeIndex globalBase pp
- !canon = V.unsafeIndex canonOf pp
- let inner !pos
- | pos >= e = pure ()
- | otherwise = do
- lg <- VUM.unsafeRead localGid pos
- let !g = base + VU.unsafeIndex canon lg
- !row = VU.unsafeIndex sortedRows pos
- VUM.unsafeWrite gidAt pos g
- VUM.unsafeWrite rtgM row g
- c <- VUM.unsafeRead counts g
- VUM.unsafeWrite counts g (c + 1)
- inner (pos + 1)
- inner s
- scanPos (pp + 1)
- scanPos 0
- -- offsets = exclusive prefix sum of counts.
- offsM <- VUM.new (nGroups + 1)
- let scan !k !acc
- | k > nGroups = pure ()
- | otherwise = do
- VUM.unsafeWrite offsM k acc
- c <- if k < nGroups then VUM.unsafeRead counts k else pure 0
- scan (k + 1) (acc + c)
- scan 0 0
- -- valueIndices: place each sorted position's row at its group's cursor.
- -- Iterating sorted positions in order keeps rows in original order within a
- -- group (the partition counting sort and grouping both preserve it).
- cursor <- VUM.new (max 1 nGroups)
- forM_ [0 .. nGroups - 1] $ \k -> VUM.unsafeRead offsM k >>= VUM.unsafeWrite cursor k
- visM <- VUM.new (max 1 n)
- let placeVis !pos
- | pos >= n = pure ()
- | otherwise = do
- g <- VUM.unsafeRead gidAt pos
- let !row = VU.unsafeIndex sortedRows pos
- c <- VUM.unsafeRead cursor g
- VUM.unsafeWrite visM c row
- VUM.unsafeWrite cursor g (c + 1)
- placeVis (pos + 1)
- placeVis 0
- rtg <- VU.unsafeFreeze rtgM
- offs <- VU.unsafeFreeze offsM
- vis <- VU.unsafeFreeze visM
- pure (rtg, vis, offs)
-
--------------------------------------------------------------------------------
--- Thread fan-out (plain forkIO + MVar join, no sparks)
--------------------------------------------------------------------------------
-
--- | Run each action on its own thread; rethrow the first failure (in order).
-forkJoin_ :: [IO ()] -> IO ()
-forkJoin_ actions = do
- vars <- mapM spawn actions
- results <- mapM takeMVar vars
- mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
- where
- spawn act = do
- var <- newEmptyMVar
- _ <- forkIO (try act >>= putMVar var)
- pure var
diff --git a/src/DataFrame/Internal/Hash.hs b/src/DataFrame/Internal/Hash.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Hash.hs
+++ /dev/null
@@ -1,128 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE CPP #-}
-{-# LANGUAGE MagicHash #-}
-
-{- |
-A poor-man's hash used by 'DataFrame.Internal.Grouping' to bucket rows
-without depending on the @hashable@ package.
-
-Each value is folded into an 'Int' accumulator with an FxHash-style step
-(rotate, xor, multiply). It is intentionally small and not cryptographically
-strong — it only needs to spread group-key tuples well enough that
-'Data.IntMap' bucketing produces sensible groups.
--}
-module DataFrame.Internal.Hash (
- fnvOffset,
- nullSalt,
- mixInt,
- mixDouble,
- mixBool,
- mixChar,
- mixText,
- mixBytes,
- mixShow,
-) where
-
-import Data.Bits (rotateL, unsafeShiftL, unsafeShiftR, xor)
-import Data.Char (ord)
-import qualified Data.Text as T
-import qualified Data.Text.Array as A
-#if MIN_VERSION_text(2,1,0)
-import Data.Array.Byte (ByteArray (ByteArray))
-#else
-import Data.Text.Array (Array (ByteArray))
-#endif
-import Data.Text.Internal (Text (Text))
-import GHC.Exts (Int (I#), indexWord8Array#, indexWord8ArrayAsWord64#)
-import GHC.Word (Word64 (W64#), Word8 (W8#))
-
-{- | FNV-1a 64-bit offset basis (used as the initial accumulator).
-The literal is unsigned and exceeds 'Int' range, so we round-trip through
-'Word64' to get the well-defined two's-complement bit pattern.
--}
-fnvOffset :: Int
-fnvOffset = fromIntegral (0xcbf29ce484222325 :: Word64)
-
--- | FNV-1a 64-bit prime.
-fnvPrime :: Int
-fnvPrime = 0x00000100000001b3
-
-{- | Sentinel mixed in for a /null/ slot of a nullable column, so that a
-@Nothing@ does not hash to the same value as a present @Just x@ that happens to
-store the same underlying bits (notably @Just 0@). A fixed distinctive constant
-(the 64-bit golden-ratio mix constant) keeps null hashing deterministic; a real
-value equal to it collides only as rarely as any other hash collision.
--}
-nullSalt :: Int
-nullSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word64)
-
-{- | Mix an 'Int' into the accumulator.
-
-An FxHash-style step (rotate the accumulator, xor the value, multiply by a large
-odd constant). The rotate diffuses each value's bits across all positions before
-the next is folded in, so small/adjacent integers — common as group keys — do
-not produce the structured collisions that a plain @(acc `xor` x) * prime@ does
-once several columns are combined. Grouping trusts hash equality, so this
-robustness is what keeps distinct rows in distinct groups.
--}
-mixInt :: Int -> Int -> Int
-mixInt acc x = (rotateL acc 13 `xor` x) * fnvPrime
-{-# INLINE mixInt #-}
-
-{- | Mix a 'Double' into the accumulator. Loses sub-millisecond precision
-but matches the bucketing the old hashable-based code used.
--}
-mixDouble :: Int -> Double -> Int
-mixDouble acc d = mixInt acc (floor (d * 1000))
-{-# INLINE mixDouble #-}
-
-mixBool :: Int -> Bool -> Int
-mixBool acc b = mixInt acc (if b then 1 else 0)
-{-# INLINE mixBool #-}
-
-mixChar :: Int -> Char -> Int
-mixChar acc = mixInt acc . ord
-{-# INLINE mixChar #-}
-
-{- | Mix a 'T.Text' value into the accumulator over its raw UTF-8 bytes,
-eight at a time. Reading a whole 'Word64' per step (rather than decoding and
-mixing one codepoint at a time) cuts the multiply count ~8x on long keys while
-staying collision-equivalent: UTF-8 is injective, so equal 'T.Text's mix to the
-same value and distinct ones almost never collide. The trailing @len `mod` 8@
-bytes are folded in individually.
--}
-mixText :: Int -> T.Text -> Int
-mixText !acc (Text arr off len) = mixBytes acc arr off len
-{-# INLINE mixText #-}
-
-{- | Mix a raw UTF-8 byte slice @[off, off+len)@ of a 'Data.Text.Array.Array'
-into the accumulator, eight bytes at a time. The shared kernel behind
-'mixText' and the packed-text hash path, so the two never drift.
--}
-mixBytes :: Int -> A.Array -> Int -> Int -> Int
-mixBytes !acc arr off len = goBytes (goWords acc off) wordsEnd
- where
- !(ByteArray ba) = arr
- !nWords = len `unsafeShiftR` 3
- !wordsEnd = off + (nWords `unsafeShiftL` 3)
- !end = off + len
- goWords !h !i
- | i >= wordsEnd = h
- | otherwise =
- let !(I# i#) = i
- !w = fromIntegral (W64# (indexWord8ArrayAsWord64# ba i#)) :: Int
- in goWords (mixInt h w) (i + 8)
- goBytes !h !i
- | i >= end = h
- | otherwise =
- let !(I# i#) = i
- !b = fromIntegral (W8# (indexWord8Array# ba i#)) :: Int
- in goBytes (mixInt h b) (i + 1)
-{-# INLINE mixBytes #-}
-
-{- | Fallback for arbitrary 'Show'-able values. Slower but covers types
-without a dedicated combinator (e.g. 'Day', 'UTCTime').
--}
-mixShow :: (Show a) => Int -> a -> Int
-mixShow acc = mixText acc . T.pack . show
-{-# INLINE mixShow #-}
diff --git a/src/DataFrame/Internal/HashTable.hs b/src/DataFrame/Internal/HashTable.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/HashTable.hs
+++ /dev/null
@@ -1,113 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-
-{- |
-A flat, unboxed, open-addressing (linear-probe) hash table that maps a row's
-key-hash to a /dense group id/, verifying the real key on every hash hit.
-
-The table is three parallel unboxed 'VUM.MVector's keyed by hash slot:
-
- * @htHash@ — the stored hash at each slot.
- * @htGroup@ — the dense group id stored at each slot (@-1@ marks an empty
- slot, since real group ids are @>= 0@).
- * @htRep@ — the representative row index of that group, used to re-verify
- the real key columns on a hash hit and so reject collisions.
-
-It is 'PrimMonad'-polymorphic: it runs in 'Control.Monad.ST.ST' for the current
-single-threaded 'DataFrame.Internal.Grouping.groupBy' and can run in 'IO' inside
-a per-worker partition once grouping is parallelised. The lookup-or-insert loop
-('htInsert') trusts the caller-supplied @eqRow@ predicate to compare the key
-columns of two rows by index, fixing the hash-only bucketing that previously
-merged colliding keys.
--}
-module DataFrame.Internal.HashTable (
- HashTable (..),
- newHashTable,
- htInsert,
- nextPow2Above,
-) where
-
-import Control.Monad.Primitive (PrimMonad, PrimState)
-import Data.Bits ((.&.))
-import qualified Data.Vector.Unboxed.Mutable as VUM
-
-{- | An open-addressing linear-probe table. @htMask@ is @capacity - 1@ (capacity
-is a power of two) and maps a hash to its home slot.
--}
-data HashTable s = HashTable
- { htHash :: !(VUM.MVector s Int)
- , htGroup :: !(VUM.MVector s Int)
- , htRep :: !(VUM.MVector s Int)
- , htMask :: !Int
- }
-
-{- | Smallest power of two strictly greater than @n@, at least 2. Sizes the
-table so the load factor stays below ~0.5 even when every row is a distinct
-group.
--}
-nextPow2Above :: Int -> Int
-nextPow2Above n = go 2
- where
- go !p
- | p > n = p
- | otherwise = go (p * 2)
-{-# INLINE nextPow2Above #-}
-
-{- | Allocate an empty table able to hold up to @n@ distinct groups while
-keeping the load factor under ~0.5 (capacity @= nextPow2Above (2*n)@). All
-group slots start empty (@-1@).
--}
-newHashTable :: (PrimMonad m) => Int -> m (HashTable (PrimState m))
-newHashTable n = do
- let !cap = nextPow2Above (2 * max 1 n)
- h <- VUM.unsafeNew cap
- g <- VUM.replicate cap (-1)
- r <- VUM.unsafeNew cap
- pure (HashTable h g r (cap - 1))
-{-# INLINE newHashTable #-}
-
-{- | Look up @row@ (with precomputed @hash@) in the table, returning its dense
-group id. On an empty slot the row starts a new group: the caller's
-@nextGroup@ thunk supplies the next dense id, and the row is recorded as that
-group's representative. On a stored-hash match the real key is re-verified with
-@eqRow rep row@ before the existing id is returned; a mismatch is a hash
-collision and probing continues. The returned 'Bool' is 'True' when a new group
-was created, letting the caller bump its group counter without a second read.
--}
-htInsert ::
- (PrimMonad m) =>
- HashTable (PrimState m) ->
- -- | @eqRow a b@: do rows @a@ and @b@ have equal key columns?
- (Int -> Int -> Bool) ->
- -- | Next dense group id to assign if this row starts a new group.
- Int ->
- -- | Row index being inserted.
- Int ->
- -- | Precomputed hash of the row's key.
- Int ->
- m (Int, Bool)
-htInsert ht eqRow nextGroup row hash = go (hash .&. mask)
- where
- !mask = htMask ht
- !hs = htHash ht
- !gs = htGroup ht
- !rs = htRep ht
- go !slot = do
- g <- VUM.unsafeRead gs slot
- if g < 0
- then do
- VUM.unsafeWrite hs slot hash
- VUM.unsafeWrite gs slot nextGroup
- VUM.unsafeWrite rs slot row
- pure (nextGroup, True)
- else do
- h <- VUM.unsafeRead hs slot
- if h == hash
- then do
- rep <- VUM.unsafeRead rs slot
- if eqRow rep row
- then pure (g, False)
- else go ((slot + 1) .&. mask)
- else go ((slot + 1) .&. mask)
-{-# INLINE htInsert #-}
diff --git a/src/DataFrame/Internal/Interpreter.hs b/src/DataFrame/Internal/Interpreter.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Interpreter.hs
+++ /dev/null
@@ -1,1103 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE UndecidableInstances #-}
-{-# OPTIONS_GHC -Wno-orphans #-}
-
-module DataFrame.Internal.Interpreter (
- -- * New core API
- Value (..),
- Ctx (..),
- eval,
- materialize,
-
- -- * Backward-compatible API
- interpret,
- interpretAggregation,
- AggregationResult (..),
-) where
-
-import Data.Bifunctor (first)
-import qualified Data.Map as M
-import qualified Data.Text as T
-import Data.Type.Equality (TestEquality (testEquality), type (:~:) (Refl))
-import qualified Data.Vector as V
-import qualified Data.Vector.Generic as VG
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import DataFrame.Errors
-import DataFrame.Internal.Column
-import DataFrame.Internal.DataFrame
-import DataFrame.Internal.Expression
-import qualified DataFrame.Internal.Grouping as G
-import DataFrame.Internal.Types
-import Type.Reflection (
- Typeable,
- typeRep,
- )
-
-import Data.Int (Int16, Int32, Int64, Int8)
-
--- Specializations for common aggregation types to avoid dictionary overhead.
--- foldLinearGroups: mean accumulator
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Double -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Float -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Int -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Int8 -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Int16 -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Int32 -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (MeanAcc -> Int64 -> MeanAcc) ->
- MeanAcc ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
--- foldLinearGroups: count accumulator
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Double -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Float -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Int -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Int8 -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Int16 -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Int32 -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int -> Int64 -> Int) ->
- Int ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
--- foldLinearGroups: sum/min/max (acc == elem)
-{-# SPECIALIZE foldLinearGroups ::
- (Double -> Double -> Double) ->
- Double ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Float -> Float -> Float) ->
- Float ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int8 -> Int8 -> Int8) ->
- Int8 ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int16 -> Int16 -> Int16) ->
- Int16 ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int32 -> Int32 -> Int32) ->
- Int32 ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE foldLinearGroups ::
- (Int64 -> Int64 -> Int64) ->
- Int64 ->
- Column ->
- VU.Vector Int ->
- Int ->
- Either DataFrameException Column
- #-}
-
--- mapColumn: finalize
-{-# SPECIALIZE mapColumn ::
- (MeanAcc -> Double) -> Column -> Either DataFrameException Column
- #-}
-{-# SPECIALIZE mapColumn ::
- (Double -> Double) -> Column -> Either DataFrameException Column
- #-}
-{-# SPECIALIZE mapColumn ::
- (Float -> Float) -> Column -> Either DataFrameException Column
- #-}
-{-# SPECIALIZE mapColumn ::
- (Int -> Int) -> Column -> Either DataFrameException Column
- #-}
-
--- zipWithColumns: binary ops
-{-# SPECIALIZE zipWithColumns ::
- (Double -> Double -> Double) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Float -> Float -> Float) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int -> Int -> Int) -> Column -> Column -> Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int8 -> Int8 -> Int8) -> Column -> Column -> Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int16 -> Int16 -> Int16) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int32 -> Int32 -> Int32) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int64 -> Int64 -> Int64) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
--- Bool-returning binary comparators (hot path for Expr Bool used in
--- DecisionTree splits)
-{-# SPECIALIZE zipWithColumns ::
- (Double -> Double -> Bool) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Float -> Float -> Bool) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Int -> Int -> Bool) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-{-# SPECIALIZE zipWithColumns ::
- (Bool -> Bool -> Bool) ->
- Column ->
- Column ->
- Either DataFrameException Column
- #-}
-
--- Bool-mapping unary ops (e.g. 'not')
-{-# SPECIALIZE mapColumn ::
- (Bool -> Bool) -> Column -> Either DataFrameException Column
- #-}
-
--------------------------------------------------------------------------------
--- Value: the unified result type
--------------------------------------------------------------------------------
-
-{- | The result of interpreting an expression. Keeps literals as scalars
-until the point where a concrete column is needed, avoiding premature
-broadcast allocations.
--}
-data Value a where
- -- | A single value, not yet broadcast to any length.
- Scalar :: (Columnable a) => !a -> Value a
- {- | A flat column (one element per row in the flat case, or one
- element per group after aggregation).
- -}
- Flat :: (Columnable a) => !Column -> Value a
- {- | A grouped column: one 'Column' slice per group. Only produced
- when interpreting inside a 'GroupCtx'.
- -}
- Group :: (Columnable a) => !(V.Vector Column) -> Value a
-
-instance (Show a) => Show (Value a) where
- show (Scalar v) = show v
- show (Flat v) = show v
- show (Group v) = show v
-
--- | The interpretation context.
-data Ctx
- = FlatCtx DataFrame
- | GroupCtx GroupedDataFrame
-
--------------------------------------------------------------------------------
--- Materialisation
--------------------------------------------------------------------------------
-
-{- | Force a 'Value' into a flat 'Column' of the given length. Scalars
-are broadcast; flat columns are returned as-is.
--}
-materialize :: forall a. (Columnable a) => Int -> Value a -> Column
-materialize n (Scalar v) = broadcastScalar @a n v
-materialize _ (Flat c) = c
-materialize _ (Group _) =
- error "materialize: cannot flatten a grouped value to a single column"
-
-{- | Replicate a scalar to a column of length @n@, choosing the most
-efficient representation.
--}
-broadcastScalar :: forall a. (Columnable a) => Int -> a -> Column
-broadcastScalar n v = case sUnbox @a of
- STrue -> fromUnboxedVector (VU.replicate n v)
- SFalse -> fromVector (V.replicate n v)
-
--------------------------------------------------------------------------------
--- Lifting: the core combinators
--------------------------------------------------------------------------------
-
--- | Apply a pure function to a 'Value'.
-liftValue ::
- (Columnable b, Columnable a) =>
- (b -> a) -> Value b -> Either DataFrameException (Value a)
-liftValue f (Scalar v) = Right (Scalar (f v))
-liftValue f (Flat col) = Flat <$> mapColumn f col
-liftValue f (Group gs) = Group <$> V.mapM (mapColumn f) gs
-{-# INLINEABLE liftValue #-}
-
-{- | Apply a binary function to two 'Value's. When one side is a
-'Scalar' the operation degenerates to a 'liftValue' — this is how the
-old @Binary op (Lit l) right@ special cases are recovered without
-explicit pattern matches in the evaluator.
--}
-liftValue2 ::
- (Columnable c, Columnable b, Columnable a) =>
- (c -> b -> a) ->
- Value c ->
- Value b ->
- Either DataFrameException (Value a)
-liftValue2 f (Scalar l) (Scalar r) = Right (Scalar (f l r))
-liftValue2 f (Scalar l) v = liftValue (f l) v
-liftValue2 f v (Scalar r) = liftValue (`f` r) v
-liftValue2 f (Flat l) (Flat r) = Flat <$> zipWithColumns f l r
-liftValue2 f (Group ls) (Group rs)
- | V.length ls == V.length rs =
- Group <$> V.zipWithM (zipWithColumns f) ls rs
--- Shape mismatches: aggregated vs. non-aggregated.
-liftValue2 _ (Flat _) (Group _) =
- Left $ AggregatedAndNonAggregatedException "aggregated" "non-aggregated"
-liftValue2 _ (Group _) (Flat _) =
- Left $ AggregatedAndNonAggregatedException "non-aggregated" "aggregated"
-liftValue2 _ (Group _) (Group _) =
- Left $ InternalException "Group count mismatch in binary operation"
-{-# INLINEABLE liftValue2 #-}
-
--- | Branch on a boolean 'Value', selecting from two same-typed 'Value's.
-branchValue ::
- forall a.
- (Columnable a) =>
- Value Bool ->
- Value a ->
- Value a ->
- Either DataFrameException (Value a)
-branchValue (Scalar True) l _ = Right l
-branchValue (Scalar False) _ r = Right r
-branchValue cond (Scalar l) (Scalar r) =
- liftValue (\c -> if c then l else r) cond
-branchValue cond (Scalar l) r =
- liftValue2 (\c rv -> if c then l else rv) cond r
-branchValue cond l (Scalar r) =
- liftValue2 (\c lv -> if c then lv else r) cond l
-branchValue (Flat cc) (Flat lc) (Flat rc) =
- Flat <$> branchColumn @a cc lc rc
-branchValue (Group cgs) (Group lgs) (Group rgs)
- | V.length cgs == V.length lgs
- && V.length lgs == V.length rgs =
- Group
- <$> V.generateM
- (V.length cgs)
- ( \i ->
- branchColumn @a (cgs V.! i) (lgs V.! i) (rgs V.! i)
- )
-branchValue _ _ _ =
- Left $
- AggregatedAndNonAggregatedException
- "if-then-else branches"
- "mismatched shapes"
-{-# INLINEABLE branchValue #-}
-
-{- | Low-level column branch: given a boolean column and two same-typed
-columns, produce the element-wise selection.
--}
-branchColumn ::
- forall a.
- (Columnable a) =>
- Column ->
- Column ->
- Column ->
- Either DataFrameException Column
-branchColumn cc lc rc = do
- cs <- toVector @Bool @V.Vector cc
- ls <- toVector @a @V.Vector lc
- rs <- toVector @a @V.Vector rc
- pure $
- fromVector @a $
- V.zipWith3 (\c l r -> if c then l else r) cs ls rs
-
--------------------------------------------------------------------------------
--- Error enrichment
--------------------------------------------------------------------------------
-
-{- | Wrap an interpretation step so that any 'TypeMismatchException' gets
-annotated with the expression that was being evaluated.
--}
-addContext ::
- (Show a) => Expr a -> Either DataFrameException b -> Either DataFrameException b
-addContext expr = first (enrichError (show expr))
-
-enrichError :: String -> DataFrameException -> DataFrameException
-enrichError loc (TypeMismatchException ctx) =
- TypeMismatchException
- ctx
- { callingFunctionName =
- callingFunctionName ctx <|+> Just "eval"
- , errorColumnName =
- errorColumnName ctx <|+> Just loc
- }
- where
- -- Prefer the existing value; fall back to the new one.
- Nothing <|+> b = b
- a <|+> _ = a
-enrichError _ e = e
-
--------------------------------------------------------------------------------
--- Group slicing
--------------------------------------------------------------------------------
-
-{- | Given a flat column and grouping metadata, produce one 'Column' per
-group. Each result column is an O(1) slice into a sorted copy of the
-input — the sort happens once, not per-group.
--}
-sliceGroups :: Column -> VU.Vector Int -> VU.Vector Int -> V.Vector Column
-sliceGroups col os indices = case col of
- PackedText _ _ -> sliceGroups (materializePacked col) os indices
- BoxedColumn bm vec ->
- let !sorted =
- V.generate
- (VU.length indices)
- ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`))
- in V.generate nGroups $ \i ->
- BoxedColumn
- (fmap (bitmapSlice (start i) (len i)) bm)
- (V.unsafeSlice (start i) (len i) sorted)
- UnboxedColumn bm vec ->
- let !sorted = VU.unsafeBackpermute vec indices
- in V.generate nGroups $ \i ->
- UnboxedColumn
- (fmap (bitmapSlice (start i) (len i)) bm)
- (VU.unsafeSlice (start i) (len i) sorted)
- where
- !nGroups = VU.length os - 1
- start i = os `VU.unsafeIndex` i
- len i = os `VU.unsafeIndex` (i + 1) - start i
-{-# INLINE sliceGroups #-}
-
-numGroups :: GroupedDataFrame -> Int
-numGroups gdf = VU.length (offsets gdf) - 1
-
--- | Build the inverse of a permutation vector.
-invertPermutation :: VU.Vector Int -> VU.Vector Int
-invertPermutation perm = VU.create $ do
- let !n = VU.length perm
- inv <- VUM.new n
- VU.imapM_ (flip (VUM.unsafeWrite inv)) perm
- return inv
-{-# INLINE invertPermutation #-}
-
--------------------------------------------------------------------------------
--- 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, Read a) =>
- (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 Nothing (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
- UnboxedColumn (Just bm) (v :: VU.Vector c) ->
- case sFloating @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Double))
- else onResult (Left "null")
- )
- SFalse -> case sIntegral @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Double))
- else onResult (Left "null")
- )
- SFalse -> castMismatch @c @b
- BoxedColumn _ _ -> tryParseWith @Double onResult col
- PackedText _ _ -> promoteToDoubleWith onResult (materializePacked col)
-
-promoteToFloatWith ::
- forall b.
- (Columnable b) =>
- (Either String Float -> b) -> Column -> Either DataFrameException Column
-promoteToFloatWith onResult col = case col of
- UnboxedColumn Nothing (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
- UnboxedColumn (Just bm) (v :: VU.Vector c) ->
- case sFloating @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (realToFrac (VU.unsafeIndex v i) :: Float))
- else onResult (Left "null")
- )
- SFalse -> case sIntegral @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Float))
- else onResult (Left "null")
- )
- SFalse -> castMismatch @c @b
- BoxedColumn _ _ -> tryParseWith @Float onResult col
- PackedText _ _ -> promoteToFloatWith onResult (materializePacked col)
-
-promoteToIntWith ::
- forall b.
- (Columnable b) =>
- (Either String Int -> b) -> Column -> Either DataFrameException Column
-promoteToIntWith onResult col = case col of
- UnboxedColumn Nothing (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
- UnboxedColumn (Just bm) (v :: VU.Vector c) ->
- case sFloating @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (round (realToFrac (VU.unsafeIndex v i) :: Double)))
- else onResult (Left "null")
- )
- SFalse -> case sIntegral @c of
- STrue ->
- Right $
- fromVector @b
- ( V.generate (VU.length v) $ \i ->
- if bitmapTestBit bm i
- then onResult (Right (fromIntegral (VU.unsafeIndex v i) :: Int))
- else onResult (Left "null")
- )
- SFalse -> castMismatch @c @b
- BoxedColumn _ _ -> tryParseWith @Int onResult col
- PackedText _ _ -> promoteToIntWith onResult (materializePacked 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)
- _ -> case reads (show s) of
- [(x, "")] -> f (Right x)
- _ -> f (Left s)
-
-tryParseWith ::
- forall a b.
- (Columnable a, Columnable b, Read a) =>
- (Either String a -> b) -> Column -> Either DataFrameException Column
-tryParseWith onResult col = case col of
- PackedText _ _ -> tryParseWith onResult (materializePacked col)
- BoxedColumn bm (v :: V.Vector c) ->
- case testEquality (typeRep @c) (typeRep @String) of
- Just Refl -> case bm of
- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult) v
- Just bitmap ->
- Right $
- fromVector @b $
- V.imap
- ( \i x ->
- if bitmapTestBit bitmap i then parseWith onResult x else onResult (Left "null")
- )
- v
- Nothing ->
- case testEquality (typeRep @c) (typeRep @T.Text) of
- Just Refl -> case bm of
- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . T.unpack) v
- Just bitmap ->
- Right $
- fromVector @b $
- V.imap
- ( \i x ->
- if bitmapTestBit bitmap i
- then parseWith onResult (T.unpack x)
- else onResult (Left "null")
- )
- v
- Nothing -> castMismatch @c @b
- UnboxedColumn bm (v :: VU.Vector c) -> case bm of
- Nothing -> Right $ fromVector @b $ V.map (parseWith onResult . show) (V.convert v)
- Just bitmap ->
- Right $
- fromVector @b $
- V.imap
- ( \i x ->
- if bitmapTestBit bitmap i
- then parseWith onResult (show x)
- else onResult (Left "null")
- )
- (V.convert v)
-
-{- | 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 Nothing (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 ->
- case testEquality (typeRep @b) (typeRep @(Maybe (Maybe c))) of
- Just _ -> Just $ Right $ fromVector @(Maybe c) wrapped
- Nothing -> Nothing
- BoxedColumn Nothing (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
- _ -> 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
--------------------------------------------------------------------------------
-
-{- | Evaluate an expression in a given context, producing a 'Value'.
-This single function replaces both the old @interpret@ (flat) and
-@interpretAggregation@ (grouped) code paths.
--}
-eval ::
- forall a.
- (Columnable a) =>
- Ctx -> Expr a -> Either DataFrameException (Value a)
--- Leaves -----------------------------------------------------------------
-
-eval _ (Lit v) = Right (Scalar v)
-eval (FlatCtx df) (Col name) =
- case getColumn name df of
- Nothing ->
- Left $ ColumnsNotFoundException [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 $
- ColumnsNotFoundException
- [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 $
- ColumnsNotFoundException [name] "" (M.keys $ columnIndices df)
- Just c -> Flat <$> promoteColumnWith onResult c
-eval (GroupCtx gdf) (CastWith name _tag onResult) =
- case getColumn name (fullDataframe gdf) of
- Nothing ->
- Left $
- ColumnsNotFoundException
- [name]
- ""
- (M.keys $ columnIndices $ fullDataframe 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 (inner :: Expr b)) = addContext expr $ do
- v <- eval @b ctx inner
- liftValue (unaryFn op) v
-
--- Binary -----------------------------------------------------------------
-
-eval ctx expr@(Binary op (left :: Expr c) (right :: Expr b)) =
- addContext expr $ do
- l <- eval @c ctx left
- r <- eval @b ctx right
- liftValue2 (binaryFn op) l r
-
--- If ---------------------------------------------------------------------
-
-eval ctx expr@(If cond l r) = addContext expr $ do
- c <- eval @Bool ctx cond
- lv <- eval @a ctx l
- rv <- eval @a ctx r
- branchValue c lv rv
-
--- Over (window function) -------------------------------------------------
-
-eval (FlatCtx df) expr@(Over keys inner) = addContext expr $ do
- let gdf = G.groupBy keys df
- v <- eval (GroupCtx gdf) inner
- case v of
- Scalar s ->
- Right (Scalar s)
- Flat groupCol ->
- -- Scalar agg (mean, sum, median): one value per group.
- -- Broadcast via rowToGroup: row i gets value at group rowToGroup[i].
- Right (Flat (atIndicesStable (rowToGroup gdf) groupCol))
- Group groupCols -> do
- -- Concatenate in sorted order, then unsort to original row order.
- sorted <- V.fold1M' concatColumns groupCols
- let inv = invertPermutation (valueIndices gdf)
- Right (Flat (atIndicesStable inv sorted))
-eval (GroupCtx _) expr@(Over _ _) =
- addContext expr $
- Left
- ( InternalException
- "Over (window function) is not supported inside a grouped context"
- )
--- Fast path: FoldAgg (seeded) on a bare Col in GroupCtx.
--- Avoids the O(n) backpermute in sliceGroups by folding directly over
--- permuted indices. Only matches when inner is exactly (Col name).
-
-eval (GroupCtx gdf) expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) (Col name :: Expr b)) =
- addContext expr $
- case getColumn name (fullDataframe gdf) of
- Nothing ->
- Left $
- ColumnsNotFoundException
- [name]
- ""
- (M.keys $ columnIndices $ fullDataframe gdf)
- Just col ->
- Flat <$> foldLinearGroups @b @a f seed col (rowToGroup gdf) (numGroups gdf)
--- Fast path: FoldAgg (seedless) on a bare Col in GroupCtx.
-
-eval (GroupCtx gdf) expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) (Col name :: Expr b)) =
- addContext expr $
- case testEquality (typeRep @a) (typeRep @b) of
- Nothing ->
- Left $
- InternalException
- "Type mismatch in seedless fold: \
- \accumulator and element types must match"
- Just Refl ->
- case getColumn name (fullDataframe gdf) of
- Nothing ->
- Left $
- ColumnsNotFoundException
- [name]
- ""
- (M.keys $ columnIndices $ fullDataframe gdf)
- Just col ->
- Flat <$> foldl1DirectGroups @b f col (valueIndices gdf) (offsets gdf)
--- Fast path: MergeAgg on a bare Col in GroupCtx.
-
-eval
- (GroupCtx gdf)
- expr@( Agg
- (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))
- (Col name :: Expr b)
- ) =
- addContext expr $
- case getColumn name (fullDataframe gdf) of
- Nothing ->
- Left $
- ColumnsNotFoundException
- [name]
- ""
- (M.keys $ columnIndices $ fullDataframe gdf)
- Just col ->
- Flat
- <$> ( foldLinearGroups @b step seed col (rowToGroup gdf) (numGroups gdf)
- >>= mapColumn finalize
- )
--- Aggregation: CollectAgg ------------------------------------------------
-
-eval ctx expr@(Agg (CollectAgg _ (f :: v b -> a)) inner) =
- addContext expr $ do
- v <- eval @b ctx inner
- case v of
- Scalar _ ->
- Left $
- InternalException
- "Cannot apply a collection aggregation to a scalar"
- Flat col ->
- Scalar <$> applyCollect @v @b @a f col
- Group gs ->
- Flat . fromVector
- <$> V.mapM (applyCollect @v @b @a f) gs
-
--- Aggregation: FoldAgg with seed -----------------------------------------
-
-eval ctx expr@(Agg (FoldAgg _ (Just seed) (f :: a -> b -> a)) inner) =
- addContext expr $ do
- v <- eval @b ctx inner
- case v of
- Scalar x -> Right (broadcastFold ctx seed f x)
- Flat col ->
- Scalar <$> foldlColumn @b @a f seed col
- Group gs ->
- Flat . fromVector
- <$> V.mapM (foldlColumn @b @a f seed) gs
-
--- Aggregation: MergeAgg --------------------------------------------------
-
-eval
- ctx
- expr@( Agg
- (MergeAgg _ seed (step :: acc -> b -> acc) _ (finalize :: acc -> a))
- (inner :: Expr b)
- ) =
- addContext expr $ do
- v <- eval @b ctx inner
- case v of
- Scalar x -> case broadcastFold ctx seed step x of
- Scalar acc -> Right (Scalar (finalize acc))
- Flat col -> Flat <$> mapColumn @acc @a finalize col
- Group _ ->
- Left
- ( InternalException
- "broadcastFold unexpectedly produced a Group value"
- )
- Flat col ->
- Scalar . finalize <$> foldlColumn @b step seed col
- Group gs ->
- Flat . fromVector
- <$> V.mapM (fmap finalize . foldlColumn @b step seed) gs
-
--- Aggregation: FoldAgg without seed (fold1) ------------------------------
-
-eval ctx expr@(Agg (FoldAgg _ Nothing (f :: a -> b -> a)) inner) =
- addContext expr $
- case testEquality (typeRep @a) (typeRep @b) of
- Nothing ->
- Left $
- InternalException
- "Type mismatch in seedless fold: \
- \accumulator and element types must match"
- Just Refl -> do
- v <- eval @b ctx inner
- case v of
- Scalar _ ->
- Left $
- InternalException
- "fold1 requires at least one element"
- Flat col ->
- Scalar <$> foldl1Column @a f col
- Group gs ->
- Flat . fromVector
- <$> V.mapM (foldl1Column @a f) gs
-
-broadcastFold ::
- forall acc b.
- (Columnable acc) =>
- Ctx -> acc -> (acc -> b -> acc) -> b -> Value acc
-broadcastFold (FlatCtx df) seed step x =
- let n = fst (dataframeDimensions df)
- in Scalar (iterateStep n step seed x)
-broadcastFold (GroupCtx gdf) seed step x =
- let offs = offsets gdf
- ng = VU.length offs - 1
- results =
- V.generate ng $ \i ->
- let sz = offs VU.! (i + 1) - offs VU.! i
- in iterateStep sz step seed x
- in Flat (fromVector results)
-
-iterateStep :: Int -> (acc -> b -> acc) -> acc -> b -> acc
-iterateStep n step = go n
- where
- go 0 !acc _ = acc
- go k !acc x = go (k - 1) (step acc x) x
-
-{- | Apply a 'CollectAgg' function to a single column, extracting the
-appropriate vector type and applying the aggregation function.
--}
-applyCollect ::
- forall v b a.
- (VG.Vector v b, Typeable v, Columnable b, Columnable a) =>
- (v b -> a) -> Column -> Either DataFrameException a
-applyCollect f col = f <$> toVector @b @v col
-
-{- | Result of interpreting an expression in a grouped context.
-Retained for backward compatibility with 'aggregate' and friends.
--}
-data AggregationResult a
- = UnAggregated Column
- | Aggregated (TypedColumn a)
-
-{- | Interpret an expression against a flat 'DataFrame', producing a
-typed column. This is the original top-level entry point; internally
-it calls 'eval' and materialises the result.
-
-NOTE: unlike the old implementation, 'Lit' values are no longer
-eagerly broadcast. The broadcast happens here, at the boundary,
-via 'materialize'.
--}
-interpret ::
- forall a.
- (Columnable a) =>
- DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
-interpret df expr = do
- v <- eval (FlatCtx df) expr
- pure $ TColumn $ materialize @a (fst (dataframeDimensions df)) v
-
-{- | Interpret an expression against a 'GroupedDataFrame',
-distinguishing aggregated results from bare column references.
-Internally calls 'eval'.
--}
-interpretAggregation ::
- forall a.
- (Columnable a) =>
- GroupedDataFrame ->
- Expr a ->
- Either DataFrameException (AggregationResult a)
-interpretAggregation gdf expr = do
- v <- eval (GroupCtx gdf) expr
- case v of
- Scalar a ->
- Right $
- Aggregated $
- TColumn $
- broadcastScalar @a (numGroups gdf) a
- Flat col ->
- Right $ Aggregated $ TColumn col
- Group _ ->
- -- The Column payload is intentionally unused — the only
- -- call-site ('aggregate') immediately throws
- -- 'UnaggregatedException' on this constructor.
- Right $ UnAggregated $ BoxedColumn @T.Text Nothing V.empty
diff --git a/src/DataFrame/Internal/Nullable.hs b/src/DataFrame/Internal/Nullable.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Nullable.hs
+++ /dev/null
@@ -1,500 +0,0 @@
-{-# 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,
- widenCmpOp,
- 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 Int32 Float where widen1 = fromIntegral; widen2 = id
-instance NumericWidenOp Float Int32 where
- widen1 = id
- widen2 = fromIntegral
-instance NumericWidenOp Int32 Double where widen1 = fromIntegral; widen2 = id
-instance NumericWidenOp Double Int32 where
- widen1 = id
- widen2 = fromIntegral
-instance NumericWidenOp Int64 Float where widen1 = fromIntegral; widen2 = id
-instance NumericWidenOp Float Int64 where
- widen1 = id
- widen2 = fromIntegral
-instance NumericWidenOp Int64 Double where widen1 = fromIntegral; widen2 = id
-instance NumericWidenOp Double Int64 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)
-
--- | Apply a comparison function after widening both operands to their common type.
-widenCmpOp ::
- forall a b.
- (NumericWidenOp a b) =>
- (Promote a b -> Promote a b -> Bool) ->
- a ->
- b ->
- Bool
-widenCmpOp f x y = f (widen1 @a @b x) (widen2 @a @b y)
-
--- | Result type of a widening binary operator, accounting for nullable wrappers.
-type WidenResult a b = NullLift2Result a b (Promote (BaseType a) (BaseType b))
-
--- ---------------------------------------------------------------------------
--- 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/PackedText.hs b/src/DataFrame/Internal/PackedText.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/PackedText.hs
+++ /dev/null
@@ -1,169 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-
-{- | Packed-text payload + byte-slice primitives. A 'PackedTextData' shares a
-single UTF-8 byte buffer across all rows of a string column, with @n+1@ row
-offsets, so no per-row 'Data.Text.Text' header is materialized at freeze.
-'Data.Text.Text' is produced only on demand (display, typed extraction) via
-the same decode path that the boxed-Text builder used.
-
-A gathered/joined/sorted result keeps sharing that buffer: instead of copying
-bytes it carries a @ptSel@ selection vector that reindexes the base rows, so a
-permuted or row-exploded column stays a 'PackedText' (shared buffer + permuted
-indices) rather than materializing back to boxed 'Data.Text.Text'.
--}
-module DataFrame.Internal.PackedText (
- PackedTextData (..),
- mkPackedContiguous,
- packedGather,
- packedTake,
- packedRowOffsetVec,
- packedLength,
- packedSlice,
- packedIndexText,
- sliceEqBytes,
- sliceCmpBytes,
-) where
-
-import qualified Data.Text as T
-import qualified Data.Text.Array as A
-import qualified Data.Vector.Unboxed as VU
-
-import Data.Ord (comparing)
-import Data.Text.Internal (Text (Text))
-import DataFrame.Internal.Utf8 (isValidUtf8Slice, lenientDecodeSlice)
-
-{- | A shared UTF-8 byte buffer plus @n+1@ row offsets (base row @r@ spans bytes
-@[offsets!r, offsets!(r+1))@). Validity lives in the enclosing column's
-@Maybe Bitmap@, mirroring 'BoxedColumn'/'UnboxedColumn'.
-
-@ptSel@ is an optional selection layer: when @Nothing@ the column is the
-contiguous base (row @i@ == base row @i@). When @Just sel@, logical row @i@ is
-base row @sel!i@; this is how a gather/join/sort result shares the buffer
-without copying bytes. Out-of-range entries in @sel@ (e.g. a join @-1@
-sentinel) decode to the empty slice and are masked by the column bitmap.
--}
-data PackedTextData = PackedTextData
- { ptBytes :: {-# UNPACK #-} !A.Array
- , ptOffsets :: {-# UNPACK #-} !(VU.Vector Int)
- , ptSel :: !(Maybe (VU.Vector Int))
- }
-
--- | Build a contiguous packed payload (no selection): the freeze-path shape.
-mkPackedContiguous :: A.Array -> VU.Vector Int -> PackedTextData
-mkPackedContiguous arr offs = PackedTextData arr offs Nothing
-{-# INLINE mkPackedContiguous #-}
-
-{- | Reindex a packed payload by a selection vector, sharing the byte buffer
-and base offsets. Logical row @i@ becomes base row @indices!i@. A negative or
-out-of-range index decodes to the empty slice (callers mask it with a bitmap).
-Composes with an existing selection so a gather of a gather still shares the
-buffer.
--}
-packedGather :: VU.Vector Int -> PackedTextData -> PackedTextData
-packedGather indices (PackedTextData arr offs msel) =
- let !base = VU.length offs - 1
- clamp r = if r >= 0 && r < base then r else -1
- sel' = case msel of
- Nothing -> VU.map clamp indices
- Just s ->
- VU.map
- (\i -> if i >= 0 && i < VU.length s then clamp (VU.unsafeIndex s i) else -1)
- indices
- in PackedTextData arr offs (Just sel')
-{-# INLINE packedGather #-}
-
-{- | Take the first @k@ logical rows, sharing the byte buffer. With a selection
-layer the selection is sliced to @k@ entries; without one a base-row selection
-@[0 .. k-1]@ is installed (slicing the contiguous offsets would still leave the
-trailing bytes addressable, but a short selection caps 'packedLength' to @k@).
-O(k), no byte copy or decode — the fix for cheap @take@/display on a 1e7-row
-packed column.
--}
-packedTake :: Int -> PackedTextData -> PackedTextData
-packedTake k (PackedTextData arr offs msel) =
- let !base = VU.length offs - 1
- !k' = max 0 k
- in case msel of
- Just s -> PackedTextData arr offs (Just (VU.take k' s))
- Nothing -> PackedTextData arr offs (Just (VU.enumFromN 0 (min k' base)))
-{-# INLINE packedTake #-}
-
--- | Map a logical row index to its base row, honoring any selection layer.
-baseRow :: PackedTextData -> Int -> Int
-baseRow (PackedTextData _ _ Nothing) i = i
-baseRow (PackedTextData _ _ (Just sel)) i = VU.unsafeIndex sel i
-{-# INLINE baseRow #-}
-
--- | Row count: @length sel@ when selected, else @length offsets - 1@.
-packedLength :: PackedTextData -> Int
-packedLength (PackedTextData _ offs Nothing) = VU.length offs - 1
-packedLength (PackedTextData _ _ (Just sel)) = VU.length sel
-{-# INLINE packedLength #-}
-
--- | Raw byte slice for logical row @i@: @(buffer, offset, length)@. The hot accessor.
-packedSlice :: PackedTextData -> Int -> (A.Array, Int, Int)
-packedSlice p@(PackedTextData arr offs _) i =
- let !r = baseRow p i
- in if r < 0
- then (arr, 0, 0)
- else
- let o = VU.unsafeIndex offs r in (arr, o, VU.unsafeIndex offs (r + 1) - o)
-{-# INLINE packedSlice #-}
-
-{- | The shared buffer + contiguous @n+1@ offsets when the payload is the
-unselected base. A selected (gathered) payload has non-contiguous rows that a
-single offset vector cannot express, so this returns @Nothing@ and the caller
-decodes per-row via 'packedIndexText'. Lets the boxed-Text fallback take the
-fast contiguous 'sliceTextVector' path when possible.
--}
-packedRowOffsetVec :: PackedTextData -> Maybe (A.Array, VU.Vector Int)
-packedRowOffsetVec (PackedTextData arr offs Nothing) = Just (arr, offs)
-packedRowOffsetVec _ = Nothing
-{-# INLINE packedRowOffsetVec #-}
-
-{- | On-demand single 'Data.Text.Text' for row @i@, using the same
-validate-or-lenient decode as 'sliceTextVector' so output is bit-identical.
--}
-packedIndexText :: PackedTextData -> Int -> T.Text
-packedIndexText p i =
- let (arr, o, l) = packedSlice p i
- in decodeField arr o l
-{-# INLINE packedIndexText #-}
-
--- Decode one field exactly as 'sliceTextVector' does per row.
-decodeField :: A.Array -> Int -> Int -> T.Text
-decodeField arr o l
- | l == 0 = T.empty
- | isValidUtf8Slice arr o l = Text arr o l
- | otherwise = lenientDecodeSlice arr o l
-{-# INLINE decodeField #-}
-
-{- | Byte-wise equality of two slices. UTF-8 is injective on valid scalar
-sequences and lenient decode is deterministic, so this agrees with
-@Text@'s '==' on the decoded values.
--}
-sliceEqBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Bool
-sliceEqBytes a ao al b bo bl
- | al /= bl = False
- | otherwise = go 0
- where
- go !k
- | k >= al = True
- | A.unsafeIndex a (ao + k) == A.unsafeIndex b (bo + k) = go (k + 1)
- | otherwise = False
-{-# INLINE sliceEqBytes #-}
-
-{- | Unsigned byte-lexicographic comparison (memcmp semantics). For
-well-formed UTF-8 this matches 'Data.Text.compare' exactly, since UTF-8
-byte order equals codepoint order for all valid scalars.
--}
-sliceCmpBytes :: A.Array -> Int -> Int -> A.Array -> Int -> Int -> Ordering
-sliceCmpBytes a ao al b bo bl = go 0
- where
- !m = min al bl
- go !k
- | k >= m = compare al bl
- | otherwise = case comparing id (A.unsafeIndex a (ao + k)) (A.unsafeIndex b (bo + k)) of
- EQ -> go (k + 1)
- r -> r
-{-# INLINE sliceCmpBytes #-}
diff --git a/src/DataFrame/Internal/ParRadixSort.hs b/src/DataFrame/Internal/ParRadixSort.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/ParRadixSort.hs
+++ /dev/null
@@ -1,294 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-
-{- |
-Parallel stable sort of row indices by the ascending unsigned order of a
-per-row 'Int' hash. Shared by the join build-side 'CompactIndex' construction
-(@DataFrame.Operations.Join@), which previously paid a single-threaded
-comparison sort over the whole build side — the dominant serial cost of a large
-inner join (the @1e7 x 1e7@ big-inner case).
-
-@parSortByHash n hashes@ returns @(sortedHashes, sortedIndices)@ where
-@sortedIndices@ lists @[0, n)@ in ascending 'sortKey' order of their hash, ties
-broken by ascending original index (stable), and @sortedHashes[k] ==
-hashes[sortedIndices[k]]@. Bit-for-bit identical to the old stable merge sort's
-output ordering, so equal-hash rows stay contiguous (the run scan in
-'buildCompactIndex' depends on this) and within a run keep original-row order.
-
-Strategy (mirrors "DataFrame.Internal.GroupingPar"): a counting sort buckets
-rows by the top @log2 p@ bits of their unsigned key into @p@ partitions laid out
-in ascending key order; @caps@ 'forkIO' workers then LSD-radix-sort each
-partition by the full 56 remaining low bits. Because partitions are already in
-global key order and each per-partition sort is stable, concatenating them
-reproduces the global stable order with no merge step. A sequential LSD radix
-sort is used below 'parSortThreshold' or on a single capability.
--}
-module DataFrame.Internal.ParRadixSort (
- parSortByHash,
- parSortThreshold,
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import Control.Monad (forM_, when)
-import Data.Bits (countLeadingZeros, unsafeShiftR, (.&.))
-import Data.IORef (atomicModifyIORef', newIORef)
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import Data.Word (Word64)
-import DataFrame.Internal.RadixRank (sortKey)
-import System.IO.Unsafe (unsafePerformIO)
-
-{- | Below this many rows the partition/fork overhead is not worth it; the
-caller's sequential LSD radix path is used instead.
--}
-parSortThreshold :: Int
-parSortThreshold = 500000
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
-{- | Top-bits partition index of a hash: the high @64 - shift@ bits of its
-unsigned 'sortKey'. Ascending partition order equals ascending key order.
--}
-partIx :: Int -> Int -> Int
-partIx shift h = fromIntegral ((fromIntegral (sortKey h) :: Word64) `unsafeShiftR` shift)
-{-# INLINE partIx #-}
-
--- | Number of partitions: a power of two, at least @4 * caps@, floored at 256.
-numPartitionsFor :: Int -> Int
-numPartitionsFor caps = go 1
- where
- target = max 256 (4 * caps)
- go p
- | p >= target = p
- | otherwise = go (p * 2)
-
--- | @floor (log2 x)@ for a power-of-two @x@.
-intLog2 :: Int -> Int
-intLog2 x = 63 - countLeadingZeros x
-{-# INLINE intLog2 #-}
-
-{- | Parallel stable sort of @[0, n)@ by ascending unsigned hash order. See the
-module header for the ordering contract.
--}
-parSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)
-parSortByHash n hashes
- | n <= 1 =
- (hashes, VU.enumFromN 0 n)
- | n < parSortThreshold || capabilities <= 1 =
- seqSortByHash n hashes
- | otherwise = unsafePerformIO (parSortByHashIO n hashes)
-{-# NOINLINE parSortByHash #-}
-
--------------------------------------------------------------------------------
--- Sequential LSD radix sort (also the per-partition worker kernel)
--------------------------------------------------------------------------------
-
-{- | Stable LSD radix sort of @[0, n)@ by ascending 'sortKey' of their hash, 8
-bits per pass over the full 64-bit key. Returns @(sortedHashes, sortedIndices)@.
--}
-seqSortByHash :: Int -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)
-seqSortByHash n hashes = unsafePerformIO $ do
- keysA <- VUM.new n
- orderA <- VUM.new n
- let seed !i
- | i >= n = pure ()
- | otherwise = do
- VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes i))
- VUM.unsafeWrite orderA i i
- seed (i + 1)
- seed 0
- keysB <- VUM.new n
- orderB <- VUM.new n
- radixPasses n keysA orderA keysB orderB
- order <- VU.unsafeFreeze orderA
- pure (VU.unsafeBackpermute hashes order, order)
-
-{- | Run all eight stable 8-bit LSD passes, ping-ponging between the two
-key/order buffer pairs so the sorted order lands back in @(keysA, orderA)@.
-@keysA[i]@ must already hold @sortKey (hash of orderA[i])@ on entry.
--}
-radixPasses ::
- Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- IO ()
-radixPasses n keysA orderA keysB orderB = do
- counts <- VUM.new 256
- let pass ::
- Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- IO ()
- pass !shiftBits !srcK !srcO !dstK !dstO = do
- VUM.set counts 0
- let count !i
- | i >= n = pure ()
- | otherwise = do
- k <- VUM.unsafeRead srcK i
- let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
- VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)
- count (i + 1)
- count 0
- let scan !b !acc
- | b >= 256 = pure ()
- | otherwise = do
- c <- VUM.unsafeRead counts b
- VUM.unsafeWrite counts b acc
- scan (b + 1) (acc + c)
- scan 0 0
- let place !i
- | i >= n = pure ()
- | otherwise = do
- k <- VUM.unsafeRead srcK i
- o <- VUM.unsafeRead srcO i
- let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
- pos <- VUM.unsafeRead counts b
- VUM.unsafeWrite counts b (pos + 1)
- VUM.unsafeWrite dstK pos k
- VUM.unsafeWrite dstO pos o
- place (i + 1)
- place 0
- pass 0 keysA orderA keysB orderB
- pass 8 keysB orderB keysA orderA
- pass 16 keysA orderA keysB orderB
- pass 24 keysB orderB keysA orderA
- pass 32 keysA orderA keysB orderB
- pass 40 keysB orderB keysA orderA
- pass 48 keysA orderA keysB orderB
- pass 56 keysB orderB keysA orderA
-
--------------------------------------------------------------------------------
--- Parallel path: counting-sort partition, then per-partition sort in parallel
--------------------------------------------------------------------------------
-
-parSortByHashIO :: Int -> VU.Vector Int -> IO (VU.Vector Int, VU.Vector Int)
-parSortByHashIO n hashes = do
- caps <- getNumCapabilities
- let !p = numPartitionsFor caps
- !shift = 64 - intLog2 p
- -- Phase 1: counting sort of row indices into ascending-key partitions.
- (partStart, partRows) <- partitionRows n hashes p shift
- -- Phase 2: stable-sort each partition by full key, in parallel. Each worker
- -- owns disjoint [partStart[pp], partStart[pp+1]) output ranges, so the
- -- single shared output buffers are written race-free.
- outOrder <- VUM.new n
- outKeys <- VUM.new n
- sortPartitions caps p partStart partRows hashes outOrder outKeys
- order <- VU.unsafeFreeze outOrder
- pure (VU.unsafeBackpermute hashes order, order)
-
-{- | Bucket every row index into its top-bits partition by a counting sort.
-Returns the exclusive prefix sum @partStart@ (length @p+1@, @partStart[p] == n@)
-and the row indices laid out partition-by-partition in ascending key order.
--}
-partitionRows ::
- Int -> VU.Vector Int -> Int -> Int -> IO (VU.Vector Int, VU.Vector Int)
-partitionRows n hashes p shift = do
- counts <- VUM.replicate (p + 1) (0 :: Int)
- let countLoop !i
- | i >= n = pure ()
- | otherwise = do
- let !pp = partIx shift (VU.unsafeIndex hashes i)
- c <- VUM.unsafeRead counts pp
- VUM.unsafeWrite counts pp (c + 1)
- countLoop (i + 1)
- countLoop 0
- partStartM <- VUM.new (p + 1)
- let scan !k !acc
- | k > p = pure ()
- | otherwise = do
- VUM.unsafeWrite partStartM k acc
- c <- if k < p then VUM.unsafeRead counts k else pure 0
- scan (k + 1) (acc + c)
- scan 0 0
- cursor <- VUM.new p
- forM_ [0 .. p - 1] $ \k -> VUM.unsafeRead partStartM k >>= VUM.unsafeWrite cursor k
- rowsM <- VUM.new (max 1 n)
- let place !i
- | i >= n = pure ()
- | otherwise = do
- let !pp = partIx shift (VU.unsafeIndex hashes i)
- pos <- VUM.unsafeRead cursor pp
- VUM.unsafeWrite rowsM pos i
- VUM.unsafeWrite cursor pp (pos + 1)
- place (i + 1)
- place 0
- partStart <- VU.unsafeFreeze partStartM
- partRows <- VU.unsafeFreeze rowsM
- pure (partStart, partRows)
-
-{- | Stable-sort each partition by full key, writing sorted original indices
-into @outOrder@ and their hashes into @outKeys@ at the partition's slot range.
-Forks @caps@ workers that pull partition indices off a shared atomic counter.
-Within a partition the counting sort already left rows in ascending original
-order, so the LSD radix sort's stability reproduces the global @(key, row)@
-order. Partitions below two elements are already sorted (counting sort kept
-original order) and are copied directly.
--}
-sortPartitions ::
- Int ->
- Int ->
- VU.Vector Int ->
- VU.Vector Int ->
- VU.Vector Int ->
- VUM.IOVector Int ->
- VUM.IOVector Int ->
- IO ()
-sortPartitions caps p partStart partRows hashes outOrder outKeys = do
- next <- newIORef 0
- let sortOne !pp = do
- let !s = VU.unsafeIndex partStart pp
- !e = VU.unsafeIndex partStart (pp + 1)
- !sz = e - s
- when (sz > 0) $
- if sz == 1
- then do
- let !r = VU.unsafeIndex partRows s
- VUM.unsafeWrite outOrder s r
- VUM.unsafeWrite outKeys s (VU.unsafeIndex hashes r)
- else do
- keysA <- VUM.new sz
- orderA <- VUM.new sz
- let seed !i
- | i >= sz = pure ()
- | otherwise = do
- let !r = VU.unsafeIndex partRows (s + i)
- VUM.unsafeWrite keysA i (sortKey (VU.unsafeIndex hashes r))
- VUM.unsafeWrite orderA i r
- seed (i + 1)
- seed 0
- keysB <- VUM.new sz
- orderB <- VUM.new sz
- radixPasses sz keysA orderA keysB orderB
- let emit !i
- | i >= sz = pure ()
- | otherwise = do
- o <- VUM.unsafeRead orderA i
- VUM.unsafeWrite outOrder (s + i) o
- VUM.unsafeWrite outKeys (s + i) (VU.unsafeIndex hashes o)
- emit (i + 1)
- emit 0
- worker = do
- i <- atomicModifyIORef' next (\j -> (j + 1, j))
- when (i < p) $ sortOne i >> worker
- forkJoin_ (replicate caps worker)
-
--- | Run each action on its own thread; rethrow the first failure (in order).
-forkJoin_ :: [IO ()] -> IO ()
-forkJoin_ actions = do
- vars <- mapM spawn actions
- results <- mapM takeMVar vars
- mapM_ (either (throwIO :: SomeException -> IO ()) pure) results
- where
- spawn act = do
- var <- newEmptyMVar
- _ <- forkIO (try act >>= putMVar var)
- pure var
diff --git a/src/DataFrame/Internal/RadixRank.hs b/src/DataFrame/Internal/RadixRank.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/RadixRank.hs
+++ /dev/null
@@ -1,114 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-
-{- |
-Stable rank of a set of group representatives by the ascending unsigned order of
-their hash. Shared by the sequential ('DataFrame.Internal.Grouping') and parallel
-('DataFrame.Internal.GroupingPar') group-by canonical-ordering steps so they
-stay bit-for-bit identical.
-
-@rankByHash readHash ng@ returns @rank@ with @rank[gid] = position@ of group
-@gid@ when groups are ordered by ascending unsigned 'sortKey' of @readHash gid@.
-A stable LSD radix sort (8 bits per pass, 8 passes) keeps groups with equal
-hash in their original @gid@ order; callers number @gid@s so that this matches
-the @repRow@ tie-break of the old comparison sort. @O(ng)@, no boxed tuples or
-comparison closures — the lever for the @1e7@-distinct-group case (Q10).
--}
-module DataFrame.Internal.RadixRank (
- rankByHash,
- sortKey,
-) where
-
-import Control.Monad (when)
-import Control.Monad.Primitive (PrimMonad)
-import Data.Bits (unsafeShiftR, (.&.))
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import Data.Word (Word64)
-
-{- | Unsigned sort key of a hash: ascending 'Word64' order of @sortKey h@ equals
-ascending signed-'Int' order of @h@. Reinterpreted back to 'Int' for the
-byte-wise radix passes (the @.&. 0xff@ byte mask makes the arithmetic shift's
-sign extension irrelevant).
--}
-sortKey :: Int -> Int
-sortKey h = fromIntegral (fromIntegral h + 0x8000000000000000 :: Word64)
-{-# INLINE sortKey #-}
-
--- | See the module header. @readHash@ supplies the hash of local group @gid@.
-rankByHash ::
- forall m. (PrimMonad m) => (Int -> m Int) -> Int -> m (VU.Vector Int)
-rankByHash readHash ng = do
- rankM <- VUM.new (max 1 ng)
- if ng <= 1
- then when (ng == 1) (VUM.unsafeWrite rankM 0 0)
- else do
- keysA <- VUM.new ng
- orderA <- VUM.new ng
- let seed !i
- | i >= ng = pure ()
- | otherwise = do
- h <- readHash i
- VUM.unsafeWrite keysA i (sortKey h)
- VUM.unsafeWrite orderA i i
- seed (i + 1)
- seed 0
- keysB <- VUM.new ng
- orderB <- VUM.new ng
- counts <- VUM.new 256
- let pass ::
- Int ->
- VUM.MVector (VUM.PrimState m) Int ->
- VUM.MVector (VUM.PrimState m) Int ->
- VUM.MVector (VUM.PrimState m) Int ->
- VUM.MVector (VUM.PrimState m) Int ->
- m ()
- pass !shiftBits !srcK !srcO !dstK !dstO = do
- VUM.set counts 0
- let count !i
- | i >= ng = pure ()
- | otherwise = do
- k <- VUM.unsafeRead srcK i
- let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
- VUM.unsafeRead counts b >>= VUM.unsafeWrite counts b . (+ 1)
- count (i + 1)
- count 0
- let scan !b !acc
- | b >= 256 = pure ()
- | otherwise = do
- c <- VUM.unsafeRead counts b
- VUM.unsafeWrite counts b acc
- scan (b + 1) (acc + c)
- scan 0 0
- let place !i
- | i >= ng = pure ()
- | otherwise = do
- k <- VUM.unsafeRead srcK i
- o <- VUM.unsafeRead srcO i
- let !b = (k `unsafeShiftR` shiftBits) .&. 0xff
- pos <- VUM.unsafeRead counts b
- VUM.unsafeWrite counts b (pos + 1)
- VUM.unsafeWrite dstK pos k
- VUM.unsafeWrite dstO pos o
- place (i + 1)
- place 0
- -- 8 stable passes over the 64-bit key; ping-pong so the final
- -- sorted order lands back in (keysA, orderA).
- pass 0 keysA orderA keysB orderB
- pass 8 keysB orderB keysA orderA
- pass 16 keysA orderA keysB orderB
- pass 24 keysB orderB keysA orderA
- pass 32 keysA orderA keysB orderB
- pass 40 keysB orderB keysA orderA
- pass 48 keysA orderA keysB orderB
- pass 56 keysB orderB keysA orderA
- -- orderA[rank] = gid; invert to rank[gid] = rank.
- let inv !r
- | r >= ng = pure ()
- | otherwise = do
- g <- VUM.unsafeRead orderA r
- VUM.unsafeWrite rankM g r
- inv (r + 1)
- inv 0
- VU.unsafeFreeze rankM
-{-# INLINEABLE rankByHash #-}
diff --git a/src/DataFrame/Internal/Row.hs b/src/DataFrame/Internal/Row.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Row.hs
+++ /dev/null
@@ -1,209 +0,0 @@
-{-# LANGUAGE ExistentialQuantification #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE InstanceSigs #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-module DataFrame.Internal.Row where
-
-import qualified Data.List as L
-import qualified Data.Map as M
-import qualified Data.Text as T
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-
-import Control.Exception (throw)
-import Data.Function (on)
-import Data.Maybe (catMaybes, fromMaybe, isNothing, mapMaybe)
-import Data.Type.Equality (TestEquality (..))
-import Data.Typeable (type (:~:) (..))
-import DataFrame.Errors (DataFrameException (..))
-import DataFrame.Internal.Column
-import DataFrame.Internal.DataFrame
-import DataFrame.Internal.Expression (Expr (..))
-import DataFrame.Internal.PackedText (packedIndexText, packedLength)
-import Type.Reflection (typeOf, typeRep)
-
-data Any where
- Value :: (Columnable a) => a -> Any
- -- Saves us the extra indirection we get from making Value (Maybe a)
- -- and having to unpack it again to check for nulls.
- -- Instead, we just have Null as a separate constructor.
- Null :: Any
-
-instance Eq Any where
- (==) :: Any -> Any -> Bool
- (Value a) == (Value b) = fromMaybe False $ do
- Refl <- testEquality (typeOf a) (typeOf b)
- return $ a == b
- Null == Null = True
- _ == _ = False
-
-instance Show Any where
- show :: Any -> String
- show (Value a) = T.unpack (showValue a)
- show Null = "null"
-
-showValue :: forall a. (Columnable a) => a -> T.Text
-showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of
- Just Refl -> v
- Nothing -> case testEquality (typeRep @a) (typeRep @String) of
- Just Refl -> T.pack v
- Nothing -> (T.pack . show) v
-
--- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.
-toAny :: forall a. (Columnable a) => a -> Any
-toAny = Value
-
--- | Unwraps a value from an \Any\ type. A 'Null' cell yields 'Nothing'.
-fromAny :: forall a. (Columnable a) => Any -> Maybe a
-fromAny Null = Nothing
-fromAny (Value (v :: b)) = do
- Refl <- testEquality (typeRep @a) (typeRep @b)
- pure v
-
-{- | Wrap a column cell into an 'Any', honouring the column's null bitmap: a slot
-marked invalid becomes 'Null', any other slot becomes a 'Value'. Only needs
-@Columnable a@ (the stored element type), not @Columnable (Maybe a)@.
--}
-cellAny :: (Columnable a) => Maybe Bitmap -> Int -> a -> Any
-cellAny Nothing _ x = Value x
-cellAny (Just bm) i x = if bitmapTestBit bm i then Value x else Null
-
-type Row = V.Vector Any
-
-(!?) :: [a] -> Int -> Maybe a
-(!?) [] _ = Nothing
-(!?) (x : _) 0 = Just x
-(!?) (_x : xs) n = (!?) xs (n - 1)
-
-{- | Reconstruct column @i@ from a list of rows. The element type is taken from
-the first non-'Null' cell; cells of a different type are skipped. If any cell is
-'Null' the result is a nullable column (built with 'fromMaybeVec', which needs
-only @Columnable a@), so a round-trip through 'toRowList'/'fromRows' preserves
-nulls.
--}
-mkColumnFromRow :: Int -> [[Any]] -> Column
-mkColumnFromRow i rows =
- let cells = mapMaybe (!? i) rows
- in case L.find isValue cells of
- Nothing -> fromList ([] :: [T.Text])
- Just (Value (_ :: a)) ->
- let collect Null = Just (Nothing :: Maybe a)
- collect (Value (v' :: b)) =
- case testEquality (typeRep @a) (typeRep @b) of
- Just Refl -> Just (Just v')
- Nothing -> Nothing
- maybes = mapMaybe collect cells
- in if any isNothing maybes
- then fromMaybeVec (V.fromList maybes)
- else fromList (catMaybes maybes)
- Just Null -> fromList ([] :: [T.Text]) -- unreachable: find isValue
- where
- isValue (Value _) = True
- isValue Null = False
-
-{- | Converts the entire dataframe to a list of rows.
-
-Each row contains all columns in the dataframe, ordered by their column indices.
-The rows are returned in their natural order (from index 0 to n-1).
-
-==== __Examples__
-
->>> toRowList df
-[[("name", "Alice"), ("age", 25), ...], [("name", "Bob"), ("age", 30), ...], ...]
-
-==== __Performance note__
-
-This function materializes all rows into a list, which may be memory-intensive
-for large dataframes. Consider using 'toRowVector' if you need random access
-or streaming operations.
--}
-toRowList :: DataFrame -> [[(T.Text, Any)]]
-toRowList df =
- let
- names = map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df))
- in
- map
- (zip names . V.toList . mkRowRep df names)
- [0 .. (fst (dataframeDimensions df) - 1)]
-
-{- | Converts the dataframe to a vector of rows with only the specified columns.
-
-Each row will contain only the columns named in the @names@ parameter.
-This is useful when you only need a subset of columns or want to control
-the column order in the resulting rows.
-
-==== __Parameters__
-
-[@names@] List of column names to include in each row. The order of names
- determines the order of fields in the resulting rows.
-
-[@df@] The dataframe to convert.
-
-==== __Examples__
-
->>> toRowVector ["name", "age"] df
-Vector of rows with only name and age fields
-
->>> toRowVector [] df -- Empty column list
-Vector of empty rows (one per dataframe row)
--}
-toRowVector :: [T.Text] -> DataFrame -> V.Vector Row
-toRowVector names df = V.generate (fst (dataframeDimensions df)) (mkRowRep df names)
-
-{- | Given a row gets the value associated with a field.
-
-==== __Examples__
-
->>> map (rowValue (F.col @Int "age")) (toRowList df)
-[25,30, ...]
--}
-rowValue :: forall a. Expr a -> [(T.Text, Any)] -> Maybe a
-rowValue (Col name) row = lookup name row >>= fromAny @a
-rowValue _ _ = error "Can only get rowValue of column reference"
-
-mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row
-mkRowFromArgs names df i = V.map get (V.fromList names)
- where
- get name = case getColumn name df of
- Nothing ->
- throw $
- ColumnsNotFoundException
- [name]
- "[INTERNAL] mkRowFromArgs"
- (M.keys $ columnIndices df)
- Just (BoxedColumn bm column) -> cellAny bm i (column V.! i)
- Just (UnboxedColumn bm column) -> cellAny bm i (column VU.! i)
- Just (PackedText bm p) -> cellAny bm i (packedIndexText p i)
-
--- This function will return the items in the order that is specified
--- by the user. For example, if the dataframe consists of the columns
--- "Age", "Pclass", "Name", and the user asks for ["Name", "Age"],
--- this will order the values in the order ["Mr Smith", 50]
-mkRowRep :: DataFrame -> [T.Text] -> Int -> Row
-mkRowRep df names i = V.generate (L.length names) (\index -> get (names' V.! index))
- where
- names' = V.fromList names
- throwError name =
- error $
- "Column "
- ++ T.unpack name
- ++ " has less items than "
- ++ "the other columns at index "
- ++ show i
- get name = case getColumn name df of
- Just (BoxedColumn bm c) -> case c V.!? i of
- Just e -> cellAny bm i e
- Nothing -> throwError name
- Just (UnboxedColumn bm c) -> case c VU.!? i of
- Just e -> cellAny bm i e
- Nothing -> throwError name
- Just (PackedText bm p)
- | i < packedLength p -> cellAny bm i (packedIndexText p i)
- | otherwise -> throwError name
- Nothing ->
- throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
diff --git a/src/DataFrame/Internal/RowHash.hs b/src/DataFrame/Internal/RowHash.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/RowHash.hs
+++ /dev/null
@@ -1,232 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-{- | Row-hash kernels with a parallel driver.
-
-The per-row key hash is the sole input to grouping and the join build/probe.
-For a single wide pass over many rows (notably a 1e7-row text/factor join key)
-the hashing is the dominant cost and is embarrassingly parallel: each row's hash
-depends only on that row's own bytes, so hashing disjoint row ranges into
-disjoint slots of one shared vector is race-free and produces a result
-/bit-for-bit identical/ to the sequential single-pass hash.
-
-'hashRowRange' is the shared per-range kernel (used sequentially and by every
-worker); 'computeRowHashesIO' forks one worker per capability over contiguous
-row ranges above 'parRowHashThreshold', else runs the range once. The mixing per
-column type mirrors the grouping hash exactly so grouping and joins agree.
--}
-module DataFrame.Internal.RowHash (
- computeRowHashesIO,
- hashRowRange,
- parRowHashThreshold,
-) where
-
-import Control.Concurrent (forkIO, getNumCapabilities)
-import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
-import Control.Exception (SomeException, throwIO, try)
-import qualified Data.Text as T
-import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
-import qualified Data.Vector as V
-import qualified Data.Vector.Unboxed as VU
-import qualified Data.Vector.Unboxed.Mutable as VUM
-import System.IO.Unsafe (unsafePerformIO)
-import Type.Reflection (typeRep)
-
-import DataFrame.Internal.Column (Bitmap, Column (..), bitmapTestBit)
-import DataFrame.Internal.Hash (
- fnvOffset,
- mixBytes,
- mixDouble,
- mixInt,
- mixShow,
- mixText,
- nullSalt,
- )
-import DataFrame.Internal.PackedText (
- PackedTextData (..),
- packedSlice,
- )
-import DataFrame.Internal.Types (
- SBool (..),
- sFloating,
- sIntegral,
- )
-
-{- | At least this many rows make the fork/coordination overhead of the parallel
-hash worth it. Below it the sequential single range is used. Matches the
-grouping/join parallel thresholds so the whole pipeline switches together.
--}
-parRowHashThreshold :: Int
-parRowHashThreshold = 200000
-
-capabilities :: Int
-capabilities = unsafePerformIO getNumCapabilities
-{-# NOINLINE capabilities #-}
-
-{- | Compute the per-row key hash over the (already selected) key columns of an
-@n@-row frame. Forks one worker per capability over contiguous row ranges when
-the row count justifies it (>= 'parRowHashThreshold' and more than one
-capability); otherwise hashes the single full range. The output is identical for
-any capability count: each row's hash is a pure function of its own bytes and
-workers own disjoint row ranges.
--}
-computeRowHashesIO :: Int -> [Column] -> IO (VU.Vector Int)
-computeRowHashesIO n selected = do
- mv <- VUM.unsafeNew (max 1 n)
- let runRange lo hi = hashRowRange mv lo hi selected
- if n >= parRowHashThreshold && capabilities > 1
- then do
- let !caps = capabilities
- !per = (n + caps - 1) `div` caps
- spawn w = do
- var <- newEmptyMVar
- let !lo = min n (w * per)
- !hi = min n (lo + per)
- _ <- forkIO (try (runRange lo hi) >>= putMVar var)
- pure var
- vars <- mapM spawn [0 .. caps - 1]
- rs <- mapM takeMVar vars
- mapM_ (either (throwIO @SomeException) pure) rs
- else runRange 0 n
- VU.unsafeFreeze (VUM.slice 0 n mv)
-
-{- | Mix every selected column over the row range @[lo, hi)@ into @mv@, seeding
-each slot with 'fnvOffset' first. The seeding and per-column mixing must match
-'DataFrame.Operations.Aggregation.computeRowHashes' byte-for-byte so grouping
-and joins bucket identically.
--}
-hashRowRange :: VUM.IOVector Int -> Int -> Int -> [Column] -> IO ()
-hashRowRange mv lo hi cols = do
- seedRange mv lo hi
- mapM_ (mixColumnRange mv lo hi) cols
-
-seedRange :: VUM.IOVector Int -> Int -> Int -> IO ()
-seedRange mv lo hi = go lo
- where
- go !i
- | i >= hi = pure ()
- | otherwise = VUM.unsafeWrite mv i fnvOffset >> go (i + 1)
-
-{- | Fold one column's values over @[lo, hi)@ into the running hashes. The branch
-structure mirrors the sequential grouping hash: typed unboxed fast paths, then a
-'mixShow' fallback, with the null bitmap mixing 'nullSalt'.
--}
-mixColumnRange :: VUM.IOVector Int -> Int -> Int -> Column -> IO ()
-mixColumnRange mv lo hi = \case
- UnboxedColumn ubm (v :: VU.Vector a) ->
- case testEquality (typeRep @a) (typeRep @Int) of
- Just Refl -> unboxedRange mv lo hi ubm mixInt v
- Nothing ->
- case testEquality (typeRep @a) (typeRep @Double) of
- Just Refl -> unboxedRange mv lo hi ubm mixDouble v
- Nothing ->
- case sIntegral @a of
- STrue ->
- unboxedRange mv lo hi ubm (\h d -> mixInt h (fromIntegral @a @Int d)) v
- SFalse ->
- case sFloating @a of
- STrue ->
- unboxedRange mv lo hi ubm (\h d -> mixDouble h (realToFrac d :: Double)) v
- SFalse ->
- unboxedRange mv lo hi ubm mixShow v
- BoxedColumn bm (v :: V.Vector a) ->
- case testEquality (typeRep @a) (typeRep @T.Text) of
- Just Refl -> boxedRange mv lo hi bm mixText v
- Nothing -> boxedRange mv lo hi bm mixShow v
- PackedText bm p -> packedRange mv lo hi bm p
-
-{- | Mix an unboxed column's range, mixing 'nullSalt' at null slots. @INLINE@d to
-specialise on the element type and mixing function per call site.
--}
-unboxedRange ::
- (VU.Unbox a) =>
- VUM.IOVector Int ->
- Int ->
- Int ->
- Maybe Bitmap ->
- (Int -> a -> Int) ->
- VU.Vector a ->
- IO ()
-unboxedRange mv lo hi ubm mix v = go lo
- where
- go !i
- | i >= hi = pure ()
- | otherwise = do
- h <- VUM.unsafeRead mv i
- let !h' = case ubm of
- Just bm | not (bitmapTestBit bm i) -> mixInt h nullSalt
- _ -> mix h (VU.unsafeIndex v i)
- VUM.unsafeWrite mv i h'
- go (i + 1)
-{-# INLINE unboxedRange #-}
-
-boxedRange ::
- VUM.IOVector Int ->
- Int ->
- Int ->
- Maybe Bitmap ->
- (Int -> a -> Int) ->
- V.Vector a ->
- IO ()
-boxedRange mv lo hi bm mix v = go lo
- where
- go !i
- | i >= hi = pure ()
- | otherwise = do
- h <- VUM.unsafeRead mv i
- let !h' = case bm of
- Just bm' | not (bitmapTestBit bm' i) -> mixInt h nullSalt
- _ -> mix h (V.unsafeIndex v i)
- VUM.unsafeWrite mv i h'
- go (i + 1)
-{-# INLINE boxedRange #-}
-
-{- | Mix a packed-text column's range over its raw UTF-8 byte slices. The
-contiguous (unselected) payload is the hot path: hoist the byte buffer and the
-@n+1@ offset vector out of the loop and index them directly, so each row mixes
-@[offs!i, offs!(i+1))@ with no per-row selection 'Maybe' test or 'packedSlice'
-tuple. A selected payload (a gather/join result) falls back to 'packedSlice'.
--}
-packedRange ::
- VUM.IOVector Int ->
- Int ->
- Int ->
- Maybe Bitmap ->
- PackedTextData ->
- IO ()
-packedRange mv lo hi bm p =
- case ptSel p of
- Nothing -> contiguous (ptBytes p) (ptOffsets p)
- Just _ -> selected
- where
- valid i = case bm of
- Just bm' -> bitmapTestBit bm' i
- Nothing -> True
- contiguous !arr !offs = go lo
- where
- go !i
- | i >= hi = pure ()
- | otherwise = do
- h <- VUM.unsafeRead mv i
- let !o = VU.unsafeIndex offs i
- !l = VU.unsafeIndex offs (i + 1) - o
- !h' = if valid i then mixBytes h arr o l else mixInt h nullSalt
- VUM.unsafeWrite mv i h'
- go (i + 1)
- selected = go lo
- where
- go !i
- | i >= hi = pure ()
- | otherwise = do
- h <- VUM.unsafeRead mv i
- let !h' =
- if valid i
- then let (arr, o, l) = packedSlice p i in mixBytes h arr o l
- else mixInt h nullSalt
- VUM.unsafeWrite mv i h'
- go (i + 1)
-{-# INLINE packedRange #-}
diff --git a/src/DataFrame/Internal/Simplify.hs b/src/DataFrame/Internal/Simplify.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Simplify.hs
+++ /dev/null
@@ -1,422 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE MultiWayIf #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE PatternSynonyms #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-
-module DataFrame.Internal.Simplify (
- simplify,
- simplifyPredicatePair,
-
- -- * Path-condition entailment (for fitted-tree pruning)
- PredFact,
- factTrue,
- factFalse,
- entails,
-) where
-
-import Control.Monad (guard)
-import Data.Maybe (fromMaybe)
-import Data.Type.Equality (testEquality, (:~:) (Refl))
-import Type.Reflection (eqTypeRep, typeRep, (:~~:) (HRefl), pattern App)
-
-import DataFrame.Internal.Column (Columnable)
-import DataFrame.Internal.Expression (
- BinaryOp,
- Expr (..),
- UnaryOp (unaryName),
- eqExpr,
- normalize,
- )
-import DataFrame.Operators (
- NullAnd,
- NullEq,
- NullGeq,
- NullGt,
- NullLeq,
- NullLt,
- NullNeq,
- NullOr,
- (.==.),
- )
-
-simplify :: forall a. (Columnable a) => Expr a -> Expr a
-simplify e
- | isBoolish @a = fixpoint (10 :: Int) e
- | otherwise = e
- where
- fixpoint 0 x = x
- fixpoint n x = let x' = simplifyB x in if eqExpr x x' then x else fixpoint (n - 1) x'
-
-isBoolish :: forall a. (Columnable a) => Bool
-isBoolish =
- case ( testEquality (typeRep @a) (typeRep @Bool)
- , testEquality (typeRep @a) (typeRep @(Maybe Bool))
- ) of
- (Just Refl, _) -> True
- (_, Just Refl) -> True
- _ -> False
-
-data Conn = ConnAnd | ConnOr
-
-connOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Conn
-connOf _
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullAnd) = Just ConnAnd
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullOr) = Just ConnOr
- | otherwise = Nothing
-
-simplifyB :: forall a. (Columnable a) => Expr a -> Expr a
-simplifyB expr = case expr of
- Binary (op :: op c b a) l r
- | Just conn <- connOf op
- , Just Refl <- testEquality (typeRep @c) (typeRep @a)
- , Just Refl <- testEquality (typeRep @b) (typeRep @a) ->
- let l' = simplifyB l; r' = simplifyB r
- in fromMaybe (Binary op l' r') (combine conn l' r')
- | otherwise -> expr
- Unary (op :: op b a) inner
- | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)
- , Just Refl <- testEquality (typeRep @b) (typeRep @Bool)
- , unaryName op == "not" ->
- simplifyNot op (simplifyB inner)
- | otherwise -> expr
- If c t f ->
- let c' = simplify c
- t' = simplifyB t
- f' = simplifyB f
- in case asBoolLit c' of
- Just True -> t'
- Just False -> f'
- Nothing
- | eqExpr t' f' -> t'
- | Just Refl <- testEquality (typeRep @a) (typeRep @Bool)
- , asBoolLit t' == Just True
- , asBoolLit f' == Just False ->
- c'
- | otherwise -> If c' t' f'
- _ -> expr
-
-simplifyNot :: (UnaryOp op) => op Bool Bool -> Expr Bool -> Expr Bool
-simplifyNot op inner = case asBoolLit inner of
- Just b -> Lit (not b)
- Nothing -> case inner of
- Unary (op2 :: op2 b2 Bool) inner2
- | unaryName op2 == "not"
- , Just Refl <- testEquality (typeRep @b2) (typeRep @Bool) ->
- inner2
- _ -> Unary op inner
-
-combine :: (Columnable a) => Conn -> Expr a -> Expr a -> Maybe (Expr a)
-combine ConnAnd = combineAnd
-combine ConnOr = combineOr
-
-asBoolLit :: forall a. (Columnable a) => Expr a -> Maybe Bool
-asBoolLit (Lit v) =
- case testEquality (typeRep @a) (typeRep @Bool) of
- Just Refl -> Just v
- Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of
- Just Refl -> v
- Nothing -> Nothing
-asBoolLit _ = Nothing
-
-{- | Polymorphic boolean literal: @Lit b@ for @Expr Bool@, @Lit (Just b)@ for
-@Expr (Maybe Bool)@.
--}
-litBoolish :: forall a. (Columnable a) => Bool -> Maybe (Expr a)
-litBoolish v =
- case testEquality (typeRep @a) (typeRep @Bool) of
- Just Refl -> Just (Lit v)
- Nothing -> case testEquality (typeRep @a) (typeRep @(Maybe Bool)) of
- Just Refl -> Just (Lit (Just v))
- Nothing -> Nothing
-
-combineAnd :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)
-combineAnd l r
- | eqExpr l r = Just l
- | asBoolLit l == Just False = litBoolish False
- | asBoolLit r == Just False = litBoolish False
- | asBoolLit l == Just True = Just r
- | asBoolLit r == Just True = Just l
- | absorbs ConnOr l r = Just l
- | absorbs ConnOr r l = Just r
- | otherwise = simplifyPredicatePair True l r
-
-combineOr :: (Columnable a) => Expr a -> Expr a -> Maybe (Expr a)
-combineOr l r
- | eqExpr l r = Just l
- | asBoolLit l == Just True = litBoolish True
- | asBoolLit r == Just True = litBoolish True
- | asBoolLit l == Just False = Just r
- | asBoolLit r == Just False = Just l
- | absorbs ConnAnd l r = Just l
- | absorbs ConnAnd r l = Just r
- | otherwise = simplifyPredicatePair False l r
-
-absorbs :: (Columnable a) => Conn -> Expr a -> Expr a -> Bool
-absorbs conn x (Binary (op :: op c b a) ya yb)
- | Just c' <- connOf op
- , sameConn conn c'
- , Just Refl <- testEquality (typeRep @c) (typeRep @a)
- , Just Refl <- testEquality (typeRep @b) (typeRep @a) =
- eqExpr x ya || eqExpr x yb
-absorbs _ _ _ = False
-
-sameConn :: Conn -> Conn -> Bool
-sameConn ConnAnd ConnAnd = True
-sameConn ConnOr ConnOr = True
-sameConn _ _ = False
-
-data Cmp = CLt | CLeq | CGt | CGeq | CEq | CNeq deriving (Eq)
-
-data NullK = Total | FalseOnNull | UnknownOnNull deriving (Eq)
-
-data Atom = Atom
- { aCmp :: Cmp
- , aThr :: !Double
- , aKey :: String
- , aNull :: NullK
- , aIntegral :: Bool
- }
-
-cmpOf :: forall op c b r. (BinaryOp op) => op c b r -> Maybe Cmp
-cmpOf _
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLt) = Just CLt
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullLeq) = Just CLeq
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGt) = Just CGt
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullGeq) = Just CGeq
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullEq) = Just CEq
- | Just HRefl <- eqTypeRep (typeRep @op) (typeRep @NullNeq) = Just CNeq
- | otherwise = Nothing
-
-isLower, isUpper :: Cmp -> Bool
-isLower c = c == CGt || c == CGeq
-isUpper c = c == CLt || c == CLeq
-
--- | True if @x@ is a @Maybe _@ type.
-isMaybeTy :: forall x. (Columnable x) => Bool
-isMaybeTy = case typeRep @x of
- App con _ -> case eqTypeRep con (typeRep @Maybe) of Just HRefl -> True; _ -> False
- _ -> False
-
-litDouble :: forall b. (Columnable b) => Expr b -> Maybe Double
-litDouble (Lit v) =
- case testEquality (typeRep @b) (typeRep @Double) of
- Just Refl -> Just v
- Nothing -> case testEquality (typeRep @b) (typeRep @Int) of
- Just Refl -> Just (fromIntegral v)
- Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Double)) of
- Just Refl -> v
- Nothing -> case testEquality (typeRep @b) (typeRep @(Maybe Int)) of
- Just Refl -> fromIntegral <$> v
- Nothing -> Nothing
-litDouble _ = Nothing
-
-{- | True for a column lifted from an integral type (never NaN): @toDouble (col …)@
-or a column whose type is itself integral.
--}
-integralColE :: forall c. (Columnable c) => Expr c -> Bool
-integralColE (Unary op _) = unaryName op == "toDouble"
-integralColE _ =
- or
- [ matches @Int
- , matches @(Maybe Int)
- ]
- where
- matches :: forall t. (Columnable t) => Bool
- matches = case testEquality (typeRep @c) (typeRep @t) of Just Refl -> True; _ -> False
-
-atomOf :: forall a. (Columnable a) => Expr a -> Maybe Atom
-atomOf (Unary fm (Binary (op :: op c b r) (colE :: Expr c) litE))
- | unaryName fm == "fromMaybe"
- , Just cmp <- cmpOf op
- , Just t <- litDouble litE =
- Just (Atom cmp t (show (normalize colE)) FalseOnNull (integralColE colE))
-atomOf (Binary (op :: op c b a) (colE :: Expr c) litE)
- | Just cmp <- cmpOf op
- , Just t <- litDouble litE =
- let nk = if isMaybeTy @c then UnknownOnNull else Total
- in Just (Atom cmp t (show (normalize colE)) nk (integralColE colE))
-atomOf _ = Nothing
-
-simplifyPredicatePair ::
- forall a. (Columnable a) => Bool -> Expr a -> Expr a -> Maybe (Expr a)
-simplifyPredicatePair isAnd a b = do
- atomA <- atomOf a
- atomB <- atomOf b
- guard (aKey atomA == aKey atomB)
- let nk = aNull atomA
- integral = aIntegral atomA
- if isAnd
- then andAtoms a atomA b atomB nk integral
- else orAtoms a atomA b atomB nk integral
-
--- | Contradiction folds to a literal False unless null-rows make it unknown.
-litFalseGated :: (Columnable a) => NullK -> Maybe (Expr a)
-litFalseGated UnknownOnNull = Nothing
-litFalseGated _ = litBoolish False
-
-{- | Tautology to literal True is sound only for total (never-null) atoms; the
-exhaustive-cover form additionally needs a non-NaN (integral) column.
--}
-litTrueTotal :: (Columnable a) => NullK -> Maybe (Expr a)
-litTrueTotal Total = litBoolish True
-litTrueTotal _ = Nothing
-
-andAtoms ::
- (Columnable a) =>
- Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)
-andAtoms a atomA b atomB nk _ =
- let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB
- in if
- | isLower cA, isLower cB, cA == cB -> Just (if tA >= tB then a else b)
- | isUpper cA, isUpper cB, cA == cB -> Just (if tA <= tB then a else b)
- | isLower cA, isUpper cB -> lu cA tA cB tB
- | isUpper cA, isLower cB -> lu cB tB cA tA
- | cA == CEq, cB == CEq -> if tA == tB then Just a else litFalseGated nk
- | cA == CEq, cB == CNeq -> if tA == tB then litFalseGated nk else Just a
- | cA == CNeq, cB == CEq -> if tA == tB then litFalseGated nk else Just b
- | cA == CEq -> if satisfies tA cB tB then Just a else litFalseGated nk
- | cB == CEq -> if satisfies tB cA tA then Just b else litFalseGated nk
- | cA == CNeq, cB == CNeq -> Nothing
- | cA == CNeq -> if outside tA cB tB then Just b else Nothing
- | cB == CNeq -> if outside tB cA tA then Just a else Nothing
- | otherwise -> Nothing
- where
- lu lc lo uc hi
- | lo > hi = litFalseGated nk
- | lo == hi, lc == CGeq, uc == CLeq = pointEq a lo
- | lo == hi = litFalseGated nk
- | otherwise = Nothing
-
-orAtoms ::
- (Columnable a) =>
- Expr a -> Atom -> Expr a -> Atom -> NullK -> Bool -> Maybe (Expr a)
-orAtoms a atomA b atomB nk integral =
- let cA = aCmp atomA; tA = aThr atomA; cB = aCmp atomB; tB = aThr atomB
- in if
- | isLower cA, isLower cB, cA == cB -> Just (if tA <= tB then a else b)
- | isUpper cA, isUpper cB, cA == cB -> Just (if tA >= tB then a else b)
- | isUpper cA
- , isLower cB
- , nk == Total
- , integral
- , covers cB tB cA tA ->
- litTrueTotal nk
- | isLower cA
- , isUpper cB
- , nk == Total
- , integral
- , covers cA tA cB tB ->
- litTrueTotal nk
- | cA == CNeq, cB == CNeq -> if tA == tB then Just a else litTrueTotal nk
- | cA == CEq, cB == CNeq -> if tA == tB then litTrueTotal nk else Just b
- | cA == CNeq, cB == CEq -> if tA == tB then litTrueTotal nk else Just a
- | cA == CEq, cB == CEq -> if tA == tB then Just a else Nothing
- | otherwise -> Nothing
-
-{- | Build @col == t@ for the point-collapse rule; only strict @Expr Bool@ over a
-@Double@ column (otherwise bail).
--}
-pointEq :: forall a. (Columnable a) => Expr a -> Double -> Maybe (Expr a)
-pointEq atom lo = case testEquality (typeRep @a) (typeRep @Bool) of
- Just Refl -> (\colE -> colE .==. Lit lo) <$> recoverColD atom
- Nothing -> Nothing
-
-recoverColD :: Expr x -> Maybe (Expr Double)
-recoverColD (Binary _ (colE :: Expr c) _) =
- case testEquality (typeRep @c) (typeRep @Double) of
- Just Refl -> Just colE
- _ -> Nothing
-recoverColD (Unary _ inner) = recoverColD inner
-recoverColD _ = Nothing
-
-covers :: Cmp -> Double -> Cmp -> Double -> Bool
-covers lowerCmp lo upperCmp hi =
- lo < hi || (lo == hi && (lowerCmp == CGeq || upperCmp == CLeq))
-
-satisfies :: Double -> Cmp -> Double -> Bool
-satisfies t CGt tb = t > tb
-satisfies t CGeq tb = t >= tb
-satisfies t CLt tb = t < tb
-satisfies t CLeq tb = t <= tb
-satisfies _ _ _ = False
-
-outside :: Double -> Cmp -> Double -> Bool
-outside t CGt tb = t <= tb
-outside t CGeq tb = t < tb
-outside t CLt tb = t >= tb
-outside t CLeq tb = t > tb
-outside _ _ _ = False
-
--- ---------------------------------------------------------------------------
--- Path-condition entailment for fitted-tree pruning.
--- ---------------------------------------------------------------------------
-
--- | A known same-column threshold fact accumulated along a tree path.
-data PredFact = PredFact !String !Cmp !Double
-
--- | The fact a branch's true edge establishes (the condition holds).
-factTrue :: Expr Bool -> Maybe PredFact
-factTrue e = (\a -> PredFact (aKey a) (aCmp a) (aThr a)) <$> atomOf e
-
-{- | The fact a branch's false edge establishes (the negated condition). Only
-sound for non-NaN (integral) columns — a NaN row takes the false edge too,
-so @¬(x>t)@ is not a clean @x<=t@ bound for floats.
--}
-factFalse :: Expr Bool -> Maybe PredFact
-factFalse e = do
- a <- atomOf e
- guard (aIntegral a && aNull a == Total)
- nc <- negCmp (aCmp a)
- pure (PredFact (aKey a) nc (aThr a))
-
-negCmp :: Cmp -> Maybe Cmp
-negCmp CLt = Just CGeq
-negCmp CLeq = Just CGt
-negCmp CGt = Just CLeq
-negCmp CGeq = Just CLt
-negCmp _ = Nothing
-
-{- | @entails facts cond@: 'Just' 'True' when the path facts force @cond@ true,
-'Just' 'False' when they force it false, 'Nothing' when undecided.
--}
-entails :: [PredFact] -> Expr Bool -> Maybe Bool
-entails facts cond = do
- a <- atomOf cond
- let decisions =
- [ d
- | PredFact fk fc ft <- facts
- , fk == aKey a
- , Just d <- [factImplies (fc, ft) (aCmp a, aThr a)]
- ]
- case decisions of
- (d : _) -> Just d
- [] -> Nothing
-
-{- | Does the fact's solution set sit inside @cond@ ('Just' 'True'), disjoint
-from it ('Just' 'False'), or neither ('Nothing')? Boundary strictness is
-honoured: e.g. @x<=t@ does NOT entail @x=t ∧ x<=t@ is not empty.
--}
-factImplies :: (Cmp, Double) -> (Cmp, Double) -> Maybe Bool
-factImplies (fc, ft) (cc, tc)
- | isLower fc, isLower cc, subset = Just True
- | isUpper fc, isUpper cc, subset = Just True
- | isLower fc, isUpper cc, disjointAtEq = Just False
- | isUpper fc, isLower cc, disjointBelow = Just False
- | otherwise = Nothing
- where
- fIncl = fc == CGeq || fc == CLeq
- cIncl = cc == CGeq || cc == CLeq
- -- same-direction containment: strictly tighter, or equal threshold where the
- -- fact's boundary inclusivity is no stronger than the condition's.
- subset =
- (if isLower fc then ft > tc else ft < tc)
- || (ft == tc && (not fIncl || cIncl))
- -- lower fact ∩ upper cond empty: fact starts above cond's top, or they meet
- -- at a point that is not in both.
- disjointAtEq = ft > tc || (ft == tc && not (fIncl && cIncl))
- -- upper fact ∩ lower cond empty (mirror).
- disjointBelow = ft < tc || (ft == tc && not (fIncl && cIncl))
diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Types.hs
+++ /dev/null
@@ -1,161 +0,0 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
-{-# LANGUAGE ConstraintKinds #-}
-{-# LANGUAGE DataKinds #-}
-{-# LANGUAGE DeriveTraversable #-}
-{-# LANGUAGE ExistentialQuantification #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE PolyKinds #-}
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE UndecidableInstances #-}
-
-module DataFrame.Internal.Types where
-
-import Data.Int (Int16, Int32, Int64, Int8)
-import Data.Kind (Constraint, Type)
-import Data.Typeable (Typeable)
-import qualified Data.Vector.Unboxed as VU
-import Data.Word (Word16, Word32, Word64, Word8)
-
-type Columnable' a = (Typeable a, Show a, Eq a)
-
-{- | Inline replacement for @Data.These.These@ to keep @dataframe-core@ free
-of the @these@ package dependency. Only the three constructors and the
-derived classes are used internally.
--}
-data These a b = This a | That b | These a b
- deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)
-
-{- | A type with column representations used to select the
-"right" representation when specializing the `toColumn` function.
--}
-data Rep
- = RBoxed
- | RUnboxed
- | RNullableBoxed
-
--- | Type-level if statement.
-type family If (cond :: Bool) (yes :: k) (no :: k) :: k where
- If 'True yes _ = yes
- If 'False _ no = no
-
--- | All unboxable types (according to the `vector` package).
-type family Unboxable (a :: Type) :: Bool where
- Unboxable Int = 'True
- Unboxable Int8 = 'True
- Unboxable Int16 = 'True
- Unboxable Int32 = 'True
- Unboxable Int64 = 'True
- Unboxable Word = 'True
- Unboxable Word8 = 'True
- Unboxable Word16 = 'True
- Unboxable Word32 = 'True
- Unboxable Word64 = 'True
- Unboxable Char = 'True
- Unboxable Bool = 'True
- Unboxable Double = 'True
- Unboxable Float = 'True
- Unboxable _ = 'False
-
-type family Numeric (a :: Type) :: Bool where
- Numeric Integer = 'True
- Numeric Int = 'True
- Numeric Int8 = 'True
- Numeric Int16 = 'True
- Numeric Int32 = 'True
- Numeric Int64 = 'True
- Numeric Word = 'True
- Numeric Word8 = 'True
- Numeric Word16 = 'True
- Numeric Word32 = 'True
- Numeric Word64 = 'True
- Numeric Double = 'True
- Numeric Float = 'True
- Numeric _ = 'False
-
--- | Compute the column representation tag for any 'a'.
-type family KindOf a :: Rep where
- KindOf (Maybe a) = 'RNullableBoxed
- KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed
-
--- | Type-level boolean for constraint/type comparison.
-data SBool (b :: Bool) where
- STrue :: SBool 'True
- SFalse :: SBool 'False
-
--- | The runtime witness for our type-level branching.
-class SBoolI (b :: Bool) where
- sbool :: SBool b
-
-instance SBoolI 'True where sbool = STrue
-instance SBoolI 'False where sbool = SFalse
-
--- | Type-level function to determine whether or not a type is unboxa
-sUnbox :: forall a. (SBoolI (Unboxable a)) => SBool (Unboxable a)
-sUnbox = sbool @(Unboxable a)
-
-sNumeric :: forall a. (SBoolI (Numeric a)) => SBool (Numeric a)
-sNumeric = sbool @(Numeric a)
-
-type family When (flag :: Bool) (c :: Constraint) :: Constraint where
- When 'True c = c
- When 'False c = () -- empty constraint
-
-type UnboxIf a = When (Unboxable a) (VU.Unbox a)
-
-type family IntegralTypes (a :: Type) :: Bool where
- IntegralTypes Integer = 'True
- IntegralTypes Int = 'True
- IntegralTypes Int8 = 'True
- IntegralTypes Int16 = 'True
- IntegralTypes Int32 = 'True
- IntegralTypes Int64 = 'True
- IntegralTypes Word = 'True
- IntegralTypes Word8 = 'True
- IntegralTypes Word16 = 'True
- IntegralTypes Word32 = 'True
- IntegralTypes Word64 = 'True
- IntegralTypes _ = 'False
-
-sIntegral :: forall a. (SBoolI (IntegralTypes a)) => SBool (IntegralTypes a)
-sIntegral = sbool @(IntegralTypes a)
-
-type IntegralIf a = When (IntegralTypes a) (Integral a)
-
-type family FloatingTypes (a :: Type) :: Bool where
- FloatingTypes Float = 'True
- FloatingTypes Double = 'True
- FloatingTypes _ = 'False
-
-sFloating :: forall a. (SBoolI (FloatingTypes a)) => SBool (FloatingTypes a)
-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/Internal/Utf8.hs b/src/DataFrame/Internal/Utf8.hs
deleted file mode 100644
--- a/src/DataFrame/Internal/Utf8.hs
+++ /dev/null
@@ -1,98 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-
-{- | UTF-8 validation and @decodeUtf8Lenient@-parity slice decoding used by
-'DataFrame.Internal.ColumnBuilder' to turn shared byte buffers into 'Text'.
--}
-module DataFrame.Internal.Utf8 (
- isValidUtf8Slice,
- isUtf8Boundary,
- lenientDecodeSlice,
- sliceTextVector,
-) where
-
-import qualified Data.Text as T
-import qualified Data.Text.Array as A
-import qualified Data.Vector as VB
-import qualified Data.Vector.Mutable as VBM
-import qualified Data.Vector.Unboxed as VU
-
-import Data.Text.Internal (Text (..))
-import Data.Text.Internal.Encoding.Utf8 (
- DecoderResult (..),
- utf8DecodeContinue,
- utf8DecodeStart,
- )
-import Data.Text.Internal.Validate (isValidUtf8ByteArray)
-import Data.Word (Word8)
-
--- | Whether @len@ bytes starting at @off@ are well-formed UTF-8.
-isValidUtf8Slice :: A.Array -> Int -> Int -> Bool
-isValidUtf8Slice = isValidUtf8ByteArray
-{-# INLINE isValidUtf8Slice #-}
-
-{- | Whether a byte may start a code point (i.e. is not a continuation
-byte). Field slices of a valid buffer are themselves valid iff every
-field starts on a boundary.
--}
-isUtf8Boundary :: Word8 -> Bool
-isUtf8Boundary w = w < 0x80 || w >= 0xC0
-{-# INLINE isUtf8Boundary #-}
-
-{- | Decode a byte slice exactly like @decodeUtf8Lenient@: greedy decode at
-each position; any byte that cannot begin a complete, valid sequence within
-the slice becomes one U+FFFD and decoding resumes at the next byte.
--}
-lenientDecodeSlice :: A.Array -> Int -> Int -> T.Text
-lenientDecodeSlice arr off len = T.pack (go off)
- where
- !end = off + len
- go !i
- | i >= end = []
- | otherwise = case tryDecode i of
- Just (c, i') -> c : go i'
- Nothing -> '\xFFFD' : go (i + 1)
- tryDecode !i = loop (utf8DecodeStart (A.unsafeIndex arr i)) (i + 1)
- where
- loop (Accept c) !j = Just (c, j)
- loop Reject _ = Nothing
- loop (Incomplete st cp) !j
- | j >= end = Nothing
- | otherwise = loop (utf8DecodeContinue (A.unsafeIndex arr j) st cp) (j + 1)
-
-{- | Slice forced 'Text' values off a shared array; row @i@ spans bytes
-@[offs!i, offs!(i+1))@. The offsets need not start at byte 0, so a row
-sub-range of a larger offset vector slices independently (parallel text
-merging uses this). Fast path: validate the spanned bytes once and check
-every field starts on a code-point boundary. Slow path: per-field
-validation with lenient decoding of invalid fields.
--}
-sliceTextVector :: A.Array -> VU.Vector Int -> VB.Vector T.Text
-sliceTextVector arr offs = VB.create $ do
- mv <- VBM.unsafeNew n
- let fill dec = go 0
- where
- go !i
- | i >= n = pure ()
- | otherwise = do
- let o = VU.unsafeIndex offs i
- !t = dec o (VU.unsafeIndex offs (i + 1) - o)
- VBM.unsafeWrite mv i t
- go (i + 1)
- if fast then fill mkSlice else fill decodeField
- pure mv
- where
- n = VU.length offs - 1
- base = VU.unsafeIndex offs 0
- used = VU.unsafeIndex offs n
- boundariesOk !i
- | i >= n = True
- | otherwise =
- let o = VU.unsafeIndex offs i
- in (o >= used || isUtf8Boundary (A.unsafeIndex arr o))
- && boundariesOk (i + 1)
- fast = isValidUtf8Slice arr base (used - base) && boundariesOk 0
- mkSlice o l = if l == 0 then T.empty else Text arr o l
- decodeField o l
- | l == 0 = T.empty
- | isValidUtf8Slice arr o l = Text arr o l
- | otherwise = lenientDecodeSlice arr o l
diff --git a/src/DataFrame/Operators.hs b/src/DataFrame/Operators.hs
deleted file mode 100644
--- a/src/DataFrame/Operators.hs
+++ /dev/null
@@ -1,425 +0,0 @@
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE TypeOperators #-}
-
-module DataFrame.Operators where
-
-import Data.Function ((&))
-import qualified Data.Text as T
-import DataFrame.Internal.Column (Columnable)
-import DataFrame.Internal.Expression (
- BinUDF (MkBinaryOp),
- BinaryOp (
- binaryCommutative,
- binaryFn,
- binaryName,
- binaryPrecedence,
- binarySymbol
- ),
- Expr (Binary, Col, If, Lit, Unary),
- NamedExpr,
- UExpr (UExpr),
- UnUDF (MkUnaryOp),
- )
-import DataFrame.Internal.Nullable (
- BaseType,
- DivWidenOp,
- NullCmpResult,
- NullLift2Op (applyNull2),
- NullableCmpOp (nullCmpOp),
- NumericWidenOp,
- WidenResult,
- WidenResultDiv,
- divArithOp,
- widenArithOp,
- widenCmpOp,
- )
-import DataFrame.Internal.Types (Promote, PromoteDiv)
-
-infixr 8 .^^, .^^., .^, .^.
-infixl 7 .*, ./, .*., ./.
-infixl 6 .+, .-, .+., .-.
-infix 4 .==, .==., .<, .<., .<=, .<=., .>=, .>=., .>, .>., ./=, ./=.
-infixr 3 .&&, .&&.
-infixr 2 .||, .||.
-infixr 0 .=
-
-(|>) :: a -> (a -> b) -> b
-(|>) = (&)
-
-as :: (Columnable a) => Expr a -> T.Text -> NamedExpr
-as expr colName = (colName, UExpr expr)
-
-name :: (Show a) => Expr a -> T.Text
-name (Col n) = n
-name other =
- error $
- "You must call `name` on a column reference. Not the expression: " ++ show other
-
-col :: (Columnable a) => T.Text -> Expr a
-col = Col
-
-ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
-ifThenElse = If
-
-lit :: (Columnable a) => a -> Expr a
-lit = Lit
-
-(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr
-(.=) = flip as
-
-liftDecorated ::
- (Columnable a, Columnable b) =>
- (a -> b) -> T.Text -> Maybe T.Text -> Expr a -> Expr b
-liftDecorated f opName rep = Unary (MkUnaryOp f opName rep)
-
-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 opName rep comm prec =
- Binary (MkBinaryOp f opName rep comm prec)
-
-data NullEq a b c where
- NullEq ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Eq (Promote (BaseType a) (BaseType b))
- ) =>
- NullEq a b (NullCmpResult a b)
-
-data NullNeq a b c where
- NullNeq ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Eq (Promote (BaseType a) (BaseType b))
- ) =>
- NullNeq a b (NullCmpResult a b)
-
-data NullLt a b c where
- NullLt ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- NullLt a b (NullCmpResult a b)
-
-data NullGt a b c where
- NullGt ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- NullGt a b (NullCmpResult a b)
-
-data NullLeq a b c where
- NullLeq ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- NullLeq a b (NullCmpResult a b)
-
-data NullGeq a b c where
- NullGeq ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- NullGeq a b (NullCmpResult a b)
-
-data NullAnd a b c where
- NullAnd ::
- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
- NullAnd a b (NullCmpResult a b)
-
-data NullOr a b c where
- NullOr ::
- (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>
- NullOr a b (NullCmpResult a b)
-
-instance BinaryOp NullEq where
- binaryFn NullEq = applyNull2 (widenCmpOp (==))
- binaryName NullEq = "eq"
- binarySymbol NullEq = Just ".=="
- binaryCommutative NullEq = True
- binaryPrecedence NullEq = 4
-instance BinaryOp NullNeq where
- binaryFn NullNeq = applyNull2 (widenCmpOp (/=))
- binaryName NullNeq = "neq"
- binarySymbol NullNeq = Just "./="
- binaryCommutative NullNeq = True
- binaryPrecedence NullNeq = 4
-instance BinaryOp NullLt where
- binaryFn NullLt = applyNull2 (widenCmpOp (<))
- binaryName NullLt = "lt"
- binarySymbol NullLt = Just ".<"
- binaryPrecedence NullLt = 4
-instance BinaryOp NullGt where
- binaryFn NullGt = applyNull2 (widenCmpOp (>))
- binaryName NullGt = "gt"
- binarySymbol NullGt = Just ".>"
- binaryPrecedence NullGt = 4
-instance BinaryOp NullLeq where
- binaryFn NullLeq = applyNull2 (widenCmpOp (<=))
- binaryName NullLeq = "leq"
- binarySymbol NullLeq = Just ".<="
- binaryPrecedence NullLeq = 4
-instance BinaryOp NullGeq where
- binaryFn NullGeq = applyNull2 (widenCmpOp (>=))
- binaryName NullGeq = "geq"
- binarySymbol NullGeq = Just ".>="
- binaryPrecedence NullGeq = 4
-instance BinaryOp NullAnd where
- binaryFn NullAnd = nullCmpOp (&&)
- binaryName NullAnd = "nulland"
- binarySymbol NullAnd = Just ".&&"
- binaryCommutative NullAnd = True
- binaryPrecedence NullAnd = 3
-instance BinaryOp NullOr where
- binaryFn NullOr = nullCmpOp (||)
- binaryName NullOr = "nullor"
- binarySymbol NullOr = Just ".||"
- binaryCommutative NullOr = True
- binaryPrecedence NullOr = 2
-
-(.==.) ::
- (Columnable a, Eq a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(.==.) = lift2Decorated (==) "eq" (Just ".==.") True 4
-
-(./=.) ::
- (Columnable a, Eq a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(./=.) = lift2Decorated (/=) "neq" (Just "./=.") True 4
-
-(.<.) ::
- (Columnable a, Ord a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(.<.) = lift2Decorated (<) "lt" (Just ".<.") False 4
-
-(.>.) ::
- (Columnable a, Ord a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(.>.) = lift2Decorated (>) "gt" (Just ".>.") False 4
-
-(.<=.) ::
- (Columnable a, Ord a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(.<=.) = lift2Decorated (<=) "leq" (Just ".<=.") False 4
-
-(.>=.) ::
- (Columnable a, Ord a) =>
- Expr a ->
- Expr a ->
- Expr Bool
-(.>=.) = lift2Decorated (>=) "geq" (Just ".>=.") False 4
-
-(.+.) :: (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. Widens numeric operands to their common type,
-so @Expr Double .== Expr Int@ typechecks. Returns @Maybe Bool@ when either
-operand is nullable.
--}
-(.==) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Eq (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(.==) = Binary NullEq
-
--- | Nullable-aware inequality. Widens numeric operands to their common type.
-(./=) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Eq (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(./=) = Binary NullNeq
-
--- | Nullable-aware less-than. Widens numeric operands to their common type.
-(.<) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(.<) = Binary NullLt
-
--- | Nullable-aware greater-than. Widens numeric operands to their common type.
-(.>) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(.>) = Binary NullGt
-
-{- | Nullable-aware less-than-or-equal. Widens numeric operands to their
-common type, so @Expr Double .<= Expr Int@ typechecks.
--}
-(.<=) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(.<=) = Binary NullLeq
-
--- | Nullable-aware greater-than-or-equal. Widens numeric operands to their common type.
-(.>=) ::
- ( NumericWidenOp (BaseType a) (BaseType b)
- , NullLift2Op a b Bool (NullCmpResult a b)
- , Ord (Promote (BaseType a) (BaseType b))
- ) =>
- Expr a ->
- Expr b ->
- Expr (NullCmpResult a b)
-(.>=) = Binary NullGeq
-
-(.&&.) :: 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)
-(.&&) = Binary NullAnd
-
--- | 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)
-(.||) = Binary NullOr
-
-(.^^) ::
- ( 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/Typed/Freeze.hs b/src/DataFrame/Typed/Freeze.hs
--- a/src/DataFrame/Typed/Freeze.hs
+++ b/src/DataFrame/Typed/Freeze.hs
@@ -9,16 +9,22 @@
-- * Safe boundary
freeze,
freezeWithError,
+ freezeOrThrow,
-- * Escape hatches
thaw,
unsafeFreeze,
+
+ -- * Frame coercion
+ ToDataFrame (..),
) where
+import Control.Exception (throwIO)
import qualified Data.Text as T
import Type.Reflection (SomeTypeRep)
import Data.List (stripPrefix)
+import DataFrame.Errors (DataFrameException (InternalException))
import qualified DataFrame.Internal.Column as C
import DataFrame.Internal.DataFrame (columnNames)
import qualified DataFrame.Internal.DataFrame as D
@@ -43,11 +49,28 @@
Left err -> Left err
Right _ -> Right (TDF df)
+{- | Validate and wrap like 'freezeWithError', but throw a 'DataFrameException'
+in 'IO' on mismatch. The throwing boundary used by the typed readers
+(@readCsv@ \/ @readParquet@).
+-}
+freezeOrThrow ::
+ forall cols. (KnownSchema cols) => D.DataFrame -> IO (TypedDataFrame cols)
+freezeOrThrow = either (throwIO . InternalException) pure . freezeWithError @cols
+
{- | Unwrap a typed DataFrame back to the untyped representation.
Always safe; discards type information.
-}
thaw :: TypedDataFrame cols -> D.DataFrame
thaw (TDF df) = df
+
+class ToDataFrame f where
+ toDataFrame :: f -> D.DataFrame
+
+instance ToDataFrame D.DataFrame where
+ toDataFrame = id
+
+instance ToDataFrame (TypedDataFrame cols) where
+ toDataFrame = thaw
{- | Wrap an untyped DataFrame without any validation.
Used internally after delegation where the library guarantees schema correctness.
diff --git a/src/DataFrame/Typed/Schema.hs b/src/DataFrame/Typed/Schema.hs
--- a/src/DataFrame/Typed/Schema.hs
+++ b/src/DataFrame/Typed/Schema.hs
@@ -20,6 +20,7 @@
HasName,
RemoveColumn,
Impute,
+ SetColumnType,
SubsetSchema,
ExcludeSchema,
RenameInSchema,
@@ -32,6 +33,12 @@
AssertPresent,
AssertAllPresent,
AssertKeyTypesMatch,
+ AssertDisjoint,
+ AssertAllColumnsHaveType,
+ AssertRealColumn,
+ AllColumnsReal,
+ AllDouble,
+ IsRealType,
IsElem,
-- * Maybe-stripping families
@@ -60,9 +67,12 @@
AllKnownSymbol (..),
) where
+import Data.Int (Int16, Int32, Int64, Int8)
import Data.Kind (Constraint, Type)
import Data.Proxy (Proxy (..))
import qualified Data.Text as T
+import qualified Data.Vector.Unboxed as VU
+import Data.Word (Word16, Word32, Word64, Word8)
import GHC.TypeLits
import Type.Reflection (SomeTypeRep, Typeable, someTypeRep)
@@ -96,6 +106,13 @@
Impute name (col ': rest) = col ': Impute name rest
Impute name '[] = '[]
+type family SetColumnType (name :: Symbol) (b :: Type) (cols :: [Type]) :: [Type] where
+ SetColumnType name b (Column name _ ': rest) = Column name b ': rest
+ SetColumnType name b (col ': rest) = col ': SetColumnType name b rest
+ SetColumnType name b '[] =
+ TypeError
+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")
+
-- | Add type to the end of a list.
type family Snoc (xs :: [k]) (x :: k) :: [k] where
Snoc '[] x = '[x]
@@ -256,6 +273,94 @@
':<>: 'Text " in the left table but "
':<>: 'ShowType r
':<>: 'Text " in the right table"
+ )
+
+type family AssertDisjoint (left :: [Type]) (right :: [Type]) :: Constraint where
+ AssertDisjoint left right =
+ AssertDisjointHelper (SharedNames left right) left right
+
+type family
+ AssertDisjointHelper (shared :: [Symbol]) (left :: [Type]) (right :: [Type]) ::
+ Constraint
+ where
+ AssertDisjointHelper '[] left right = ()
+ AssertDisjointHelper (n ': ns) left right =
+ TypeError
+ ( 'Text "Cannot horizontally merge: column '"
+ ':<>: 'Text n
+ ':<>: 'Text "' appears in both schemas"
+ )
+
+type family
+ AssertAllColumnsHaveType (names :: [Symbol]) (a :: Type) (cols :: [Type]) ::
+ Constraint
+ where
+ AssertAllColumnsHaveType '[] a cols = ()
+ AssertAllColumnsHaveType (n ': ns) a cols =
+ ( SafeLookup n cols ~ a
+ , AssertPresent n cols
+ , AssertAllColumnsHaveType ns a cols
+ )
+
+-- | Is @a@ a real, unboxed numeric type — i.e. a valid numeric-column element?
+type family IsRealType (a :: Type) :: Bool where
+ IsRealType Int = 'True
+ IsRealType Int8 = 'True
+ IsRealType Int16 = 'True
+ IsRealType Int32 = 'True
+ IsRealType Int64 = 'True
+ IsRealType Word = 'True
+ IsRealType Word8 = 'True
+ IsRealType Word16 = 'True
+ IsRealType Word32 = 'True
+ IsRealType Word64 = 'True
+ IsRealType Double = 'True
+ IsRealType Float = 'True
+ IsRealType _ = 'False
+
+{- | Emit a readable compile error when the column named @name@ is not a
+real-number type, naming the calling function @fn@, the column, and the type it
+actually has. Used by the numeric extractors so a wrong column type reads as a
+repairable message rather than a bare @No instance for Real …@.
+-}
+type family AssertRealColumn (fn :: Symbol) (name :: Symbol) (a :: Type) :: Constraint where
+ AssertRealColumn fn name a = AssertRealColumnGo fn name a (IsRealType a)
+
+type family
+ AssertRealColumnGo (fn :: Symbol) (name :: Symbol) (a :: Type) (isReal :: Bool) ::
+ Constraint
+ where
+ AssertRealColumnGo fn name a 'True = ()
+ AssertRealColumnGo fn name a 'False =
+ TypeError
+ ( 'Text fn
+ ':<>: 'Text ": expected a real number column for '"
+ ':<>: 'Text name
+ ':<>: 'Text "' but instead you gave "
+ ':<>: 'ShowType a
+ )
+
+{- | Constraint that every column in the schema is a real (numeric), unboxed
+type. Lets the whole-frame matrix extractors ('toDoubleMatrix' and friends) be
+total — a non-numeric or nullable column is a compile error (with the offending
+column named, via 'AssertRealColumn'), not a runtime 'Left'.
+-}
+type family AllColumnsReal (fn :: Symbol) (cols :: [Type]) :: Constraint where
+ AllColumnsReal fn '[] = ()
+ AllColumnsReal fn (Column n a ': rest) =
+ (AssertRealColumn fn n a, Real a, VU.Unbox a, AllColumnsReal fn rest)
+
+-- TODO: mchavinda - we can generalist to AllX
+type family AllDouble (cols :: [Type]) :: Constraint where
+ AllDouble '[] = ()
+ AllDouble (Column n Double ': rest) = AllDouble rest
+ AllDouble (Column n a ': rest) =
+ TypeError
+ ( 'Text "Column '"
+ ':<>: 'Text n
+ ':<>: 'Text "' must be Double for this model, but is "
+ ':<>: 'ShowType a
+ ':$$: 'Text "Convert it (toDouble) or drop it before fitting."
)
{- | Strip 'Maybe' from all columns. Used by 'filterAllJust'.
diff --git a/src/DataFrame/Typed/Types.hs b/src/DataFrame/Typed/Types.hs
--- a/src/DataFrame/Typed/Types.hs
+++ b/src/DataFrame/Typed/Types.hs
@@ -2,10 +2,12 @@
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
-{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module DataFrame.Typed.Types (
@@ -18,6 +20,10 @@
-- * Typed expressions (schema-validated)
TExpr (..),
+ -- * Lifting untyped expressions into the typed world
+ AsTExpr,
+ ToTExpr (..),
+
-- * Typed sort orders
TSortOrder (..),
@@ -70,6 +76,24 @@
Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API.
-}
newtype TExpr (cols :: [Type]) a = TExpr {unTExpr :: Expr a}
+
+-- | Shows the underlying expression; the schema phantom is type-level only.
+instance (Show a) => Show (TExpr cols a) where
+ showsPrec d (TExpr e) = showsPrec d e
+
+{- | The typed counterpart of an untyped expression type for schema @cols@:
+@AsTExpr cols (Expr r) = TExpr cols r@. Lets a result type follow the frame —
+an @Expr@ over a plain frame becomes a @TExpr@ over a typed one.
+-}
+type family AsTExpr (cols :: [Type]) (e :: Type) :: Type where
+ AsTExpr cols (Expr r) = TExpr cols r
+
+-- | Lift an untyped expression into its 'TExpr' for schema @cols@.
+class ToTExpr (cols :: [Type]) e where
+ toTExpr :: e -> AsTExpr cols e
+
+instance ToTExpr cols (Expr r) where
+ toTExpr = TExpr
-- | A typed sort order validated against schema @cols@.
data TSortOrder (cols :: [Type]) where