source: GenericIOSQLite.cxx @ da65757

Revision da65757, 12.4 KB checked in by Hal Finkel <hfinkel@…>, 9 years ago (diff)

add license to all files
  • Property mode set to 100644
Line 
1/*
2 *                    Copyright (C) 2015, UChicago Argonne, LLC
3 *                               All Rights Reserved
4 *
5 *                               Generic IO (ANL-15-066)
6 *                     Hal Finkel, Argonne National Laboratory
7 *
8 *                              OPEN SOURCE LICENSE
9 *
10 * Under the terms of Contract No. DE-AC02-06CH11357 with UChicago Argonne,
11 * LLC, the U.S. Government retains certain rights in this software.
12 *
13 * Redistribution and use in source and binary forms, with or without
14 * modification, are permitted provided that the following conditions are met:
15 *
16 *   1. Redistributions of source code must retain the above copyright notice,
17 *      this list of conditions and the following disclaimer.
18 *
19 *   2. Redistributions in binary form must reproduce the above copyright
20 *      notice, this list of conditions and the following disclaimer in the
21 *      documentation and/or other materials provided with the distribution.
22 *
23 *   3. Neither the names of UChicago Argonne, LLC or the Department of Energy
24 *      nor the names of its contributors may be used to endorse or promote
25 *      products derived from this software without specific prior written
26 *      permission.
27 *
28 * *****************************************************************************
29 *
30 *                                  DISCLAIMER
31 * THE SOFTWARE IS SUPPLIED “AS IS” WITHOUT WARRANTY OF ANY KIND.  NEITHER THE
32 * UNTED STATES GOVERNMENT, NOR THE UNITED STATES DEPARTMENT OF ENERGY, NOR
33 * UCHICAGO ARGONNE, LLC, NOR ANY OF THEIR EMPLOYEES, MAKES ANY WARRANTY,
34 * EXPRESS OR IMPLIED, OR ASSUMES ANY LEGAL LIABILITY OR RESPONSIBILITY FOR THE
35 * ACCURACY, COMPLETENESS, OR USEFULNESS OF ANY INFORMATION, DATA, APPARATUS,
36 * PRODUCT, OR PROCESS DISCLOSED, OR REPRESENTS THAT ITS USE WOULD NOT INFRINGE
37 * PRIVATELY OWNED RIGHTS.
38 *
39 * *****************************************************************************
40 */
41
42#include "GenericIO.h"
43#include <cstring>
44#include <vector>
45#include <string>
46#include <limits>
47#include <sstream>
48#include <sqlite3ext.h>
49SQLITE_EXTENSION_INIT1
50
51using namespace std;
52using namespace gio;
53
54class PrinterBase {
55public:
56  virtual ~PrinterBase() {}
57  virtual void print(sqlite3_context *cxt, size_t i) = 0;
58};
59
60template <class T>
61class Printer : public PrinterBase {
62public:
63  Printer(GenericIO &G, size_t MNE, const string &N)
64    : Data(MNE + G.requestedExtraSpace()/sizeof(T)) {
65    G.addVariable(N, Data, true);
66  }
67
68  virtual void print(sqlite3_context *cxt, size_t i) {
69    T d = Data[i];
70    if (!numeric_limits<T>::is_integer) {
71      sqlite3_result_double(cxt, double(d));
72    } else {
73      if (sizeof(T) > 4) {
74        sqlite3_result_int64(cxt, (sqlite3_int64) d);
75      } else {
76        sqlite3_result_int(cxt, int(d));
77      }
78    }
79  }
80
81protected:
82  vector<T> Data;
83};
84
85template <typename T>
86PrinterBase *addPrinter(GenericIO::VariableInfo &V,
87                GenericIO &GIO, size_t MNE) {
88  if (sizeof(T) != V.Size)
89    return 0;
90
91  if (V.IsFloat == numeric_limits<T>::is_integer)
92    return 0;
93  if (V.IsSigned != numeric_limits<T>::is_signed)
94    return 0;
95
96  return new Printer<T>(GIO, MNE, V.Name);
97}
98
99// See: Michael Owens. Query Anything with SQLite.
100// Dr. Dobb's Journal. Nov. 2007.
101// http://www.drdobbs.com/database/query-anything-with-sqlite/202802959
102// See: http://www.sqlite.org/c3ref/load_extension.html
103// See: http://www.sqlite.org/vtab.html
104
105struct gio_vtab {
106  gio_vtab(const string &FN) : GIO(FN), MaxNumElems(0) {
107    memset(&vtab, 0, sizeof(sqlite3_vtab));
108
109    GIO.openAndReadHeader(false);
110
111    int NR = GIO.readNRanks();
112    for (int i = 0; i < NR; ++i) {
113      size_t NElem = GIO.readNumElems(i);
114      MaxNumElems = max(MaxNumElems, NElem);
115    }
116  }
117
118  sqlite3_vtab vtab;
119  GenericIO    GIO;
120  size_t       MaxNumElems;
121};
122
123struct gio_cursor {
124  gio_cursor(const GenericIO &G, size_t MaxNElem) : GIO(G) {
125    memset(&cursor, 0, sizeof(sqlite3_vtab_cursor));
126
127    vector<GenericIO::VariableInfo> VI;
128    GIO.getVariableInfo(VI);
129
130    Vars.resize(VI.size());
131    for (size_t i = 0; i < VI.size(); ++i) {
132      PrinterBase *P = 0;
133
134#define ADD_PRINTER(T) \
135      if (!P) P = addPrinter<T>(VI[i], GIO, MaxNElem)
136      ADD_PRINTER(float);
137      ADD_PRINTER(double);
138      ADD_PRINTER(unsigned char);
139      ADD_PRINTER(signed char);
140      ADD_PRINTER(int16_t);
141      ADD_PRINTER(uint16_t);
142      ADD_PRINTER(int32_t);
143      ADD_PRINTER(uint32_t);
144      ADD_PRINTER(int64_t);
145      ADD_PRINTER(uint64_t);
146#undef ADD_PRINTER
147
148      if (P) Printers.push_back(P);
149    }
150  }
151
152  sqlite3_vtab_cursor    cursor;
153  int                    Rank;
154  size_t                 Pos;
155  uint64_t               AbsPos;
156  GenericIO              GIO;
157  vector< vector<char> > Vars;
158  vector<PrinterBase *>  Printers;
159  vector<bool>           RankMask;
160};
161
162static
163int gio_connect(sqlite3* db, void *, int argc, const char *const *argv,
164               sqlite3_vtab **ppVTab, char **pzErr) {
165  if (argc < 4) {
166    *pzErr = sqlite3_mprintf("%s", "No input file specified!");
167    return SQLITE_ERROR;
168  }
169
170  *pzErr = 0;
171  try {
172    string FileName(argv[3]);
173
174    gio_vtab *tab = new gio_vtab(FileName);
175    if (!tab)
176      return SQLITE_ERROR;
177
178    vector<GenericIO::VariableInfo> VI;
179    tab->GIO.getVariableInfo(VI);
180
181    string Schema = "CREATE TABLE x(_rank INTEGER";
182    if (VI.size())
183      Schema += ", ";
184
185    for (size_t i = 0; i < VI.size(); ++i) {
186      if (VI[i].Size > 8)
187        continue;
188
189      Schema += VI[i].Name + (VI[i].IsFloat ? " REAL" : " INTEGER");
190      if (i != VI.size() - 1)
191        Schema += ", ";
192    }
193
194    Schema += ")";
195    if (sqlite3_declare_vtab(db, Schema.c_str()) != SQLITE_OK) {
196      *pzErr = sqlite3_mprintf("Could not declare schema: %s", Schema.c_str());
197      delete tab;
198      return SQLITE_ERROR;
199    }
200
201    *ppVTab = &tab->vtab;
202  } catch(exception &e) {
203     *pzErr = sqlite3_mprintf("%s", e.what());
204     return SQLITE_ERROR;
205  }
206
207  return SQLITE_OK;
208}
209
210static
211int gio_create(sqlite3* db, void *pAux, int argc, const char *const *argv,
212               sqlite3_vtab **ppVTab, char **pzErr) {
213  return gio_connect(db, pAux, argc, argv, ppVTab, pzErr);
214}
215
216static
217int gio_disconnect(sqlite3_vtab *pVTab) {
218  gio_vtab *tab = (gio_vtab *) pVTab;
219
220  delete tab;
221  return SQLITE_OK;
222}
223
224static
225int gio_destroy(sqlite3_vtab *pVTab) {
226  return gio_disconnect(pVTab);
227}
228
229static
230int gio_best_index(sqlite3_vtab *pVTab, sqlite3_index_info *pInfo) {
231  gio_vtab *tab = (gio_vtab *) pVTab;
232
233  int IndexIndex = 0;
234  stringstream IndexName;
235  for (int i = 0; i < pInfo->nConstraint; ++i) {
236    if (!pInfo->aConstraint[i].usable)
237      continue;
238
239    char opcode = 0;
240    switch (pInfo->aConstraint[i].op) {
241    case SQLITE_INDEX_CONSTRAINT_EQ:
242      opcode = 'E';
243      break;
244    case SQLITE_INDEX_CONSTRAINT_GT:
245      opcode = 'g';
246      break;
247    case SQLITE_INDEX_CONSTRAINT_LT:
248      opcode = 'l';
249      break;
250    case SQLITE_INDEX_CONSTRAINT_GE:
251      opcode = 'G';
252      break;
253    case SQLITE_INDEX_CONSTRAINT_LE:
254      opcode = 'L';
255      break;
256    }
257
258    if (!opcode)
259      continue;
260
261    // For now, we only support column 0 (_rank)
262    // TODO: support for x, y, z, etc.
263    if (pInfo->aConstraint[i].iColumn != 0)
264      continue;
265
266    if (!IndexName.str().empty())
267      IndexName << ";";
268    IndexName << pInfo->aConstraint[i].iColumn << ":" << opcode;
269
270    pInfo->aConstraintUsage[i].argvIndex = 1 + IndexIndex++;
271    pInfo->aConstraintUsage[i].omit = 1;
272  }
273
274  if (!IndexName.str().empty()) {
275    pInfo->idxNum = 1;
276    pInfo->idxStr = sqlite3_mprintf("%s", IndexName.str().c_str());
277    if (!pInfo->idxStr)
278      return SQLITE_NOMEM;
279
280    pInfo->needToFreeIdxStr = 1;
281
282    // This is a crude estimate, without the actual range information,
283    // this is the best that we can do.
284    if (tab->GIO.readNRanks() > 0)
285      pInfo->estimatedCost = (double) tab->GIO.readNumElems(0);
286  } else {
287    uint64_t TotalNumElems = tab->GIO.readTotalNumElems();
288    if (TotalNumElems == (uint64_t) -1)
289      TotalNumElems = ((double) tab->GIO.readNumElems(0))*
290                        tab->GIO.readNRanks();
291    pInfo->estimatedCost = (double) TotalNumElems;
292  }
293
294  if (pInfo->nOrderBy == 1 && pInfo->aOrderBy[0].iColumn == 0 &&
295      pInfo->aOrderBy[0].desc == 0 /* ASC */)
296    pInfo->orderByConsumed = 1;
297
298  return SQLITE_OK;
299}
300
301static
302int gio_open(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor) {
303  gio_vtab   *tab = (gio_vtab *) pVTab;
304  gio_cursor *cur = new gio_cursor(tab->GIO, tab->MaxNumElems);
305  if (!cur)
306    return SQLITE_ERROR;
307
308  cur->cursor.pVtab = pVTab;
309  *ppCursor = &cur->cursor;
310
311  return SQLITE_OK;
312}
313
314static
315int gio_close(sqlite3_vtab_cursor* pCursor) {
316  gio_cursor *cur = (gio_cursor *) pCursor;
317
318  delete cur;
319  return SQLITE_OK;
320}
321
322static
323int gio_filter(sqlite3_vtab_cursor* pCursor, int idxNum, const char *idxStr,
324               int argc, sqlite3_value **argv) {
325  gio_cursor *cur = (gio_cursor *) pCursor;
326  gio_vtab   *tab = (gio_vtab *) cur->cursor.pVtab;
327
328  cur->RankMask.resize(cur->GIO.readNRanks(), true);
329  if (idxNum == 1) {
330    stringstream ss(idxStr);
331    string IndexSpec;
332    for (int arg = 0; getline(ss, IndexSpec, ';'); ++arg) {
333      if (arg >= argc)
334        return SQLITE_ERROR;
335
336      stringstream tt(IndexSpec);
337      int Col; char Colon, Op;
338      tt >> Col >> Colon >> Op;
339      if (Colon != ':')
340        return SQLITE_ERROR;
341
342      // TODO: Support for other columns.
343      if (Col != 0)
344        return SQLITE_ERROR;
345
346      int Rank = sqlite3_value_int(argv[arg]);
347      for (int i = 0; i < (int) cur->RankMask.size(); ++i) {
348        bool Incl = false;
349        if (Op == 'E')
350          Incl = (i == Rank);
351        else if (Op == 'g')
352          Incl = (i > Rank);
353        else if (Op == 'G')
354          Incl = (i >= Rank);
355        else if (Op == 'l')
356          Incl = (i < Rank);
357        else if (Op == 'L')
358          Incl = (i <= Rank);
359
360        cur->RankMask[i] = cur->RankMask[i] && Incl;
361      }
362    }
363  }
364
365  cur->Pos = 0;
366  cur->AbsPos = 0;
367
368  cur->Rank = -1;
369  do {
370    try {
371      ++cur->Rank;
372      if (cur->Rank == cur->GIO.readNRanks()) {
373        break;
374      }
375
376      if (cur->RankMask[cur->Rank])
377        cur->GIO.readData(cur->Rank, false);
378      else
379        cur->AbsPos += cur->GIO.readNumElems(cur->Rank);
380    } catch (exception &e) {
381        cerr << e.what() << endl;
382        return SQLITE_IOERR;
383    }
384  } while (!cur->RankMask[cur->Rank] ||
385           cur->GIO.readNumElems(cur->Rank) == 0);
386
387  return SQLITE_OK;
388}
389
390static
391int gio_next(sqlite3_vtab_cursor* pCursor) {
392  gio_cursor *cur = (gio_cursor *) pCursor;
393
394  if (cur->Pos == cur->GIO.readNumElems(cur->Rank) - 1) {
395    do {
396      try {
397        ++cur->Rank;
398        if (cur->Rank == cur->GIO.readNRanks()) {
399          break;
400        }
401
402        if (cur->RankMask[cur->Rank])
403          cur->GIO.readData(cur->Rank, false);
404        else
405          cur->AbsPos += cur->GIO.readNumElems(cur->Rank);
406      } catch (exception &e) {
407        cerr << e.what() << endl;
408        return SQLITE_IOERR;
409      }
410    } while (!cur->RankMask[cur->Rank] ||
411             cur->GIO.readNumElems(cur->Rank) == 0);
412
413    cur->Pos = 0;
414  } else {
415    ++cur->Pos;
416  }
417
418  ++cur->AbsPos;
419
420  return SQLITE_OK;
421}
422
423static
424int gio_eof(sqlite3_vtab_cursor* pCursor) {
425  gio_cursor *cur = (gio_cursor *) pCursor;
426
427  if (cur->Rank == cur->GIO.readNRanks())
428    return 1;
429
430  if (cur->Rank == cur->GIO.readNRanks() - 1 &&
431      cur->Pos >= (cur->GIO.readNumElems(cur->Rank) - 1))
432    return 1;
433
434  return 0;
435}
436
437static
438int gio_column(sqlite3_vtab_cursor* pCursor, sqlite3_context* cxt, int n) {
439  gio_cursor *cur = (gio_cursor *) pCursor;
440
441  if (n == 0)
442    sqlite3_result_int(cxt, cur->Rank);
443  else if ((n-1) >= cur->Printers.size())
444    sqlite3_result_null(cxt);
445  else
446    cur->Printers[n-1]->print(cxt, cur->Pos);
447
448  return SQLITE_OK;
449}
450
451static
452int gio_rowid(sqlite3_vtab_cursor* pCursor, sqlite3_int64 *pRowid) {
453  gio_cursor *cur = (gio_cursor *) pCursor;
454
455  *pRowid = (sqlite3_int64) cur->AbsPos;
456  return SQLITE_OK;
457}
458
459extern "C"
460int sq3_register_virtual_gio(sqlite3 * db) {
461  static const struct sqlite3_module gio_module = {
462    2,                         /* iVersion */
463    gio_create,
464    gio_connect,
465    gio_best_index,
466    gio_disconnect,
467    gio_destroy,
468    gio_open,
469    gio_close,
470    gio_filter,
471    gio_next,
472    gio_eof,
473    gio_column,
474    gio_rowid,
475    0,
476    0,
477    0,
478    0,
479    0,
480    0,
481    0,
482    0,
483    0,
484    0
485  };
486
487  return sqlite3_create_module( db, "GenericIO", &gio_module, NULL );
488}
489
490extern "C"
491int sqlite3_extension_init(sqlite3 *db, char **,
492                           const sqlite3_api_routines *pApi) {
493  SQLITE_EXTENSION_INIT2(pApi);
494  sq3_register_virtual_gio(db);
495  return SQLITE_OK;
496}
497
Note: See TracBrowser for help on using the repository browser.