feat: Correctly stub and read logical type time unit.

Also adds test file generated from duckdb SQL similar to #147

Author: mchav

src/DataFrame/IO/Parquet.hs

@@ -1,5 +1,7 @@
+{-# LANGUAGE NumericUnderscores #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
 module DataFrame.IO.Parquet where
@@ -14,6 +16,8 @@ import qualified Data.List as L
 import qualified Data.Map as M
 import qualified Data.Text as T
 import Data.Text.Encoding
+import Data.Time
+import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
 import Data.Word
 import qualified DataFrame.Internal.Column as DI
 import DataFrame.Internal.DataFrame (DataFrame)
@@ -28,6 +32,7 @@ import DataFrame.IO.Parquet.Thrift
 import DataFrame.IO.Parquet.Types
 import System.Directory (doesDirectoryExist)
 
+import qualified Data.Vector.Unboxed as VU
 import System.FilePath ((</>))
 
 {- | Read a parquet file from path and load it into a dataframe.
@@ -93,13 +98,15 @@ readParquet path = do
             let schemaTail = drop 1 (schema fileMetadata)
             let colPath = columnPathInSchema (columnMetaData colChunk)
             let (maxDef, maxRep) = levelsForPath schemaTail colPath
+            let lType = logicalType (schemaTail !! colIdx)
             column <-
                 processColumnPages
                     (maxDef, maxRep)
                     pages
                     (columnType metadata)
                     primaryEncoding
                     maybeTypeLength
+                    lType
 
             modifyIORef colMap (M.insertWith DI.concatColumnsEither colName column)
 
@@ -172,8 +179,9 @@ processColumnPages ::
     ParquetType ->
     ParquetEncoding ->
     Maybe Int32 ->
+    LogicalType ->
     IO DI.Column
-processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do
+processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength lType = do
     let dictPages = filter isDictionaryPage pages
     let dataPages = filter isDataPage pages
 
@@ -206,10 +214,10 @@ processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do
                                  in pure (toMaybeBool maxDef defLvls vals)
                             PINT32 ->
                                 let (vals, _) = readNInt32 nPresent afterLvls
-                                 in pure (toMaybeInt32 maxDef defLvls vals)
+                                 in pure (applyLogicalType lType $ toMaybeInt32 maxDef defLvls vals)
                             PINT64 ->
                                 let (vals, _) = readNInt64 nPresent afterLvls
-                                 in pure (toMaybeInt64 maxDef defLvls vals)
+                                 in pure (applyLogicalType lType $ toMaybeInt64 maxDef defLvls vals)
                             PINT96 ->
                                 let (vals, _) = readNInt96Times nPresent afterLvls
                                  in pure (toMaybeUTCTime maxDef defLvls vals)
@@ -258,10 +266,10 @@ processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do
                                  in pure (toMaybeBool maxDef defLvls vals)
                             PINT32 ->
                                 let (vals, _) = readNInt32 nPresent afterLvls
-                                 in pure (toMaybeInt32 maxDef defLvls vals)
+                                 in pure (applyLogicalType lType $ toMaybeInt32 maxDef defLvls vals)
                             PINT64 ->
                                 let (vals, _) = readNInt64 nPresent afterLvls
-                                 in pure (toMaybeInt64 maxDef defLvls vals)
+                                 in pure (applyLogicalType lType $ toMaybeInt64 maxDef defLvls vals)
                             PINT96 ->
                                 let (vals, _) = readNInt96Times nPresent afterLvls
                                  in pure (toMaybeUTCTime maxDef defLvls vals)
@@ -296,3 +304,37 @@ processColumnPages (maxDef, maxRep) pages pType _ maybeTypeLength = do
         (c : cs) ->
             pure $
                 L.foldl' (\l r -> fromRight (error "concat failed") (DI.concatColumns l r)) c cs
+
+applyLogicalType :: LogicalType -> DI.Column -> DI.Column
+applyLogicalType (TimestampType isUTC unit) col =
+    fromRight col $
+        DI.mapColumn
+            (microsecondsToUTCTime . (* (1_000_000 `div` unitDivisor unit)))
+            col
+applyLogicalType (DecimalType precision scale) col
+    | precision <= 9 = case DI.toVector @Int32 @VU.Vector col of
+        Right xs ->
+            DI.fromUnboxedVector $
+                VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs
+        Left _ -> col
+    | precision <= 18 = case DI.toVector @Int64 @VU.Vector col of
+        Right xs ->
+            DI.fromUnboxedVector $
+                VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs
+        Left _ -> col
+    | otherwise = col
+applyLogicalType _ col = col
+
+microsecondsToUTCTime :: Int64 -> UTCTime
+microsecondsToUTCTime us =
+    posixSecondsToUTCTime (fromIntegral us / 1_000_000)
+
+unitDivisor :: TimeUnit -> Int64
+unitDivisor MILLISECONDS = 1_000
+unitDivisor MICROSECONDS = 1_000_000
+unitDivisor NANOSECONDS = 1_000_000_000
+unitDivisor TIME_UNIT_UNKNOWN = 1
+
+applyScale :: Int32 -> Int32 -> Double
+applyScale scale rawValue =
+    fromIntegral rawValue / (10 ^ scale)

src/DataFrame/IO/Parquet/Thrift.hs

@@ -1139,7 +1139,7 @@ readTimeType isAdjustedToUTC unit buf pos lastFieldId = do
         Nothing -> return (TimeType{isAdjustedToUTC = isAdjustedToUTC, unit = unit})
         Just (elemType, identifier) -> case identifier of
             1 -> do
-                isAdjustedToUTC' <- (== compactBooleanTrue) <$> readAndAdvance pos buf
+                let isAdjustedToUTC' = elemType == toTType compactBooleanTrue
                 readTimeType isAdjustedToUTC' unit buf pos identifier
             2 -> do
                 unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0
@@ -1156,16 +1156,15 @@ readTimestampType ::
 readTimestampType isAdjustedToUTC unit buf pos lastFieldId = do
     fieldContents <- readField buf pos lastFieldId
     case fieldContents of
-        Nothing -> return (TimestampType isAdjustedToUTC unit)
+        Nothing -> return (TimestampType{isAdjustedToUTC = isAdjustedToUTC, unit = unit})
         Just (elemType, identifier) -> case identifier of
             1 -> do
-                isAdjustedToUTC' <- (== compactBooleanTrue) <$> readNoAdvance pos buf
-                readTimestampType False unit buf pos identifier
+                let isAdjustedToUTC' = elemType == toTType compactBooleanTrue
+                readTimestampType isAdjustedToUTC' unit buf pos identifier
             2 -> do
-                _ <- readField buf pos 0
                 unit' <- readUnit TIME_UNIT_UNKNOWN buf pos 0
                 readTimestampType isAdjustedToUTC unit' buf pos identifier
-            _ -> error $ "UNKNOWN field ID for TimestampType" ++ show identifier
+            _ -> error $ "UNKNOWN field ID for TimestampType " ++ show identifier
 
 readUnit :: TimeUnit -> BS.ByteString -> IORef Int -> Int16 -> IO TimeUnit
 readUnit unit buf pos lastFieldId = do
@@ -1174,9 +1173,12 @@ readUnit unit buf pos lastFieldId = do
         Nothing -> return unit
         Just (elemType, identifier) -> case identifier of
             1 -> do
+                _ <- readField buf pos 0
                 readUnit MILLISECONDS buf pos identifier
             2 -> do
+                _ <- readField buf pos 0
                 readUnit MICROSECONDS buf pos identifier
             3 -> do
+                _ <- readField buf pos 0
                 readUnit NANOSECONDS buf pos identifier
             n -> error $ "Unknown time unit: " ++ show n