loader.cpp

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>

#include "khash.h"
#include "utils.h"

#define PERSON_FIELD_ID 0
#define PERSON_FIELD_BIRTHDAY 4
#define PERSON_FIELD_LOCATION 8
#define KNOWS_FIELD_PERSON 0
#define KNOWS_FIELD_FRIEND 1
#define INTEREST_FIELD_PERSON 0
#define INTEREST_FIELD_INTEREST 1

// hash map needs long keys (large person ids), but unsigned int is enough for person offsets
KHASH_MAP_INIT_INT64(pht, unsigned int)
khash_t(pht) *person_offsets;

FILE   *person_out = nullptr;
FILE   *interest_out = nullptr;
FILE   *knows_out = nullptr;
Person *person_map = nullptr;
Person *person = nullptr;

unsigned long person_id = 0;
unsigned long person_id_prev = 0;
unsigned long knows_id = 0;

// person offset can be smaller, we do not have so many
unsigned int person_offset = 0;
unsigned long knows_offset = 0;
unsigned long interest_offset = 0;

void person_line_handler(unsigned char nfields, char** tokens) {
	(void)nfields;
	int ret;
	khiter_t k;
	Person row{};

	row.person_id = std::strtoul(tokens[PERSON_FIELD_ID], nullptr, 10);
	row.birthday = birthday_to_short(tokens[PERSON_FIELD_BIRTHDAY]);
	row.location = static_cast<unsigned short>(std::atoi(tokens[PERSON_FIELD_LOCATION]));

	// add mapping person id -> offset to hash table
	k = kh_put(pht, person_offsets, row.person_id, &ret);
	kh_value(person_offsets, k) = person_offset;
	// write binary person record to file
	std::fwrite(&row, sizeof(Person), 1, person_out);
	person_offset++;
}

void updatePerson() {
	person = &person_map[kh_value(person_offsets,
		kh_get(pht, person_offsets, person_id))];
	person_id_prev = person_id;
}

void knows_line_handler(unsigned char nfields, char** tokens) {
	(void)nfields;
	unsigned int knows_person_offset;
	person_id = std::strtoul(tokens[KNOWS_FIELD_PERSON], nullptr, 10);
	unsigned long knows_id = std::strtoul(tokens[KNOWS_FIELD_FRIEND], nullptr, 10);

	if (person_id != person_id_prev) {
		updatePerson();
		person->knows_first = knows_offset;
		person->knows_n = 0;
	}
	// lookup other person and write offset
	knows_person_offset = kh_value(person_offsets,
			kh_get(pht, person_offsets, knows_id));
	std::fwrite(&knows_person_offset, sizeof(unsigned int), 1, knows_out);
	knows_offset++;
	person->knows_n++;
}

void interest_line_handler(unsigned char nfields, char** tokens) {
	(void)nfields;
	unsigned short interest_id;
	person_id = std::strtoul(tokens[INTEREST_FIELD_PERSON], nullptr, 10);
	interest_id = static_cast<unsigned short>(std::atoi(tokens[INTEREST_FIELD_INTEREST]));
	if (person_id != person_id_prev) {
		updatePerson();
		person->interests_first = interest_offset;
		person->interest_n = 0;
	}
	std::fwrite(&interest_id, sizeof(unsigned short), 1, interest_out);
	interest_offset++;
	person->interest_n++;
}

int main(int argc, char *argv[]) {
	if (argc < 3) {
		std::fprintf(stderr, "Usage: [csv input path] [output path]\n");
		std::exit(-1);
	}

	std::string person_input_file    = makepath(argv[1], "person", "csv");
	std::string interest_input_file  = makepath(argv[1], "interest", "csv");
	std::string knows_input_file     = makepath(argv[1], "knows", "csv");
	std::string person_output_file   = makepath(argv[2], "person", "bin");
	std::string interest_output_file = makepath(argv[2], "interest", "bin");
	std::string knows_output_file    = makepath(argv[2], "knows", "bin");

	int person_map_fd;
	struct stat st;

	if (stat(argv[2], &st) == -1) {
		if (mkdir(argv[2], 0700) != 0) {
			std::fprintf(stderr, "Unable to create output directory %s\n", argv[2]);
			std::exit(-1);
		}
	}

	// first pass person, parse person, write to binary and store bin offset in hash table
	person_offsets = kh_init(pht);
	person_out = open_binout(person_output_file.c_str());
	parse_csv(person_input_file.c_str(), &person_line_handler);

	// mmap person.bin binary file for updates
	std::fclose(person_out);
	person_map_fd = open(person_output_file.c_str(), O_RDWR);
	person_map = static_cast<Person *>(mmap(0, person_offset * sizeof(Person), PROT_READ | PROT_WRITE,
		MAP_SHARED, person_map_fd, 0));
	if (person_map_fd < 0 || person_map == MAP_FAILED) {
		std::fprintf(stderr, "Failed to map person binary.\n");
		std::exit(-1);
	}
	close(person_map_fd);

	// pass through interest and friends, write to binary, set offsets in person
	person_id = 0;
	interest_out = open_binout(interest_output_file.c_str());
	parse_csv(interest_input_file.c_str(), &interest_line_handler);
	std::fclose(interest_out);

	person_id = 0;
	knows_out = open_binout(knows_output_file.c_str());
	parse_csv(knows_input_file.c_str(), &knows_line_handler);
	std::fclose(knows_out);
	munmap(person_map, person_offset * sizeof(Person));
	kh_destroy(pht, person_offsets);

	return 0;
}