feat: Declare parquet should respect nullable columns.

Author: mchav

data/sharded/part-0.parquet

@@ -41,9 +41,11 @@ import Data.Time
 import qualified Data.Vector as V
 import qualified Data.Vector.Unboxed as VU
 import Data.Word
+import qualified Data.Set as S
 import qualified DataFrame.IO.CSV as CSV
 import qualified DataFrame.IO.Parquet as Parquet
 import DataFrame.IO.Parquet.Thrift
+import DataFrame.IO.Parquet.Types (columnNullCount)
 import DataFrame.Internal.Nullable (
     BaseType,
     NullLift1Op (applyNull1),
@@ -537,33 +539,39 @@ declareColumnsFromCsvFile path = do
 declareColumnsFromParquetFile :: String -> DecsQ
 declareColumnsFromParquetFile path = do
     isDir <- liftIO $ doesDirectoryExist path
-
     let pat = if isDir then path </> "*.parquet" else path
-
     matches <- liftIO $ glob pat
-
     files <- liftIO $ filterM (fmap Prelude.not . doesDirectoryExist) matches
-    df <-
-        liftIO $
-            foldM
-                ( \acc p -> do
-                    (metadata, _) <- liftIO (Parquet.readMetadataFromPath p)
-                    let d = schemaToEmptyDataFrame (schema metadata)
-                    pure $ acc <> d
-                )
-                DataFrame.Internal.DataFrame.empty
-                files
+    metas <- liftIO $ mapM (fmap fst . Parquet.readMetadataFromPath) files
+    let nullableCols :: S.Set T.Text
+        nullableCols = S.fromList
+            [ T.pack (last colPath)
+            | meta  <- metas
+            , rg    <- rowGroups meta
+            , cc    <- rowGroupColumns rg
+            , let cm      = columnMetaData cc
+                  colPath = columnPathInSchema cm
+            , Prelude.not (null colPath)
+            , columnNullCount (columnStatistics cm) > 0
+            ]
+    let df = foldl (\acc meta -> acc <> schemaToEmptyDataFrame nullableCols (schema meta))
+                   DataFrame.Internal.DataFrame.empty
+                   metas
     declareColumns df
 
-schemaToEmptyDataFrame :: [SchemaElement] -> DataFrame
-schemaToEmptyDataFrame elems =
+schemaToEmptyDataFrame :: S.Set T.Text -> [SchemaElement] -> DataFrame
+schemaToEmptyDataFrame nullableCols elems =
     let leafElems = filter (\e -> numChildren e == 0) elems
-     in fromNamedColumns (map schemaElemToColumn leafElems)
+     in fromNamedColumns (map (schemaElemToColumn nullableCols) leafElems)
 
-schemaElemToColumn :: SchemaElement -> (T.Text, Column)
-schemaElemToColumn elem =
-    let name = elementName elem
-     in (name, emptyColumnForType (elementType elem))
+schemaElemToColumn :: S.Set T.Text -> SchemaElement -> (T.Text, Column)
+schemaElemToColumn nullableCols elem =
+    let name   = elementName elem
+        isNull = name `S.member` nullableCols
+        col    = if isNull
+                     then emptyNullableColumnForType (elementType elem)
+                     else emptyColumnForType         (elementType elem)
+     in (name, col)
 
 emptyColumnForType :: TType -> Column
 emptyColumnForType = \case
@@ -578,6 +586,19 @@ emptyColumnForType = \case
     STRING -> fromList @T.Text []
     other -> error $ "Unsupported parquet type for column: " <> show other
 
+emptyNullableColumnForType :: TType -> Column
+emptyNullableColumnForType = \case
+    BOOL   -> fromList @(Maybe Bool)   []
+    BYTE   -> fromList @(Maybe Word8)  []
+    I16    -> fromList @(Maybe Int16)  []
+    I32    -> fromList @(Maybe Int32)  []
+    I64    -> fromList @(Maybe Int64)  []
+    I96    -> fromList @(Maybe Int64)  []
+    FLOAT  -> fromList @(Maybe Float)  []
+    DOUBLE -> fromList @(Maybe Double) []
+    STRING -> fromList @(Maybe T.Text) []
+    other  -> error $ "Unsupported parquet type for column: " <> show other
+
 declareColumnsFromCsvWithOpts :: CSV.ReadOptions -> String -> DecsQ
 declareColumnsFromCsvWithOpts opts path = do
     df <- liftIO (CSV.readSeparated opts path)

data/sharded/part-1.parquet

@@ -10,9 +10,15 @@ import qualified DataFrame.IO.Parquet as DP
 
 import qualified Data.ByteString as BS
 import Data.Int
+import qualified Data.Set as S
 import Data.Text (Text)
+import qualified Data.Text as T
 import Data.Time
 import Data.Word
+import DataFrame.Internal.Column (hasMissing)
+import DataFrame.Internal.DataFrame (unsafeGetColumn)
+import DataFrame.IO.Parquet.Thrift (columnMetaData, columnPathInSchema, columnStatistics, rowGroupColumns, rowGroups, schema)
+import DataFrame.IO.Parquet.Types (columnNullCount)
 import DataFrame.Internal.Binary (
     littleEndianWord32,
     littleEndianWord64,
@@ -1687,6 +1693,46 @@ nationDictMalformed =
             (D.readParquet "./tests/data/nation.dict-malformed.parquet")
         )
 
+shardedNullableSchema :: Test
+shardedNullableSchema =
+    TestCase $ do
+        metas <- mapM (fmap fst . DP.readMetadataFromPath)
+            ["data/sharded/part-0.parquet", "data/sharded/part-1.parquet"]
+        let nullableCols = S.fromList
+                [ last (map T.pack colPath)
+                | meta  <- metas
+                , rg    <- rowGroups meta
+                , cc    <- rowGroupColumns rg
+                , let cm      = columnMetaData cc
+                      colPath = columnPathInSchema cm
+                , not (null colPath)
+                , columnNullCount (columnStatistics cm) > 0
+                ]
+            df = foldl (\acc meta -> acc <> F.schemaToEmptyDataFrame nullableCols (schema meta))
+                       D.empty
+                       metas
+        assertBool "id should be nullable"    (hasMissing (unsafeGetColumn "id"    df))
+        assertBool "name should be nullable"  (hasMissing (unsafeGetColumn "name"  df))
+        assertBool "score should be nullable" (hasMissing (unsafeGetColumn "score" df))
+
+singleShardNoNulls :: Test
+singleShardNoNulls =
+    TestCase $ do
+        (meta, _) <- DP.readMetadataFromPath "data/sharded/part-0.parquet"
+        let nullableCols = S.fromList
+                [ last (map T.pack colPath)
+                | rg    <- rowGroups meta
+                , cc    <- rowGroupColumns rg
+                , let cm      = columnMetaData cc
+                      colPath = columnPathInSchema cm
+                , not (null colPath)
+                , columnNullCount (columnStatistics cm) > 0
+                ]
+            df = F.schemaToEmptyDataFrame nullableCols (schema meta)
+        assertBool "id should NOT be nullable"    (not (hasMissing (unsafeGetColumn "id"    df)))
+        assertBool "name should NOT be nullable"  (not (hasMissing (unsafeGetColumn "name"  df)))
+        assertBool "score should NOT be nullable" (not (hasMissing (unsafeGetColumn "score" df)))
+
 tests :: [Test]
 tests =
     [ allTypesPlain
@@ -1780,4 +1826,7 @@ tests =
     , unknownLogicalType
     , -- Group 12: malformed files
       nationDictMalformed
+    , -- Group 13: metadata-based null detection
+      shardedNullableSchema
+    , singleShardNoNulls
     ]