dataframe-huggingface (empty) → 1.0.0.0
raw patch · 4 files changed
+439/−0 lines, 4 filesdep +Globdep +aesondep +base
Dependencies added: Glob, aeson, base, bytestring, dataframe-core, dataframe-huggingface, dataframe-lazy, dataframe-operations, dataframe-parquet, dataframe-parsing, directory, filepath, http-conduit, temporary, text
Files
- LICENSE +20/−0
- dataframe-huggingface.cabal +57/−0
- src/DataFrame/IO/HuggingFace.hs +307/−0
- test/Main.hs +55/−0
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2026 Michael Chavinda++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ dataframe-huggingface.cabal view
@@ -0,0 +1,57 @@+cabal-version: 2.4+name: dataframe-huggingface+version: 1.0.0.0++synopsis: Read Parquet datasets from HuggingFace into dataframes.+description:+ @DataFrame.IO.HuggingFace@ — resolve and download Parquet files from+ HuggingFace (@hf://@) datasets and load them as 'DataFrame's, eagerly+ or via the lazy/streaming engine. This is the home for the heavier+ @aeson@ and @http-conduit@ dependencies that network access requires,+ keeping @dataframe-parquet@ and @dataframe-lazy@ free of them.++bug-reports: https://github.com/mchav/dataframe/issues+license: MIT+license-file: LICENSE+author: Michael Chavinda+maintainer: mschavinda@gmail.com+copyright: (c) 2024-2026 Michael Chavinda+category: Data+tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2++common warnings+ ghc-options:+ -Wincomplete-patterns+ -Wincomplete-uni-patterns+ -Wunused-imports+ -Wunused-local-binds+ -Wunused-packages++library+ import: warnings+ exposed-modules: DataFrame.IO.HuggingFace+ build-depends: base >= 4 && < 5,+ aeson >= 0.11.0.0 && < 3,+ bytestring >= 0.11 && < 0.13,+ dataframe-core ^>= 1.1,+ dataframe-lazy ^>= 1.1,+ dataframe-operations ^>= 1.1.1,+ dataframe-parquet ^>= 1.1,+ dataframe-parsing ^>= 1.0.2,+ directory >= 1.3.0.0 && < 2,+ filepath >= 1.4 && < 2,+ Glob >= 0.10 && < 1,+ http-conduit >= 2.3 && < 3,+ temporary >= 1.3 && < 2,+ text >= 2.1 && < 3+ hs-source-dirs: src+ default-language: Haskell2010++test-suite tests+ import: warnings+ type: exitcode-stdio-1.0+ main-is: Main.hs+ hs-source-dirs: test+ build-depends: base >= 4 && < 5,+ dataframe-huggingface+ default-language: Haskell2010
+ src/DataFrame/IO/HuggingFace.hs view
@@ -0,0 +1,307 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}++{- | Read Parquet datasets from HuggingFace (@hf://@) into 'DataFrame's.++A @hf://@ URI has the form @hf:\/\/datasets\/{owner}\/{dataset}\/{glob}@.+When the glob contains wildcards the dataset's file list is resolved via the+HuggingFace datasets-server API; otherwise the file is fetched directly from+the repo. Resolved files are downloaded to a temporary location and then read+with the local "DataFrame.IO.Parquet" reader.++This module is the home for the heavier @aeson@ and @http-conduit@+dependencies, keeping @dataframe-parquet@ and @dataframe-lazy@ free of them.+-}+module DataFrame.IO.HuggingFace (+ -- * Eager readers+ readParquet,+ readParquetWithOpts,+ readParquetFiles,+ readParquetFilesWithOpts,++ -- * Lazy / streaming reader+ scanParquet,++ -- * Read options (re-exported from "DataFrame.IO.Parquet")+ ParquetReadOptions (..),+ defaultParquetReadOptions,++ -- * URI helpers+ HFRef (..),+ HFParquetFile (..),+ isHFUri,+ parseHFUri,+ hasGlob,+ directHFUrl,+ hfUrlRepoPath,++ -- * Resolution and download+ getHFToken,+ resolveHFUrls,+ downloadHFFiles,+ fetchHFParquetFiles,+) where++import Control.Exception (try)+import Control.Monad (forM, when)+import Data.Aeson (FromJSON (..), eitherDecodeStrict, withObject, (.:))+import qualified Data.ByteString as BS+import qualified Data.List as L+import qualified Data.Text as T+import Data.Text.Encoding (encodeUtf8)+import DataFrame.IO.Parquet (+ ParquetReadOptions (..),+ defaultParquetReadOptions,+ )+import qualified DataFrame.IO.Parquet as Parquet+import DataFrame.Internal.DataFrame (DataFrame)+import DataFrame.Internal.Schema (Schema)+import qualified DataFrame.Lazy as Lazy+import DataFrame.Operations.Merge ()+import qualified DataFrame.Operations.Subset as DS+import Network.HTTP.Simple (+ getResponseBody,+ getResponseStatusCode,+ httpBS,+ parseRequest,+ setRequestHeader,+ )+import System.Directory (getHomeDirectory, getTemporaryDirectory)+import System.Environment (lookupEnv)+import System.FilePath ((</>))+import System.FilePath.Glob (compile, match)+import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory)++-- Eager readers -----------------------------------------------------------++{- | Read a Parquet dataset from a @hf://@ URI (or a local path) into a+'DataFrame'.++==== __Example__+@+ghci> readParquet "hf://datasets/owner/dataset/data/train-*.parquet"+@+-}+readParquet :: FilePath -> IO DataFrame+readParquet = readParquetWithOpts defaultParquetReadOptions++-- | Read a Parquet dataset from a @hf://@ URI (or local path) with options.+readParquetWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame+readParquetWithOpts opts path+ | isHFUri path = readDownloaded opts =<< fetchHFParquetFiles path+ | otherwise = Parquet.readParquetWithOpts opts path++{- | Read multiple Parquet files from a @hf://@ URI (or a local directory/glob).++For multi-file reads, @rowRange@ is applied once after concatenation (global+range semantics).+-}+readParquetFiles :: FilePath -> IO DataFrame+readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions++-- | Read multiple Parquet files from a @hf://@ URI (or local dir/glob) with options.+readParquetFilesWithOpts :: ParquetReadOptions -> FilePath -> IO DataFrame+readParquetFilesWithOpts opts path+ | isHFUri path = readDownloaded opts =<< fetchHFParquetFiles path+ | otherwise = Parquet.readParquetFilesWithOpts opts path++{- | Read a list of downloaded local files, concatenate them, and apply the+global row range once over the result. -}+readDownloaded :: ParquetReadOptions -> [FilePath] -> IO DataFrame+readDownloaded opts files = do+ let optsNoRange = opts{rowRange = Nothing}+ dfs <- mapM (Parquet.readParquetWithOpts optsNoRange) files+ pure (applyRowRange opts (mconcat dfs))++applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame+applyRowRange opts df =+ maybe df (`DS.range` df) (rowRange opts)++-- Lazy / streaming reader -------------------------------------------------++{- | Build a lazy 'Lazy.LazyDataFrame' over a @hf://@ dataset (or local path).++The HuggingFace files are resolved and downloaded to a temporary directory up+front, then scanned lazily from local disk — so the query runs in constant+memory but the whole dataset is fetched before scanning begins.+-}+scanParquet :: Schema -> T.Text -> IO Lazy.LazyDataFrame+scanParquet schema uri+ | isHFUri (T.unpack uri) = do+ dir <- downloadToDir (T.unpack uri)+ pure (Lazy.scanParquet schema (T.pack dir))+ | otherwise = pure (Lazy.scanParquet schema uri)++-- | Resolve a @hf://@ URI and download all matching files into a fresh temp dir.+downloadToDir :: FilePath -> IO FilePath+downloadToDir uri = do+ (mToken, files) <- resolveFiles uri+ tmp <- getCanonicalTemporaryDirectory+ dir <- createTempDirectory tmp "hf-parquet"+ _ <- downloadHFFilesTo dir mToken files+ pure dir++-- HuggingFace resolution --------------------------------------------------++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)++-- | Download files into the system temp directory, returning the local paths.+downloadHFFiles :: Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath]+downloadHFFiles mToken files = do+ tmpDir <- getTemporaryDirectory+ downloadHFFilesTo tmpDir mToken files++-- | Download files into @destDir@, returning the local paths.+downloadHFFilesTo :: FilePath -> Maybe BS.ByteString -> [HFParquetFile] -> IO [FilePath]+downloadHFFilesTo destDir mToken files =+ forM files $ \f -> do+ -- Derive a collision-resistant 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 = destDir </> 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++{- | Resolve a @hf://@ URI to the token and the list of files to fetch, without+downloading. -}+resolveFiles :: FilePath -> IO (Maybe BS.ByteString, [HFParquetFile])+resolveFiles uri = do+ ref <- case parseHFUri uri of+ Left err -> ioError (userError err)+ Right r -> pure r+ mToken <- getHFToken+ files <-+ if hasGlob (hfGlob ref)+ then do+ hfFiles <- resolveHFUrls mToken ref+ when (null hfFiles) $+ ioError $ userError $ "No parquet files found for " ++ uri+ pure hfFiles+ else do+ let url = directHFUrl ref+ let filename = last $ T.splitOn "/" (hfGlob ref)+ pure [HFParquetFile url "" "" filename]+ pure (mToken, files)++-- | Resolve and download all matching files for a @hf://@ URI to temp paths.+fetchHFParquetFiles :: FilePath -> IO [FilePath]+fetchHFParquetFiles uri = do+ (mToken, files) <- resolveFiles uri+ downloadHFFiles mToken files
+ test/Main.hs view
@@ -0,0 +1,55 @@+{-# LANGUAGE OverloadedStrings #-}++-- Pure unit tests for the HuggingFace URI helpers. The network paths+-- (resolve/download) need live HTTP and are not exercised here.+module Main (main) where++import Control.Monad (unless)+import DataFrame.IO.HuggingFace+import System.Exit (exitFailure)++check :: String -> Bool -> IO Bool+check name ok = do+ unless ok $ putStrLn ("FAIL: " ++ name)+ pure ok++main :: IO ()+main = do+ let glob = "hf://datasets/owner/ds/data/train-*.parquet"+ noGlob = "hf://datasets/owner/ds/data/train.parquet"+ bare = "hf://datasets/owner/ds"+ results <-+ sequence+ [ check "isHFUri hf" (isHFUri glob)+ , check "isHFUri local" (not (isHFUri "/tmp/local.parquet"))+ , check "parseHFUri ref" $+ case parseHFUri glob of+ Right r ->+ hfOwner r == "owner"+ && hfDataset r == "ds"+ && hfGlob r == "data/train-*.parquet"+ Left _ -> False+ , check "parseHFUri bare is Left" $+ case parseHFUri bare of+ Left _ -> True+ Right _ -> False+ , check "hasGlob wildcard" (hasGlob "data/train-*.parquet")+ , check "hasGlob none" (not (hasGlob "data/train.parquet"))+ , check "directHFUrl" $+ case parseHFUri noGlob of+ Right r ->+ directHFUrl r+ == "https://huggingface.co/datasets/owner/ds/resolve/main/data/train.parquet"+ Left _ -> False+ , check "hfUrlRepoPath" $+ hfUrlRepoPath+ ( HFParquetFile+ "https://huggingface.co/datasets/o/d/resolve/main/data/train-0.parquet"+ ""+ ""+ ""+ )+ == "data/train-0.parquet"+ ]+ unless (and results) exitFailure+ putStrLn "dataframe-huggingface: all pure-helper tests passed"