Skip to content
Open
Show file tree
Hide file tree
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
10 changes: 8 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,19 @@ g++ -o mycode mycode.cpp -L/path/to/install/dir -lcnpy -lz --std=c++11

# Description:

There are two functions for writing data: `npy_save` and `npz_save`.
There are three functions for writing data: `npy_save`, `npz_save`, and `npz_save_compressed`.

- `npy_save` saves a single array to a .npy file.
- `npz_save` saves arrays to a .npz file. Accepts an optional `compress` parameter (default: false) to enable zlib deflate compression.
- `npz_save_compressed` convenience wrapper that calls `npz_save` with compression enabled (equivalent to `numpy.savez_compressed`).

There are 3 functions for reading:
- `npy_load` will load a .npy file.
- `npz_load(fname)` will load a .npz and return a dictionary of NpyArray structues.
- `npz_load(fname)` will load a .npz and return a dictionary of NpyArray structures.
- `npz_load(fname,varname)` will load and return the NpyArray for data varname from the specified .npz file.

Both compressed and uncompressed .npz files are supported for reading, including files created by `numpy.savez_compressed()` which use ZIP64 format.

The data structure for loaded data is below.
Data is accessed via the `data<T>()`-method, which returns a pointer of the specified type (which must match the underlying datatype of the data).
The array shape and word size are read from the npy header.
Expand Down
97 changes: 70 additions & 27 deletions cnpy.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,12 @@
#include<stdexcept>
#include <regex>

template<typename T> static T read_unaligned(const void* source) {
T value;
memcpy(&value, source, sizeof(T));
return value;
}

