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;
}

}  // namespace

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;
}