Indexing certain frames

javelin-frames/benchmarks/Main.hs

@@ -0,0 +1,45 @@
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE TypeFamilies #-}
+
+import           Control.DeepSeq    ( NFData, rnf )
+import           Control.Exception  ( evaluate )
+import           Criterion.Main     ( bench, nf, defaultMain )
+
+import           Data.Frame ( Column, Frameable, Indexable, Row, Frame )
+import qualified Data.Frame as Frame
+import qualified Data.Vector as Vector
+
+import           GHC.Generics ( Generic )
+
+
+data Bench t
+    = MkBench { field1 :: Column t Int
+              , field2 :: Column t Int
+              , field3 :: Column t Int
+              , field4 :: Column t Int
+              , field5 :: Column t Int
+              , field6 :: Column t Int
+              }
+    deriving (Generic, Frameable)
+
+instance NFData (Row Bench)
+instance NFData (Frame Bench)
+
+instance Indexable Bench where
+    type Key Bench = Int
+
+    index = field1
+
+
+main :: IO ()
+main = do
+    let rs = Vector.fromList [MkBench ix 0 0 0 0 0 | ix <- [0::Int .. 100_000]]
+        fr = Frame.fromRows rs
+    evaluate $ rnf rs
+    evaluate $ rnf fr
+    defaultMain
+        [ bench "fromRows" $ nf (Frame.fromRows) rs
+        , bench "toRows"   $ nf (Frame.toRows) fr
+        , bench "lookup"   $ nf (Frame.lookup 100) fr 
+        , bench "at"       $ nf (Frame.at (100, field5)) fr 
+        ]

javelin-frames/javelin-frames.cabal

@@ -60,3 +60,14 @@ test-suite javelin-frames-test
                       tasty-hedgehog,
                       vector
 
+benchmark bench-frames
+    import:           common
+    type:             exitcode-stdio-1.0
+    ghc-options:      -rtsopts
+    hs-source-dirs:   benchmarks
+    main-is:          Main.hs
+    build-depends:    base >=4.15.0.0 && <4.22,
+                      criterion ^>=1.6,
+                      deepseq,
+                      javelin-frames,
+                      vector

javelin-frames/src/Data/Frame.hs

@@ -76,7 +76,7 @@ module Data.Frame (
     mapFrame, filterFrame, zipFramesWith, foldlFrame,
 
     -- * Indexing operations
-    Indexable(Key, index), lookup
+    Indexable(Key, index), lookup, at
 ) where
 
 import Data.Functor.Identity (Identity(..))
@@ -128,16 +128,32 @@ class GToRows tI tV where
 
 instance GToRows (Rec0 a) (Rec0 (Vector a)) where
     gtoRows = Data.Vector.map K1 . unK1
+    {-# INLINEABLE gtoRows #-}
 
 instance (GToRows tI1 tV1, GToRows tI2 tV2) 
     => GToRows (tI1 :*: tI2) (tV1 :*: tV2) where
     gtoRows (xs :*: ys) = Data.Vector.zipWith (:*:) (gtoRows xs) (gtoRows ys)
+    {-# INLINEABLE gtoRows #-}
 
 instance (GToRows tI tV) => GToRows (M1 i c tI) (M1 i c tV) where
     -- gtoRows :: M1 i c tV a -> Vector (M1 i c tI a)
     gtoRows = Data.Vector.map M1 . gtoRows . unM1
 
 
+class GILookup tI tV where
+    gilookup :: Int -> tV a -> Maybe (tI a)
+
+instance GILookup (Rec0 a) (Rec0 (Vector a)) where
+    gilookup ix vs = K1 <$> (unK1 vs) Data.Vector.!? ix
+
+instance (GILookup tI1 tV1, GILookup tI2 tV2)
+    => GILookup (tI1 :*: tI2) (tV1 :*: tV2) where
+        gilookup ix (xs :*: ys) = liftA2 (:*:) (gilookup ix xs) (gilookup ix ys)
+
+instance (GILookup tI tV) => GILookup (M1 i c tI) (M1 i c tV) where
+    gilookup ix = fmap M1 . gilookup ix . unM1
+
+
 -- | Typeclass that endows any record type @t@ with the ability to be packaged
 -- as a dataframe.
 --
@@ -170,6 +186,20 @@ class Frameable t where
     toRows = Data.Vector.map to . gtoRows . from
 
 
+    -- | Look up a row from the frame by integer index
+    ilookup :: Int -> Frame t -> Maybe (Row t)
+
+    default ilookup :: ( Generic (t Identity)
+                       , Generic (t Vector)
+                       , GILookup (Rep (Row t)) (Rep (Frame t))
+                       )
+                    => Int
+                    -> Frame t
+                    -> Maybe (Row t)
+    ilookup ix = fmap to . gilookup ix . from
+
+
+
 -- | Map a function over each row individually.
 mapFrame :: (Frameable t1, Frameable t2)
          => (Row t1 -> Row t2)
@@ -241,6 +271,18 @@ lookup :: (Indexable t)
        -> Frame t
        -> Maybe (Row t)
 lookup key fr 
-    = do
-        mix <- Data.Vector.findIndex (==key) (index fr)
-        pure $ (toRows fr) Data.Vector.! mix
+    = Data.Vector.findIndex (==key) (index fr) 
+    >>= flip ilookup fr
+
+
+-- | Lookup an element of a frame by row and column.
+--
+-- This is much more efficient than looking up an entire row 
+-- using `lookup`, and then selecting a specific field from a row.
+at :: (Indexable t)
+   => (Key t, Frame t -> Vector a)
+   -> Frame t
+   -> Maybe a
+at (row, col) fr 
+    = Data.Vector.findIndex (==row) (index fr)
+    >>= \ix -> (col fr) Data.Vector.!? ix