Function ParseODirectHeader

dali/util/numpy.cc:200–248

}

void ParseODirectHeader(HeaderData &parsed_header, InputStream *src, size_t o_direct_alignm,
                        size_t o_direct_read_len_alignm) {
  const size_t token_len = 6 + 1 + 1 + 2;
  size_t token_read_len = align_up(token_len + 1, o_direct_read_len_alignm);
  auto token_mem =
      mm::alloc_raw_shared<char, mm::memory_kind::host>(token_read_len, o_direct_alignm);
  char *token = token_mem.get();
  auto file = dynamic_cast<ODirectFileStream *>(src);
  DALI_ENFORCE(
      file,
      "Could not read the numpy file header: expected file stream opened with O_DIRECT flag.");
  int64_t nread = file->ReadAt(token, token_read_len, 0);
  DALI_ENFORCE(nread <= static_cast<Index>(token_read_len) &&
                   nread >= static_cast<Index>(std::min(src->Size(), token_read_len)),
               make_string("Can not read header: ",
                           static_cast<Index>(std::min(src->Size(), token_read_len)), " <= ", nread,
                           " <= ", token_read_len));
  auto char_tmp = token[token_len];
  token[token_len] = '\0';

  CheckNpyVersion(token);
  auto header_len = GetHeaderLen(token);

  // The header_len can have up to 2**16 - 1 bytes. We do not support V2 headers
  // (with up to 4GB - 4 byte header len), as those are used by numpy to save structured
  // arrays (where dtype can be different for each column and the columns have arbitrary names).
  // Parsing such a dtype in the header will fail.
  // https://numpy.org/neps/nep-0001-npy-format.html
  size_t aligned_token_header_len =
      align_up(token_len + header_len + 1, std::max(o_direct_alignm, o_direct_read_len_alignm));
  // if header_len goes beyond the previously allocated and read memory reallocate and read again
  // otherwise reuse
  if (token_read_len != aligned_token_header_len) {
    token_mem = mm::alloc_raw_shared<char, mm::memory_kind::host>(aligned_token_header_len,
                                                                  o_direct_alignm);
    nread = file->ReadAt(token_mem.get(), aligned_token_header_len, 0);
    DALI_ENFORCE(nread <= static_cast<Index>(aligned_token_header_len) &&
                     nread >= static_cast<Index>(std::min(src->Size(), aligned_token_header_len)),
                 make_string("Can not read header: ",
                             static_cast<Index>(std::min(src->Size(), aligned_token_header_len)),
                             " <= ", nread, " <= ", aligned_token_header_len));
  } else {
    // restore overriden character
    token[token_len] = char_tmp;
  }
  char *header = token_mem.get() + token_len;
  header[header_len] = '\0';
  ParseHeaderItself(parsed_header, header, header_len);
}

void ParseHeader(HeaderData &parsed_header, InputStream *src) {
  // check if the file is actually a numpy file