header.cc 2.5 KB
Newer Older
D
"init"  
dzhwinter 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/fluid/recordio/header.h"
16 17 18

#include <string>

Y
Yu Yang 已提交
19
#include "paddle/fluid/platform/enforce.h"
D
"init"  
dzhwinter 已提交
20 21 22 23 24 25 26 27 28 29 30 31 32

namespace paddle {
namespace recordio {

Header::Header()
    : num_records_(0),
      checksum_(0),
      compressor_(Compressor::kNoCompress),
      compress_size_(0) {}

Header::Header(uint32_t num, uint32_t sum, Compressor c, uint32_t cs)
    : num_records_(num), checksum_(sum), compressor_(c), compress_size_(cs) {}

Y
Yu Yang 已提交
33
bool Header::Parse(std::istream& is) {
Y
Yu Yang 已提交
34
  uint32_t magic;
F
fengjiayi 已提交
35 36
  is.read(reinterpret_cast<char*>(&magic), sizeof(uint32_t));
  size_t read_size = is.gcount();
Y
Yu Yang 已提交
37 38 39
  if (read_size < sizeof(uint32_t)) {
    return false;
  }
Y
Yu Yang 已提交
40 41
  PADDLE_ENFORCE_EQ(magic, kMagicNumber);

Y
Yu Yang 已提交
42 43 44 45
  is.read(reinterpret_cast<char*>(&num_records_), sizeof(uint32_t))
      .read(reinterpret_cast<char*>(&checksum_), sizeof(uint32_t))
      .read(reinterpret_cast<char*>(&compressor_), sizeof(uint32_t))
      .read(reinterpret_cast<char*>(&compress_size_), sizeof(uint32_t));
Y
Yu Yang 已提交
46
  return true;
D
"init"  
dzhwinter 已提交
47 48
}

Y
Yu Yang 已提交
49
void Header::Write(std::ostream& os) const {
Y
Yu Yang 已提交
50 51
  os.write(reinterpret_cast<const char*>(&kMagicNumber), sizeof(uint32_t))
      .write(reinterpret_cast<const char*>(&num_records_), sizeof(uint32_t))
Y
Yu Yang 已提交
52 53 54
      .write(reinterpret_cast<const char*>(&checksum_), sizeof(uint32_t))
      .write(reinterpret_cast<const char*>(&compressor_), sizeof(uint32_t))
      .write(reinterpret_cast<const char*>(&compress_size_), sizeof(uint32_t));
D
"init"  
dzhwinter 已提交
55 56 57
}

std::ostream& operator<<(std::ostream& os, Header h) {
Y
Yu Yang 已提交
58 59
  os << "Header: " << h.NumRecords() << ", " << h.Checksum() << ", "
     << static_cast<uint32_t>(h.CompressType()) << ", " << h.CompressSize();
D
"init"  
dzhwinter 已提交
60 61 62 63 64 65 66 67
  return os;
}

bool operator==(Header l, Header r) {
  return l.NumRecords() == r.NumRecords() && l.Checksum() == r.Checksum() &&
         l.CompressType() == r.CompressType() &&
         l.CompressSize() == r.CompressSize();
}
68 69 70

}  // namespace recordio
}  // namespace paddle