feat: Add ability to read parquet with just the schema.

Author: mchav

app/Synthesis.hs

@@ -9,7 +9,7 @@ import qualified DataFrame.Functions as F
 
 import Data.Char
 import DataFrame.DecisionTree
-import DataFrame.Operators
+import DataFrame.Operators hiding (name)
 
 import System.Random
 

src/DataFrame/Functions.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes #-}
@@ -12,7 +13,7 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE UndecidableInstances #-}
 
-module DataFrame.Functions where
+module DataFrame.Functions (module DataFrame.Functions, module DataFrame.Operators) where
 
 import DataFrame.Internal.Column
 import DataFrame.Internal.DataFrame (
@@ -21,22 +22,26 @@ import DataFrame.Internal.DataFrame (
  )
 import DataFrame.Internal.Expression hiding (normalize)
 import DataFrame.Internal.Statistics
+import DataFrame.Operations.Core
 
 import Control.Applicative
 import Control.Monad
 import Control.Monad.IO.Class
 import qualified Data.Char as Char
 import Data.Function
 import Data.Functor
+import Data.Int
 import qualified Data.List as L
 import qualified Data.Map as M
 import qualified Data.Maybe as Maybe
 import qualified Data.Text as T
 import Data.Time
 import qualified Data.Vector as V
 import qualified Data.Vector.Unboxed as VU
+import Data.Word
 import qualified DataFrame.IO.CSV as CSV
 import qualified DataFrame.IO.Parquet as Parquet
+import DataFrame.IO.Parquet.Thrift
 import DataFrame.Operators
 import Debug.Trace (trace)
 import Language.Haskell.TH
@@ -45,21 +50,6 @@ import Text.Regex.TDFA
 import Prelude hiding (maximum, minimum)
 import Prelude as P
 
-name :: (Show a) => Expr a -> T.Text
-name (Col n) = n
-name other =
-    error $
-        "You must call `name` on a column reference. Not the expression: " ++ show other
-
-col :: (Columnable a) => T.Text -> Expr a
-col = Col
-
-ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
-ifThenElse = If
-
-lit :: (Columnable a) => a -> Expr a
-lit = Lit
-
 lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b
 lift f =
     Unary (MkUnaryOp{unaryFn = f, unaryName = "unaryUdf", unarySymbol = Nothing})
@@ -450,12 +440,35 @@ declareColumnsFromCsvFile path = do
             (CSV.readSeparated (CSV.defaultReadOptions{CSV.numColumns = Just 100}) path)
     declareColumns df
 
--- TODO: We don't have to read the whole file, we can just read the schema.
 declareColumnsFromParquetFile :: String -> DecsQ
 declareColumnsFromParquetFile path = do
-    df <- liftIO (Parquet.readParquet path)
+    metadata <- liftIO (Parquet.readMetadataFromPath path)
+    let df = schemaToEmptyDataFrame (schema metadata)
     declareColumns df
 
+schemaToEmptyDataFrame :: [SchemaElement] -> DataFrame
+schemaToEmptyDataFrame elems =
+    let leafElems = filter (\e -> numChildren e == 0) elems
+     in fromNamedColumns (map schemaElemToColumn leafElems)
+
+schemaElemToColumn :: SchemaElement -> (T.Text, Column)
+schemaElemToColumn elem =
+    let name = elementName elem
+     in (name, emptyColumnForType (elementType elem))
+
+emptyColumnForType :: TType -> Column
+emptyColumnForType = \case
+    BOOL -> fromList @Bool []
+    BYTE -> fromList @Word8 []
+    I16 -> fromList @Int16 []
+    I32 -> fromList @Int32 []
+    I64 -> fromList @Int64 []
+    I96 -> fromList @Int64 []
+    FLOAT -> fromList @Float []
+    DOUBLE -> fromList @Double []
+    STRING -> fromList @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)

src/DataFrame/Operators.hs

@@ -15,7 +15,7 @@ import DataFrame.Internal.Expression (
         binaryPrecedence,
         binarySymbol
     ),
-    Expr (Binary, Lit),
+    Expr (Binary, Col, If, Lit),
     NamedExpr,
     UExpr (UExpr),
  )
@@ -32,6 +32,21 @@ infixr 0 .=
 as :: (Columnable a) => Expr a -> T.Text -> NamedExpr
 as expr name = (name, UExpr expr)
 
+name :: (Show a) => Expr a -> T.Text
+name (Col n) = n
+name other =
+    error $
+        "You must call `name` on a column reference. Not the expression: " ++ show other
+
+col :: (Columnable a) => T.Text -> Expr a
+col = Col
+
+ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
+ifThenElse = If
+
+lit :: (Columnable a) => a -> Expr a
+lit = Lit
+
 (.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr
 (.=) = flip as