diff --git a/dataframe-parquet.cabal b/dataframe-parquet.cabal --- a/dataframe-parquet.cabal +++ b/dataframe-parquet.cabal @@ -1,15 +1,16 @@ cabal-version: 2.4 name: dataframe-parquet -version: 1.0.1.2 +version: 1.1.0.0 synopsis: Parquet reader and writer for the dataframe ecosystem. description: @DataFrame.IO.Parquet@ — pure-Haskell Parquet 2.0 reader and writer (with snappy, zstd, gzip codecs, dictionary decoding, nested - list/repeated columns, predicate pushdown, and HTTP range reads). - Heavy package — pulls in @pinch@, @streamly@, compression codecs, - and @http-conduit@. Most users want @dataframe-csv@ instead unless - they specifically need Parquet. + list/repeated columns, and predicate pushdown). + Heavy package — pulls in @pinch@, @streamly@, and compression codecs. + Reading directly from HuggingFace (@hf://@) datasets lives in the + separate @dataframe-huggingface@ package. Most users want + @dataframe-csv@ instead unless they specifically need Parquet. bug-reports: https://github.com/mchav/dataframe/issues license: MIT @@ -45,7 +46,6 @@ DataFrame.IO.Parquet.Utils DataFrame.IO.Utils.RandomAccess build-depends: base >= 4 && < 5, - aeson >= 0.11.0.0 && < 3, bytestring >= 0.11 && < 0.13, containers >= 0.6.7 && < 0.9, dataframe-core ^>= 1.1, @@ -54,12 +54,11 @@ directory >= 1.3.0.0 && < 2, filepath >= 1.4 && < 2, Glob >= 0.10 && < 1, - http-conduit >= 2.3 && < 3, pinch >= 0.5 && < 1, snappy-hs ^>= 0.1, streamly-bytestring >= 0.2.0 && < 0.4, streamly-core >= 0.2.3 && < 0.4, - text >= 2.0 && < 3, + text >= 2.1 && < 3, time >= 1.12 && < 2, vector ^>= 0.13, zlib >= 0.5 && < 1, 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 @@ -9,10 +9,9 @@ module DataFrame.IO.Parquet where -import Control.Exception (throw, try) +import Control.Exception (throw) import Control.Monad import Control.Monad.IO.Class (MonadIO (..)) -import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:)) import Data.Bits (Bits (shiftL), (.|.)) import qualified Data.ByteString as BS import Data.Either (fromRight) @@ -26,7 +25,6 @@ import qualified Data.List as L import qualified Data.Map as Map import qualified Data.Text as T -import Data.Text.Encoding (encodeUtf8) import Data.Time (UTCTime) import Data.Time.Clock.POSIX (posixSecondsToUTCTime) import qualified Data.Vector as Vector @@ -82,23 +80,11 @@ import DataFrame.Internal.Expression (Expr, getColumns) import DataFrame.Operations.Merge () import qualified DataFrame.Operations.Subset as DS -import Network.HTTP.Simple ( - getResponseBody, - getResponseStatusCode, - httpBS, - parseRequest, - setRequestHeader, - ) import qualified Pinch import qualified Streamly.Data.Stream as Stream -import System.Directory ( - doesDirectoryExist, - getHomeDirectory, - getTemporaryDirectory, - ) -import System.Environment (lookupEnv) +import System.Directory (doesDirectoryExist) import System.FilePath (()) -import System.FilePath.Glob (compile, glob, match) +import System.FilePath.Glob (glob) import System.IO (IOMode (ReadMode)) -- Options ----------------------------------------------------------------- @@ -175,13 +161,7 @@ are auto-included for decoding, then projected back to the requested output columns. -} readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame -readParquetWithOpts opts path - | isHFUri path = do - paths <- fetchHFParquetFiles path - let optsNoRange = opts{rowRange = Nothing} - dfs <- mapM (_readParquetWithOpts Nothing optsNoRange) paths - pure (applyRowRange opts (mconcat dfs)) - | otherwise = _readParquetWithOpts Nothing opts path +readParquetWithOpts = _readParquetWithOpts Nothing -- | Internal entry point used by tests to force non-seekable mode. _readParquetWithOpts :: @@ -212,29 +192,23 @@ @ -} readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame -readParquetFilesWithOpts opts path - | isHFUri path = do - files <- fetchHFParquetFiles path - let optsWithoutRowRange = opts{rowRange = Nothing} - dfs <- mapM (_readParquetWithOpts Nothing optsWithoutRowRange) files - pure (applyRowRange opts (mconcat dfs)) - | otherwise = do - isDir <- doesDirectoryExist path +readParquetFilesWithOpts opts path = do + isDir <- doesDirectoryExist path - let pat = if isDir then path "*.parquet" else path + let pat = if isDir then path "*.parquet" else path - matches <- glob pat + matches <- glob pat - files <- filterM (fmap not . doesDirectoryExist) matches + files <- filterM (fmap not . doesDirectoryExist) matches - case files of - [] -> - error $ - "readParquetFiles: no parquet files found for " ++ path - _ -> do - let optsWithoutRowRange = opts{rowRange = Nothing} - dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files - pure (applyRowRange opts (mconcat dfs)) + case files of + [] -> + error $ + "readParquetFiles: no parquet files found for " ++ path + _ -> do + let optsWithoutRowRange = opts{rowRange = Nothing} + dfs <- mapM (readParquetWithOpts optsWithoutRowRange) files + pure (applyRowRange opts (mconcat dfs)) -- Core parsing pipeline --------------------------------------------------- @@ -586,155 +560,3 @@ microsecondsToUTCTime :: Int64 -> UTCTime microsecondsToUTCTime us = posixSecondsToUTCTime (fromIntegral us / 1_000_000) - --- HuggingFace support ----------------------------------------------------- - -data HFRef = HFRef - { hfOwner :: T.Text - , hfDataset :: T.Text - , hfGlob :: T.Text - } - -data HFParquetFile = HFParquetFile - { hfpUrl :: T.Text - , hfpConfig :: T.Text - , hfpSplit :: T.Text - , hfpFilename :: T.Text - } - deriving (Show) - -instance FromJSON HFParquetFile where - parseJSON = withObject "HFParquetFile" $ \o -> - HFParquetFile - <$> o .: "url" - <*> o .: "config" - <*> o .: "split" - <*> o .: "filename" - -newtype HFParquetResponse = HFParquetResponse {hfParquetFiles :: [HFParquetFile]} - -instance FromJSON HFParquetResponse where - parseJSON = withObject "HFParquetResponse" $ \o -> - HFParquetResponse <$> o .: "parquet_files" - -isHFUri :: FilePath -> Bool -isHFUri = L.isPrefixOf "hf://" - -parseHFUri :: FilePath -> Either String HFRef -parseHFUri path = - let stripped = drop (length ("hf://datasets/" :: String)) path - in case T.splitOn "/" (T.pack stripped) of - (owner : dataset : rest) - | not (null rest) -> - Right $ HFRef owner dataset (T.intercalate "/" rest) - _ -> - Left $ "Invalid hf:// URI (expected hf://datasets/owner/dataset/glob): " ++ path - -getHFToken :: IO (Maybe BS.ByteString) -getHFToken = do - envToken <- lookupEnv "HF_TOKEN" - case envToken of - Just t -> pure (Just (encodeUtf8 (T.pack t))) - Nothing -> do - home <- getHomeDirectory - let tokenPath = home ".cache" "huggingface" "token" - result <- try (BS.readFile tokenPath) :: IO (Either IOError BS.ByteString) - case result of - Right bs -> pure (Just (BS.takeWhile (/= 10) bs)) - Left _ -> pure Nothing - -{- | Extract the repo-relative path from a HuggingFace download URL. -URL format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/{ref}/{path} -Returns the {path} portion (e.g. "data/train-00000-of-00001.parquet"). --} -hfUrlRepoPath :: HFParquetFile -> String -hfUrlRepoPath f = - case T.breakOn "/resolve/" (hfpUrl f) of - (_, rest) - | not (T.null rest) -> - -- Drop "/resolve/", then drop the ref component (up to and including "/") - T.unpack $ T.drop 1 $ T.dropWhile (/= '/') $ T.drop (T.length "/resolve/") rest - _ -> - T.unpack (hfpConfig f) T.unpack (hfpSplit f) T.unpack (hfpFilename f) - -matchesGlob :: T.Text -> HFParquetFile -> Bool -matchesGlob g f = match (compile (T.unpack g)) (hfUrlRepoPath f) - -resolveHFUrls :: Maybe BS.ByteString -> HFRef -> IO [HFParquetFile] -resolveHFUrls mToken ref = do - let dataset = hfOwner ref <> "/" <> hfDataset ref - let apiUrl = "https://datasets-server.huggingface.co/parquet?dataset=" ++ T.unpack dataset - req0 <- parseRequest apiUrl - let req = case mToken of - Nothing -> req0 - Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 - resp <- httpBS req - let status = getResponseStatusCode resp - when (status /= 200) $ - ioError $ - userError $ - "HuggingFace API returned status " - ++ show status - ++ " for dataset " - ++ T.unpack dataset - case eitherDecodeStrict (getResponseBody resp) of - Left err -> ioError $ userError $ "Failed to parse HF API response: " ++ err - Right hfResp -> pure $ filter (matchesGlob (hfGlob ref)) (hfParquetFiles hfResp) - -downloadHFFiles :: Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath] -downloadHFFiles mToken files = do - tmpDir <- getTemporaryDirectory - forM files $ \f -> do - -- Derive a collision-resistant temp name from the URL path components - let fname = case (hfpConfig f, hfpSplit f) of - (c, s) | T.null c && T.null s -> T.unpack (hfpFilename f) - (c, s) -> T.unpack c <> "_" <> T.unpack s <> "_" <> T.unpack (hfpFilename f) - let destPath = tmpDir fname - req0 <- parseRequest (T.unpack (hfpUrl f)) - let req = case mToken of - Nothing -> req0 - Just tok -> setRequestHeader "Authorization" ["Bearer " <> tok] req0 - resp <- httpBS req - let status = getResponseStatusCode resp - when (status /= 200) $ - ioError $ - userError $ - "Failed to download " ++ T.unpack (hfpUrl f) ++ " (HTTP " ++ show status ++ ")" - BS.writeFile destPath (getResponseBody resp) - pure destPath - --- | True when the path contains glob wildcard characters. -hasGlob :: T.Text -> Bool -hasGlob = T.any (\c -> c == '*' || c == '?' || c == '[') - -{- | Build the direct HF repo download URL for a path with no wildcards. -Format: https://huggingface.co/datasets/{owner}/{dataset}/resolve/main/{path} --} -directHFUrl :: HFRef -> T.Text -directHFUrl ref = - "https://huggingface.co/datasets/" - <> hfOwner ref - <> "/" - <> hfDataset ref - <> "/resolve/main/" - <> hfGlob ref - -fetchHFParquetFiles :: FilePath -> IO [FilePath] -fetchHFParquetFiles uri = do - ref <- case parseHFUri uri of - Left err -> ioError (userError err) - Right r -> pure r - mToken <- getHFToken - if hasGlob (hfGlob ref) - then do - hfFiles <- resolveHFUrls mToken ref - when (null hfFiles) $ - ioError $ - userError $ - "No parquet files found for " ++ uri - downloadHFFiles mToken hfFiles - else do - -- Direct repo file download — no datasets-server needed - let url = directHFUrl ref - let filename = last $ T.splitOn "/" (hfGlob ref) - downloadHFFiles mToken [HFParquetFile url "" "" filename]