dataframe-core (empty) → 1.0.0.0
raw patch · 21 files changed
+6332/−0 lines, 21 filesdep +basedep +bytestringdep +containers
Dependencies added: base, bytestring, containers, random, text, time, vector
Files
- LICENSE +20/−0
- dataframe-core.cabal +61/−0
- src/DataFrame/Display/Terminal/Colours.hs +14/−0
- src/DataFrame/Display/Terminal/PrettyPrint.hs +104/−0
- src/DataFrame/Errors.hs +187/−0
- src/DataFrame/Internal/Column.hs +1707/−0
- src/DataFrame/Internal/DataFrame.hs +371/−0
- src/DataFrame/Internal/Expression.hs +397/−0
- src/DataFrame/Internal/Grouping.hs +194/−0
- src/DataFrame/Internal/Hash.hs +62/−0
- src/DataFrame/Internal/Interpreter.hs +1064/−0
- src/DataFrame/Internal/Nullable.hs +500/−0
- src/DataFrame/Internal/Row.hs +175/−0
- src/DataFrame/Internal/Types.hs +161/−0
- src/DataFrame/Operators.hs +329/−0
- src/DataFrame/Typed/Freeze.hs +98/−0
- src/DataFrame/Typed/Generic.hs +205/−0
- src/DataFrame/Typed/Record.hs +105/−0
- src/DataFrame/Typed/Schema.hs +441/−0
- src/DataFrame/Typed/Types.hs +114/−0
- src/DataFrame/Typed/Util.hs +23/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2026 Michael Chavinda++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ dataframe-core.cabal view
@@ -0,0 +1,61 @@+cabal-version: 2.4+name: dataframe-core+version: 1.0.0.0++synopsis: Core data structures for the dataframe library.+description:+ Minimal interchange-format types for the @dataframe@ ecosystem:+ 'Column', 'DataFrame', 'Bitmap', the untyped expression/interpreter,+ and the typed-schema phantom layer. Contains no Template Haskell and+ no file I/O. Lightweight dependency footprint (base, vector,+ containers, time, random, bytestring, text) so other packages can+ exchange dataframes by-value without pulling in the full+ @dataframe@ package.++bug-reports: https://github.com/mchav/dataframe/issues+license: MIT+license-file: LICENSE+author: Michael Chavinda+maintainer: mschavinda@gmail.com+copyright: (c) 2024-2025 Michael Chavinda+category: Data+tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2++common warnings+ ghc-options:+ -Wincomplete-patterns+ -Wincomplete-uni-patterns+ -Wunused-imports+ -Wunused-local-binds++library+ import: warnings+ exposed-modules:+ DataFrame.Errors+ DataFrame.Operators+ DataFrame.Display.Terminal.Colours+ DataFrame.Display.Terminal.PrettyPrint+ DataFrame.Internal.Column+ DataFrame.Internal.DataFrame+ DataFrame.Internal.Expression+ DataFrame.Internal.Grouping+ DataFrame.Internal.Hash+ DataFrame.Internal.Interpreter+ DataFrame.Internal.Nullable+ DataFrame.Internal.Row+ DataFrame.Internal.Types+ DataFrame.Typed.Freeze+ DataFrame.Typed.Generic+ DataFrame.Typed.Record+ DataFrame.Typed.Schema+ DataFrame.Typed.Types+ DataFrame.Typed.Util+ build-depends: base >= 4 && < 5,+ bytestring >= 0.11 && < 0.13,+ containers >= 0.6.7 && < 0.9,+ random >= 1 && < 2,+ text >= 2.0 && < 3,+ time >= 1.12 && < 2,+ vector ^>= 0.13+ hs-source-dirs: src+ default-language: Haskell2010
+ src/DataFrame/Display/Terminal/Colours.hs view
@@ -0,0 +1,14 @@+module DataFrame.Display.Terminal.Colours where++-- terminal color functions+red :: String -> String+red s = "\ESC[31m" ++ s ++ "\ESC[0m"++green :: String -> String+green s = "\ESC[32m" ++ s ++ "\ESC[0m"++brightGreen :: String -> String+brightGreen s = "\ESC[92m" ++ s ++ "\ESC[0m"++brightBlue :: String -> String+brightBlue s = "\ESC[94m" ++ s ++ "\ESC[0m"
+ src/DataFrame/Display/Terminal/PrettyPrint.hs view
@@ -0,0 +1,104 @@+{-# 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+ consolidatedHeader =+ if isMarkdown+ then zipWith (\h t -> h <> "<br>" <> t) header types+ else header+ cs = map (\h -> ColDesc center h left) consolidatedHeader+ nRows = case columns 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+ types+ columns+ 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) columns+ 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 types+ : separator+ : rowLines+ in T.unlines outputLines
+ src/DataFrame/Errors.hs view
@@ -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)+ ]
+ src/DataFrame/Internal/Column.hs view
@@ -0,0 +1,1707 @@+{-# 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.Types+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++{- | 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 _ _ = False++-- | Return the 'Maybe Bitmap' from a column.+columnBitmap :: Column -> Maybe Bitmap+columnBitmap (BoxedColumn bm _) = bm+columnBitmap (UnboxedColumn bm _) = bm++-- ---------------------------------------------------------------------------+-- 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 _ = 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 _ = 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++{- | 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)+ 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"++{- | 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+ 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` ()++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 ++ "]"++{- | 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+ )+ (==) _ _ = False++{- | 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+ 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+ 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+{-# 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+{-# 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)+{-# 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)+{-# 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)+{-# 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+ 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+ 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+ 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+ 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+ 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+ 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+ 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+ 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 (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+ (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++-- | 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 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 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+ (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 (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)++concatColumnsEither :: Column -> Column -> Column+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))++-- | 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 _ _ _ =+ 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+ 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+ 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+ 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+ 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))
+ src/DataFrame/Internal/DataFrame.hs view
@@ -0,0 +1,371 @@+{-# 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 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++ -- 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 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)++stripJust :: T.Text -> T.Text+stripJust = fromMaybe "null" . T.stripPrefix "Just "
+ src/DataFrame/Internal/Expression.hs view
@@ -0,0 +1,397 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Internal.Expression where++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)++data UnaryOp a b = MkUnaryOp+ { unaryFn :: a -> b+ , unaryName :: T.Text+ , unarySymbol :: Maybe T.Text+ }++data BinaryOp a b c = MkBinaryOp+ { binaryFn :: a -> b -> c+ , binaryName :: T.Text+ , binarySymbol :: Maybe T.Text+ , binaryCommutative :: Bool+ , binaryPrecedence :: Int+ }++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 ::+ (Columnable a, Columnable b) => UnaryOp b a -> Expr b -> Expr a+ Binary ::+ (Columnable c, Columnable b, Columnable a) =>+ BinaryOp 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)++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) ++ ")"
+ src/DataFrame/Internal/Grouping.hs view
@@ -0,0 +1,194 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Internal.Grouping (+ groupBy,+ buildRowToGroup,+ changingPoints,+) where++import qualified Data.IntMap.Strict as IM+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 (runST)+import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))+import DataFrame.Errors+import DataFrame.Internal.Column (+ Column (..),+ bitmapTestBit,+ )+import DataFrame.Internal.DataFrame (DataFrame (..), GroupedDataFrame (..))+import DataFrame.Internal.Hash+import DataFrame.Internal.Types+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.+-}+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+ | otherwise =+ let !vis = VU.map fst valIndices+ !os = changingPoints valIndices+ !n = nRows df+ in Grouped+ df+ names+ vis+ os+ (buildRowToGroup n vis os)+ where+ indicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` names) (columnIndices df)+ -- valIndices is a vector of (rowIdx, rowHash) sorted by rowHash so that+ -- runs of equal-hash rows are adjacent. Hashes are computed with the+ -- in-tree FNV combinators from "DataFrame.Internal.Hash", and bucketed+ -- with 'Data.IntMap.Strict' to keep the dep set minimal (no hashable,+ -- no vector-algorithms).+ valIndices = runST $ do+ let n = nRows df+ mh <- VUM.replicate n fnvOffset+ let selectedCols = map (columns df V.!) indicesToGroup+ forM_ selectedCols $ \case+ UnboxedColumn _ (v :: VU.Vector a) ->+ case testEquality (typeRep @a) (typeRep @Int) of+ Just Refl ->+ VU.imapM_+ ( \i x -> do+ !h <- VUM.unsafeRead mh i+ VUM.unsafeWrite mh i (mixInt h x)+ )+ v+ Nothing ->+ case testEquality (typeRep @a) (typeRep @Double) of+ Just Refl ->+ VU.imapM_+ ( \i d -> do+ !h <- VUM.unsafeRead mh i+ VUM.unsafeWrite mh i (mixDouble h d)+ )+ v+ Nothing ->+ case sIntegral @a of+ STrue ->+ VU.imapM_+ ( \i d -> do+ !h <- VUM.unsafeRead mh i+ VUM.unsafeWrite mh i (mixInt h (fromIntegral @a @Int d))+ )+ v+ SFalse ->+ case sFloating @a of+ STrue ->+ VU.imapM_+ ( \i d -> do+ !h <- VUM.unsafeRead mh i+ VUM.unsafeWrite mh i (mixDouble h (realToFrac d :: Double))+ )+ v+ SFalse ->+ VU.imapM_+ ( \i d -> do+ !h <- VUM.unsafeRead mh i+ VUM.unsafeWrite mh i (mixShow h d)+ )+ 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 0 -- null sentinel+ _ -> 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 0 -- null sentinel+ _ -> mixShow h d+ VUM.unsafeWrite mh i h'+ )+ v+ hashes <- VU.unsafeFreeze mh+ -- Bucket row indices by hash using an IntMap, then walk it in+ -- ascending key order to emit (rowIdx, hash) pairs grouped by+ -- hash. Each bucket's accumulated list is reversed so rows come+ -- out in the original row order.+ let buckets =+ VU.ifoldl'+ (\acc i h -> IM.insertWith (++) h [i] acc)+ IM.empty+ hashes+ ordered =+ [ (i, h)+ | (h, is) <- IM.toAscList buckets+ , i <- reverse is+ ]+ return (VU.fromList ordered)++-- 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)
+ src/DataFrame/Internal/Hash.hs view
@@ -0,0 +1,62 @@+{- |+A poor-man's hash used by 'DataFrame.Internal.Grouping' to bucket rows+without depending on the @hashable@ package.++The hash is FNV-1a-shaped: an accumulator is repeatedly @xor@ed with the+next chunk and multiplied by an FNV prime. 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,+ mixInt,+ mixDouble,+ mixBool,+ mixChar,+ mixText,+ mixShow,+) where++import Data.Bits (xor)+import Data.Char (ord)+import qualified Data.Text as T++-- | FNV-1a 64-bit offset basis (used as the initial accumulator).+fnvOffset :: Int+fnvOffset = 0xcbf29ce484222325++-- | FNV-1a 64-bit prime.+fnvPrime :: Int+fnvPrime = 0x00000100000001b3++-- | Mix an 'Int' into the accumulator.+mixInt :: Int -> Int -> Int+mixInt acc x = (acc `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, byte by byte.+mixText :: Int -> T.Text -> Int+mixText = T.foldl' (\a c -> mixInt a (ord c))+{-# INLINE mixText #-}++{- | 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 #-}
+ src/DataFrame/Internal/Interpreter.hs view
@@ -0,0 +1,1064 @@+{-# 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+ #-}++-------------------------------------------------------------------------------+-- 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++{- | 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"++-- | 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"++{- | 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+ 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++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++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++-- | 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+ 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 :: UnaryOp b a) inner) = addContext expr $ do+ v <- eval @b ctx inner+ liftValue (unaryFn op) v++-- Binary -----------------------------------------------------------------++eval ctx expr@(Binary (op :: BinaryOp c b a) left right) =+ 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
+ src/DataFrame/Internal/Nullable.hs view
@@ -0,0 +1,500 @@+{-# 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))
+ src/DataFrame/Internal/Row.hs view
@@ -0,0 +1,175 @@+{-# 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 (fromMaybe)+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 Type.Reflection (typeOf, typeRep)++data Any where+ Value :: (Columnable a) => a -> Any++instance Eq Any where+ (==) :: Any -> Any -> Bool+ (Value a) == (Value b) = fromMaybe False $ do+ Refl <- testEquality (typeOf a) (typeOf b)+ return $ a == b++instance Show Any where+ show :: Any -> String+ show (Value a) = T.unpack (showValue a)++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.+fromAny :: forall a. (Columnable a) => Any -> Maybe a+fromAny (Value (v :: b)) = do+ Refl <- testEquality (typeRep @a) (typeRep @b)+ pure v++type Row = V.Vector Any++(!?) :: [a] -> Int -> Maybe a+(!?) [] _ = Nothing+(!?) (x : _) 0 = Just x+(!?) (_x : xs) n = (!?) xs (n - 1)++mkColumnFromRow :: Int -> [[Any]] -> Column+mkColumnFromRow i rows = case rows of+ [] -> fromList ([] :: [T.Text])+ (row : _) -> case row !? i of+ Nothing -> fromList ([] :: [T.Text])+ Just (Value (v :: a)) -> fromList $ reverse $ L.foldl' addToList [v] (drop 1 rows)+ where+ addToList acc r = case r !? i of+ Nothing -> acc+ Just (Value (v' :: b)) -> case testEquality (typeRep @a) (typeRep @b) of+ Nothing -> acc+ Just Refl -> v' : acc++{- | 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 _ column) -> toAny (column V.! i)+ Just (UnboxedColumn _ column) -> toAny (column VU.! 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 _ c) -> case c V.!? i of+ Just e -> toAny e+ Nothing -> throwError name+ Just (UnboxedColumn _ c) -> case c VU.!? i of+ Just e -> toAny e+ Nothing -> throwError name+ Nothing ->+ throw $ ColumnsNotFoundException [name] "mkRowRep" (M.keys $ columnIndices df)
+ src/DataFrame/Internal/Types.hs view
@@ -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++-- | 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
+ src/DataFrame/Operators.hs view
@@ -0,0 +1,329 @@+{-# LANGUAGE FlexibleContexts #-}+{-# 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 (+ BinaryOp (+ MkBinaryOp,+ binaryCommutative,+ binaryFn,+ binaryName,+ binaryPrecedence,+ binarySymbol+ ),+ Expr (Binary, Col, If, Lit, Unary),+ NamedExpr,+ UExpr (UExpr),+ UnaryOp (MkUnaryOp, unaryFn, unaryName, unarySymbol),+ )+import DataFrame.Internal.Nullable (+ BaseType,+ DivWidenOp,+ NullCmpResult,+ NullLift2Op (applyNull2),+ NullableCmpOp (nullCmpOp),+ NumericWidenOp,+ WidenResult,+ WidenResultDiv,+ divArithOp,+ widenArithOp,+ 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{unaryFn = f, unaryName = opName, unarySymbol = 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+ { binaryFn = f+ , binaryName = opName+ , binarySymbol = rep+ , binaryCommutative = comm+ , binaryPrecedence = prec+ }+ )++(.==.) ::+ (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)+(.==) = lift2Decorated (applyNull2 (widenCmpOp (==))) "eq" (Just ".==") True 4++-- | 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)+(./=) = lift2Decorated (applyNull2 (widenCmpOp (/=))) "neq" (Just "./=") True 4++-- | 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)+(.<) = lift2Decorated (applyNull2 (widenCmpOp (<))) "lt" (Just ".<") False 4++-- | 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)+(.>) = lift2Decorated (applyNull2 (widenCmpOp (>))) "gt" (Just ".>") False 4++{- | 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)+(.<=) = lift2Decorated (applyNull2 (widenCmpOp (<=))) "leq" (Just ".<=") False 4++-- | 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)+(.>=) = lift2Decorated (applyNull2 (widenCmpOp (>=))) "geq" (Just ".>=") False 4++(.&&.) :: Expr Bool -> Expr Bool -> Expr Bool+(.&&.) = lift2Decorated (&&) "and" (Just ".&&.") True 3++(.||.) :: Expr Bool -> Expr Bool -> Expr Bool+(.||.) = lift2Decorated (||) "or" (Just ".||.") True 2++-- | Nullable-aware logical AND. Returns @Maybe Bool@ when either operand is nullable.+(.&&) ::+ (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+ Expr a ->+ Expr b ->+ Expr (NullCmpResult a b)+(.&&) = lift2Decorated (nullCmpOp (&&)) "nulland" (Just ".&&") True 3++-- | Nullable-aware logical OR. Returns @Maybe Bool@ when either operand is nullable.+(.||) ::+ (NullableCmpOp a b (NullCmpResult a b), BaseType a ~ Bool) =>+ Expr a ->+ Expr b ->+ Expr (NullCmpResult a b)+(.||) = lift2Decorated (nullCmpOp (||)) "nullor" (Just ".||") True 2++(.^^) ::+ ( Columnable (BaseType a)+ , Columnable (BaseType b)+ , Fractional (BaseType a)+ , Integral (BaseType b)+ , NumericWidenOp (BaseType a) (BaseType b)+ , NullLift2Op a b (BaseType a) a+ , Num (Promote (BaseType a) (BaseType b))+ ) =>+ Expr a -> Expr b -> Expr a+(.^^) = lift2Decorated (applyNull2 (^^)) "pow" (Just ".^^") False 8++(.^) ::+ ( Columnable (BaseType a)+ , Columnable (BaseType b)+ , Num (BaseType a)+ , Integral (BaseType b)+ , NumericWidenOp (BaseType a) (BaseType b)+ , NullLift2Op a b (BaseType a) a+ , Num (Promote (BaseType a) (BaseType b))+ ) =>+ Expr a -> Expr b -> Expr a+(.^) = lift2Decorated (applyNull2 (^)) "pow" (Just ".^") False 8++-- Same-type (non-nullable) exponentiation operators++(.^^.) ::+ (Columnable a, Columnable b, Fractional a, Integral b) =>+ Expr a -> Expr b -> Expr a+(.^^.) = lift2Decorated (^^) "pow" (Just ".^^.") False 8++(.^.) ::+ (Columnable a, Columnable b, Num a, Integral b) =>+ Expr a -> Expr b -> Expr a+(.^.) = lift2Decorated (^) "pow" (Just ".^.") False 8
+ src/DataFrame/Typed/Freeze.hs view
@@ -0,0 +1,98 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module DataFrame.Typed.Freeze (+ -- * Safe boundary+ freeze,+ freezeWithError,++ -- * Escape hatches+ thaw,+ unsafeFreeze,+) where++import qualified Data.Text as T+import Type.Reflection (SomeTypeRep)++import Data.List (stripPrefix)+import qualified DataFrame.Internal.Column as C+import DataFrame.Internal.DataFrame (columnNames)+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Typed.Schema (KnownSchema (..))+import DataFrame.Typed.Types (TypedDataFrame (..))++{- | Validate that an untyped 'DataFrame' matches the expected schema @cols@,+then wrap it. Returns 'Nothing' on mismatch.+-}+freeze ::+ forall cols. (KnownSchema cols) => D.DataFrame -> Maybe (TypedDataFrame cols)+freeze df = case validateSchema @cols df of+ Left _ -> Nothing+ Right _ -> Just (TDF df)++-- | Like 'freeze' but returns a descriptive error message on failure.+freezeWithError ::+ forall cols.+ (KnownSchema cols) =>+ D.DataFrame -> Either T.Text (TypedDataFrame cols)+freezeWithError df = case validateSchema @cols df of+ Left err -> Left err+ Right _ -> Right (TDF df)++{- | Unwrap a typed DataFrame back to the untyped representation.+Always safe; discards type information.+-}+thaw :: TypedDataFrame cols -> D.DataFrame+thaw (TDF df) = df++{- | Wrap an untyped DataFrame without any validation.+Used internally after delegation where the library guarantees schema correctness.+-}+unsafeFreeze :: D.DataFrame -> TypedDataFrame cols+unsafeFreeze = TDF++validateSchema ::+ forall cols.+ (KnownSchema cols) =>+ D.DataFrame -> Either T.Text ()+validateSchema df = mapM_ checkCol (schemaEvidence @cols)+ where+ checkCol :: (T.Text, SomeTypeRep) -> Either T.Text ()+ checkCol (name, expectedRep) = case D.getColumn name df of+ Nothing ->+ Left $+ "Column '"+ <> name+ <> "' not found in DataFrame. "+ <> "Available columns: "+ <> T.pack (show (columnNames df))+ Just col ->+ if matchesType expectedRep col+ then Right ()+ else+ Left $+ "Type mismatch on column '"+ <> name+ <> "': expected "+ <> T.pack (show expectedRep)+ <> ", got "+ <> T.pack (C.columnTypeString col)++{- | Check if a Column's element type matches the expected SomeTypeRep.+For nullable columns (those with a bitmap), @Maybe a@ in the schema matches+a column whose inner type is @a@, since we store nullable data as+@BoxedColumn (Just bm) a@ or @UnboxedColumn (Just bm) a@ rather than+@Column (Maybe a)@.+-}+matchesType :: SomeTypeRep -> C.Column -> Bool+matchesType expected col =+ let expectedStr = show expected+ colTypeStr = C.columnTypeString col+ in expectedStr == colTypeStr+ || ( C.hasMissing col -- nullable column: schema says "Maybe X", column stores "X" with a bitmap+ && Just colTypeStr == stripPrefix "Maybe " expectedStr+ )
+ src/DataFrame/Typed/Generic.hs view
@@ -0,0 +1,205 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++{- |+Module : DataFrame.Typed.Generic+License : MIT++Generic-based opt-in for record-to-schema derivation. Mirrors the Template+Haskell splice in "DataFrame.Typed.TH" but builds the schema type from a+@GHC.Generics.Generic@ instance instead of @reify@.++Use it like this:++@+data Order = Order+ { orderId :: Int64+ , region :: Text+ , amount :: Double+ } deriving (Show, Eq, Generic)++type OrderSchema = SchemaOf Order++instance HasSchema Order OrderSchema where+ toColumns = genericToColumns+ fromColumns = genericFromColumns+@++Field names are translated with the @CamelCase -> snake_case@ rule+(matching 'DataFrame.Typed.TH.camelToSnake'); use 'SchemaOfRaw' if you+want the schema to keep the record selector names verbatim — in that+case you cannot use 'genericToColumns' \/ 'genericFromColumns' and must+either hand-roll the instance or use the TH splice with a custom name+transform.+-}+module DataFrame.Typed.Generic (+ -- * Type-level schema derivation+ NameCase (..),+ SchemaOf,+ SchemaOfRaw,+ RepToSchema,+ CamelToSnake,++ -- * Value-level default methods+ genericToColumns,+ genericFromColumns,+ GHasColumns,+) where++import Data.Kind (Type)+import Data.Proxy (Proxy (..))+import qualified Data.Text as T+import qualified Data.Vector as VB+import GHC.Generics (+ C,+ D,+ Generic (..),+ K1 (..),+ M1 (..),+ Meta (..),+ S,+ type (:*:) (..),+ )+import GHC.TypeLits (+ CharToNat,+ ConsSymbol,+ KnownSymbol,+ NatToChar,+ Symbol,+ UnconsSymbol,+ symbolVal,+ type (+),+ )++import Data.Type.Bool (If, type (&&))+import Data.Type.Ord (type (<=?))++import qualified DataFrame.Internal.Column as C+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Typed.Record (requireColumn)+import DataFrame.Typed.Schema (Append)+import DataFrame.Typed.Types (Column)+import DataFrame.Typed.Util (camelToSnake)++{- | Field-name policy applied to record selectors when computing+'RepToSchema'.++* 'SnakeCase' — translate @camelCaseField@ to @\"camel_case_field\"@.+* 'IdentityCase' — keep the selector name verbatim.+-}+data NameCase = SnakeCase | IdentityCase++{- | The schema type @[Column name ty, ...]@ derived from the 'Rep' of a+record type, with the given 'NameCase' applied to each field name.+-}+type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [Type] where+ RepToSchema nc (M1 D _ f) = RepToSchema nc f+ RepToSchema nc (M1 C _ f) = RepToSchema nc f+ RepToSchema nc (a :*: b) = Append (RepToSchema nc a) (RepToSchema nc b)+ RepToSchema nc (M1 S ('MetaSel ('Just name) _ _ _) (K1 _ a)) =+ '[Column (TransformName nc name) a]++type family TransformName (nc :: NameCase) (name :: Symbol) :: Symbol where+ TransformName 'SnakeCase s = CamelToSnake s+ TransformName 'IdentityCase s = s++-- | Type-level camelCase -> snake_case. Matches 'camelToSnake' at the value level.+type family CamelToSnake (s :: Symbol) :: Symbol where+ CamelToSnake s = SnakeStart (UnconsSymbol s)++type family SnakeStart (mu :: Maybe (Char, Symbol)) :: Symbol where+ SnakeStart 'Nothing = ""+ SnakeStart ('Just '(c, r)) =+ ConsSymbol (ToLowerChar c) (SnakeRest (UnconsSymbol r))++type family SnakeRest (mu :: Maybe (Char, Symbol)) :: Symbol where+ SnakeRest 'Nothing = ""+ SnakeRest ('Just '(c, r)) =+ SnakeStep (IsUpperChar c) c (SnakeRest (UnconsSymbol r))++type family SnakeStep (up :: Bool) (c :: Char) (rest :: Symbol) :: Symbol where+ SnakeStep 'True c rest = ConsSymbol '_' (ConsSymbol (ToLowerChar c) rest)+ SnakeStep 'False c rest = ConsSymbol c rest++type family IsUpperChar (c :: Char) :: Bool where+ IsUpperChar c =+ (CharToNat 'A' <=? CharToNat c) && (CharToNat c <=? CharToNat 'Z')++type family ToLowerChar (c :: Char) :: Char where+ ToLowerChar c = If (IsUpperChar c) (NatToChar (CharToNat c + 32)) c++-- | Snake_case schema derived from @a@'s 'Generic' representation.+type SchemaOf a = RepToSchema 'SnakeCase (Rep a)++-- | Identity-cased schema derived from @a@'s 'Generic' representation.+type SchemaOfRaw a = RepToSchema 'IdentityCase (Rep a)++{- | Walks the 'Rep' tree of a record, producing or consuming a list of+named columns. Used by 'genericToColumns' \/ 'genericFromColumns'.+-}+class GHasColumns (r :: Type -> Type) where+ gToColumns :: [r p] -> [(T.Text, C.Column)]+ gFromColumns :: D.DataFrame -> Either T.Text [r p]++instance (GHasColumns f) => GHasColumns (M1 D meta f) where+ gToColumns rs = gToColumns (map unM1 rs)+ gFromColumns df = fmap (map M1) (gFromColumns df)++instance (GHasColumns f) => GHasColumns (M1 C meta f) where+ gToColumns rs = gToColumns (map unM1 rs)+ gFromColumns df = fmap (map M1) (gFromColumns df)++instance (GHasColumns a, GHasColumns b) => GHasColumns (a :*: b) where+ gToColumns rs =+ gToColumns (map (\(x :*: _) -> x) rs)+ ++ gToColumns (map (\(_ :*: y) -> y) rs)+ gFromColumns df = do+ as <- gFromColumns df+ bs <- gFromColumns df+ pure (zipWith (:*:) as bs)++instance+ (KnownSymbol name, C.Columnable a) =>+ GHasColumns+ ( M1+ S+ ('MetaSel ('Just name) su ss ds)+ (K1 i a)+ )+ where+ gToColumns rs =+ let colName = T.pack (camelToSnake (symbolVal (Proxy @name)))+ vals = map (unK1 . unM1) rs+ in [(colName, C.fromList vals)]+ gFromColumns df = do+ let colName = T.pack (camelToSnake (symbolVal (Proxy @name)))+ v <- requireColumn @a colName df+ pure (map (M1 . K1) (VB.toList v))++{- | Default implementation of 'DataFrame.Typed.Record.toColumns' for any+@Generic@ record. Field names are translated with @camelCase -> snake_case@.++@+instance HasSchema Order (SchemaOf Order) where+ toColumns = genericToColumns+ fromColumns = genericFromColumns+@+-}+genericToColumns ::+ forall a. (Generic a, GHasColumns (Rep a)) => [a] -> [(T.Text, C.Column)]+genericToColumns = gToColumns . map from++{- | Default implementation of 'DataFrame.Typed.Record.fromColumns' for any+@Generic@ record.+-}+genericFromColumns ::+ forall a. (Generic a, GHasColumns (Rep a)) => D.DataFrame -> Either T.Text [a]+genericFromColumns df = fmap (map to) (gFromColumns df)
+ src/DataFrame/Typed/Record.hs view
@@ -0,0 +1,105 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}++{- |+Module : DataFrame.Typed.Record+License : MIT++Bridge a Haskell record type to a typed dataframe schema. Instances are+typically generated by the @deriveSchemaFromType@ Template Haskell splice+(see "DataFrame.Typed.TH"), but can be written by hand as well, or+plugged into the @Generic@-derived defaults from "DataFrame.Typed.Generic".+-}+module DataFrame.Typed.Record (+ -- * Class+ HasSchema (..),++ -- * Untyped helpers+ fromRecords,+ toRecords,++ -- * Typed helpers+ fromRecordsTyped,+ toRecordsTyped,++ -- * Helpers used by generated code+ requireColumn,+) where++import Data.Kind (Type)+import qualified Data.Text as T+import qualified Data.Vector as VB++import qualified DataFrame.Internal.Column as C+import DataFrame.Internal.DataFrame (fromNamedColumns)+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Typed.Types (TypedDataFrame (..))++{- | Bridge a Haskell record type @a@ to a typed-dataframe schema.++The schema is exposed as an associated type family 'Schema' so that+instances can pick it up from a 'GHC.Generics.Rep' computation (see+'DataFrame.Typed.Generic.SchemaOf') or from an explicit list emitted by+'DataFrame.Typed.TH.deriveSchemaFromType'.++@toColumns@ explodes a list of records into a list of named columns.+@fromColumns@ reconstructs the records from a 'D.DataFrame', returning+@Left err@ if a column is missing or has the wrong type.+-}+class HasSchema a where+ type Schema a :: [Type]+ toColumns :: [a] -> [(T.Text, C.Column)]+ fromColumns :: D.DataFrame -> Either T.Text [a]++{- | Build an untyped 'D.DataFrame' from a list of records.++@+data Order = Order { orderId :: Int64, region :: Text, amount :: Double }+\$(deriveSchemaFromType ''Order)++xs :: [Order]+xs = [Order 1 "us" 10.0, Order 2 "eu" 20.0]++df :: DataFrame+df = fromRecords xs+@+-}+fromRecords :: (HasSchema a) => [a] -> D.DataFrame+fromRecords = fromNamedColumns . toColumns++{- | Parse a list of records out of an untyped 'D.DataFrame'.++Returns @Left err@ on schema mismatch (missing column, wrong type).+-}+toRecords :: (HasSchema a) => D.DataFrame -> Either T.Text [a]+toRecords = fromColumns++-- | Like 'fromRecords' but returns a 'TypedDataFrame' tagged with the schema.+fromRecordsTyped :: forall a. (HasSchema a) => [a] -> TypedDataFrame (Schema a)+fromRecordsTyped = TDF . fromRecords++-- | Like 'toRecords' but accepts a 'TypedDataFrame'.+toRecordsTyped ::+ forall a. (HasSchema a) => TypedDataFrame (Schema a) -> Either T.Text [a]+toRecordsTyped (TDF df) = fromColumns df++{- | Extract a column as a boxed vector by name, returning a 'T.Text' error+on missing column or type mismatch.++Used by code generated by 'DataFrame.Typed.TH.deriveSchemaFromType'.+-}+requireColumn ::+ forall a.+ (C.Columnable a) => T.Text -> D.DataFrame -> Either T.Text (VB.Vector a)+requireColumn name df = case D.getColumn name df of+ Nothing ->+ Left $ "Column '" <> name <> "' not found in DataFrame"+ Just col -> case C.toVector col of+ Right v -> Right v+ Left e ->+ Left $+ "Column '" <> name <> "': " <> T.pack (show e)
+ src/DataFrame/Typed/Schema.hs view
@@ -0,0 +1,441 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-}++module DataFrame.Typed.Schema (+ -- * Type families for schema manipulation+ Lookup,+ SafeLookup,+ HasName,+ RemoveColumn,+ Impute,+ SubsetSchema,+ ExcludeSchema,+ RenameInSchema,+ RenameManyInSchema,+ Append,+ Snoc,+ Reverse,+ ColumnNames,+ AssertAbsent,+ AssertPresent,+ AssertAllPresent,+ IsElem,++ -- * Maybe-stripping families+ StripAllMaybe,+ StripMaybeAt,++ -- * Join schema families+ SharedNames,+ UniqueLeft,+ InnerJoinSchema,+ LeftJoinSchema,+ RightJoinSchema,+ FullOuterJoinSchema,+ WrapMaybe,+ WrapMaybeColumns,+ CollidingColumns,++ -- * GroupBy helpers+ GroupKeyColumns,++ -- * KnownSchema class+ KnownSchema (..),++ -- * Helpers+ AllKnownSymbol (..),+) where++import Data.Kind (Constraint, Type)+import Data.Proxy (Proxy (..))+import qualified Data.Text as T+import GHC.TypeLits+import Type.Reflection (SomeTypeRep, Typeable, someTypeRep)++import DataFrame.Internal.Column (Columnable)+import DataFrame.Internal.Types (These)+import DataFrame.Typed.Types (Column)++-- | Look up the element type of a column by name.+type family Lookup (name :: Symbol) (cols :: [Type]) :: Type where+ Lookup name (Column name a ': _) = a+ Lookup name (Column _ _ ': rest) = Lookup name rest+ Lookup name '[] =+ TypeError+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")++{- | Like 'Lookup', but returns a harmless fallback ('Int') instead of+'TypeError' when the column is not found. Use together with+'AssertPresent' so the error fires exactly once.+-}+type family SafeLookup (name :: Symbol) (cols :: [Type]) :: Type where+ SafeLookup name (Column name a ': _) = a+ SafeLookup name (Column _ _ ': rest) = SafeLookup name rest+ SafeLookup name '[] = Int++-- | Unwrap a Maybe from a type after we impute values.+type family Impute (name :: Symbol) (cols :: [Type]) :: [Type] where+ Impute name (Column name (Maybe a) ': rest) = Column name a ': rest+ Impute name (Column name _ ': rest) =+ TypeError+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' is not of kind Maybe *")+ Impute name (col ': rest) = col ': Impute name rest+ Impute name '[] = '[]++-- | Add type to the end of a list.+type family Snoc (xs :: [k]) (x :: k) :: [k] where+ Snoc '[] x = '[x]+ Snoc (y ': ys) x = y ': Snoc ys x++-- | Check whether a column name exists in a schema (type-level Bool).+type family HasName (name :: Symbol) (cols :: [Type]) :: Bool where+ HasName name (Column name _ ': _) = 'True+ HasName name (Column _ _ ': rest) = HasName name rest+ HasName name '[] = 'False++-- | Remove a column by name from a schema.+type family RemoveColumn (name :: Symbol) (cols :: [Type]) :: [Type] where+ RemoveColumn name (Column name _ ': rest) = rest+ RemoveColumn name (col ': rest) = col ': RemoveColumn name rest+ RemoveColumn name '[] = '[]++-- | Select a subset of columns by a list of names.+type family SubsetSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where+ SubsetSchema '[] cols = '[]+ SubsetSchema (n ': ns) cols = Column n (Lookup n cols) ': SubsetSchema ns cols++-- | Exclude columns by a list of names.+type family ExcludeSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where+ ExcludeSchema names '[] = '[]+ ExcludeSchema names (Column n a ': rest) =+ ExcludeSchemaHelper (IsElem n names) n a names rest++type family+ ExcludeSchemaHelper+ (found :: Bool)+ (n :: Symbol)+ (a :: Type)+ (names :: [Symbol])+ (rest :: [Type]) ::+ [Type]+ where+ ExcludeSchemaHelper 'True n a names rest = ExcludeSchema names rest+ ExcludeSchemaHelper 'False n a names rest =+ Column n a ': ExcludeSchema names rest++-- | Type-level elem for Symbols+type family IsElem (x :: Symbol) (xs :: [Symbol]) :: Bool where+ IsElem x '[] = 'False+ IsElem x (x ': _) = 'True+ IsElem x (_ ': xs) = IsElem x xs++-- | Rename a column in the schema.+type family RenameInSchema (old :: Symbol) (new :: Symbol) (cols :: [Type]) :: [Type] where+ RenameInSchema old new (Column old a ': rest) = Column new a ': rest+ RenameInSchema old new (col ': rest) = col ': RenameInSchema old new rest+ RenameInSchema old new '[] =+ TypeError+ ('Text "Cannot rename: column '" ':<>: 'Text old ':<>: 'Text "' not found")++-- | Rename multiple columns.+type family RenameManyInSchema (pairs :: [(Symbol, Symbol)]) (cols :: [Type]) :: [Type] where+ RenameManyInSchema '[] cols = cols+ RenameManyInSchema ('(old, new) ': rest) cols =+ RenameManyInSchema rest (RenameInSchema old new cols)++-- | Append two type-level lists.+type family Append (xs :: [k]) (ys :: [k]) :: [k] where+ Append '[] ys = ys+ Append (x ': xs) ys = x ': Append xs ys++-- | Reverse a type-level list.+type family Reverse (xs :: [Type]) :: [Type] where+ Reverse xs = ReverseAcc xs '[]++type family ReverseAcc (xs :: [Type]) (acc :: [Type]) :: [Type] where+ ReverseAcc '[] acc = acc+ ReverseAcc (x ': xs) acc = ReverseAcc xs (x ': acc)++-- | Extract column names as a type-level list of Symbols.+type family ColumnNames (cols :: [Type]) :: [Symbol] where+ ColumnNames '[] = '[]+ ColumnNames (Column n _ ': rest) = n ': ColumnNames rest++-- | Assert that a column name is absent from the schema (for derive/insert).+type family AssertAbsent (name :: Symbol) (cols :: [Type]) :: Constraint where+ AssertAbsent name cols = AssertAbsentHelper name (HasName name cols) cols++type family+ AssertAbsentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::+ Constraint+ where+ AssertAbsentHelper name 'False cols = ()+ AssertAbsentHelper name 'True cols =+ TypeError+ ( 'Text "Column '"+ ':<>: 'Text name+ ':<>: 'Text "' already exists in schema. "+ ':<>: 'Text "Use replaceColumn to overwrite."+ )++-- | Assert that a column name is present in the schema.+type family AssertPresent (name :: Symbol) (cols :: [Type]) :: Constraint where+ AssertPresent name cols = AssertPresentHelper name (HasName name cols) cols++type family+ AssertPresentHelper (name :: Symbol) (found :: Bool) (cols :: [Type]) ::+ Constraint+ where+ AssertPresentHelper name 'True cols = ()+ AssertPresentHelper name 'False cols =+ TypeError+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")++-- | Assert that a column name is present in the schema.+type family AssertAllPresent (name :: [Symbol]) (cols :: [Type]) :: Constraint where+ AssertAllPresent (name ': rest) cols =+ AssertAllPresentHelper (HasName name cols) name rest cols+ AssertAllPresent '[] cols = ()++type family+ AssertAllPresentHelper+ (found :: Bool)+ (name :: Symbol)+ (rest :: [Symbol])+ (cols :: [Type]) ::+ Constraint+ where+ AssertAllPresentHelper 'True name rest cols = AssertAllPresent rest cols+ AssertAllPresentHelper 'False name rest cols =+ TypeError+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")++{- | Strip 'Maybe' from all columns. Used by 'filterAllJust'.++@Column "x" (Maybe Double)@ becomes @Column "x" Double@.+@Column "y" Int@ stays @Column "y" Int@.+-}+type family StripAllMaybe (cols :: [Type]) :: [Type] where+ StripAllMaybe '[] = '[]+ StripAllMaybe (Column n (Maybe a) ': rest) = Column n a ': StripAllMaybe rest+ StripAllMaybe (Column n a ': rest) = Column n a ': StripAllMaybe rest++{- | Strip 'Maybe' from a single named column. Used by 'filterJust'.++@StripMaybeAt "x" '[Column "x" (Maybe Double), Column "y" Int]@+ = @'[Column "x" Double, Column "y" Int]@+-}+type family StripMaybeAt (name :: Symbol) (cols :: [Type]) :: [Type] where+ StripMaybeAt name (Column name (Maybe a) ': rest) = Column name a ': rest+ StripMaybeAt name (Column name a ': rest) = Column name a ': rest+ StripMaybeAt name (col ': rest) = col ': StripMaybeAt name rest+ StripMaybeAt name '[] =+ TypeError+ ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")++-- | Extract column names that appear in both schemas.+type family SharedNames (left :: [Type]) (right :: [Type]) :: [Symbol] where+ SharedNames '[] right = '[]+ SharedNames (Column n _ ': rest) right =+ SharedNamesHelper (HasName n right) n rest right++type family+ SharedNamesHelper+ (found :: Bool)+ (n :: Symbol)+ (rest :: [Type])+ (right :: [Type]) ::+ [Symbol]+ where+ SharedNamesHelper 'True n rest right = n ': SharedNames rest right+ SharedNamesHelper 'False n rest right = SharedNames rest right++-- | Columns from @left@ whose names do NOT appear in @right@.+type family UniqueLeft (left :: [Type]) (rightNames :: [Symbol]) :: [Type] where+ UniqueLeft '[] _ = '[]+ UniqueLeft (Column n a ': rest) rn =+ UniqueLeftHelper (IsElem n rn) n a rest rn++type family+ UniqueLeftHelper+ (found :: Bool)+ (n :: Symbol)+ (a :: Type)+ (rest :: [Type])+ (rn :: [Symbol]) ::+ [Type]+ where+ UniqueLeftHelper 'True n a rest rn = UniqueLeft rest rn+ UniqueLeftHelper 'False n a rest rn = Column n a ': UniqueLeft rest rn++-- | Wrap column types in Maybe.+type family WrapMaybe (cols :: [Type]) :: [Type] where+ WrapMaybe '[] = '[]+ WrapMaybe (Column n a ': rest) = Column n (Maybe a) ': WrapMaybe rest++-- | Wrap selected columns in Maybe by name list.+type family WrapMaybeColumns (names :: [Symbol]) (cols :: [Type]) :: [Type] where+ WrapMaybeColumns names '[] = '[]+ WrapMaybeColumns names (Column n a ': rest) =+ WrapMaybeColumnsHelper (IsElem n names) n a names rest++type family+ WrapMaybeColumnsHelper+ (found :: Bool)+ (n :: Symbol)+ (a :: Type)+ (names :: [Symbol])+ (rest :: [Type]) ::+ [Type]+ where+ WrapMaybeColumnsHelper 'True n a names rest =+ Column n (Maybe a) ': WrapMaybeColumns names rest+ WrapMaybeColumnsHelper 'False n a names rest =+ Column n a ': WrapMaybeColumns names rest++-- | Columns in left whose names collide with right (excluding keys).+type family CollidingColumns (left :: [Type]) (right :: [Type]) (keys :: [Symbol]) :: [Type] where+ CollidingColumns '[] _ _ = '[]+ CollidingColumns (Column n a ': rest) right keys =+ CollidingColumnsHelper1 (IsElem n keys) n a rest right keys++type family+ CollidingColumnsHelper1+ (isKey :: Bool)+ (n :: Symbol)+ (a :: Type)+ (rest :: [Type])+ (right :: [Type])+ (keys :: [Symbol]) ::+ [Type]+ where+ CollidingColumnsHelper1 'True n a rest right keys =+ CollidingColumns rest right keys+ CollidingColumnsHelper1 'False n a rest right keys =+ CollidingColumnsHelper2 (HasName n right) n a rest right keys++type family+ CollidingColumnsHelper2+ (inRight :: Bool)+ (n :: Symbol)+ (a :: Type)+ (rest :: [Type])+ (right :: [Type])+ (keys :: [Symbol]) ::+ [Type]+ where+ CollidingColumnsHelper2 'True n a rest right keys =+ Column n (These a (Lookup n right)) ': CollidingColumns rest right keys+ CollidingColumnsHelper2 'False n a rest right keys =+ CollidingColumns rest right keys++-- | Inner join result schema.+type family InnerJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+ InnerJoinSchema keys left right =+ Append+ (SubsetSchema keys left)+ ( Append+ (UniqueLeft left (Append keys (ColumnNames right)))+ ( Append+ (UniqueLeft right (Append keys (ColumnNames left)))+ (CollidingColumns left right keys)+ )+ )++-- | Left join result schema.+type family LeftJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+ LeftJoinSchema keys left right =+ Append+ (SubsetSchema keys left)+ ( Append+ (UniqueLeft left (Append keys (ColumnNames right)))+ ( Append+ (WrapMaybe (UniqueLeft right (Append keys (ColumnNames left))))+ (CollidingColumns left right keys)+ )+ )++-- | Right join result schema.+type family RightJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) :: [Type] where+ RightJoinSchema keys left right =+ Append+ (SubsetSchema keys right)+ ( Append+ (WrapMaybe (UniqueLeft left (Append keys (ColumnNames right))))+ ( Append+ (UniqueLeft right (Append keys (ColumnNames left)))+ (CollidingColumns left right keys)+ )+ )++-- | Full outer join result schema.+type family+ FullOuterJoinSchema (keys :: [Symbol]) (left :: [Type]) (right :: [Type]) ::+ [Type]+ where+ FullOuterJoinSchema keys left right =+ Append+ (WrapMaybe (SubsetSchema keys left))+ ( Append+ (WrapMaybe (UniqueLeft left (Append keys (ColumnNames right))))+ ( Append+ (WrapMaybe (UniqueLeft right (Append keys (ColumnNames left))))+ (CollidingColumns left right keys)+ )+ )++-- | Extract Column entries from a schema whose names appear in @keys@.+type family GroupKeyColumns (keys :: [Symbol]) (cols :: [Type]) :: [Type] where+ GroupKeyColumns keys '[] = '[]+ GroupKeyColumns keys (Column n a ': rest) =+ GroupKeyColumnsHelper (IsElem n keys) n a keys rest++type family+ GroupKeyColumnsHelper+ (found :: Bool)+ (n :: Symbol)+ (a :: Type)+ (keys :: [Symbol])+ (rest :: [Type]) ::+ [Type]+ where+ GroupKeyColumnsHelper 'True n a keys rest =+ Column n a ': GroupKeyColumns keys rest+ GroupKeyColumnsHelper 'False n a keys rest = GroupKeyColumns keys rest++-- | Provides runtime evidence of a schema: a list of (name, TypeRep) pairs.+class KnownSchema (cols :: [Type]) where+ schemaEvidence :: [(T.Text, SomeTypeRep)]++instance KnownSchema '[] where+ schemaEvidence = []++instance+ (KnownSymbol name, Typeable a, Columnable a, KnownSchema rest) =>+ KnownSchema (Column name a ': rest)+ where+ schemaEvidence =+ (T.pack (symbolVal (Proxy @name)), someTypeRep (Proxy @a))+ : schemaEvidence @rest++-- | A class that provides a list of 'Text' values for a type-level list of Symbols.+class AllKnownSymbol (names :: [Symbol]) where+ symbolVals :: [T.Text]++instance AllKnownSymbol '[] where+ symbolVals = []++instance (KnownSymbol n, AllKnownSymbol ns) => AllKnownSymbol (n ': ns) where+ symbolVals = T.pack (symbolVal (Proxy @n)) : symbolVals @ns
+ src/DataFrame/Typed/Types.hs view
@@ -0,0 +1,114 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeOperators #-}++module DataFrame.Typed.Types (+ -- * Core phantom-typed wrapper+ TypedDataFrame (..),++ -- * Column phantom type (no constructors)+ Column,++ -- * Typed expressions (schema-validated)+ TExpr (..),++ -- * Typed sort orders+ TSortOrder (..),++ -- * Grouped typed dataframe+ TypedGrouped (..),++ -- * Typed aggregation builder (Option B)+ TAgg (..),+ taggToNamedExprs,++ -- * Re-export These+ These (..),+) where++import Data.Kind (Type)+import GHC.TypeLits (Symbol)++import qualified Data.Text as T+import DataFrame.Internal.Column (Columnable)+import qualified DataFrame.Internal.DataFrame as D+import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..))+import DataFrame.Internal.Types (These (..))++{- | A phantom-typed wrapper over the untyped 'DataFrame'.++The type parameter @cols@ is a type-level list of @Column name ty@ entries+that tracks the schema at compile time. All operations delegate to the+untyped core at runtime and update the phantom type at compile time.+-}+newtype TypedDataFrame (cols :: [Type]) = TDF {unTDF :: D.DataFrame}++instance Show (TypedDataFrame cols) where+ show (TDF df) = show df++instance Eq (TypedDataFrame cols) where+ (TDF a) == (TDF b) = a == b++{- | A phantom type that pairs a type-level column name ('Symbol')+with its element type. Has no value-level constructors — used+purely at the type level to describe schemas.+-}+data Column (name :: Symbol) (a :: Type)++{- | A typed expression validated against schema @cols@, producing values of type @a@.++Unlike the untyped 'Expr a', a 'TExpr' can only be constructed through+type-safe combinators ('col', 'lit', arithmetic operations) that verify+column references exist in the schema with the correct type.++Use 'unTExpr' to extract the underlying 'Expr' for delegation to the untyped API.+-}+newtype TExpr (cols :: [Type]) a = TExpr {unTExpr :: Expr a}++-- | A typed sort order validated against schema @cols@.+data TSortOrder (cols :: [Type]) where+ Asc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols+ Desc :: (Columnable a, Ord a) => TExpr cols a -> TSortOrder cols++-- | A phantom-typed wrapper over 'GroupedDataFrame'.+newtype TypedGrouped (keys :: [Symbol]) (cols :: [Type])+ = TGD {unTGD :: D.GroupedDataFrame}++{- | Internal aggregation chain. Each cons prepends a 'Column' to the+@aggs@ phantom list. End users never construct this directly — they+compose 'DataFrame.Typed.Aggregate.as' entries with @(.)@ and let+'DataFrame.Typed.Aggregate.aggregate' apply the composition to+'TAggNil'.++@+as \@\"total\" (F.sum salary)+ . as \@\"avg_age\" (F.mean age)+@+-}+data TAgg (keys :: [Symbol]) (cols :: [Type]) (aggs :: [Type]) where+ TAggNil :: TAgg keys cols '[]+ TAggCons ::+ (Columnable a) =>+ -- | column name+ T.Text ->+ -- | typed aggregation expression+ TExpr cols a ->+ -- | rest+ TAgg keys cols aggs ->+ TAgg keys cols (Column name a ': aggs)++{- | Extract the runtime 'NamedExpr' list from a 'TAgg', in+declaration order (reversed from the cons-built order).+-}+taggToNamedExprs :: TAgg keys cols aggs -> [NamedExpr]+taggToNamedExprs = reverse . go+ where+ go :: TAgg keys cols aggs -> [NamedExpr]+ go TAggNil = []+ go (TAggCons name (TExpr expr) rest) = (name, UExpr expr) : go rest
+ src/DataFrame/Typed/Util.hs view
@@ -0,0 +1,23 @@+module DataFrame.Typed.Util (+ camelToSnake,+) where++import Data.Char (isUpper, toLower)++{- | Convert a camelCase identifier to snake_case.++>>> camelToSnake "orderId"+"order_id"+>>> camelToSnake "amountUS"+"amount_u_s"+>>> camelToSnake "region"+"region"+-}+camelToSnake :: String -> String+camelToSnake [] = []+camelToSnake (c : cs) = toLower c : go cs+ where+ go [] = []+ go (x : xs)+ | isUpper x = '_' : toLower x : go xs+ | otherwise = x : go xs