mariadb-columnstore-engine/tools/idb_comp/main.cpp

#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <cerrno>
//#define NDEBUG
#include <cassert>
#include <cstring>

#include <lzo/lzoconf.h>
#include <lzo/lzo1x.h>

#include <iostream>
#include <iomanip>
#include <sstream>
#include <inttypes.h>
#include <cstdlib>
using namespace std;

#include <boost/scoped_array.hpp>

namespace
{
size_t exp_buf_size;
unsigned vflg;

const string lzo_rctos(int r)
{
    switch (r)
    {
        case LZO_E_INPUT_NOT_CONSUMED:
            return "LZO_E_INPUT_NOT_CONSUMED";

        default:
            break;
    }

    return "Unknown Error!";
}

void usage()
{
    cout << "usage: idb_comp [-e size][-h] file(s)" << endl
         << "\t-e size size (in KB) of expansion buffer" << endl
         << "\t-h      display this help" << endl;
}

int doit(const string& infile)
{
    int rc = 0;
    int ifd = open(infile.c_str(), O_RDONLY | O_LARGEFILE | O_DIRECT);

    if (ifd < 0)
    {
        cerr << infile << ": open: " << strerror(errno) << endl;
        return 1;
    }

    struct stat istatbuf;

    fstat(ifd, &istatbuf);

    string outname(infile);

    string::size_type ptr;

    ptr = outname.find_last_of('.');

    if (ptr != string::npos)
        outname.erase(ptr);

    ptr = outname.find_last_of('/');

    if (ptr != string::npos)
        outname.erase(0, ptr + 1);

    outname = "./" + outname + ".cmp";

    int ofd = open(outname.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_LARGEFILE | O_DIRECT, 0644);

    if (ofd < 0)
    {
        cerr << outname << ": open: " << strerror(errno) << endl;
        close(ifd);
        return 1;
    }

    lzo_init();

    ssize_t nread = -1;
    ssize_t nwritten = -1;
    lzo_bytep ibuf;
    lzo_bytep tibuf;
    lzo_bytep cbuf;
    lzo_bytep tcbuf;
    lzo_bytep wkmem;
    lzo_uint32 ibuf_len = 0;
    lzo_uint cbuf_len = 0;

    ibuf_len = 512 * 1024 * 8;
    tibuf = new lzo_byte[ibuf_len + 4095];

    if ((ptrdiff_t)tibuf & 0xfffULL)
        ibuf = (lzo_bytep)((ptrdiff_t)tibuf & 0xfffffffffffff000ULL) + 4096;
    else
        ibuf = tibuf;

    cbuf_len = 512 * 1024 * 8;
    tcbuf = new lzo_byte[cbuf_len + 4095 + exp_buf_size * 1024];

    if ((ptrdiff_t)tcbuf & 0xfff)
        cbuf = (lzo_bytep)((ptrdiff_t)tcbuf & 0xfffffffffffff000ULL) + 4096;
    else
        cbuf = tcbuf;

    wkmem = new lzo_byte[LZO1X_1_15_MEM_COMPRESS];

    int r = LZO_E_OK;

    const int TOTAL_HDR_LEN = 4096 * 2;
    char* hdrbuf = new char[TOTAL_HDR_LEN + 4095];
    memset(hdrbuf, 0, TOTAL_HDR_LEN + 4095);
    char* hdrbufp = 0;

    if ((ptrdiff_t)hdrbuf & 0xfff)
        hdrbufp = (char*)((ptrdiff_t)hdrbuf & 0xfffffffffffff000ULL) + 4096;
    else
        hdrbufp = hdrbuf;

    struct compHdr
    {
        uint64_t ptrs[512];
    };

    idbassert(sizeof(compHdr) <= 4096);

    compHdr* hdrptr1 = (compHdr*)hdrbufp;
    compHdr* hdrptr  = hdrptr1 + 1; // advance to 2nd hdr to store compression ptrs
    lseek(ofd, TOTAL_HDR_LEN, SEEK_SET);

    nread = read(ifd, ibuf, ibuf_len);

    int idx = 0;
    off_t cmpoff = TOTAL_HDR_LEN;

    while (nread > 0)
    {
        cbuf_len = 512 * 1024 * 8;
        memset(cbuf, 0, cbuf_len);
        r = lzo1x_1_15_compress(ibuf, nread, cbuf, &cbuf_len, wkmem);

        if (r != LZO_E_OK)
        {
            cerr << "compression failed!: " << r << endl;
            rc = 1;
            goto out;
        }

        if (cbuf_len > (unsigned)nread)
        {
            cerr << "WARNING: expansion detected! (output grew by " << (cbuf_len - nread) << " bytes)" << endl;
            idbassert((cbuf_len - nread) <= exp_buf_size * 1024);
        }

        if (cbuf_len & 0xfff)
            cbuf_len = (cbuf_len & 0xfffffffffffff000ULL) + 4096;

        //cbuf_len = 512 * 1024 * 8;
        nwritten = write(ofd, cbuf, cbuf_len);

        if (nwritten < 0 || (unsigned)nwritten != cbuf_len)
        {
            cerr << outname << ": write: " << strerror(errno) << " (" << nwritten << ')' << endl;
            rc = 1;
            goto out;
        }

        if (vflg > 0)
        {
            lzo_bytep tbuf;
            lzo_uint tbuflen = 4 * 1024 * 1024 + 4;
            boost::scoped_array<lzo_byte> tbuf_sa(new lzo_byte[tbuflen]);
            tbuf = tbuf_sa.get();
            cout << "idx: " << idx << " off: " << cmpoff << " size: " << cbuf_len;
            r = lzo1x_decompress(cbuf, cbuf_len, tbuf, &tbuflen, 0);
            cout << " r: " << lzo_rctos(r) << " size: " << tbuflen << endl;
        }

        hdrptr->ptrs[idx] = cmpoff;
        idx++;
        cmpoff += cbuf_len;

        nread = read(ifd, ibuf, ibuf_len);
    }

    if (nread < 0)
    {
        cerr << infile << ": read: " << strerror(errno) << endl;
        rc = 1;
        goto out;
    }

    hdrptr->ptrs[idx] = cmpoff;
    idbassert(idx <= 64);

    // Fill in meta-data information in first header
    hdrptr1->ptrs[0] = 0xfdc119a384d0778eULL;
    hdrptr1->ptrs[1] = 1;
    hdrptr1->ptrs[2] = 1;

    nwritten = pwrite(ofd, hdrbufp, TOTAL_HDR_LEN, 0);
    idbassert(nwritten == TOTAL_HDR_LEN);

out:
    delete [] wkmem;
    delete [] tcbuf;
    delete [] tibuf;
    fsync(ofd);
    struct stat ostatbuf;
    fstat(ofd, &ostatbuf);
    idbassert(ostatbuf.st_size == (signed)hdrptr->ptrs[idx]);
    delete [] hdrbuf;
    cout << infile << ": Input Size: " << istatbuf.st_size
         << " Output size: " << ostatbuf.st_size
         << " Compression: " << (100LL - (ostatbuf.st_size * 100LL / istatbuf.st_size)) << '%' << endl;
    close(ofd);
    close(ifd);

    return rc;
}

}

int main(int argc, char** argv)
{
    opterr = 0;
    int c;
    exp_buf_size = 128;
    vflg = 0;

    while ((c = getopt(argc, argv, "e:vh")) != -1)
        switch (c)
        {
            case 'e':
                exp_buf_size = atoi(optarg);
                break;

            case 'v':
                vflg++;
                break;

            case 'h':
            default:
                usage();
                return (c == 'h' ? 0 : 1);
                break;
        }

    if ((argc - optind) < 1)
    {
        usage();
        return 1;
    }

    int rc = 0;

    for (int i = optind; i < argc; i++)
        if (doit(argv[i]))
            rc = 1;

    return rc;
}