feat: safeRead now defaults reading columns to Maybe a

There is an analogue for Parquet as well.

Author: mchav

dataframe-persistent/dataframe-persistent.cabal

@@ -35,13 +35,11 @@ library
     exposed-modules: DataFrame.IO.Persistent,
                     DataFrame.IO.Persistent.TH
     build-depends:    base >= 4 && <5,
-                      bytestring >= 0.11 && < 0.13,
                       containers >= 0.6.7 && < 0.9,
                       dataframe ^>= 1,
                       persistent >= 2.14 && < 3,
                       template-haskell >= 2.0 && < 3,
                       text >= 2.0 && < 3,
-                      time >= 1.12 && < 2,
                       transformers >= 0.5 && < 0.7,
                       vector ^>= 0.13
     hs-source-dirs:   src

dataframe-persistent/src/DataFrame/IO/Persistent.hs

@@ -41,19 +41,16 @@ module DataFrame.IO.Persistent (
 import Control.Monad (forM)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Trans.Reader (ReaderT)
-import Data.ByteString (ByteString)
 import qualified Data.Map.Strict as M
 import Data.Proxy (Proxy (..))
 import Data.Text (Text)
 import qualified Data.Text as T
-import Data.Time (Day, TimeOfDay, UTCTime)
 import qualified Data.Vector as V
 import qualified DataFrame.Internal.Column as DFCol
 import DataFrame.Internal.DataFrame (DataFrame (..))
 import qualified DataFrame.Internal.DataFrame as DF
 import Database.Persist
 import Database.Persist.Sql hiding (Column)
-import Database.Persist.Types (fieldHaskell, getEntityFields, unFieldNameHS)
 import Unsafe.Coerce (unsafeCoerce)
 
 -- | Get number of rows in a DataFrame

dataframe-persistent/src/DataFrame/IO/Persistent/TH.hs

@@ -14,21 +14,16 @@ module DataFrame.IO.Persistent.TH (
     derivePersistentDataFrame,
 ) where
 
-import Control.Monad (forM, when)
+import Control.Monad (forM)
 import Data.Char
-import Data.List (foldl')
-import Data.Text (Text)
 import qualified Data.Text as T
 import qualified Data.Vector as V
 import DataFrame.Functions (col)
 import DataFrame.IO.Persistent
-import qualified DataFrame.Internal.Column as DFCol
 import DataFrame.Internal.Expression
 import Database.Persist
-import Database.Persist.Sql (fromSqlKey)
-import Database.Persist.TH
 import Language.Haskell.TH
-import Language.Haskell.TH.Syntax (Lift, lift)
+import Language.Haskell.TH.Syntax (lift)
 
 import Debug.Trace (trace)
 

dataframe.cabal

@@ -20,6 +20,9 @@ extra-source-files: cbits/process_csv.h
                     cbits/arrow_abi.h
                     cbits/dataframe_arrow.h
                     cbits/rts_init.c
+                    tests/data/typing/texts.txt
+                    tests/data/typing/texts_with_empties.txt
+                    tests/data/typing/texts_with_empties_and_nullish.txt
                     data/titanic/*.csv
                     data/sharded/*.parquet
                     tests/data/*.csv
@@ -274,6 +277,7 @@ test-suite tests
                    Operations.Typing,
                    LazyParquet,
                    Parquet,
+                   ParquetTestData,
                    Properties,
                    Monad
     build-depends:  base >= 4 && < 5,

scripts/presubmit.sh

@@ -8,14 +8,14 @@ set -e
 cabal build
 cabal test
 
-cd ../dataframe-persistent
+cd ./dataframe-persistent
 
 cabal build
 
 cd ../dataframe-hasktorch
 
 cabal build
 
-cd examples
+cd ../examples
 
 cabal build all

src/DataFrame/IO/CSV.hs

@@ -220,7 +220,7 @@ defaultReadOptions =
     ReadOptions
         { headerSpec = UseFirstRow
         , typeSpec = InferFromSample 100
-        , safeRead = True
+        , safeRead = False
         , dateFormat = "%Y-%m-%d"
         , columnSeparator = ','
         , numColumns = Nothing
@@ -408,11 +408,14 @@ freezeBuilderColumn (BuilderBS _ _) =
         "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn"
 
 finalizeBuilderColumn :: ReadOptions -> BuilderColumn -> IO Column
-finalizeBuilderColumn opts (BuilderBS gv validRef) = do
-    vec <- freezePagedVector gv
-    valid <- freezePagedUnboxedVector validRef
-    return $! inferColumnFromBS opts vec valid
-finalizeBuilderColumn _ bc = freezeBuilderColumn bc
+finalizeBuilderColumn opts bc = do
+    col <- case bc of
+        BuilderBS gv validRef -> do
+            vec <- freezePagedVector gv
+            valid <- freezePagedUnboxedVector validRef
+            return $! inferColumnFromBS opts vec valid
+        _ -> freezeBuilderColumn bc
+    return $! if safeRead opts then ensureOptional col else col
 
 inferColumnFromBS ::
     ReadOptions -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column

src/DataFrame/IO/Parquet.hs

@@ -19,10 +19,11 @@ import qualified Data.Text as T
 import Data.Text.Encoding
 import Data.Time
 import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
+import qualified Data.Vector as V
 import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))
 import DataFrame.Internal.Binary (littleEndianWord32)
 import qualified DataFrame.Internal.Column as DI
-import DataFrame.Internal.DataFrame (DataFrame)
+import DataFrame.Internal.DataFrame (DataFrame, columns)
 import DataFrame.Internal.Expression (Expr, getColumns)
 import qualified DataFrame.Operations.Core as DI
 import DataFrame.Operations.Merge ()
@@ -75,6 +76,8 @@ data ParquetReadOptions = ParquetReadOptions
     -- ^ Optional row filter expression applied before projection.
     , rowRange :: Maybe (Int, Int)
     -- ^ Optional row slice @(start, end)@ with start-inclusive/end-exclusive semantics.
+    , safeColumns :: Bool
+    -- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema.
     }
     deriving (Eq, Show)
 
@@ -87,6 +90,7 @@ ParquetReadOptions
     { selectedColumns = Nothing
     , predicate = Nothing
     , rowRange = Nothing
+    , safeColumns = False
     }
 @
 -}
@@ -96,6 +100,7 @@ defaultParquetReadOptions =
         { selectedColumns = Nothing
         , predicate = Nothing
         , rowRange = Nothing
+        , safeColumns = False
         }
 
 -- Public API --------------------------------------------------------------
@@ -349,9 +354,15 @@ applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame
 applyPredicate opts df =
     maybe df (`DS.filterWhere` df) (predicate opts)
 
+applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame
+applySafeRead opts df
+    | safeColumns opts = df{columns = V.map DI.ensureOptional (columns df)}
+    | otherwise = df
+
 applyReadOptions :: ParquetReadOptions -> DataFrame -> DataFrame
 applyReadOptions opts =
-    applyRowRange opts
+    applySafeRead opts
+        . applyRowRange opts
         . applySelectedColumns opts
         . applyPredicate opts
 

src/DataFrame/Internal/Column.hs

@@ -872,6 +872,12 @@ freezeColumn' nulls (MUnboxedColumn col)
                         )
 {-# INLINE freezeColumn' #-}
 
+-- | Promote a non-nullable column to OptionalColumn. No-op when already optional.
+ensureOptional :: Column -> Column
+ensureOptional (BoxedColumn col) = OptionalColumn (VB.map Just col)
+ensureOptional (UnboxedColumn col) = OptionalColumn (VB.generate (VU.length col) (Just . (col `VU.unsafeIndex`)))
+ensureOptional c = c
+
 -- | Fills the end of a column, up to n, with Nothing. Does nothing if column has length greater than n.
 expandColumn :: Int -> Column -> Column
 expandColumn n (OptionalColumn col) = OptionalColumn $ col <> VB.replicate (n - VG.length col) Nothing
@@ -1350,7 +1356,4 @@ toUnboxedVector column =
                         , errorColumnName = Nothing
                         }
                     )
-{-# SPECIALIZE toUnboxedVector ::
-    Column -> Either DataFrameException (VU.Vector Double)
-    #-}
 {-# INLINE toUnboxedVector #-}

src/DataFrame/Lazy/IO/CSV.hs

@@ -30,6 +30,7 @@ import DataFrame.Internal.Column (
     Column (..),
     MutableColumn (..),
     columnLength,
+    ensureOptional,
     freezeColumn',
     writeColumn,
  )
@@ -60,7 +61,7 @@ defaultOptions =
     ReadOptions
         { hasHeader = True
         , inferTypes = True
-        , safeRead = True
+        , safeRead = False
         , rowRange = Nothing
         , seekPos = Nothing
         , totalRows = Nothing
@@ -218,7 +219,8 @@ freezeColumn ::
     IO Column
 freezeColumn mutableCols nulls opts colIndex = do
     col <- VM.unsafeRead mutableCols colIndex
-    freezeColumn' (nulls V.! colIndex) col
+    frozen <- freezeColumn' (nulls V.! colIndex) col
+    return $! if safeRead opts then ensureOptional frozen else frozen
 {-# INLINE freezeColumn #-}
 
 -- ---------------------------------------------------------------------------

src/DataFrame/Operations/Typing.hs

@@ -15,7 +15,7 @@ import Data.Maybe (fromMaybe)
 import qualified Data.Proxy as P
 import Data.Time
 import Data.Type.Equality (TestEquality (..))
-import DataFrame.Internal.Column (Column (..), fromVector)
+import DataFrame.Internal.Column (Column (..), ensureOptional, fromVector)
 import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)
 import DataFrame.Internal.Parsing
 import DataFrame.Internal.Schema
@@ -78,14 +78,16 @@ parseFromExamples opts cols =
         examples = V.map converter (V.take (sampleSize opts) cols)
         asMaybeText = V.map converter cols
         dfmt = parseDateFormat opts
+        result =
+            case makeParsingAssumption dfmt examples of
+                BoolAssumption -> handleBoolAssumption asMaybeText
+                IntAssumption -> handleIntAssumption asMaybeText
+                DoubleAssumption -> handleDoubleAssumption asMaybeText
+                TextAssumption -> handleTextAssumption asMaybeText
+                DateAssumption -> handleDateAssumption dfmt asMaybeText
+                NoAssumption -> handleNoAssumption dfmt asMaybeText
      in
-        case makeParsingAssumption dfmt examples of
-            BoolAssumption -> handleBoolAssumption asMaybeText
-            IntAssumption -> handleIntAssumption asMaybeText
-            DoubleAssumption -> handleDoubleAssumption asMaybeText
-            TextAssumption -> handleTextAssumption asMaybeText
-            DateAssumption -> handleDateAssumption dfmt asMaybeText
-            NoAssumption -> handleNoAssumption dfmt asMaybeText
+        if parseSafe opts then ensureOptional result else result
 
 handleBoolAssumption :: V.Vector (Maybe T.Text) -> Column
 handleBoolAssumption asMaybeText