data_utils.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: generate inputs and targets for the DLRM benchmark
#
# Utility function(s) to download and pre-process public data sets
#   - Criteo Kaggle Display Advertising Challenge Dataset
#     https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset
#   - Criteo Terabyte Dataset
#     https://labs.criteo.com/2013/12/download-terabyte-click-logs
#
# After downloading dataset, run:
#   getCriteoAdData(
#       datafile="<path-to-train.txt>",
#       o_filename=kaggleAdDisplayChallenge_processed.npz,
#       max_ind_range=-1,
#       sub_sample_rate=0.0,
#       days=7,
#       data_split='train',
#       randomize='total',
#       criteo_kaggle=True,
#       memory_map=False
#   )
#   getCriteoAdData(
#       datafile="<path-to-day_{0,...,23}>",
#       o_filename=terabyte_processed.npz,
#       max_ind_range=-1,
#       sub_sample_rate=0.0,
#       days=24,
#       data_split='train',
#       randomize='total',
#       criteo_kaggle=False,
#       memory_map=False
#   )

from __future__ import absolute_import, division, print_function, unicode_literals

import sys
# import os
from os import path
from multiprocessing import Process, Manager
# import io
# from io import StringIO
# import collections as coll

import numpy as np
import pdb

def convertUStringToDistinctIntsDict(mat, convertDicts, counts):
    # Converts matrix of unicode strings into distinct integers.
    #
    # Inputs:
    #     mat (np.array): array of unicode strings to convert
    #     convertDicts (list): dictionary for each column
    #     counts (list): number of different categories in each column
    #
    # Outputs:
    #     out (np.array): array of output integers
    #     convertDicts (list): dictionary for each column
    #     counts (list): number of different categories in each column

    # check if convertDicts and counts match correct length of mat
    if len(convertDicts) != mat.shape[1] or len(counts) != mat.shape[1]:
        print("Length of convertDicts or counts does not match input shape")
        print("Generating convertDicts and counts...")

        convertDicts = [{} for _ in range(mat.shape[1])]
        counts = [0 for _ in range(mat.shape[1])]

    # initialize output
    out = np.zeros(mat.shape)

    for j in range(mat.shape[1]):
        for i in range(mat.shape[0]):
            # add to convertDict and increment count
            if mat[i, j] not in convertDicts[j]:
                convertDicts[j][mat[i, j]] = counts[j]
                counts[j] += 1
            out[i, j] = convertDicts[j][mat[i, j]]

    return out, convertDicts, counts


def convertUStringToDistinctIntsUnique(mat, mat_uni, counts):
    # mat is an array of 0,...,# samples, with each being 26 categorical features

    # check if mat_unique and counts match correct length of mat
    if len(mat_uni) != mat.shape[1] or len(counts) != mat.shape[1]:
        print("Length of mat_unique or counts does not match input shape")
        print("Generating mat_unique and counts...")

        mat_uni = [np.array([]) for _ in range(mat.shape[1])]
        counts = [0 for _ in range(mat.shape[1])]

    # initialize output
    out = np.zeros(mat.shape)
    ind_map = [np.array([]) for _ in range(mat.shape[1])]

    # find out and assign unique ids to features
    for j in range(mat.shape[1]):
        m = mat_uni[j].size
        mat_concat = np.concatenate((mat_uni[j], mat[:, j]))
        mat_uni[j], ind_map[j] = np.unique(mat_concat, return_inverse=True)
        out[:, j] = ind_map[j][m:]
        counts[j] = mat_uni[j].size

    return out, mat_uni, counts


def processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, pre_comp_counts):
    # Process Kaggle Display Advertising Challenge or Terabyte Dataset
    # by converting unicode strings in X_cat to integers and
    # converting negative integer values in X_int.
    #
    # Loads data in the form "{kaggle|terabyte}_day_i.npz" where i is the day.
    #
    # Inputs:
    #   d_path (str): path for {kaggle|terabyte}_day_i.npz files
    #   i (int): splits in the dataset (typically 0 to 7 or 0 to 24)

    # process data if not all files exist
    filename_i = npzfile + "_{0}_processed.npz".format(i)

    if path.exists(filename_i):
        print("Using existing " + filename_i, end="\n")
    else:
        print("Not existing " + filename_i)
        with np.load(npzfile + "_{0}.npz".format(i)) as data:
            # categorical features
            '''
            # Approach 1a: using empty dictionaries
            X_cat, convertDicts, counts = convertUStringToDistinctIntsDict(
                data["X_cat"], convertDicts, counts
            )
            '''
            '''
            # Approach 1b: using empty np.unique
            X_cat, convertDicts, counts = convertUStringToDistinctIntsUnique(
                data["X_cat"], convertDicts, counts
            )
            '''
            # Approach 2a: using pre-computed dictionaries
            X_cat_t = np.zeros(data["X_cat_t"].shape)
            for j in range(26):
                for k, x in enumerate(data["X_cat_t"][j, :]):
                    X_cat_t[j, k] = convertDicts[j][x]
            # continuous features
            X_int = data["X_int"]
            X_int[X_int < 0] = 0
            # targets
            y = data["y"]

        np.savez_compressed(
            filename_i,
            # X_cat = X_cat,
            X_cat=np.transpose(X_cat_t),  # transpose of the data
            X_int=X_int,
            y=y,
        )
        print("Processed " + filename_i, end="\n")
    # sanity check (applicable only if counts have been pre-computed & are re-computed)
    # for j in range(26):
    #    if pre_comp_counts[j] != counts[j]:
    #        sys.exit("ERROR: Sanity check on counts has failed")
    # print("\nSanity check on counts passed")

    return


