Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion src/datasets/features/pdf.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from ..download.download_config import DownloadConfig
from ..table import array_cast
from ..utils.file_utils import is_local_path, xopen
from ..utils.py_utils import string_to_dict
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
Expand Down Expand Up @@ -216,6 +216,38 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr
storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
return array_cast(storage, self.pa_type)

def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
"""Embed PDF files into the Arrow array.

Args:
storage (`pa.StructArray`):
PyArrow array to embed.

Returns:
`pa.StructArray`: Array in the PDF arrow storage type, that is
`pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
"""

@no_op_if_value_is_null
def path_to_bytes(path):
with xopen(path, "rb") as f:
bytes_ = f.read()
return bytes_

bytes_array = pa.array(
[
(path_to_bytes(x["path"]) if x["bytes"] is None else x["bytes"]) if x is not None else None
for x in storage.to_pylist()
],
type=pa.binary(),
)
path_array = pa.array(
[os.path.basename(path) if path is not None else None for path in storage.field("path").to_pylist()],
type=pa.string(),
)
storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null())
return array_cast(storage, self.pa_type)


def encode_pdfplumber_pdf(pdf: "pdfplumber.pdf.PDF") -> dict:
"""
Expand Down
Loading