refactor: Move mutable column to its own GADT.

Author: mchav

src/DataFrame/IO/CSV.hs

@@ -24,7 +24,7 @@ import Control.Applicative ((<$>), (<|>), (<*>), (<*), (*>), many)
 import Control.Monad (forM_, zipWithM_, unless, void)
 import Data.Attoparsec.Text
 import Data.Char
-import DataFrame.Internal.Column (Column(..), freezeColumn', writeColumn, columnLength)
+import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength)
 import DataFrame.Internal.DataFrame (DataFrame(..))
 import DataFrame.Internal.Parsing
 import DataFrame.Operations.Typing
@@ -107,13 +107,13 @@ readSeparated c opts path = do
             }
 {-# INLINE readSeparated #-}
 
-getInitialDataVectors :: Int -> VM.IOVector Column -> [T.Text] -> IO ()
+getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO ()
 getInitialDataVectors n mCol xs = do
     forM_ (zip [0..] xs) $ \(i, x) -> do
         col <- case inferValueType x of
-                "Int" -> MutableUnboxedColumn <$>  ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c)
-                "Double" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c)
-                _ -> MutableBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c)
+                "Int" -> MUnboxedColumn <$>  ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c)
+                "Double" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c)
+                _ -> MBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c)
         VM.unsafeWrite mCol i col
 {-# INLINE getInitialDataVectors #-}
 
@@ -128,7 +128,7 @@ inferValueType s = let
 {-# INLINE inferValueType #-}
 
 -- | Reads rows from the handle and stores values in mutable vectors.
-fillColumns :: Int -> Char -> VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Handle -> IO ()
+fillColumns :: Int -> Char -> VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> Handle -> IO ()
 fillColumns n c mutableCols nullIndices handle = do
     input <- newIORef (mempty :: T.Text)
     forM_ [1..n] $ \i -> do
@@ -148,7 +148,7 @@ fillColumns n c mutableCols nullIndices handle = do
 {-# INLINE fillColumns #-}
 
 -- | Writes a value into the appropriate column, resizing the vector if necessary.
-writeValue :: VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO ()
+writeValue :: VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO ()
 writeValue mutableCols nullIndices count colIndex value = do
     col <- VM.unsafeRead mutableCols colIndex
     res <- writeColumn count value col
@@ -157,7 +157,7 @@ writeValue mutableCols nullIndices count colIndex value = do
 {-# INLINE writeValue #-}
 
 -- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.
-freezeColumn :: VM.IOVector Column -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column
+freezeColumn :: VM.IOVector MutableColumn -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column
 freezeColumn mutableCols nulls opts colIndex = do
     col <- VM.unsafeRead mutableCols colIndex
     freezeColumn' (nulls V.! colIndex) col

src/DataFrame/Internal/Column.hs

@@ -54,9 +54,10 @@ data Column where
   GroupedBoxedColumn :: Columnable a => VB.Vector (VB.Vector a) -> Column
   GroupedUnboxedColumn :: (Columnable a, VU.Unbox a) => VB.Vector (VU.Vector a) -> Column
   GroupedOptionalColumn :: (Columnable a) => VB.Vector (VB.Vector (Maybe a)) -> Column
-  -- These are used purely for I/O, not to store live data.
-  MutableBoxedColumn :: Columnable a => VBM.IOVector a -> Column
-  MutableUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> Column
+
+data MutableColumn where
+  MBoxedColumn :: Columnable a => VBM.IOVector a -> MutableColumn
+  MUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> MutableColumn
 
 -- | A TypedColumn is a wrapper around our type-erased column.
 -- It is used to type check expressions on columns.
@@ -515,14 +516,14 @@ zipWithColumns f left right = let
 {-# INLINE zipWithColumns #-}
 
 -- Functions for mutable columns (intended for IO).
-writeColumn :: Int -> T.Text -> Column -> IO (Either T.Text Bool)
-writeColumn i value (MutableBoxedColumn (col :: VBM.IOVector a)) = let
+writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool)
+writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) = let
   in case testEquality (typeRep @a) (typeRep @T.Text) of
       Just Refl -> (if isNullish value
                     then VBM.unsafeWrite col i "" >> return (Left $! value)
                     else VBM.unsafeWrite col i value >> return (Right True))
       Nothing -> return (Left value)
-writeColumn i value (MutableUnboxedColumn (col :: VUM.IOVector a)) =
+writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) =
   case testEquality (typeRep @a) (typeRep @Int) of
       Just Refl -> case readInt value of
         Just v -> VUM.unsafeWrite col i v >> return (Right True)
@@ -535,12 +536,12 @@ writeColumn i value (MutableUnboxedColumn (col :: VUM.IOVector a)) =
 writeColumn _ _ _ = error "Cannot write to immutable column"
 {-# INLINE writeColumn #-}
 
-freezeColumn' :: [(Int, T.Text)] -> Column -> IO Column
-freezeColumn' nulls (MutableBoxedColumn col)
+freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column
+freezeColumn' nulls (MBoxedColumn col)
   | null nulls = BoxedColumn <$> VB.unsafeFreeze col
   | all (isNullish . snd) nulls = OptionalColumn . VB.imap (\i v -> if i `elem` map fst nulls then Nothing else Just v) <$> VB.unsafeFreeze col
   | otherwise  = BoxedColumn . VB.imap (\i v -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right v) <$> VB.unsafeFreeze col
-freezeColumn' nulls (MutableUnboxedColumn col)
+freezeColumn' nulls (MUnboxedColumn col)
   | null nulls = UnboxedColumn <$> VU.unsafeFreeze col
   | all (isNullish . snd) nulls = VU.unsafeFreeze col >>= \c -> return $ OptionalColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Nothing else Just (c VU.! i))
   | otherwise  = VU.unsafeFreeze col >>= \c -> return $ BoxedColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right (c VU.! i))

src/DataFrame/Lazy/IO/CSV.hs

@@ -24,7 +24,7 @@ import Control.Applicative ((<$>), (<|>), (<*>), (<*), (*>), many)
 import Control.Monad (forM_, zipWithM_, unless, when, void, replicateM_)
 import Data.Attoparsec.Text
 import Data.Char
-import DataFrame.Internal.Column (Column(..), freezeColumn', writeColumn, columnLength)
+import DataFrame.Internal.Column (Column(..), MutableColumn(..), freezeColumn', writeColumn, columnLength)
 import DataFrame.Internal.DataFrame (DataFrame(..))
 import DataFrame.Internal.Parsing
 import DataFrame.Operations.Typing
@@ -122,13 +122,13 @@ readSeparated c opts path = do
             }, (pos, unconsumed, r + 1))
 {-# INLINE readSeparated #-}
 
-getInitialDataVectors :: Int -> VM.IOVector Column -> [T.Text] -> IO ()
+getInitialDataVectors :: Int -> VM.IOVector MutableColumn -> [T.Text] -> IO ()
 getInitialDataVectors n mCol xs = do
     forM_ (zip [0..] xs) $ \(i, x) -> do
         col <- case inferValueType x of
-                "Int" -> MutableUnboxedColumn <$>  ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c)
-                "Double" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c)
-                _ -> MutableBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c)
+                "Int" -> MUnboxedColumn <$>  ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c)
+                "Double" -> MUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c)
+                _ -> MBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c)
         VM.unsafeWrite mCol i col
 {-# INLINE getInitialDataVectors #-}
 
@@ -154,7 +154,7 @@ readSingleLine c unused handle = parseWith (TIO.hGetChunk handle) (parseRow c) u
                   return (row, unconsumed)
 
 -- | Reads rows from the handle and stores values in mutable vectors.
-fillColumns :: Int -> Char -> VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> T.Text -> Handle -> IO (T.Text, Int)
+fillColumns :: Int -> Char -> VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> T.Text -> Handle -> IO (T.Text, Int)
 fillColumns n c mutableCols nullIndices unused handle = do
     input <- newIORef unused
     rowsRead <- newIORef (0 :: Int)
@@ -179,7 +179,7 @@ fillColumns n c mutableCols nullIndices unused handle = do
 {-# INLINE fillColumns #-}
 
 -- | Writes a value into the appropriate column, resizing the vector if necessary.
-writeValue :: VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO ()
+writeValue :: VM.IOVector MutableColumn -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO ()
 writeValue mutableCols nullIndices count colIndex value = do
     col <- VM.unsafeRead mutableCols colIndex
     res <- writeColumn count value col
@@ -188,7 +188,7 @@ writeValue mutableCols nullIndices count colIndex value = do
 {-# INLINE writeValue #-}
 
 -- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.
-freezeColumn :: VM.IOVector Column -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column
+freezeColumn :: VM.IOVector MutableColumn -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO Column
 freezeColumn mutableCols nulls opts colIndex = do
     col <- VM.unsafeRead mutableCols colIndex
     freezeColumn' (nulls V.! colIndex) col