diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: dataframe-parquet -version: 1.6.0.0 +version: 1.6.1.0 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer @@ -60,8 +60,8 @@ build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, containers >= 0.6.7 && < 0.10, - dataframe-core >= 2.6 && < 2.7, - dataframe-operations >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, + dataframe-operations >= 2.7 && < 2.8, dataframe-parsing >= 2.2 && < 2.3, primitive >= 0.7 && < 0.11, directory >= 1.3.0.0 && < 2, @@ -85,8 +85,8 @@ hs-source-dirs: tests build-depends: base >= 4 && < 5, bytestring >= 0.11 && < 0.14, - dataframe-core >= 2.6 && < 2.7, - dataframe-operations >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, + dataframe-operations >= 2.7 && < 2.8, dataframe-parquet, directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, @@ -101,7 +101,7 @@ other-modules: DataFrame10GB hs-source-dirs: stress build-depends: base >= 4 && < 5, - dataframe-core >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, dataframe-parquet, filepath >= 1.4 && < 2, temporary >= 1.3 && < 1.5, @@ -120,7 +120,7 @@ build-depends: base >= 4 && < 5, criterion >= 1 && < 2, deepseq >= 1.4 && < 2, - dataframe-core >= 2.6 && < 2.7, + dataframe-core >= 2.7 && < 2.8, dataframe-parquet, directory >= 1.3 && < 2, filepath >= 1.4 && < 2, diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs --- a/src/DataFrame/IO/Parquet.hs +++ b/src/DataFrame/IO/Parquet.hs @@ -175,8 +175,10 @@ import DataFrame.Internal.Control.Concurrent (forkJoin_) import DataFrame.Internal.DataFrame (DataFrame (..)) import DataFrame.Internal.Expression ( - BinaryOp (binaryName), + BinOpId (..), + BinaryOp (binaryId), Expr (..), + builtinOf, getColumns, ) import DataFrame.Operations.Merge () @@ -419,7 +421,8 @@ chunks = unField (rg_columns rg) go :: Expr Bool -> Bool go (Binary op l r) - | binaryName op == "and" = case (l, r) of + | Just tag <- builtinOf (binaryId op) + , tag `elem` [BAnd, BNullAnd] = case (l, r) of (l' :: Expr lt, r' :: Expr rt) | Just Refl <- testEquality (typeRep @lt) (typeRep @Bool) , Just Refl <- testEquality (typeRep @rt) (typeRep @Bool) -> @@ -429,31 +432,31 @@ (Col name :: Expr lt, Lit (v :: rt)) | Just Refl <- testEquality (typeRep @lt) (typeRep @Int) , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) -> - decide name (binaryName op) v + maybe True (\t -> decide name t v) (builtinOf (binaryId op)) (Lit (v :: lt), Col name :: Expr rt) | Just Refl <- testEquality (typeRep @lt) (typeRep @Int) , Just Refl <- testEquality (typeRep @rt) (typeRep @Int) -> - decide name (flipped (binaryName op)) v + maybe True (\t -> decide name (flipped t) v) (builtinOf (binaryId op)) _ -> True go _ = True - flipped :: T.Text -> T.Text + flipped :: BinOpId -> BinOpId flipped n = case n of - "lt" -> "gt" - "gt" -> "lt" - "leq" -> "geq" - "geq" -> "leq" + BLt -> BGt + BGt -> BLt + BLeq -> BGeq + BGeq -> BLeq other -> other -- @column v@ can hold for some row in [lo, hi]. - decide :: T.Text -> T.Text -> Int -> Bool + decide :: T.Text -> BinOpId -> Int -> Bool decide name cmp v = case columnRange name of Nothing -> True Just (lo, hi) -> case cmp of - "eq" -> lo <= v && v <= hi - "neq" -> not (lo == v && hi == v) - "lt" -> lo < v - "leq" -> lo <= v - "gt" -> hi > v - "geq" -> hi >= v + BEq -> lo <= v && v <= hi + BNeq -> not (lo == v && hi == v) + BLt -> lo < v + BLeq -> lo <= v + BGt -> hi > v + BGeq -> hi >= v _ -> True columnRange :: T.Text -> Maybe (Int, Int) columnRange name = do @@ -896,7 +899,7 @@ applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame applyPredicate opts df = - maybe df (`DS.filterWhere` df) (predicate opts) + maybe df (`DS.filter` df) (predicate opts) applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame applySafeRead opts df diff --git a/src/DataFrame/Typed/IO/Parquet.hs b/src/DataFrame/Typed/IO/Parquet.hs --- a/src/DataFrame/Typed/IO/Parquet.hs +++ b/src/DataFrame/Typed/IO/Parquet.hs @@ -8,15 +8,22 @@ readParquetWithError, readParquetWithOpts, readParquetFiles, + readParquetFilesWithOpts, + writeParquet, + writeParquetWithOptions, ) where import Control.Applicative ((<|>)) import Control.Exception (SomeException, try) import qualified Data.Text as T -import DataFrame.IO.Parquet (ParquetReadOptions (..), defaultParquetReadOptions) +import DataFrame.IO.Parquet ( + ParquetReadOptions (..), + ParquetWriteOptions, + defaultParquetReadOptions, + ) import qualified DataFrame.IO.Parquet as Parquet -import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError) +import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError, thaw) import DataFrame.Typed.Schema (KnownSchema, schemaColumnNames) import DataFrame.Typed.Types (TypedDataFrame) @@ -58,8 +65,23 @@ -- | Read a directory\/glob of Parquet files into a typed DataFrame. readParquetFiles :: forall cols. (KnownSchema cols) => FilePath -> IO (TypedDataFrame cols) -readParquetFiles path = - Parquet.readParquetFilesWithOpts - (schemaOptions @cols defaultParquetReadOptions) - path +readParquetFiles = readParquetFilesWithOpts @cols defaultParquetReadOptions + +-- | Read a directory\/glob of Parquet files with custom options. +readParquetFilesWithOpts :: + forall cols. + (KnownSchema cols) => + ParquetReadOptions -> FilePath -> IO (TypedDataFrame cols) +readParquetFilesWithOpts opts path = + Parquet.readParquetFilesWithOpts (schemaOptions @cols opts) path >>= freezeOrThrow @cols + +-- | Write a typed DataFrame to a Parquet file. +writeParquet :: FilePath -> TypedDataFrame cols -> IO () +writeParquet path = Parquet.writeParquet path . thaw + +-- | Write a typed DataFrame to a Parquet file with custom options. +writeParquetWithOptions :: + ParquetWriteOptions -> FilePath -> TypedDataFrame cols -> IO () +writeParquetWithOptions opts path = + Parquet.writeParquetWithOptions opts path . thaw