__init__.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. import struct
  2. import thrift
  3. import logging
  4. from ttypes import FileMetaData
  5. from thrift.protocol import TCompactProtocol
  6. from thrift.transport import TTransport
  7. logger = logging.getLogger("parquet")
  8. def _check_header_magic_bytes(fo):
  9. """Returns true if the file-like obj has the PAR1 magic bytes at the header"""
  10. fo.seek(0, 0)
  11. magic = fo.read(4)
  12. return magic == 'PAR1'
  13. def _check_footer_magic_bytes(fo):
  14. """Returns true if the file-like obj has the PAR1 magic bytes at the footer"""
  15. fo.seek(-4, 2) # seek to four bytes from the end of the file
  16. magic = fo.read(4)
  17. return magic == 'PAR1'
  18. def _get_footer_size(fo):
  19. """Readers the footer size in bytes, which is serialized as little endian"""
  20. fo.seek(-8, 2)
  21. tup = struct.unpack("<i", fo.read(4))
  22. return tup[0]
  23. def _read_footer(fo):
  24. """Reads the footer from the given file object, returning a FileMetaData object. This method
  25. assumes that the fo references a valid parquet file"""
  26. footer_size = _get_footer_size(fo)
  27. logger.debug("Footer size in bytes: %s", footer_size)
  28. fo.seek(-(8+footer_size), 2) # seek to beginning of footer
  29. tin = TTransport.TFileObjectTransport(fo)
  30. pin = TCompactProtocol.TCompactProtocol(tin)
  31. fmd = FileMetaData()
  32. fmd.read(pin)
  33. return fmd
  34. def read_footer(filename):
  35. with open(filename, 'rb') as fo:
  36. if not _check_header_magic_bytes(fo) or not _check_footer_magic_bytes(fo):
  37. raise ParquetFormatException("%s is not a valid parquet file (missing magic bytes".format(filename))
  38. return _read_footer(fo)