dataframe-1.3.0.0: src/DataFrame/IO/Parquet.hs
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module DataFrame.IO.Parquet where
import Control.Exception (throw, try)
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)
import Data.Functor ((<&>))
import Data.Int (Int32, Int64)
import Data.List (foldl', transpose)
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
import qualified Data.Vector.Unboxed as VU
import DataFrame.Errors (DataFrameException (ColumnsNotFoundException))
import DataFrame.IO.Parquet.Page (
PageDecoder,
UnboxedPageDecoder,
boolDecoder,
byteArrayDecoder,
doubleDecoder,
fixedLenByteArrayDecoder,
floatDecoder,
int32Decoder,
int64Decoder,
int96Decoder,
readPages,
)
import DataFrame.IO.Parquet.Seeking (
FileBufferedOrSeekable,
ForceNonSeekable,
withFileBufferedOrSeekable,
)
import DataFrame.IO.Parquet.Thrift (
ColumnChunk (..),
DecimalType (..),
FileMetadata (..),
LogicalType (..),
RowGroup (..),
ThriftType (..),
TimeUnit (..),
TimestampType (..),
unField,
)
import DataFrame.IO.Parquet.Utils (
ColumnDescription (..),
foldNonNullable,
foldNonNullableUnboxed,
foldNullable,
foldNullableUnboxed,
foldRepeated,
foldRepeatedUnboxed,
generateColumnDescriptions,
getColumnNames,
)
import DataFrame.IO.Utils.RandomAccess (
RandomAccess (..),
ReaderIO (runReaderIO),
)
import DataFrame.Internal.Column (Column, Columnable)
import qualified DataFrame.Internal.Column as DI
import DataFrame.Internal.DataFrame (DataFrame (..))
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.FilePath ((</>))
import System.FilePath.Glob (compile, glob, match)
import System.IO (IOMode (ReadMode))
-- Options -----------------------------------------------------------------
{- | Options for reading Parquet data.
These options are applied in this order:
1. predicate filtering
2. column projection
3. row range
4. safe column promotion
Column selection for @selectedColumns@ uses leaf column names only.
-}
data ParquetReadOptions = ParquetReadOptions
{ selectedColumns :: Maybe [T.Text]
{- ^ Columns to keep in the final dataframe. If set, only these columns are returned.
Predicate-referenced columns are read automatically when needed and projected out after filtering.
-}
, predicate :: Maybe (Expr Bool)
-- ^ Optional row filter expression applied before projection.
, rowRange :: Maybe (Int, Int)
-- ^ Optional row slice @(start, end)@ with start-inclusive/end-exclusive semantics.
, safeColumns :: Bool
-- ^ When True, every column is promoted to OptionalColumn after read, regardless of nullability in the schema.
}
deriving (Show)
{- | Default Parquet read options.
Equivalent to:
@
ParquetReadOptions
{ selectedColumns = Nothing
, predicate = Nothing
, rowRange = Nothing
, safeColumns = False
}
@
-}
defaultParquetReadOptions :: ParquetReadOptions
defaultParquetReadOptions =
ParquetReadOptions
{ selectedColumns = Nothing
, predicate = Nothing
, rowRange = Nothing
, safeColumns = False
}
-- Public API --------------------------------------------------------------
{- | Read a parquet file from path and load it into a dataframe.
==== __Example__
@
ghci> D.readParquet ".\/data\/mtcars.parquet"
@
-}
readParquet :: FilePath -> IO DataFrame
readParquet = readParquetWithOpts defaultParquetReadOptions
{- | Read a Parquet file using explicit read options.
==== __Example__
@
ghci> D.readParquetWithOpts
ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 10)})
ghci| "./tests/data/alltypes_plain.parquet"
@
When @selectedColumns@ is set and @predicate@ references other columns, those predicate columns
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
-- | Internal entry point used by tests to force non-seekable mode.
_readParquetWithOpts ::
ForceNonSeekable -> ParquetReadOptions -> FilePath -> IO DataFrame
_readParquetWithOpts extraConfig opts path =
withFileBufferedOrSeekable extraConfig path ReadMode $ \file ->
runReaderIO (parseParquetWithOpts opts) file
{- | Read Parquet files from a directory or glob path.
This is equivalent to calling 'readParquetFilesWithOpts' with 'defaultParquetReadOptions'.
-}
readParquetFiles :: FilePath -> IO DataFrame
readParquetFiles = readParquetFilesWithOpts defaultParquetReadOptions
{- | Read multiple Parquet files (directory or glob) using explicit options.
If @path@ is a directory, all non-directory entries are read.
If @path@ is a glob, matching files are read.
For multi-file reads, @rowRange@ is applied once after concatenation (global range semantics).
==== __Example__
@
ghci> D.readParquetFilesWithOpts
ghci| (D.defaultParquetReadOptions{D.selectedColumns = Just ["id"], D.rowRange = Just (0, 5)})
ghci| "./tests/data/alltypes_plain*.parquet"
@
-}
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
let pat = if isDir then path </> "*.parquet" else path
matches <- glob pat
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))
-- Core parsing pipeline ---------------------------------------------------
{- | Parse a Parquet file via the 'RandomAccess' handle, applying all
read options. This is the central parsing entry point used by
'_readParquetWithOpts'.
-}
parseParquetWithOpts ::
(RandomAccess m, MonadIO m) =>
ParquetReadOptions ->
m DataFrame
parseParquetWithOpts opts = do
metadata <- parseFileMetadata
let schemaElems = unField metadata.schema
allNames = getColumnNames (drop 1 schemaElems)
leafNames = L.nub (map (last . T.splitOn ".") allNames)
predicateColumns = maybe [] (L.nub . getColumns) (predicate opts)
selectedColumnsForRead = case selectedColumns opts of
Nothing -> Nothing
Just selected -> Just (L.nub (selected ++ predicateColumns))
-- TODO: When selectedColumnsForRead is Just, pass the set of required
-- column indices into the chunk parsers so that RandomAccess reads are
-- skipped for columns not in the selection, rather than decoding all
-- columns and projecting afterward.
-- TODO: When rowRange is set, compute cumulative row offsets from
-- rg_num_rows in each RowGroup and skip any group whose row interval does
-- not overlap the requested range, avoiding all decoding for those groups.
-- TODO: When predicate is set, inspect cmd_statistics min/max values for
-- predicate-referenced columns in each RowGroup and skip groups where
-- statistics prove the predicate cannot be satisfied.
-- Validate selected columns
case selectedColumnsForRead of
Nothing -> pure ()
Just requested ->
let missing = requested L.\\ leafNames
in unless (L.null missing) $
liftIO $
throw
( ColumnsNotFoundException
missing
"readParquetWithOpts"
leafNames
)
let descriptions = generateColumnDescriptions schemaElems
chunks = columnChunksForAll metadata
nCols = length chunks
nDescs = length descriptions
unless (nCols == nDescs) $
error $
"Column count mismatch: got "
<> show nCols
<> " columns but schema implied "
<> show nDescs
<> " columns"
-- Some files omit the top-level num_rows field; fall back to summing row-group counts.
let topLevelRows = fromIntegral . unField $ metadata.num_rows :: Int
rgRows =
sum $ map (fromIntegral . unField . rg_num_rows) (unField metadata.row_groups) ::
Int
vectorLength = if topLevelRows > 0 then topLevelRows else rgRows
rawCols <- zipWithM (parseColumnChunks vectorLength) chunks descriptions
let finalCols = zipWith applyDescLogicalType descriptions rawCols
indices = Map.fromList $ zip allNames [0 ..]
dimensions = (vectorLength, length finalCols)
let df =
DataFrame
(Vector.fromListN (length finalCols) finalCols)
indices
dimensions
Map.empty
return (applyReadOptions opts df)
{- | Parse the file-level Thrift metadata from the Parquet file footer.
Validates the trailing 4-byte magic marker (\"PAR1\") before decoding.
-}
parseFileMetadata :: (RandomAccess m) => m FileMetadata
parseFileMetadata = do
footerBytes <- readSuffix 8
let magic = BS.drop 4 footerBytes
when (magic /= "PAR1") $
error
( "Not a valid Parquet file: expected magic bytes \"PAR1\", got "
++ show magic
)
let size = getMetadataSize footerBytes
rawMetadata <- readSuffix (size + 8) <&> BS.take size
case Pinch.decode Pinch.compactProtocol rawMetadata of
Left e -> error $ "Failed to parse Parquet metadata: " ++ show e
Right metadata -> return metadata
where
getMetadataSize footer =
let sizes :: [Int]
sizes = map (fromIntegral . BS.index footer) [0 .. 3]
in foldl' (.|.) 0 $ zipWith shiftL sizes [0, 8 .. 24]
-- | Read the file metadata from a Parquet file at the given path.
readMetadataFromPath :: FilePath -> IO FileMetadata
readMetadataFromPath path =
withFileBufferedOrSeekable Nothing path ReadMode $
runReaderIO parseFileMetadata
-- | Read only the file metadata from an open 'FileBufferedOrSeekable' handle.
readMetadataFromHandle :: FileBufferedOrSeekable -> IO FileMetadata
readMetadataFromHandle = runReaderIO parseFileMetadata
-- | Collect column chunks per column (transposed across all row groups).
columnChunksForAll :: FileMetadata -> [[ColumnChunk]]
columnChunksForAll =
transpose . map (unField . rg_columns) . unField . row_groups
-- | Dispatch a column's chunks to the correct decoder path.
parseColumnChunks ::
(RandomAccess m, MonadIO m) =>
Int ->
[ColumnChunk] ->
ColumnDescription ->
m Column
parseColumnChunks totalRows chunks description
| description.maxRepetitionLevel == 0 && description.maxDefinitionLevel == 0 =
getNonNullableColumn totalRows description chunks
| description.maxRepetitionLevel == 0 =
getNullableColumn totalRows description chunks
| otherwise =
getRepeatedColumn description chunks
-- | Decode a required (non-nullable, non-repeated) column.
getNonNullableColumn ::
forall m.
(RandomAccess m, MonadIO m) =>
Int ->
ColumnDescription ->
[ColumnChunk] ->
m Column
getNonNullableColumn totalRows description chunks =
case description.colElementType of
Just (BOOLEAN _) -> unboxedGo boolDecoder
Just (INT32 _) -> unboxedGo int32Decoder
Just (INT64 _) -> unboxedGo int64Decoder
Just (INT96 _) -> go int96Decoder
Just (FLOAT _) -> unboxedGo floatDecoder
Just (DOUBLE _) -> unboxedGo doubleDecoder
Just (BYTE_ARRAY _) -> go byteArrayDecoder
Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of
Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"
Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))
Nothing -> error "Column has no Parquet type"
where
go ::
forall a.
(Columnable a) =>
PageDecoder a ->
m Column
go decoder =
foldNonNullable totalRows $
(\(vs, _, _) -> vs)
<$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)
unboxedGo ::
forall a.
(Columnable a, VU.Unbox a) =>
UnboxedPageDecoder a ->
m Column
unboxedGo decoder =
foldNonNullableUnboxed totalRows $
(\(vs, _, _) -> vs)
<$> Stream.unfoldEach
(readPages description decoder)
(Stream.fromList chunks)
-- | Decode an optional (nullable) column.
getNullableColumn ::
forall m.
(RandomAccess m, MonadIO m) =>
Int ->
ColumnDescription ->
[ColumnChunk] ->
m Column
getNullableColumn totalRows description chunks =
case description.colElementType of
Just (BOOLEAN _) -> unboxedGo boolDecoder
Just (INT32 _) -> unboxedGo int32Decoder
Just (INT64 _) -> unboxedGo int64Decoder
Just (INT96 _) -> go int96Decoder
Just (FLOAT _) -> unboxedGo floatDecoder
Just (DOUBLE _) -> unboxedGo doubleDecoder
Just (BYTE_ARRAY _) -> go byteArrayDecoder
Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of
Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"
Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))
Nothing -> error "Column has no Parquet type"
where
maxDef :: Int
maxDef = fromIntegral description.maxDefinitionLevel
go ::
forall a.
(Columnable a) =>
PageDecoder a ->
m Column
go decoder =
foldNullable maxDef totalRows $
(\(vs, ds, _) -> (vs, ds))
<$> Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)
unboxedGo ::
forall a.
(Columnable a, VU.Unbox a) =>
UnboxedPageDecoder a ->
m Column
unboxedGo decoder =
foldNullableUnboxed maxDef totalRows $
(\(vs, ds, _) -> (vs, ds))
<$> Stream.unfoldEach
(readPages description decoder)
(Stream.fromList chunks)
-- | Decode a repeated (list/nested) column.
getRepeatedColumn ::
forall m.
(RandomAccess m, MonadIO m) =>
ColumnDescription ->
[ColumnChunk] ->
m Column
getRepeatedColumn description chunks =
case description.colElementType of
Just (BOOLEAN _) -> unboxedGo boolDecoder
Just (INT32 _) -> unboxedGo int32Decoder
Just (INT64 _) -> unboxedGo int64Decoder
Just (INT96 _) -> go int96Decoder
Just (FLOAT _) -> unboxedGo floatDecoder
Just (DOUBLE _) -> unboxedGo doubleDecoder
Just (BYTE_ARRAY _) -> go byteArrayDecoder
Just (FIXED_LEN_BYTE_ARRAY _) -> case description.typeLength of
Nothing -> error "FIXED_LEN_BYTE_ARRAY requires type_length to be set"
Just tl -> go (fixedLenByteArrayDecoder (fromIntegral tl))
Nothing -> error "Column has no Parquet type"
where
maxRep :: Int
maxRep = fromIntegral description.maxRepetitionLevel
maxDef :: Int
maxDef = fromIntegral description.maxDefinitionLevel
go ::
forall a.
( Columnable a
, Columnable (Maybe [Maybe a])
, Columnable (Maybe [Maybe [Maybe a]])
, Columnable (Maybe [Maybe [Maybe [Maybe a]]])
) =>
PageDecoder a ->
m Column
go decoder =
foldRepeated maxRep maxDef $
Stream.unfoldEach (readPages description decoder) (Stream.fromList chunks)
unboxedGo ::
forall a.
( VU.Unbox a
, Columnable a
, Columnable (Maybe [Maybe a])
, Columnable (Maybe [Maybe [Maybe a]])
, Columnable (Maybe [Maybe [Maybe [Maybe a]]])
) =>
UnboxedPageDecoder a ->
m Column
unboxedGo decoder =
foldRepeatedUnboxed maxRep maxDef $
Stream.unfoldEach
(readPages description decoder)
(Stream.fromList chunks)
-- Options application -----------------------------------------------------
applyRowRange :: ParquetReadOptions -> DataFrame -> DataFrame
applyRowRange opts df =
maybe df (`DS.range` df) (rowRange opts)
applySelectedColumns :: ParquetReadOptions -> DataFrame -> DataFrame
applySelectedColumns opts df =
maybe df (`DS.select` df) (selectedColumns opts)
applyPredicate :: ParquetReadOptions -> DataFrame -> DataFrame
applyPredicate opts df =
maybe df (`DS.filterWhere` df) (predicate opts)
applySafeRead :: ParquetReadOptions -> DataFrame -> DataFrame
applySafeRead opts df
| safeColumns opts = df{columns = Vector.map DI.ensureOptional (columns df)}
| otherwise = df
applyReadOptions :: ParquetReadOptions -> DataFrame -> DataFrame
applyReadOptions opts =
applySafeRead opts
. applyRowRange opts
. applySelectedColumns opts
. applyPredicate opts
-- Logical type conversion -------------------------------------------------
{- | Apply a column-description's logical type annotation to convert raw
decoded values (e.g. millisecond integers → 'UTCTime').
-}
applyDescLogicalType :: ColumnDescription -> DI.Column -> DI.Column
applyDescLogicalType desc = applyLogicalType (colLogicalType desc)
applyLogicalType :: Maybe LogicalType -> DI.Column -> DI.Column
applyLogicalType (Just (LT_TIMESTAMP f)) col =
let ts = unField f
unit = unField ts.timestamp_unit
divisor = case unit of
MILLIS _ -> 1_000
MICROS _ -> 1_000_000
NANOS _ -> 1_000_000_000
in fromRight col $
DI.mapColumn
(microsecondsToUTCTime . (* (1_000_000 `div` divisor)))
col
applyLogicalType (Just (LT_DECIMAL f)) col =
let dt = unField f
scale = unField dt.decimal_scale
precision = unField dt.decimal_precision
in if precision <= 9
then case DI.toVector @Int32 @VU.Vector col of
Right xs ->
DI.fromUnboxedVector $
VU.map (\raw -> fromIntegral @Int32 @Double raw / 10 ^ scale) xs
Left _ -> col
else
if precision <= 18
then case DI.toVector @Int64 @VU.Vector col of
Right xs ->
DI.fromUnboxedVector $
VU.map (\raw -> fromIntegral @Int64 @Double raw / 10 ^ scale) xs
Left _ -> col
else col
applyLogicalType _ col = col
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]