char cnpy::BigEndianTest() {
int x = 1;
return (((char *)&x)[0]) ? '<' : '>';
Expand Down Expand Up @@ -63,7 +69,7 @@ void cnpy::parse_npy_header(unsigned char* buffer,size_t& word_size, std::vector
//std::string magic_string(buffer,6);
uint8_t major_version = *reinterpret_cast<uint8_t*>(buffer+6);
uint8_t minor_version = *reinterpret_cast<uint8_t*>(buffer+7);
uint16_t header_len = *reinterpret_cast<uint16_t*>(buffer+8);
uint16_t header_len = read_unaligned<uint16_t>(buffer+8);
std::string header(reinterpret_cast<char*>(buffer+9),header_len);

size_t loc1, loc2;
Expand Down Expand Up @@ -93,12 +99,13 @@ void cnpy::parse_npy_header(unsigned char* buffer,size_t& word_size, std::vector
bool littleEndian = (header[loc1] == '<' || header[loc1] == '|' ? true : false);
assert(littleEndian);

//char type = header[loc1+1];
char type = header[loc1+1];
//assert(type == map_type(T));

std::string str_ws = header.substr(loc1+2);
loc2 = str_ws.find("'");
word_size = atoi(str_ws.substr(0,loc2).c_str());
if(type == 'U') word_size *= 4;
}

void cnpy::parse_npy_header(FILE* fp, size_t& word_size, std::vector<size_t>& shape, bool& fortran_order) {
Expand Down Expand Up @@ -144,12 +151,13 @@ void cnpy::parse_npy_header(FILE* fp, size_t& word_size, std::vector<size_t>& sh
bool littleEndian = (header[loc1] == '<' || header[loc1] == '|' ? true : false);
assert(littleEndian);

//char type = header[loc1+1];
char type = header[loc1+1];
//assert(type == map_type(T));

std::string str_ws = header.substr(loc1+2);
loc2 = str_ws.find("'");
word_size = atoi(str_ws.substr(0,loc2).c_str());
if(type == 'U') word_size *= 4;
}

void cnpy::parse_zip_footer(FILE* fp, uint16_t& nrecs, size_t& global_header_size, size_t& global_header_offset)
Expand All @@ -161,13 +169,13 @@ void cnpy::parse_zip_footer(FILE* fp, uint16_t& nrecs, size_t& global_header_siz
throw std::runtime_error("parse_zip_footer: failed fread");

uint16_t disk_no, disk_start, nrecs_on_disk, comment_len;
disk_no = *(uint16_t*) &footer[4];
disk_start = *(uint16_t*) &footer[6];
nrecs_on_disk = *(uint16_t*) &footer[8];
nrecs = *(uint16_t*) &footer[10];
global_header_size = *(uint32_t*) &footer[12];
global_header_offset = *(uint32_t*) &footer[16];
comment_len = *(uint16_t*) &footer[20];
disk_no = read_unaligned<uint16_t>(&footer[4]);
disk_start = read_unaligned<uint16_t>(&footer[6]);
nrecs_on_disk = read_unaligned<uint16_t>(&footer[8]);
nrecs = read_unaligned<uint16_t>(&footer[10]);
global_header_size = read_unaligned<uint32_t>(&footer[12]);
global_header_offset = read_unaligned<uint32_t>(&footer[16]);
comment_len = read_unaligned<uint16_t>(&footer[20]);

assert(disk_no == 0);
assert(disk_start == 0);
Expand All @@ -188,6 +196,24 @@ cnpy::NpyArray load_the_npy_file(FILE* fp) {
return arr;
}

// Helper function to parse ZIP64 extended info from extra field
void parse_zip64_sizes(const std::vector<char>& extra_field, uint32_t& compr_bytes, uint32_t& uncompr_bytes) {
if(extra_field.size() >= 4 && (compr_bytes == 0xFFFFFFFF || uncompr_bytes == 0xFFFFFFFF)) {
uint16_t extra_id = read_unaligned<uint16_t>(&extra_field[0]);
uint16_t extra_size = read_unaligned<uint16_t>(&extra_field[2]);
if(extra_id == 0x0001 && extra_size >= 16) { // ZIP64 extended info
size_t offset = 4;
if(uncompr_bytes == 0xFFFFFFFF) {
uncompr_bytes = static_cast<uint32_t>(read_unaligned<uint64_t>(&extra_field[offset]));
offset += 8;
}
if(compr_bytes == 0xFFFFFFFF) {
compr_bytes = static_cast<uint32_t>(read_unaligned<uint64_t>(&extra_field[offset]));
}
}
}
}

cnpy::NpyArray load_the_npz_array(FILE* fp, uint32_t compr_bytes, uint32_t uncompr_bytes) {

std::vector<unsigned char> buffer_compr(compr_bytes);
Expand All @@ -205,14 +231,22 @@ cnpy::NpyArray load_the_npz_array(FILE* fp, uint32_t compr_bytes, uint32_t uncom
d_stream.avail_in = 0;
d_stream.next_in = Z_NULL;
err = inflateInit2(&d_stream, -MAX_WBITS);
if(err != Z_OK)
throw std::runtime_error("load_the_npz_array: inflateInit2 failed");

d_stream.avail_in = compr_bytes;
d_stream.next_in = &buffer_compr[0];
d_stream.avail_out = uncompr_bytes;
d_stream.next_out = &buffer_uncompr[0];

err = inflate(&d_stream, Z_FINISH);
if(err != Z_STREAM_END) {
inflateEnd(&d_stream);
throw std::runtime_error("load_the_npz_array: inflate failed");
}
err = inflateEnd(&d_stream);
if(err != Z_OK)
throw std::runtime_error("load_the_npz_array: inflateEnd failed");

std::vector<size_t> shape;
size_t word_size;
Expand Down Expand Up @@ -246,7 +280,7 @@ cnpy::npz_t cnpy::npz_load(std::string fname) {
if(local_header[2] != 0x03 || local_header[3] != 0x04) break;

//read in the variable name
uint16_t name_len = *(uint16_t*) &local_header[26];
uint16_t name_len = read_unaligned<uint16_t>(&local_header[26]);
std::string varname(name_len,' ');
size_t vname_res = fread(&varname[0],sizeof(char),name_len,fp);
if(vname_res != name_len)
Expand All @@ -256,17 +290,20 @@ cnpy::npz_t cnpy::npz_load(std::string fname) {
varname.erase(varname.end()-4,varname.end());

//read in the extra field
uint16_t extra_field_len = *(uint16_t*) &local_header[28];
uint16_t extra_field_len = read_unaligned<uint16_t>(&local_header[28]);
std::vector<char> extra_field(extra_field_len);
if(extra_field_len > 0) {
std::vector<char> buff(extra_field_len);
size_t efield_res = fread(&buff[0],sizeof(char),extra_field_len,fp);
size_t efield_res = fread(&extra_field[0],sizeof(char),extra_field_len,fp);
if(efield_res != extra_field_len)
throw std::runtime_error("npz_load: failed fread");
}

uint16_t compr_method = *reinterpret_cast<uint16_t*>(&local_header[0]+8);
uint32_t compr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+18);
uint32_t uncompr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+22);
uint16_t compr_method = read_unaligned<uint16_t>(&local_header[0]+8);
uint32_t compr_bytes = read_unaligned<uint32_t>(&local_header[0]+18);
uint32_t uncompr_bytes = read_unaligned<uint32_t>(&local_header[0]+22);

// ZIP64 support: if sizes are 0xFFFFFFFF, read from extra field
parse_zip64_sizes(extra_field, compr_bytes, uncompr_bytes);

if(compr_method == 0) {arrays[varname] = load_the_npy_file(fp);}
else {arrays[varname] = load_the_npz_array(fp,compr_bytes,uncompr_bytes);}
Expand All @@ -291,30 +328,37 @@ cnpy::NpyArray cnpy::npz_load(std::string fname, std::string varname) {
if(local_header[2] != 0x03 || local_header[3] != 0x04) break;

//read in the variable name
uint16_t name_len = *(uint16_t*) &local_header[26];
uint16_t name_len = read_unaligned<uint16_t>(&local_header[26]);
std::string vname(name_len,' ');
size_t vname_res = fread(&vname[0],sizeof(char),name_len,fp);
if(vname_res != name_len)
throw std::runtime_error("npz_load: failed fread");
vname.erase(vname.end()-4,vname.end()); //erase the lagging .npy

//read in the extra field
uint16_t extra_field_len = *(uint16_t*) &local_header[28];
fseek(fp,extra_field_len,SEEK_CUR); //skip past the extra field
uint16_t extra_field_len = read_unaligned<uint16_t>(&local_header[28]);
std::vector<char> extra_field(extra_field_len);
if(extra_field_len > 0) {
size_t efield_res = fread(&extra_field[0],sizeof(char),extra_field_len,fp);
if(efield_res != extra_field_len)
throw std::runtime_error("npz_load: failed fread");
}

uint16_t compr_method = *reinterpret_cast<uint16_t*>(&local_header[0]+8);
uint32_t compr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+18);
uint32_t uncompr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+22);
uint16_t compr_method = read_unaligned<uint16_t>(&local_header[0]+8);
uint32_t compr_bytes = read_unaligned<uint32_t>(&local_header[0]+18);
uint32_t uncompr_bytes = read_unaligned<uint32_t>(&local_header[0]+22);

// ZIP64 support: if sizes are 0xFFFFFFFF, read from extra field
parse_zip64_sizes(extra_field, compr_bytes, uncompr_bytes);

if(vname == varname) {
NpyArray array = (compr_method == 0) ? load_the_npy_file(fp) : load_the_npz_array(fp,compr_bytes,uncompr_bytes);
fclose(fp);
return array;
}
else {
//skip past the data
uint32_t size = *(uint32_t*) &local_header[22];
fseek(fp,size,SEEK_CUR);
//skip past the data (use compr_bytes for compressed data)
fseek(fp,compr_bytes,SEEK_CUR);
}
}

Expand All @@ -337,4 +381,3 @@ cnpy::NpyArray cnpy::npy_load(std::string fname) {
}



94 changes: 80 additions & 14 deletions cnpy.h
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include<typeinfo>
#include<iostream>
#include<cassert>
#include<cstring>
#include<zlib.h>
#include<map>
#include<memory>
Expand Down Expand Up @@ -130,7 +131,8 @@ namespace cnpy {
fclose(fp);
}

template<typename T> void npz_save(std::string zipname, std::string fname, const T* data, const std::vector<size_t>& shape, std::string mode = "w")
// Save arrays to NPZ file, optionally with compression
template<typename T> void npz_save(std::string zipname, std::string fname, const T* data, const std::vector<size_t>& shape, std::string mode = "w", bool compress = false)
{
//first, append a .npy to the fname
fname += ".npy";
Expand Down Expand Up @@ -165,24 +167,73 @@ namespace cnpy {
std::vector<char> npy_header = create_npy_header<T>(shape);

size_t nels = std::accumulate(shape.begin(),shape.end(),1,std::multiplies<size_t>());
size_t nbytes = nels*sizeof(T) + npy_header.size();

//get the CRC of the data to be added
uint32_t crc = crc32(0L,(uint8_t*)&npy_header[0],npy_header.size());
crc = crc32(crc,(uint8_t*)data,nels*sizeof(T));
size_t nbytes_uncompressed = nels*sizeof(T) + npy_header.size();

// Prepare data and compression parameters
std::vector<uint8_t> buffer_compressed;
size_t nbytes_on_disk;
uint16_t compression_method;
uint32_t crc;

if(compress) {
// Create uncompressed buffer (header + data)
std::vector<uint8_t> uncompressed(nbytes_uncompressed);
memcpy(&uncompressed[0], &npy_header[0], npy_header.size());
memcpy(&uncompressed[npy_header.size()], data, nels*sizeof(T));

// Get CRC of uncompressed data
crc = crc32(0L, &uncompressed[0], nbytes_uncompressed);

// Compress using zlib deflate (raw deflate, no zlib/gzip header)
uLongf max_compressed_size = compressBound(nbytes_uncompressed);
buffer_compressed.resize(max_compressed_size);

z_stream strm;
strm.zalloc = Z_NULL;
strm.zfree = Z_NULL;
strm.opaque = Z_NULL;
int ret = deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -MAX_WBITS, 8, Z_DEFAULT_STRATEGY);
if(ret != Z_OK) {
fclose(fp);
throw std::runtime_error("npz_save: deflateInit2 failed");
}

strm.avail_in = nbytes_uncompressed;
strm.next_in = &uncompressed[0];
strm.avail_out = max_compressed_size;
strm.next_out = &buffer_compressed[0];

ret = deflate(&strm, Z_FINISH);
if(ret != Z_STREAM_END) {
deflateEnd(&strm);
fclose(fp);
throw std::runtime_error("npz_save: deflate failed");
}

nbytes_on_disk = strm.total_out;
deflateEnd(&strm);
compression_method = 8; // deflate
}
else {
// No compression - CRC computed in two parts
crc = crc32(0L,(uint8_t*)&npy_header[0],npy_header.size());
crc = crc32(crc,(uint8_t*)data,nels*sizeof(T));
nbytes_on_disk = nbytes_uncompressed;
compression_method = 0; // store
}

//build the local header
std::vector<char> local_header;
local_header += "PK"; //first part of sig
local_header += (uint16_t) 0x0403; //second part of sig
local_header += (uint16_t) 20; //min version to extract
local_header += (uint16_t) 0; //general purpose bit flag
local_header += (uint16_t) 0; //compression method
local_header += (uint16_t) compression_method; //compression method
local_header += (uint16_t) 0; //file last mod time
local_header += (uint16_t) 0; //file last mod date
local_header += (uint32_t) crc; //crc
local_header += (uint32_t) nbytes; //compressed size
local_header += (uint32_t) nbytes; //uncompressed size
local_header += (uint32_t) nbytes_on_disk; //compressed size
local_header += (uint32_t) nbytes_uncompressed; //uncompressed size
local_header += (uint16_t) fname.size(); //fname length
local_header += (uint16_t) 0; //extra field length
local_header += fname;
Expand All @@ -208,13 +259,17 @@ namespace cnpy {
footer += (uint16_t) (nrecs+1); //number of records on this disk
footer += (uint16_t) (nrecs+1); //total number of records
footer += (uint32_t) global_header.size(); //nbytes of global headers
footer += (uint32_t) (global_header_offset + nbytes + local_header.size()); //offset of start of global headers, since global header now starts after newly written array
footer += (uint32_t) (global_header_offset + nbytes_on_disk + local_header.size()); //offset of start of global headers, since global header now starts after newly written array
footer += (uint16_t) 0; //zip file comment length

//write everything
fwrite(&local_header[0],sizeof(char),local_header.size(),fp);
fwrite(&npy_header[0],sizeof(char),npy_header.size(),fp);
fwrite(data,sizeof(T),nels,fp);
if(compress) {
fwrite(&buffer_compressed[0],sizeof(char),nbytes_on_disk,fp);
} else {
fwrite(&npy_header[0],sizeof(char),npy_header.size(),fp);
fwrite(data,sizeof(T),nels,fp);
}
fwrite(&global_header[0],sizeof(char),global_header.size(),fp);
fwrite(&footer[0],sizeof(char),footer.size(),fp);
fclose(fp);
Expand All @@ -226,10 +281,21 @@ namespace cnpy {
npy_save(fname, &data[0], shape, mode);
}

template<typename T> void npz_save(std::string zipname, std::string fname, const std::vector<T> data, std::string mode = "w") {
template<typename T> void npz_save(std::string zipname, std::string fname, const std::vector<T> data, std::string mode = "w", bool compress = false) {
std::vector<size_t> shape;
shape.push_back(data.size());
npz_save(zipname, fname, &data[0], shape, mode, compress);
}

// Convenience wrapper for compressed save
template<typename T> void npz_save_compressed(std::string zipname, std::string fname, const T* data, const std::vector<size_t>& shape, std::string mode = "w") {
npz_save(zipname, fname, data, shape, mode, true);
}

template<typename T> void npz_save_compressed(std::string zipname, std::string fname, const std::vector<T> data, std::string mode = "w") {
std::vector<size_t> shape;
shape.push_back(data.size());
npz_save(zipname, fname, &data[0], shape, mode);
npz_save(zipname, fname, &data[0], shape, mode, true);
}

template<typename T> std::vector<char> create_npy_header(const std::vector<size_t>& shape) {
Expand Down
Loading