contrib/pg_trgm/trgm_op.c

   1 /*
   2  * $PostgreSQL$
   3  */
   4 #include "trgm.h"
   5 #include <ctype.h>
   6 #include "utils/array.h"
   7 #include "catalog/pg_type.h"
   8 #include "tsearch/ts_locale.h"
   9
  10 PG_MODULE_MAGIC;
  11
  12 float4          trgm_limit = 0.3f;
  13
  14 PG_FUNCTION_INFO_V1(set_limit);
  15 Datum           set_limit(PG_FUNCTION_ARGS);
  16 Datum
  17 set_limit(PG_FUNCTION_ARGS)
  18 {
  19         float4          nlimit = PG_GETARG_FLOAT4(0);
  20
  21         if (nlimit < 0 || nlimit > 1.0)
  22                 elog(ERROR, "wrong limit, should be between 0 and 1");
  23         trgm_limit = nlimit;
  24         PG_RETURN_FLOAT4(trgm_limit);
  25 }
  26
  27 PG_FUNCTION_INFO_V1(show_limit);
  28 Datum           show_limit(PG_FUNCTION_ARGS);
  29 Datum
  30 show_limit(PG_FUNCTION_ARGS)
  31 {
  32         PG_RETURN_FLOAT4(trgm_limit);
  33 }
  34
  35 static int
  36 comp_trgm(const void *a, const void *b)
  37 {
  38         return CMPTRGM(a, b);
  39 }
  40
  41 static int
  42 unique_array(trgm * a, int len)
  43 {
  44         trgm       *curend,
  45                            *tmp;
  46
  47         curend = tmp = a;
  48         while (tmp - a < len)
  49                 if (CMPTRGM(tmp, curend))
  50                 {
  51                         curend++;
  52                         CPTRGM(curend, tmp);
  53                         tmp++;
  54                 }
  55                 else
  56                         tmp++;
  57
  58         return curend + 1 - a;
  59 }
  60
  61 #ifdef KEEPONLYALNUM
  62 #define iswordchr(c)    (t_isalpha(c) || t_isdigit(c))
  63 #else
  64 #define iswordchr(c)    (!t_isspace(c))
  65 #endif
  66
  67 /*
  68  * Finds first word in string, returns pointer to the word,
  69  * endword points to the character after word
  70  */
  71 static char*
  72 find_word(char *str, int lenstr, char **endword, int *charlen)
  73 {
  74         char *beginword = str;
  75
  76         while( beginword - str < lenstr && !iswordchr(beginword) )
  77                 beginword += pg_mblen(beginword);
  78
  79         if (beginword - str >= lenstr)
  80                 return NULL;
  81
  82         *endword = beginword;
  83         *charlen = 0;
  84         while( *endword - str < lenstr && iswordchr(*endword) )
  85         {
  86                 *endword += pg_mblen(*endword);
  87                 (*charlen)++;
  88         }
  89
  90         return beginword;
  91 }
  92
  93 #ifdef USE_WIDE_UPPER_LOWER
  94 static void
  95 cnt_trigram(trgm *tptr, char *str, int bytelen)
  96 {
  97         if ( bytelen == 3 )
  98         {
  99                 CPTRGM(tptr, str);
 100         }
 101         else
 102         {
 103                 pg_crc32        crc;
 104
 105                 INIT_CRC32(crc);
 106                 COMP_CRC32(crc, str, bytelen);
 107                 FIN_CRC32(crc);
 108
 109                 /*
 110                  * use only 3 upper bytes from crc, hope, it's
 111                  * good enough hashing
 112                  */
 113                 CPTRGM(tptr, &crc);
 114         }
 115 }
 116 #endif
 117
 118 /*
 119  * Adds trigramm from words (already padded).
 120  */
 121 static trgm*
 122 make_trigrams( trgm *tptr, char *str, int bytelen, int charlen )
 123 {
 124         char    *ptr = str;
 125
 126         if ( charlen < 3 )
 127                 return tptr;
 128
 129 #ifdef USE_WIDE_UPPER_LOWER
 130         if (pg_database_encoding_max_length() > 1)
 131         {
 132                 int lenfirst    = pg_mblen(str),
 133                         lenmiddle       = pg_mblen(str + lenfirst),
 134                         lenlast         = pg_mblen(str + lenfirst + lenmiddle);
 135
 136                 while( (ptr - str) + lenfirst + lenmiddle + lenlast <= bytelen )
 137                 {
 138                         cnt_trigram(tptr, ptr, lenfirst + lenmiddle + lenlast);
 139
 140                         ptr += lenfirst;
 141                         tptr++;
 142
 143                         lenfirst        = lenmiddle;
 144                         lenmiddle       = lenlast;
 145                         lenlast         = pg_mblen(ptr + lenfirst + lenmiddle);
 146                 }
 147         }
 148         else
 149 #endif
 150         {
 151                 Assert( bytelen == charlen );
 152
 153                 while (ptr - str < bytelen - 2 /* number of trigrams = strlen - 2 */ )
 154                 {
 155                         CPTRGM(tptr, ptr);
 156                         ptr++;
 157                         tptr++;
 158                 }
 159         }
 160
 161         return tptr;
 162 }
 163
 164 TRGM *
 165 generate_trgm(char *str, int slen)
 166 {
 167         TRGM       *trg;
 168         char       *buf;
 169         trgm       *tptr;
 170         int                     len,
 171                                 charlen,
 172                                 bytelen;
 173         char            *bword, *eword;
 174
 175         trg = (TRGM *) palloc(TRGMHDRSIZE + sizeof(trgm) * (slen / 2 + 1) * 3);
 176         trg->flag = ARRKEY;
 177         SET_VARSIZE(trg, TRGMHDRSIZE);
 178
 179         if (slen + LPADDING + RPADDING < 3 || slen == 0)
 180                 return trg;
 181
 182         tptr = GETARR(trg);
 183
 184         buf = palloc(sizeof(char) * (slen + 4));
 185
 186         if (LPADDING > 0)
 187         {
 188                 *buf = ' ';
 189                 if (LPADDING > 1)
 190                         *(buf + 1) = ' ';
 191         }
 192
 193         eword = str;
 194         while( (bword=find_word(eword, slen - (eword-str), &eword, &charlen)) != NULL )
 195         {
 196 #ifdef IGNORECASE
 197                 bword = lowerstr_with_len(bword, eword - bword);
 198                 bytelen = strlen(bword);
 199 #else
 200                 bytelen = eword - bword;
 201 #endif
 202
 203                 memcpy(buf + LPADDING, bword, bytelen);
 204
 205 #ifdef IGNORECASE
 206                 pfree(bword);
 207 #endif
 208                 buf[LPADDING+bytelen] = ' ';
 209                 buf[LPADDING+bytelen+1] = ' ';
 210
 211                 /*
 212                  * count trigrams
 213                  */
 214                 tptr = make_trigrams( tptr, buf, bytelen + LPADDING + RPADDING,
 215                                                                                  charlen + LPADDING + RPADDING );
 216         }
 217
 218         pfree(buf);
 219
 220         if ((len = tptr - GETARR(trg)) == 0)
 221                 return trg;
 222
 223         if (len > 0)
 224         {
 225                 qsort((void *) GETARR(trg), len, sizeof(trgm), comp_trgm);
 226                 len = unique_array(GETARR(trg), len);
 227         }
 228
 229         SET_VARSIZE(trg, CALCGTSIZE(ARRKEY, len));
 230
 231         return trg;
 232 }
 233
 234 uint32
 235 trgm2int(trgm *ptr)
 236 {
 237         uint32  val = 0;
 238
 239         val |= *( ((unsigned char*)ptr) );
 240         val <<= 8;
 241         val |= *( ((unsigned char*)ptr) + 1 );
 242         val <<= 8;
 243         val |= *( ((unsigned char*)ptr) + 2 );
 244
 245         return val;
 246 }
 247
 248 PG_FUNCTION_INFO_V1(show_trgm);
 249 Datum           show_trgm(PG_FUNCTION_ARGS);
 250 Datum
 251 show_trgm(PG_FUNCTION_ARGS)
 252 {
 253         text       *in = PG_GETARG_TEXT_P(0);
 254         TRGM       *trg;
 255         Datum      *d;
 256         ArrayType  *a;
 257         trgm       *ptr;
 258         int                     i;
 259
 260         trg = generate_trgm(VARDATA(in), VARSIZE(in) - VARHDRSZ);
 261         d = (Datum *) palloc(sizeof(Datum) * (1 + ARRNELEM(trg)));
 262
 263         for (i = 0, ptr = GETARR(trg); i < ARRNELEM(trg); i++, ptr++)
 264         {
 265                 text       *item = (text *) palloc(VARHDRSZ + Max(12, pg_database_encoding_max_length()*3) );
 266
 267                 if ( pg_database_encoding_max_length() > 1 && !ISPRINTABLETRGM(ptr) )
 268                 {
 269                         snprintf(VARDATA(item), 12, "0x%06x", trgm2int(ptr));
 270                         SET_VARSIZE(item, VARHDRSZ + strlen(VARDATA(item)));
 271                 }
 272                 else
 273                 {
 274                         SET_VARSIZE(item, VARHDRSZ + 3);
 275                         CPTRGM(VARDATA(item), ptr);
 276                 }
 277                 d[i] = PointerGetDatum(item);
 278         }
 279
 280         a = construct_array(
 281                                                 d,
 282                                                 ARRNELEM(trg),
 283                                                 TEXTOID,
 284                                                 -1,
 285                                                 false,
 286                                                 'i'
 287                 );
 288
 289         for (i = 0; i < ARRNELEM(trg); i++)
 290                 pfree(DatumGetPointer(d[i]));
 291
 292         pfree(d);
 293         pfree(trg);
 294         PG_FREE_IF_COPY(in, 0);
 295
 296         PG_RETURN_POINTER(a);
 297 }
 298
 299 float4
 300 cnt_sml(TRGM * trg1, TRGM * trg2)
 301 {
 302         trgm       *ptr1,
 303                            *ptr2;
 304         int                     count = 0;
 305         int                     len1,
 306                                 len2;
 307
 308         ptr1 = GETARR(trg1);
 309         ptr2 = GETARR(trg2);
 310
 311         len1 = ARRNELEM(trg1);
 312         len2 = ARRNELEM(trg2);
 313
 314         while (ptr1 - GETARR(trg1) < len1 && ptr2 - GETARR(trg2) < len2)
 315         {
 316                 int                     res = CMPTRGM(ptr1, ptr2);
 317
 318                 if (res < 0)
 319                         ptr1++;
 320                 else if (res > 0)
 321                         ptr2++;
 322                 else
 323                 {
 324                         ptr1++;
 325                         ptr2++;
 326                         count++;
 327                 }
 328         }
 329
 330 #ifdef DIVUNION
 331         return ((((float4) count) / ((float4) (len1 + len2 - count))));
 332 #else
 333         return (((float) count) / ((float) ((len1 > len2) ? len1 : len2)));
 334 #endif
 335
 336 }
 337
 338 PG_FUNCTION_INFO_V1(similarity);
 339 Datum           similarity(PG_FUNCTION_ARGS);
 340 Datum
 341 similarity(PG_FUNCTION_ARGS)
 342 {
 343         text       *in1 = PG_GETARG_TEXT_P(0);
 344         text       *in2 = PG_GETARG_TEXT_P(1);
 345         TRGM       *trg1,
 346                            *trg2;
 347         float4          res;
 348
 349         trg1 = generate_trgm(VARDATA(in1), VARSIZE(in1) - VARHDRSZ);
 350         trg2 = generate_trgm(VARDATA(in2), VARSIZE(in2) - VARHDRSZ);
 351
 352         res = cnt_sml(trg1, trg2);
 353
 354         pfree(trg1);
 355         pfree(trg2);
 356         PG_FREE_IF_COPY(in1, 0);
 357         PG_FREE_IF_COPY(in2, 1);
 358
 359         PG_RETURN_FLOAT4(res);
 360 }
 361
 362 PG_FUNCTION_INFO_V1(similarity_op);
 363 Datum           similarity_op(PG_FUNCTION_ARGS);
 364 Datum
 365 similarity_op(PG_FUNCTION_ARGS)
 366 {
 367         float4          res = DatumGetFloat4(DirectFunctionCall2(
 368                                                                                                                  similarity,
 369                                                                                                                  PG_GETARG_DATUM(0),
 370                                                                                                                  PG_GETARG_DATUM(1)
 371                                                                                                                  ));
 372
 373         PG_RETURN_BOOL(res >= trgm_limit);
 374 }