'''
    195841983 day 0
    199563535 day 1
    196792019 day 2
    181115208 day 3
    152115810 day 4
    172548507 day 5
    204846845 day 6
    200801003 day 7
    193772492 day 8
    198424372 day 9
    185778055 day 10
    153588700 day 11
    169003364 day 12
    194216520 day 13
    194081279 day 14
    187154596 day 15
    177984934 day 16
    163382602 day 17
    142061091 day 18
    156534237 day 19
    193627464 day 20
    192215183 day 21
    189747893 day 22
    178274637 day 23
'''


def concatCriteoAdData(
        d_path,
        d_file,
        npzfile,
        trafile,
        days,
        data_split,
        randomize,
        total_per_file,
        total_count,
        memory_map,
        o_filename
):
    # Concatenates different days and saves the result.
    #
    # Inputs:
    #   days (int): total number of days in the dataset (typically 7 or 24)
    #   d_path (str): path for {kaggle|terabyte}_day_i.npz files
    #   o_filename (str): output file name
    #
    # Output:
    #   o_file (str): output file path

    if memory_map:
        # dataset break up per fea
        # tar_fea = 1   # single target
        den_fea = 13  # 13 dense  features
        spa_fea = 26  # 26 sparse features
        # tad_fea = tar_fea + den_fea
        # tot_fea = tad_fea + spa_fea
        # create offset per file
        offset_per_file = np.array([0] + [x for x in total_per_file])
        for i in range(days):
            offset_per_file[i + 1] += offset_per_file[i]

        '''
        # Approach 1, 2 and 3 use indices, while Approach 4 does not use them
        # create indices
        indices = np.arange(total_count)
        if data_split == "none":
            if randomize == "total":
                indices = np.random.permutation(indices)
        else:
            indices = np.array_split(indices, offset_per_file[1:-1])

            # randomize train data (per day)
            if randomize == "day":  # or randomize == "total":
                for i in range(len(indices) - 1):
                    indices[i] = np.random.permutation(indices[i])
                print("Randomized indices per day ...")

            train_indices = np.concatenate(indices[:-1])
            test_indices = indices[-1]

            # randomize train data (across days)
            if randomize == "total":
                train_indices = np.random.permutation(train_indices)
                print("Randomized indices across days ...")

            indices = np.concatenate((train_indices, test_indices))
        # no reordering
        # indices = np.arange(total_count)
        '''
        '''
        # Approach 1: simple and slow (no grouping is used)
        # check if data already exists
        recreate_flag = False
        for j in range(tot_fea):
            filename_j = trafile + "_{0}_reordered.npy".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((total_count))
            for j in range(tot_fea):
                filename_j = trafile + "_{0}_reordered".format(j)
                np.save(filename_j, z)
                print("Creating " + filename_j)

            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat_t = np.transpose(data["X_cat"])
                    X_int_t = np.transpose(data["X_int"])
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                # print(filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #     + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r+')
                    if j < tar_fea:
                        fj[indices[start:end]] = y
                    elif tar_fea <= j and j < tad_fea:
                        fj[indices[start:end]] = X_int_t[j - tar_fea, :]
                    else:
                        fj[indices[start:end]] = X_cat_t[j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing " + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                size = total_per_file[i]
                X_int_t = np.zeros((den_fea, size))
                X_cat_t = np.zeros((spa_fea, size))
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                print("Creating " + filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #     + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r')
                    if j < tar_fea:
                        y = fj[start:end]
                    elif tar_fea <= j and j < tad_fea:
                        X_int_t[j - tar_fea, :] = fj[start:end]
                    else:
                        X_cat_t[j - tad_fea, :] = fj[start:end]
                    del fj

                np.savez_compressed(
                    filename_i,
                    X_cat=np.transpose(X_cat_t),  # transpose of the data
                    X_int=np.transpose(X_int_t),  # transpose of the data
                    y=y,
                )
        else:
            print("Reordered day files already exist, skipping ...")
        '''
        '''
        # Approach 2: group days
        # check if data already exists
        recreate_flag = False
        for j in range(tot_fea):
            filename_j = trafile + "_{0}_reordered.npy".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((total_count))
            for j in range(tot_fea):
                filename_j = trafile + "_{0}_reordered".format(j)
                np.save(filename_j, z)
                print("Creating " + filename_j)

            group_day = 3  # e.g. 8, 4 or 3
            group_num = days // group_day
            file_group = [i*group_day for i in range(group_num)] + [days]
            for ii in range(group_num):
                # for last may be group_size != group_num, therefore reset it below
                group_size = file_group[ii + 1] - file_group[ii]
                X_cat_t = [0]*group_size
                X_int_t = [0]*group_size
                y = [0]*group_size
                start = [0]*group_size
                end  = [0]*group_size
                for ig in range(group_size):
                    i = file_group[ii] + ig
                    filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                    # setup start and end ranges
                    start[ig] = offset_per_file[i]
                    end[ig] = offset_per_file[i + 1]
                    # print(filename_i)
                    # load a group of files
                    with np.load(filename_i) as data:
                        X_cat_t[ig] = np.transpose(data["X_cat"])
                        X_int_t[ig] = np.transpose(data["X_int"])
                        y[ig] = data["y"]
                    # sanity check
                    if total_per_file[i] != len(y[ig]):
                        sys.exit("ERROR: sanity check on number of samples failed")
                # print("start=" + str(start) + " end=" + str(end)
                #  + " diff=" + str(end[ig]-start[ig]) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r+')
                    for ig in range(group_size):
                        if j < tar_fea:
                            fj[indices[start[ig]:end[ig]]] = y[ig]
                        elif tar_fea <= j and j < tad_fea:
                            fj[indices[start[ig]:end[ig]]] = X_int_t[ig][j - tar_fea, :]
                        else:
                            fj[indices[start[ig]:end[ig]]] = X_cat_t[ig][j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing " + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for ii in range(group_num):
                # for last may be group_size != group_num, therefore reset it below
                group_size = file_group[ii + 1] - file_group[ii]
                X_cat_t= []; X_int_t = []
                for ig in range(group_size):
                    i = file_group[ii] + ig
                    X_int_t.append(np.zeros((den_fea, total_per_file[i])))
                    X_cat_t.append(np.zeros((spa_fea, total_per_file[i])))
                y = [0]*group_size
                start = [0]*group_size
                end  = [0]*group_size

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r')
                    # load a group of files
                    for ig in range(group_size):
                        i = file_group[ii] + ig
                        # setup start and end ranges
                        start[ig] = offset_per_file[i]
                        end[ig] = offset_per_file[i + 1]
                        # load data for the group of files
                        if j < tar_fea:
                            y[ig] = fj[start[ig]:end[ig]]
                        elif tar_fea <= j and j < tad_fea:
                            X_int_t[ig][j - tar_fea, :] = fj[start[ig]:end[ig]]
                        else:
                            X_cat_t[ig][j - tad_fea, :] = fj[start[ig]:end[ig]]
                    del fj

                for ig in range(group_size):
                    i = file_group[ii] + ig
                    filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                    print("Creating " + filename_i)
                    np.savez_compressed(
                        filename_i,
                        X_cat=np.transpose(X_cat_t[ig]),  # transpose of the data
                        X_int=np.transpose(X_int_t[ig]),  # transpose of the data
                        y=y[ig],
                    )
        else:
            print("Reordered day files already exist, skipping ...")
        '''
        '''
        # Approach 3: group features
        # check if data already exists
        group_fea = 5  # e.g. 8, 5 or 4
        group_num = tot_fea // group_fea
        if tot_fea % group_fea != 0:  # sanity check
            sys.exit("ERROR: the group_fea must divided tot_fea evenly.")
        recreate_flag = False
        for jn in range(group_num):
            filename_j = trafile + "_{0}_reordered{1}.npy".format(
                jn, group_fea
            )
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((group_fea, total_count))
            for jn in range(group_num):
                filename_j = trafile + "_{0}_reordered{1}".format(
                    jn, group_fea
                )
                np.save(filename_j, z)
                print("Creating " + filename_j)

            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat_t = np.transpose(data["X_cat"])
                    X_int_t = np.transpose(data["X_int"])
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                # print(filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #      + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for jn in range(group_num):
                    filename_j = trafile + "_{0}_reordered{1}.npy".format(
                        jn, group_fea
                    )
                    fj = np.load(filename_j, mmap_mode='r+')
                    for jg in range(group_fea):
                        j = jn * group_fea + jg
                        # print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
                        if j < tar_fea:
                            fj[jg, indices[start:end]] = y
                        elif tar_fea <= j and j < tad_fea:
                            fj[jg, indices[start:end]] = X_int_t[j - tar_fea, :]
                        else:
                            fj[jg, indices[start:end]] = X_cat_t[j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing" + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                size = total_per_file[i]
                X_int_t = np.zeros((den_fea, size))
                X_cat_t = np.zeros((spa_fea, size))
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                print("Creating " + filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #      + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for jn in range(group_num):
                    filename_j = trafile + "_{0}_reordered{1}.npy".format(
                        jn, group_fea
                    )
                    fj = np.load(filename_j, mmap_mode='r')
                    for jg in range(group_fea):
                        j = jn * group_fea + jg
                        # print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
                        if j < tar_fea:
                            y = fj[jg, start:end]
                        elif tar_fea <= j and j < tad_fea:
                            X_int_t[j - tar_fea, :] = fj[jg, start:end]
                        else:
                            X_cat_t[j - tad_fea, :] = fj[jg, start:end]
                    del fj

                np.savez_compressed(
                    filename_i,
                    X_cat=np.transpose(X_cat_t),  # transpose of the data
                    X_int=np.transpose(X_int_t),  # transpose of the data
                    y=y,
                )

        else:
            print("Reordered day files already exist, skipping ...")
        '''

        # Approach 4: Fisher-Yates-Rao (FYR) shuffle algorithm
        # 1st pass of FYR shuffle
        # check if data already exists
        recreate_flag = False
        for j in range(days):
            filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
            filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
            filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
            if (
                path.exists(filename_j_y)
                and path.exists(filename_j_d)
                and path.exists(filename_j_s)
            ):
                print(
                    "Using existing\n"
                    + filename_j_y + "\n"
                    + filename_j_d + "\n"
                    + filename_j_s
                )
            else:
                recreate_flag = True
        # reorder across buckets using sampling
        if recreate_flag:
            # init intermediate files (.npy appended automatically)
            for j in range(days):
                filename_j_y = npzfile + "_{0}_intermediate_y".format(j)
                filename_j_d = npzfile + "_{0}_intermediate_d".format(j)
                filename_j_s = npzfile + "_{0}_intermediate_s".format(j)
                np.save(filename_j_y, np.zeros((total_per_file[j])))
                np.save(filename_j_d, np.zeros((total_per_file[j], den_fea)))
                np.save(filename_j_s, np.zeros((total_per_file[j], spa_fea)))
            # start processing files
            total_counter = [0] * days
            for i in range(days):
                filename_i = npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat = data["X_cat"]
                    X_int = data["X_int"]
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # debug prints
                print("Reordering (1st pass) " + filename_i)

                # create buckets using sampling of random ints
                # from (discrete) uniform distribution
                buckets = []
                for _j in range(days):
                    buckets.append([])
                counter = [0] * days
                days_to_sample = days if data_split == "none" else days - 1
                if randomize == "total":
                    rand_u = np.random.randint(low=0, high=days_to_sample, size=size)
                    for k in range(size):
                        # sample and make sure elements per buckets do not overflow
                        if data_split == "none" or i < days - 1:
                            # choose bucket
                            p = rand_u[k]
                            # retry of the bucket is full
                            while total_counter[p] + counter[p] >= total_per_file[p]:
                                p = np.random.randint(low=0, high=days_to_sample)
                        else:  # preserve the last day/bucket if needed
                            p = i
                        buckets[p].append(k)
                        counter[p] += 1
                else:  # randomize is day or none
                    for k in range(size):
                        # do not sample, preserve the data in this bucket
                        p = i
                        buckets[p].append(k)
                        counter[p] += 1

                # sanity check
                if np.sum(counter) != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # debug prints
                # print(counter)
                # print(str(np.sum(counter)) + " = " + str(size))
                # print([len(x) for x in buckets])
                # print(total_counter)

                # partially feel the buckets
                for j in range(days):
                    filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
                    filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
                    filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
                    start = total_counter[j]
                    end = total_counter[j] + counter[j]
                    # target buckets
                    fj_y = np.load(filename_j_y, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_y[start:end].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_y[start:end] = y[buckets[j]]
                    del fj_y
                    # dense buckets
                    fj_d = np.load(filename_j_d, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_d[start:end, :].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_d[start:end, :] = X_int[buckets[j], :]
                    del fj_d
                    # sparse buckets
                    fj_s = np.load(filename_j_s, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_s[start:end, :].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_s[start:end, :] = X_cat[buckets[j], :]
                    del fj_s
                    # update counters for next step
                    total_counter[j] += counter[j]

        # 2nd pass of FYR shuffle
        # check if data already exists
        for j in range(days):
            filename_j = npzfile + "_{0}_reordered.npz".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # reorder within buckets
        if recreate_flag:
            for j in range(days):
                filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
                filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
                filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
                fj_y = np.load(filename_j_y)
                fj_d = np.load(filename_j_d)
                fj_s = np.load(filename_j_s)

                indices = range(total_per_file[j])
                if randomize == "day" or randomize == "total":
                    if data_split == "none" or j < days - 1:
                        indices = np.random.permutation(range(total_per_file[j]))

                filename_r = npzfile + "_{0}_reordered.npz".format(j)
                print("Reordering (2nd pass) " + filename_r)
                np.savez_compressed(
                    filename_r,
                    X_cat=fj_s[indices, :],
                    X_int=fj_d[indices, :],
                    y=fj_y[indices],
                )

        '''
        # sanity check (under no reordering norms should be zero)
        for i in range(days):
            filename_i_o = npzfile + "_{0}_processed.npz".format(i)
            print(filename_i_o)
            with np.load(filename_i_o) as data_original:
                X_cat_o = data_original["X_cat"]
                X_int_o = data_original["X_int"]
                y_o = data_original["y"]
            filename_i_r = npzfile + "_{0}_reordered.npz".format(i)
            print(filename_i_r)
            with np.load(filename_i_r) as data_reordered:
                X_cat_r = data_reordered["X_cat"]
                X_int_r = data_reordered["X_int"]
                y_r = data_reordered["y"]
            print(np.linalg.norm(y_o - y_r))
            print(np.linalg.norm(X_int_o - X_int_r))
            print(np.linalg.norm(X_cat_o - X_cat_r))
        '''

    else:
        print("Concatenating multiple days into %s.npz file" % str(d_path + o_filename))

        # load and concatenate data
        X_cat = np.zeros((total_count, 26));
        X_int = np.zeros((total_count, 13));
        y = np.zeros((total_count,)) - 1;
        loc = 0
        for i in range(days):
            filename_i = npzfile + "_{0}_processed.npz".format(i)
            with np.load(filename_i) as data:
                    this_count = len(data["y"])
                    X_cat[loc:loc+this_count,:] = data["X_cat"]
                    X_int[loc:loc+this_count,:] = data["X_int"]
                    y[loc:loc+this_count] =  data["y"]
                    loc += this_count
            #print("Loaded day:", i, "y = 1:", len(y[y == 1]), "y = 0:", len(y[y == 0]), flush=True)
            print("Loaded day:", i, "count:", loc, "/", total_count, flush=True)

        with np.load(d_path + d_file + "_fea_count.npz") as data:
            counts = data["counts"]
        print("Loaded counts!")

        np.savez_compressed(
            d_path + o_filename + ".npz",
            X_cat=X_cat,
            X_int=X_int,
            y=y,
            counts=counts,
        )

    return d_path + o_filename + ".npz"


def transformCriteoAdData(X_cat, X_int, y, days, data_split, randomize, total_per_file):
    # Transforms Criteo Kaggle or terabyte data by applying log transformation
    # on dense features and converting everything to appropriate tensors.
    #
    # Inputs:
    #     X_cat (ndarray): array of integers corresponding to preprocessed
    #                      categorical features
    #     X_int (ndarray): array of integers corresponding to dense features
    #     y (ndarray):     array of bool corresponding to labels
    #     data_split(str): flag for splitting dataset into training/validation/test
    #                      sets
    #     randomize (str): determines randomization scheme
    #         "none": no randomization
    #         "day": randomizes each day"s data (only works if split = True)
    #         "total": randomizes total dataset
    #
    # Outputs:
    #     if split:
    #         X_cat_train (tensor): sparse features for training set
    #         X_int_train (tensor): dense features for training set
    #         y_train (tensor): labels for training set
    #         X_cat_val (tensor): sparse features for validation set
    #         X_int_val (tensor): dense features for validation set
    #         y_val (tensor): labels for validation set
    #         X_cat_test (tensor): sparse features for test set
    #         X_int_test (tensor): dense features for test set
    #         y_test (tensor): labels for test set
    #     else:
    #         X_cat (tensor): sparse features
    #         X_int (tensor): dense features
    #         y (tensor): label

    # define initial set of indices
    indices = np.arange(len(y))

    # create offset per file
    offset_per_file = np.array([0] + [x for x in total_per_file])
    for i in range(days):
        offset_per_file[i + 1] += offset_per_file[i]

    # split dataset
    if data_split == 'train':
        indices = np.array_split(indices, offset_per_file[1:-1])

        # randomize train data (per day)
        if randomize == "day":  # or randomize == "total":
            for i in range(len(indices) - 1):
                indices[i] = np.random.permutation(indices[i])
            print("Randomized indices per day ...")

        train_indices = np.concatenate(indices[:-1])
        test_indices = indices[-1]
        test_indices, val_indices = np.array_split(test_indices, 2)

        print("Defined training and testing indices...")

        # randomize train data (across days)
        if randomize == "total":
            train_indices = np.random.permutation(train_indices)
            print("Randomized indices across days ...")

        # indices = np.concatenate((train_indices, test_indices))

        # create training, validation, and test sets
        X_cat_train = X_cat[train_indices]
        X_int_train = X_int[train_indices]
        y_train = y[train_indices]

        X_cat_val = X_cat[val_indices]
        X_int_val = X_int[val_indices]
        y_val = y[val_indices]

        X_cat_test = X_cat[test_indices]
        X_int_test = X_int[test_indices]
        y_test = y[test_indices]

        print("Split data according to indices...")

        X_cat_train = X_cat_train.astype(np.long)
        X_int_train = np.log(X_int_train.astype(np.float32) + 1)
        y_train = y_train.astype(np.float32)

        X_cat_val = X_cat_val.astype(np.long)
        X_int_val = np.log(X_int_val.astype(np.float32) + 1)
        y_val = y_val.astype(np.float32)

        X_cat_test = X_cat_test.astype(np.long)
        X_int_test = np.log(X_int_test.astype(np.float32) + 1)
        y_test = y_test.astype(np.float32)

        print("Converted to tensors...done!")

        return (
            X_cat_train,
            X_int_train,
            y_train,
            X_cat_val,
            X_int_val,
            y_val,
            X_cat_test,
            X_int_test,
            y_test,
        )

    else:

        # randomize data
        if randomize == "total":
            indices = np.random.permutation(indices)
            print("Randomized indices...")

        X_cat = X_cat[indices].astype(np.long)
        X_int = np.log(X_int[indices].astype(np.float32) + 1)
        y = y[indices].astype(np.float32)

        print("Converted to tensors...done!")

        return (X_cat, X_int, y, [], [], [], [], [], [])


def getCriteoAdData(
        datafile,
        o_filename,
        max_ind_range=-1,
        sub_sample_rate=0.0,
        days=7,
        data_split='train',
        randomize='total',
        criteo_kaggle=True,
        memory_map=False,
        dataset_multiprocessing=False,
):
    # Passes through entire dataset and defines dictionaries for categorical
    # features and determines the number of total categories.
    #
    # Inputs:
    #    datafile : path to downloaded raw data file
    #    o_filename (str): saves results under o_filename if filename is not ""
    #
    # Output:
    #   o_file (str): output file path

    #split the datafile into path and filename
    lstr = datafile.split("/")
    d_path = "/".join(lstr[0:-1]) + "/"
    d_file = lstr[-1].split(".")[0] if criteo_kaggle else lstr[-1]
    npzfile = d_path + ((d_file + "_day") if criteo_kaggle else d_file)
    trafile = d_path + ((d_file + "_fea") if criteo_kaggle else "fea")

    # count number of datapoints in training set
    total_file = d_path + d_file + "_day_count.npz"
    if path.exists(total_file):
        with np.load(total_file) as data:
            total_per_file = list(data["total_per_file"])
        total_count = np.sum(total_per_file)
        print("Skipping counts per file (already exist)")
    else:
        total_count = 0
        total_per_file = []
        if criteo_kaggle:
            # WARNING: The raw data consists of a single train.txt file
            # Each line in the file is a sample, consisting of 13 continuous and
            # 26 categorical features (an extra space indicates that feature is
            # missing and will be interpreted as 0).
            if path.exists(datafile):
                print("Reading data from path=%s" % (datafile))
                with open(str(datafile)) as f:
                    for _ in f:
                        total_count += 1
                total_per_file.append(total_count)
                # reset total per file due to split
                num_data_per_split, extras = divmod(total_count, days)
                total_per_file = [num_data_per_split] * days
                for j in range(extras):
                    total_per_file[j] += 1
                # split into days (simplifies code later on)
                file_id = 0
                boundary = total_per_file[file_id]
                nf = open(npzfile + "_" + str(file_id), "w")
                with open(str(datafile)) as f:
                    for j, line in enumerate(f):
                        if j == boundary:
                            nf.close()
                            file_id += 1
                            nf = open(npzfile + "_" + str(file_id), "w")
                            boundary += total_per_file[file_id]
                        nf.write(line)
                nf.close()
            else:
                sys.exit("ERROR: Criteo Kaggle Display Ad Challenge Dataset path is invalid; please download from https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset")
        else:
            # WARNING: The raw data consist of day_0.gz,... ,day_23.gz text files
            # Each line in the file is a sample, consisting of 13 continuous and
            # 26 categorical features (an extra space indicates that feature is
            # missing and will be interpreted as 0).

            lines = [ 195841983 , 199563535 , 196792019 , 181115208 , 152115810 , 172548507 , 204846845 , 200801003 , 193772492 , 198424372 , 185778055 , 153588700 , 169003364 , 194216520 , 194081279 , 187154596 , 177984934 , 163382602 , 142061091 , 156534237 , 193627464 , 192215183 , 189747893 , 178274637 ]
            for i in range(days):
                datafile_i = datafile + "_" + str(i)  # + ".gz"
                if path.exists(str(datafile_i)):
                    print("Reading data from path=%s" % (str(datafile_i)))
                    # file day_<number>
                    total_per_file_count = 0
                    #with open(str(datafile_i)) as f:
                    #    for _ in f:
                    #        total_per_file_count += 1
                    total_per_file_count = lines[i]
                    total_per_file.append(total_per_file_count)
                    total_count += total_per_file_count
                else:
                    sys.exit("ERROR: Criteo Terabyte Dataset path is invalid; please download from https://labs.criteo.com/2013/12/download-terabyte-click-logs")

    # process a file worth of data and reinitialize data
    # note that a file main contain a single or multiple splits
    def process_one_file(
            datfile,
            npzfile,
            split,
            num_data_in_split,
            dataset_multiprocessing,
            convertDictsDay=None,
            resultDay=None
    ):
        if dataset_multiprocessing:
            convertDicts_day = [{} for _ in range(26)]

        with open(str(datfile)) as f:
            y = np.zeros(num_data_in_split, dtype="i4")  # 4 byte int
            X_int = np.zeros((num_data_in_split, 13), dtype="i4")  # 4 byte int
            X_cat = np.zeros((num_data_in_split, 26), dtype="i4")  # 4 byte int
            if sub_sample_rate == 0.0:
                rand_u = 1.0
            else:
                rand_u = np.random.uniform(low=0.0, high=1.0, size=num_data_in_split)

            i = 0
            percent = 0
            for k, line in enumerate(f):
                # process a line (data point)
                line = line.split('\t')
                # set missing values to zero
                for j in range(len(line)):
                    if (line[j] == '') or (line[j] == '\n'):
                        line[j] = '0'
                # sub-sample data by dropping zero targets, if needed
                target = np.int32(line[0])
                if target == 0 and \
                   (rand_u if sub_sample_rate == 0.0 else rand_u[k]) < sub_sample_rate:
                    continue

                y[i] = target
                X_int[i] = np.array(line[1:14], dtype=np.int32)
                if max_ind_range > 0:
                    X_cat[i] = np.array(
                        list(map(lambda x: int(x, 16) % max_ind_range, line[14:])),
                        dtype=np.int32
                    )
                else:
                    X_cat[i] = np.array(
                        list(map(lambda x: int(x, 16), line[14:])),
                        dtype=np.int32
                    )

                # count uniques
                if dataset_multiprocessing:
                    for j in range(26):
                        convertDicts_day[j][X_cat[i][j]] = 1
                    # debug prints
                    if float(i)/num_data_in_split*100 > percent+1:
                        percent = int(float(i)/num_data_in_split*100)
                        print(
                            "Load %d/%d (%d%%) Split: %d  Label True: %d  Stored: %d"
                            % (
                                i,
                                num_data_in_split,
                                percent,
                                split,
                                target,
                                y[i],
                            ),
                            end="\n",
                        )
                else:
                    for j in range(26):
                        convertDicts[j][X_cat[i][j]] = 1
                    # debug prints
                    print(
                        "Load %d/%d  Split: %d  Label True: %d  Stored: %d"
                        % (
                            i,
                            num_data_in_split,
                            split,
                            target,
                            y[i],
                        ),
                        end="\r",
                    )
                i += 1

            # store num_data_in_split samples or extras at the end of file
            # count uniques
            # X_cat_t  = np.transpose(X_cat)
            # for j in range(26):
            #     for x in X_cat_t[j,:]:
            #         convertDicts[j][x] = 1
            # store parsed
            filename_s = npzfile + "_{0}.npz".format(split)
            if path.exists(filename_s):
                print("\nSkip existing " + filename_s)
            else:
                np.savez_compressed(
                    filename_s,
                    X_int=X_int[0:i, :],
                    # X_cat=X_cat[0:i, :],
                    X_cat_t=np.transpose(X_cat[0:i, :]),  # transpose of the data
                    y=y[0:i],
                )
                print("\nSaved " + npzfile + "_{0}.npz!".format(split))

        if dataset_multiprocessing:
            resultDay[split] = i
            convertDictsDay[split] = convertDicts_day
            return
        else:
            return i

    # create all splits (reuse existing files if possible)
    recreate_flag = False
    convertDicts = [{} for _ in range(26)]
    # WARNING: to get reproducable sub-sampling results you must reset the seed below
    # np.random.seed(123)
    # in this case there is a single split in each day
    for i in range(days):
        npzfile_i = npzfile + "_{0}.npz".format(i)
        npzfile_p = npzfile + "_{0}_processed.npz".format(i)
        if path.exists(npzfile_i):
            print("Skip existing " + npzfile_i)
        elif path.exists(npzfile_p):
            print("Skip existing " + npzfile_p)
        else:
            recreate_flag = True

    if recreate_flag:
        if dataset_multiprocessing:
            resultDay = Manager().dict()
            convertDictsDay = Manager().dict()
            processes = [Process(target=process_one_file,
                                 name="process_one_file:%i" % i,
                                 args=(npzfile + "_{0}".format(i),
                                       npzfile,
                                       i,
                                       total_per_file[i],
                                       dataset_multiprocessing,
                                       convertDictsDay,
                                       resultDay,
                                       )
                                 ) for i in range(0, days)]
            for process in processes:
                process.start()
            for process in processes:
                process.join()
            for day in range(days):
                total_per_file[day] = resultDay[day]
                print("Constructing convertDicts Split: {}".format(day))
                convertDicts_tmp = convertDictsDay[day]
                for i in range(26):
                    for j in convertDicts_tmp[i]:
                        convertDicts[i][j] = 1
        else:
            for i in range(days):
                total_per_file[i] = process_one_file(
                    npzfile + "_{0}".format(i),
                    npzfile,
                    i,
                    total_per_file[i],
                    dataset_multiprocessing,
                )

    # report and save total into a file
    total_count = np.sum(total_per_file)
    if not path.exists(total_file):
        np.savez_compressed(total_file, total_per_file=total_per_file)
    print("Total number of samples:", total_count)
    print("Divided into days/splits:\n", total_per_file)

    # dictionary files
    counts = np.zeros(26, dtype=np.int32)
    if recreate_flag:
        # create dictionaries
        for j in range(26):
            for i, x in enumerate(convertDicts[j]):
                convertDicts[j][x] = i
            dict_file_j = d_path + d_file + "_fea_dict_{0}.npz".format(j)
            if not path.exists(dict_file_j):
                np.savez_compressed(
                    dict_file_j,
                    unique=np.array(list(convertDicts[j]), dtype=np.int32)
                )
            counts[j] = len(convertDicts[j])
        # store (uniques and) counts
        count_file = d_path + d_file + "_fea_count.npz"
        if not path.exists(count_file):
            np.savez_compressed(count_file, counts=counts)
    else:
        # create dictionaries (from existing files)
        for j in range(26):
            with np.load(d_path + d_file + "_fea_dict_{0}.npz".format(j)) as data:
                unique = data["unique"]
            for i, x in enumerate(unique):
                convertDicts[j][x] = i
        # load (uniques and) counts
        with np.load(d_path + d_file + "_fea_count.npz") as data:
            counts = data["counts"]

    # process all splits
    if dataset_multiprocessing:
        processes = [Process(target=processCriteoAdData,
                           name="processCriteoAdData:%i" % i,
                           args=(d_path,
                                 d_file,
                                 npzfile,
                                 i,
                                 convertDicts,
                                 counts,
                                 )
                           ) for i in range(0, days)]
        for process in processes:
            process.start()
        for process in processes:
            process.join()

    else:
        for i in range(days):
            processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, counts)

    o_file = concatCriteoAdData(
        d_path,
        d_file,
        npzfile,
        trafile,
        days,
        data_split,
        randomize,
        total_per_file,
        total_count,
        memory_map,
        o_filename
    )

    return o_file


def loadDataset(
        dataset,
        max_ind_range,
        sub_sample_rate,
        randomize,
        data_split,
        raw_path="",
        pro_data="",
        memory_map=False
):
    # dataset
    if dataset == "kaggle":
        days = 7
        o_filename = "kaggleAdDisplayChallenge_processed"
    elif dataset == "terabyte":
        days = 24
        o_filename = "terabyte_processed"
    else:
        raise(ValueError("Data set option is not supported"))

    # split the datafile into path and filename
    lstr = raw_path.split("/")
    d_path = "/".join(lstr[0:-1]) + "/"
    d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
    npzfile = (d_file + "_day") if dataset == "kaggle" else d_file
    # trafile = d_path + ((d_file + "_fea") if dataset == "kaggle" else "fea")

    # check if pre-processed data is available
    data_ready = True
    if memory_map:
        for i in range(days):
            reo_data = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if not path.exists(str(reo_data)):
                data_ready = False
    else:
        if not path.exists(str(pro_data)):
            data_ready = False

    # pre-process data if needed
    # WARNNING: when memory mapping is used we get a collection of files
    if data_ready:
        print("Reading pre-processed data=%s" % (str(pro_data)))
        file = str(pro_data)
    else:
        print("Reading raw data=%s" % (str(raw_path)))
        file = getCriteoAdData(
            raw_path,
            o_filename,
            max_ind_range,
            sub_sample_rate,
            days,
            data_split,
            randomize,
            dataset == "kaggle",
            memory_map
        )

    return file, days


if __name__ == "__main__":
    ### import packages ###
    import argparse

    ### parse arguments ###
    parser = argparse.ArgumentParser(
        description="Preprocess Criteo dataset"
    )
    # model related parameters
    parser.add_argument("--max-ind-range", type=int, default=-1)
    parser.add_argument("--data-sub-sample-rate", type=float, default=0.0)  # in [0, 1]
    parser.add_argument("--data-randomize", type=str, default="total")  # or day or none
    parser.add_argument("--memory-map", action="store_true", default=False)
    parser.add_argument("--data-set", type=str, default="kaggle")  # or terabyte
    parser.add_argument("--raw-data-file", type=str, default="")
    parser.add_argument("--processed-data-file", type=str, default="")
    args = parser.parse_args()

    loadDataset(
        args.data_set,
        args.max_ind_range,
        args.data_sub_sample_rate,
        args.data_randomize,
        "train",
        args.raw_data_file,
        args.processed_data_file,
        args.memory_map
    )