diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/dataframe-huggingface.cabal b/dataframe-huggingface.cabal new file mode 100644 --- /dev/null +++ b/dataframe-huggingface.cabal @@ -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 diff --git a/src/DataFrame/IO/HuggingFace.hs b/src/DataFrame/IO/HuggingFace.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/HuggingFace.hs @@ -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 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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"