contrib/btree_gist/btree_gist.c

   1 #include "postgres.h"
   2
   3 #include "access/gist.h"
   4 #include "access/itup.h"
   5 #include "access/nbtree.h"
   6
   7 #include "utils/palloc.h"
   8 #include "utils/geo_decls.h"
   9 #include "utils/elog.h"
  10
  11 typedef int (*CMPFUNC) (const void *a, const void *b);
  12 typedef void (*BINARY_UNION) (Datum *, char *);
  13
  14 typedef struct intkey
  15 {
  16         int4            lower;
  17         int4            upper;
  18 }       INT4KEY;
  19
  20 typedef struct tskey
  21 {
  22         Timestamp       lower;
  23         Timestamp       upper;
  24 }       TSKEY;
  25
  26 /* used for sorting */
  27 typedef struct rix
  28 {
  29         int                     index;
  30         char       *r;
  31 }       RIX;
  32
  33 /*
  34 ** int4key in/out
  35 */
  36 PG_FUNCTION_INFO_V1(int4key_in);
  37 PG_FUNCTION_INFO_V1(int4key_out);
  38 Datum           int4key_in(PG_FUNCTION_ARGS);
  39 Datum           int4key_out(PG_FUNCTION_ARGS);
  40
  41 /*
  42 ** tskey in/out
  43 */
  44 PG_FUNCTION_INFO_V1(tskey_in);
  45 PG_FUNCTION_INFO_V1(tskey_out);
  46 Datum           tskey_in(PG_FUNCTION_ARGS);
  47 Datum           tskey_out(PG_FUNCTION_ARGS);
  48
  49 /*
  50 ** int4 ops
  51 */
  52 PG_FUNCTION_INFO_V1(gint4_compress);
  53 PG_FUNCTION_INFO_V1(gint4_union);
  54 PG_FUNCTION_INFO_V1(gint4_picksplit);
  55 PG_FUNCTION_INFO_V1(gint4_consistent);
  56 PG_FUNCTION_INFO_V1(gint4_penalty);
  57 PG_FUNCTION_INFO_V1(gint4_same);
  58
  59 Datum           gint4_compress(PG_FUNCTION_ARGS);
  60 Datum           gint4_union(PG_FUNCTION_ARGS);
  61 Datum           gint4_picksplit(PG_FUNCTION_ARGS);
  62 Datum           gint4_consistent(PG_FUNCTION_ARGS);
  63 Datum           gint4_penalty(PG_FUNCTION_ARGS);
  64 Datum           gint4_same(PG_FUNCTION_ARGS);
  65
  66 static void gint4_binary_union(Datum *r1, char *r2);
  67 static int      int4key_cmp(const void *a, const void *b);
  68
  69 /*
  70 ** timestamp ops
  71 */
  72 PG_FUNCTION_INFO_V1(gts_compress);
  73 PG_FUNCTION_INFO_V1(gts_union);
  74 PG_FUNCTION_INFO_V1(gts_picksplit);
  75 PG_FUNCTION_INFO_V1(gts_consistent);
  76 PG_FUNCTION_INFO_V1(gts_penalty);
  77 PG_FUNCTION_INFO_V1(gts_same);
  78
  79 Datum           gts_compress(PG_FUNCTION_ARGS);
  80 Datum           gts_union(PG_FUNCTION_ARGS);
  81 Datum           gts_picksplit(PG_FUNCTION_ARGS);
  82 Datum           gts_consistent(PG_FUNCTION_ARGS);
  83 Datum           gts_penalty(PG_FUNCTION_ARGS);
  84 Datum           gts_same(PG_FUNCTION_ARGS);
  85
  86 static void gts_binary_union(Datum *r1, char *r2);
  87 static int      tskey_cmp(const void *a, const void *b);
  88
  89 #define TimestampGetDatumFast(X) Float8GetDatumFast(X)
  90
  91 /* define for comparison */
  92 #define TSGE( ts1, ts2 ) (DatumGetBool(DirectFunctionCall2( \
  93                 timestamp_ge, \
  94                 PointerGetDatum( ts1 ), \
  95                 PointerGetDatum( ts2 ) \
  96 )))
  97 #define TSGT( ts1, ts2 ) (DatumGetBool(DirectFunctionCall2( \
  98                 timestamp_gt, \
  99                 PointerGetDatum( ts1 ), \
 100                 PointerGetDatum( ts2 ) \
 101 )))
 102 #define TSEQ( ts1, ts2 ) (DatumGetBool(DirectFunctionCall2( \
 103                 timestamp_eq, \
 104                 PointerGetDatum( ts1 ), \
 105                 PointerGetDatum( ts2 ) \
 106 )))
 107 #define TSLT( ts1, ts2 ) (DatumGetBool(DirectFunctionCall2( \
 108                 timestamp_lt, \
 109                 PointerGetDatum( ts1 ), \
 110                 PointerGetDatum( ts2 ) \
 111 )))
 112 #define TSLE( ts1, ts2 ) (DatumGetBool(DirectFunctionCall2( \
 113                 timestamp_le, \
 114                 PointerGetDatum( ts1 ), \
 115                 PointerGetDatum( ts2 ) \
 116 )))
 117
 118 /*
 119 ** Common btree-function (for all ops)
 120 */
 121 static GIST_SPLITVEC *btree_picksplit(bytea *entryvec, GIST_SPLITVEC *v,
 122                                 BINARY_UNION bu, CMPFUNC cmp);
 123
 124 PG_FUNCTION_INFO_V1(btree_decompress);
 125 Datum           btree_decompress(PG_FUNCTION_ARGS);
 126
 127 /**************************************************
 128  * int4 ops
 129  **************************************************/
 130
 131 Datum
 132 gint4_compress(PG_FUNCTION_ARGS)
 133 {
 134         GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
 135         GISTENTRY  *retval;
 136
 137         if (entry->leafkey)
 138         {
 139                 INT4KEY    *r = palloc(sizeof(INT4KEY));
 140
 141                 retval = palloc(sizeof(GISTENTRY));
 142                 r->lower = r->upper = (entry->key);
 143
 144                 gistentryinit(*retval, PointerGetDatum(r), entry->rel, entry->page,
 145                                           entry->offset, sizeof(INT4KEY), FALSE);
 146
 147         }
 148         else
 149                 retval = entry;
 150         PG_RETURN_POINTER(retval);
 151 }
 152
 153 Datum
 154 gint4_consistent(PG_FUNCTION_ARGS)
 155 {
 156         GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
 157         int4            query = PG_GETARG_INT32(1);
 158         INT4KEY    *kkk = (INT4KEY *) DatumGetPointer(entry->key);
 159         StrategyNumber strategy = (StrategyNumber) PG_GETARG_UINT16(2);
 160         bool            retval;
 161
 162         switch (strategy)
 163         {
 164                 case BTLessEqualStrategyNumber:
 165                         retval = (query >= kkk->lower);
 166                         break;
 167                 case BTLessStrategyNumber:
 168                         if (GIST_LEAF(entry))
 169                                 retval = (query > kkk->lower);
 170                         else
 171                                 retval = (query >= kkk->lower);
 172                         break;
 173                 case BTEqualStrategyNumber:
 174                         /* in leaf page kkk->lower always = kkk->upper */
 175                         if (GIST_LEAF(entry))
 176                                 retval = (query == kkk->lower);
 177                         else
 178                                 retval = (kkk->lower <= query && query <= kkk->upper);
 179                         break;
 180                 case BTGreaterStrategyNumber:
 181                         if (GIST_LEAF(entry))
 182                                 retval = (query < kkk->upper);
 183                         else
 184                                 retval = (query <= kkk->upper);
 185                         break;
 186                 case BTGreaterEqualStrategyNumber:
 187                         retval = (query <= kkk->upper);
 188                         break;
 189                 default:
 190                         retval = FALSE;
 191         }
 192         PG_RETURN_BOOL(retval);
 193 }
 194
 195 Datum
 196 gint4_union(PG_FUNCTION_ARGS)
 197 {
 198         bytea      *entryvec = (bytea *) PG_GETARG_POINTER(0);
 199         int                     i,
 200                                 numranges;
 201         INT4KEY    *cur,
 202                            *out = palloc(sizeof(INT4KEY));
 203
 204         numranges = (VARSIZE(entryvec) - VARHDRSZ) / sizeof(GISTENTRY);
 205         *(int *) PG_GETARG_POINTER(1) = sizeof(INT4KEY);
 206
 207         cur = (INT4KEY *) DatumGetPointer((((GISTENTRY *) (VARDATA(entryvec)))[0].key));
 208         out->lower = cur->lower;
 209         out->upper = cur->upper;
 210
 211         for (i = 1; i < numranges; i++)
 212         {
 213                 cur = (INT4KEY *) DatumGetPointer((((GISTENTRY *) (VARDATA(entryvec)))[i].key));
 214                 if (out->lower > cur->lower)
 215                         out->lower = cur->lower;
 216                 if (out->upper < cur->upper)
 217                         out->upper = cur->upper;
 218         }
 219
 220         PG_RETURN_POINTER(out);
 221 }
 222
 223 Datum
 224 gint4_penalty(PG_FUNCTION_ARGS)
 225 {
 226         INT4KEY    *origentry = (INT4KEY *) DatumGetPointer(((GISTENTRY *) PG_GETARG_POINTER(0))->key);
 227         INT4KEY    *newentry = (INT4KEY *) DatumGetPointer(((GISTENTRY *) PG_GETARG_POINTER(1))->key);
 228         float      *result = (float *) PG_GETARG_POINTER(2);
 229
 230         *result = Max(newentry->upper - origentry->upper, 0) +
 231                 Max(origentry->lower - newentry->lower, 0);
 232
 233         PG_RETURN_POINTER(result);
 234 }
 235
 236 Datum
 237 gint4_picksplit(PG_FUNCTION_ARGS)
 238 {
 239         PG_RETURN_POINTER(btree_picksplit(
 240                                                                           (bytea *) PG_GETARG_POINTER(0),
 241                                                                   (GIST_SPLITVEC *) PG_GETARG_POINTER(1),
 242                                                                           gint4_binary_union,
 243                                                                           int4key_cmp
 244                                                                           ));
 245 }
 246
 247 Datum
 248 gint4_same(PG_FUNCTION_ARGS)
 249 {
 250         INT4KEY    *b1 = (INT4KEY *) PG_GETARG_POINTER(0);
 251         INT4KEY    *b2 = (INT4KEY *) PG_GETARG_POINTER(1);
 252         bool       *result = (bool *) PG_GETARG_POINTER(2);
 253
 254         *result = (b1->lower == b2->lower && b1->upper == b2->upper) ? TRUE : FALSE;
 255         PG_RETURN_POINTER(result);
 256 }
 257
 258 static void
 259 gint4_binary_union(Datum *r1, char *r2)
 260 {
 261         INT4KEY    *b1;
 262         INT4KEY    *b2 = (INT4KEY *) r2;
 263
 264         if (!DatumGetPointer(*r1))
 265         {
 266                 *r1 = PointerGetDatum(palloc(sizeof(INT4KEY)));
 267                 b1 = (INT4KEY *) DatumGetPointer(*r1);
 268                 b1->upper = b2->upper;
 269                 b1->lower = b2->lower;
 270         }
 271         else
 272         {
 273                 b1 = (INT4KEY *) DatumGetPointer(*r1);
 274
 275                 b1->lower = (b1->lower > b2->lower) ?
 276                         b2->lower : b1->lower;
 277                 b1->upper = (b1->upper > b2->upper) ?
 278                         b1->upper : b2->upper;
 279         }
 280 }
 281
 282
 283 static int
 284 int4key_cmp(const void *a, const void *b)
 285 {
 286         return (((INT4KEY *) (((RIX *) a)->r))->lower - ((INT4KEY *) (((RIX *) b)->r))->lower);
 287 }
 288
 289 /**************************************************
 290  * timestamp ops
 291  **************************************************/
 292
 293 Datum
 294 gts_compress(PG_FUNCTION_ARGS)
 295 {
 296         GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
 297         GISTENTRY  *retval;
 298
 299         if (entry->leafkey)
 300         {
 301                 TSKEY      *r = (TSKEY *) palloc(sizeof(TSKEY));
 302
 303                 retval = palloc(sizeof(GISTENTRY));
 304                 r->lower = r->upper = *(Timestamp *) (entry->key);
 305                 gistentryinit(*retval, PointerGetDatum(r),
 306                                           entry->rel, entry->page,
 307                                           entry->offset, sizeof(TSKEY), FALSE);
 308         }
 309         else
 310                 retval = entry;
 311         PG_RETURN_POINTER(retval);
 312 }
 313
 314 Datum
 315 gts_consistent(PG_FUNCTION_ARGS)
 316 {
 317         GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
 318         Timestamp  *query = (Timestamp *) PG_GETARG_POINTER(1);
 319         StrategyNumber strategy = (StrategyNumber) PG_GETARG_UINT16(2);
 320         bool            retval;
 321         TSKEY      *key;
 322
 323         /*
 324          * * if entry is not leaf, use gbox_internal_consistent, * else use
 325          * gbox_leaf_consistent
 326          */
 327         if (!entry->key)
 328                 return FALSE;
 329         key = (TSKEY *) DatumGetPointer(entry->key);
 330
 331         switch (strategy)
 332         {
 333                 case BTLessEqualStrategyNumber:
 334                         retval = TSGE(query, &(key->lower));
 335                         break;
 336                 case BTLessStrategyNumber:
 337                         if (GIST_LEAF(entry))
 338                                 retval = TSGT(query, &(key->lower));
 339                         else
 340                                 retval = TSGE(query, &(key->lower));
 341                         break;
 342                 case BTEqualStrategyNumber:
 343                         /* in leaf page key->lower always = key->upper */
 344                         if (GIST_LEAF(entry))
 345                                 retval = TSEQ(query, &(key->lower));
 346                         else
 347                                 retval = (TSLE(&(key->lower), query) && TSLE(query, &(key->upper)));
 348                         break;
 349                 case BTGreaterStrategyNumber:
 350                         if (GIST_LEAF(entry))
 351                                 retval = TSLT(query, &(key->upper));
 352                         else
 353                                 retval = TSLE(query, &(key->upper));
 354                         break;
 355                 case BTGreaterEqualStrategyNumber:
 356                         retval = TSLE(query, &(key->upper));
 357                         break;
 358                 default:
 359                         retval = FALSE;
 360         }
 361         PG_RETURN_BOOL(retval);
 362 }
 363
 364 Datum
 365 gts_union(PG_FUNCTION_ARGS)
 366 {
 367         bytea      *entryvec = (bytea *) PG_GETARG_POINTER(0);
 368         int                     i,
 369                                 numranges;
 370         TSKEY      *cur,
 371                            *out = palloc(sizeof(TSKEY));
 372
 373         numranges = (VARSIZE(entryvec) - VARHDRSZ) / sizeof(GISTENTRY);
 374         *(int *) PG_GETARG_POINTER(1) = sizeof(TSKEY);
 375
 376         cur = (TSKEY *) DatumGetPointer((((GISTENTRY *) (VARDATA(entryvec)))[0].key));
 377         out->lower = cur->lower;
 378         out->upper = cur->upper;
 379
 380         for (i = 1; i < numranges; i++)
 381         {
 382                 cur = (TSKEY *) DatumGetPointer((((GISTENTRY *) (VARDATA(entryvec)))[i].key));
 383                 if (TSGT(&out->lower, &cur->lower))
 384                         out->lower = cur->lower;
 385                 if (TSLT(&out->upper, &cur->upper))
 386                         out->upper = cur->upper;
 387         }
 388
 389         PG_RETURN_POINTER(out);
 390 }
 391
 392 Datum
 393 gts_penalty(PG_FUNCTION_ARGS)
 394 {
 395         TSKEY      *origentry = (TSKEY *) DatumGetPointer(((GISTENTRY *) PG_GETARG_POINTER(0))->key);
 396         TSKEY      *newentry = (TSKEY *) DatumGetPointer(((GISTENTRY *) PG_GETARG_POINTER(1))->key);
 397         float      *result = (float *) PG_GETARG_POINTER(2);
 398         Interval   *intr;
 399
 400         intr = DatumGetIntervalP(DirectFunctionCall2(
 401                                                                                                  timestamp_mi,
 402                                                                   TimestampGetDatumFast(newentry->upper),
 403                                                            TimestampGetDatumFast(origentry->upper)));
 404
 405         /* see interval_larger */
 406         *result = Max(intr->time + intr->month * (30.0 * 86400), 0);
 407         pfree(intr);
 408
 409         intr = DatumGetIntervalP(DirectFunctionCall2(
 410                                                                                                  timestamp_mi,
 411                                                                  TimestampGetDatumFast(origentry->lower),
 412                                                                 TimestampGetDatumFast(newentry->lower)));
 413
 414         /* see interval_larger */
 415         *result += Max(intr->time + intr->month * (30.0 * 86400), 0);
 416         pfree(intr);
 417
 418         PG_RETURN_POINTER(result);
 419 }
 420
 421 Datum
 422 gts_picksplit(PG_FUNCTION_ARGS)
 423 {
 424         PG_RETURN_POINTER(btree_picksplit(
 425                                                                           (bytea *) PG_GETARG_POINTER(0),
 426                                                                   (GIST_SPLITVEC *) PG_GETARG_POINTER(1),
 427                                                                           gts_binary_union,
 428                                                                           tskey_cmp
 429                                                                           ));
 430 }
 431
 432 Datum
 433 gts_same(PG_FUNCTION_ARGS)
 434 {
 435         TSKEY      *b1 = (TSKEY *) PG_GETARG_POINTER(0);
 436         TSKEY      *b2 = (TSKEY *) PG_GETARG_POINTER(1);
 437
 438         bool       *result = (bool *) PG_GETARG_POINTER(2);
 439
 440         if (b1 && b2)
 441                 *result = (TSEQ(&(b1->lower), &(b2->lower)) && TSEQ(&(b1->upper), &(b2->upper))) ? TRUE : FALSE;
 442         else
 443                 *result = (b1 == NULL && b2 == NULL) ? TRUE : FALSE;
 444         PG_RETURN_POINTER(result);
 445 }
 446
 447 static void
 448 gts_binary_union(Datum *r1, char *r2)
 449 {
 450         TSKEY      *b1;
 451         TSKEY      *b2 = (TSKEY *) r2;
 452
 453         if (!DatumGetPointer(*r1))
 454         {
 455                 *r1 = PointerGetDatum(palloc(sizeof(TSKEY)));
 456                 b1 = (TSKEY *) DatumGetPointer(*r1);
 457                 b1->upper = b2->upper;
 458                 b1->lower = b2->lower;
 459         }
 460         else
 461         {
 462                 b1 = (TSKEY *) DatumGetPointer(*r1);
 463
 464                 b1->lower = (TSGT(&b1->lower, &b2->lower)) ?
 465                         b2->lower : b1->lower;
 466                 b1->upper = (TSGT(&b1->upper, &b2->upper)) ?
 467                         b1->upper : b2->upper;
 468         }
 469 }
 470
 471 static int
 472 tskey_cmp(const void *a, const void *b)
 473 {
 474         return DatumGetInt32(
 475                                                  DirectFunctionCall2(
 476                                                                                          timestamp_cmp,
 477                           TimestampGetDatumFast(((TSKEY *) (((RIX *) a)->r))->lower),
 478                            TimestampGetDatumFast(((TSKEY *) (((RIX *) b)->r))->lower)
 479                                                                                          )
 480                 );
 481 }
 482
 483 /**************************************************
 484  * Common btree-function (for all ops)
 485  **************************************************/
 486
 487 /*
 488 ** The GiST PickSplit method
 489 */
 490 static GIST_SPLITVEC *
 491 btree_picksplit(bytea *entryvec, GIST_SPLITVEC *v, BINARY_UNION bu, CMPFUNC cmp)
 492 {
 493         OffsetNumber i;
 494         RIX                *array;
 495         OffsetNumber maxoff;
 496         int                     nbytes;
 497
 498         maxoff = ((VARSIZE(entryvec) - VARHDRSZ) / sizeof(GISTENTRY)) - 1;
 499         nbytes = (maxoff + 2) * sizeof(OffsetNumber);
 500         v->spl_left = (OffsetNumber *) palloc(nbytes);
 501         v->spl_right = (OffsetNumber *) palloc(nbytes);
 502         v->spl_nleft = 0;
 503         v->spl_nright = 0;
 504         v->spl_ldatum = PointerGetDatum(0);
 505         v->spl_rdatum = PointerGetDatum(0);
 506         array = (RIX *) palloc(sizeof(RIX) * (maxoff + 1));
 507
 508         /* copy the data into RIXes, and sort the RIXes */
 509         for (i = FirstOffsetNumber; i <= maxoff; i = OffsetNumberNext(i))
 510         {
 511                 array[i].index = i;
 512                 array[i].r = (char *) DatumGetPointer((((GISTENTRY *) (VARDATA(entryvec)))[i].key));
 513         }
 514         qsort((void *) &array[FirstOffsetNumber], maxoff - FirstOffsetNumber + 1,
 515                   sizeof(RIX), cmp);
 516
 517         for (i = FirstOffsetNumber; i <= maxoff; i = OffsetNumberNext(i))
 518         {
 519                 if (i <= (maxoff - FirstOffsetNumber + 1) / 2)
 520                 {
 521                         v->spl_left[v->spl_nleft] = array[i].index;
 522                         v->spl_nleft++;
 523                         (*bu) (&v->spl_ldatum, array[i].r);
 524                 }
 525                 else
 526                 {
 527                         v->spl_right[v->spl_nright] = array[i].index;
 528                         v->spl_nright++;
 529                         (*bu) (&v->spl_rdatum, array[i].r);
 530                 }
 531         }
 532         pfree(array);
 533
 534         return (v);
 535 }
 536
 537 /*
 538 ** GiST DeCompress methods
 539 ** do not do anything.
 540 */
 541 Datum
 542 btree_decompress(PG_FUNCTION_ARGS)
 543 {
 544         PG_RETURN_POINTER(PG_GETARG_POINTER(0));
 545 }
 546
 547
 548 /**************************************************
 549  * In/Out for keys, not really needed
 550  **************************************************/
 551 Datum
 552 int4key_in(PG_FUNCTION_ARGS)
 553 {
 554         INT4KEY    *key = palloc(sizeof(INT4KEY));
 555
 556         if (sscanf(PG_GETARG_POINTER(0), "%d|%d", &(key->lower), &(key->upper)) != 2)
 557                 elog(ERROR, "Error in input format");
 558
 559         PG_RETURN_POINTER(key);
 560 }
 561
 562 Datum
 563 int4key_out(PG_FUNCTION_ARGS)
 564 {
 565         INT4KEY    *key = (INT4KEY *) PG_GETARG_POINTER(0);
 566         char       *str = palloc(sizeof(char) * 22);
 567
 568         sprintf(str, "%d|%d", key->lower, key->upper);
 569         PG_RETURN_POINTER(str);
 570 }
 571
 572 Datum
 573 tskey_in(PG_FUNCTION_ARGS)
 574 {
 575         elog(ERROR, "Not implemented");
 576         PG_RETURN_POINTER(NULL);
 577 }
 578
 579 Datum
 580 tskey_out(PG_FUNCTION_ARGS)
 581 {
 582         elog(ERROR, "Not implemented");
 583         PG_RETURN_POINTER(NULL);
 584 }