PolarDB-for-PostgreSQL

1
/*
2
 * contrib/btree_gist/btree_utils_num.c
3
 */
4
#include "postgres.h"
5

6
#include "btree_gist.h"
7
#include "btree_utils_num.h"
8
#include "utils/cash.h"
9
#include "utils/date.h"
10
#include "utils/timestamp.h"
11

12

13
GISTENTRY *
14
gbt_num_compress(GISTENTRY *entry, const gbtree_ninfo *tinfo)
15
{
16
	GISTENTRY  *retval;
17

18
	if (entry->leafkey)
19
	{
20
		union
21
		{
22
			int16		i2;
23
			int32		i4;
24
			int64		i8;
25
			float4		f4;
26
			float8		f8;
27
			DateADT		dt;
28
			TimeADT		tm;
29
			Timestamp	ts;
30
			Cash		ch;
31
		}			v;
32

33
		GBT_NUMKEY *r = (GBT_NUMKEY *) palloc0(tinfo->indexsize);
34
		void	   *leaf = NULL;
35

36
		switch (tinfo->t)
37
		{
38
			case gbt_t_int2:
39
				v.i2 = DatumGetInt16(entry->key);
40
				leaf = &v.i2;
41
				break;
42
			case gbt_t_int4:
43
				v.i4 = DatumGetInt32(entry->key);
44
				leaf = &v.i4;
45
				break;
46
			case gbt_t_int8:
47
				v.i8 = DatumGetInt64(entry->key);
48
				leaf = &v.i8;
49
				break;
50
			case gbt_t_oid:
51
			case gbt_t_enum:
52
				v.i4 = DatumGetObjectId(entry->key);
53
				leaf = &v.i4;
54
				break;
55
			case gbt_t_float4:
56
				v.f4 = DatumGetFloat4(entry->key);
57
				leaf = &v.f4;
58
				break;
59
			case gbt_t_float8:
60
				v.f8 = DatumGetFloat8(entry->key);
61
				leaf = &v.f8;
62
				break;
63
			case gbt_t_date:
64
				v.dt = DatumGetDateADT(entry->key);
65
				leaf = &v.dt;
66
				break;
67
			case gbt_t_time:
68
				v.tm = DatumGetTimeADT(entry->key);
69
				leaf = &v.tm;
70
				break;
71
			case gbt_t_ts:
72
				v.ts = DatumGetTimestamp(entry->key);
73
				leaf = &v.ts;
74
				break;
75
			case gbt_t_cash:
76
				v.ch = DatumGetCash(entry->key);
77
				leaf = &v.ch;
78
				break;
79
			default:
80
				leaf = DatumGetPointer(entry->key);
81
		}
82

83
		Assert(tinfo->indexsize >= 2 * tinfo->size);
84

85
		memcpy((void *) &r[0], leaf, tinfo->size);
86
		memcpy((void *) &r[tinfo->size], leaf, tinfo->size);
87
		retval = palloc(sizeof(GISTENTRY));
88
		gistentryinit(*retval, PointerGetDatum(r), entry->rel, entry->page,
89
					  entry->offset, false);
90
	}
91
	else
92
		retval = entry;
93

94
	return retval;
95
}
96

97
/*
98
 * Convert a compressed leaf item back to the original type, for index-only
99
 * scans.
100
 */
101
GISTENTRY *
102
gbt_num_fetch(GISTENTRY *entry, const gbtree_ninfo *tinfo)
103
{
104
	GISTENTRY  *retval;
105
	Datum		datum;
106

107
	Assert(tinfo->indexsize >= 2 * tinfo->size);
108

109
	/*
110
	 * Get the original Datum from the stored datum. On leaf entries, the
111
	 * lower and upper bound are the same. We just grab the lower bound and
112
	 * return it.
113
	 */
114
	switch (tinfo->t)
115
	{
116
		case gbt_t_int2:
117
			datum = Int16GetDatum(*(int16 *) entry->key);
118
			break;
119
		case gbt_t_int4:
120
			datum = Int32GetDatum(*(int32 *) entry->key);
121
			break;
122
		case gbt_t_int8:
123
			datum = Int64GetDatum(*(int64 *) entry->key);
124
			break;
125
		case gbt_t_oid:
126
		case gbt_t_enum:
127
			datum = ObjectIdGetDatum(*(Oid *) entry->key);
128
			break;
129
		case gbt_t_float4:
130
			datum = Float4GetDatum(*(float4 *) entry->key);
131
			break;
132
		case gbt_t_float8:
133
			datum = Float8GetDatum(*(float8 *) entry->key);
134
			break;
135
		case gbt_t_date:
136
			datum = DateADTGetDatum(*(DateADT *) entry->key);
137
			break;
138
		case gbt_t_time:
139
			datum = TimeADTGetDatum(*(TimeADT *) entry->key);
140
			break;
141
		case gbt_t_ts:
142
			datum = TimestampGetDatum(*(Timestamp *) entry->key);
143
			break;
144
		case gbt_t_cash:
145
			datum = CashGetDatum(*(Cash *) entry->key);
146
			break;
147
		default:
148
			datum = PointerGetDatum(entry->key);
149
	}
150

151
	retval = palloc(sizeof(GISTENTRY));
152
	gistentryinit(*retval, datum, entry->rel, entry->page, entry->offset,
153
				  false);
154
	return retval;
155
}
156

157

158

159
/*
160
** The GiST union method for numerical values
161
*/
162

163
void *
164
gbt_num_union(GBT_NUMKEY *out, const GistEntryVector *entryvec, const gbtree_ninfo *tinfo, FmgrInfo *flinfo)
165
{
166
	int			i,
167
				numranges;
168
	GBT_NUMKEY *cur;
169
	GBT_NUMKEY_R o,
170
				c;
171

172
	numranges = entryvec->n;
173
	cur = (GBT_NUMKEY *) DatumGetPointer((entryvec->vector[0].key));
174

175

176
	o.lower = &((GBT_NUMKEY *) out)[0];
177
	o.upper = &((GBT_NUMKEY *) out)[tinfo->size];
178

179
	memcpy((void *) out, (void *) cur, 2 * tinfo->size);
180

181
	for (i = 1; i < numranges; i++)
182
	{
183
		cur = (GBT_NUMKEY *) DatumGetPointer((entryvec->vector[i].key));
184
		c.lower = &cur[0];
185
		c.upper = &cur[tinfo->size];
186
		/* if out->lower > cur->lower, adopt cur as lower */
187
		if (tinfo->f_gt(o.lower, c.lower, flinfo))
188
			memcpy((void *) o.lower, (void *) c.lower, tinfo->size);
189
		/* if out->upper < cur->upper, adopt cur as upper */
190
		if (tinfo->f_lt(o.upper, c.upper, flinfo))
191
			memcpy((void *) o.upper, (void *) c.upper, tinfo->size);
192
	}
193

194
	return out;
195
}
196

197

198

199
/*
200
** The GiST same method for numerical values
201
*/
202

203
bool
204
gbt_num_same(const GBT_NUMKEY *a, const GBT_NUMKEY *b, const gbtree_ninfo *tinfo, FmgrInfo *flinfo)
205
{
206
	GBT_NUMKEY_R b1,
207
				b2;
208

209
	b1.lower = &(((GBT_NUMKEY *) a)[0]);
210
	b1.upper = &(((GBT_NUMKEY *) a)[tinfo->size]);
211
	b2.lower = &(((GBT_NUMKEY *) b)[0]);
212
	b2.upper = &(((GBT_NUMKEY *) b)[tinfo->size]);
213

214
	return (tinfo->f_eq(b1.lower, b2.lower, flinfo) &&
215
			tinfo->f_eq(b1.upper, b2.upper, flinfo));
216
}
217

218

219
void
220
gbt_num_bin_union(Datum *u, GBT_NUMKEY *e, const gbtree_ninfo *tinfo, FmgrInfo *flinfo)
221
{
222
	GBT_NUMKEY_R rd;
223

224
	rd.lower = &e[0];
225
	rd.upper = &e[tinfo->size];
226

227
	if (!DatumGetPointer(*u))
228
	{
229
		*u = PointerGetDatum(palloc0(tinfo->indexsize));
230
		memcpy((void *) &(((GBT_NUMKEY *) DatumGetPointer(*u))[0]), (void *) rd.lower, tinfo->size);
231
		memcpy((void *) &(((GBT_NUMKEY *) DatumGetPointer(*u))[tinfo->size]), (void *) rd.upper, tinfo->size);
232
	}
233
	else
234
	{
235
		GBT_NUMKEY_R ur;
236

237
		ur.lower = &(((GBT_NUMKEY *) DatumGetPointer(*u))[0]);
238
		ur.upper = &(((GBT_NUMKEY *) DatumGetPointer(*u))[tinfo->size]);
239
		if (tinfo->f_gt((void *) ur.lower, (void *) rd.lower, flinfo))
240
			memcpy((void *) ur.lower, (void *) rd.lower, tinfo->size);
241
		if (tinfo->f_lt((void *) ur.upper, (void *) rd.upper, flinfo))
242
			memcpy((void *) ur.upper, (void *) rd.upper, tinfo->size);
243
	}
244
}
245

246

247

248
/*
249
 * The GiST consistent method
250
 *
251
 * Note: we currently assume that no datatypes that use this routine are
252
 * collation-aware; so we don't bother passing collation through.
253
 */
254
bool
255
gbt_num_consistent(const GBT_NUMKEY_R *key,
256
				   const void *query,
257
				   const StrategyNumber *strategy,
258
				   bool is_leaf,
259
				   const gbtree_ninfo *tinfo,
260
				   FmgrInfo *flinfo)
261
{
262
	bool		retval;
263

264
	switch (*strategy)
265
	{
266
		case BTLessEqualStrategyNumber:
267
			retval = tinfo->f_ge(query, key->lower, flinfo);
268
			break;
269
		case BTLessStrategyNumber:
270
			if (is_leaf)
271
				retval = tinfo->f_gt(query, key->lower, flinfo);
272
			else
273
				retval = tinfo->f_ge(query, key->lower, flinfo);
274
			break;
275
		case BTEqualStrategyNumber:
276
			if (is_leaf)
277
				retval = tinfo->f_eq(query, key->lower, flinfo);
278
			else
279
				retval = (tinfo->f_le(key->lower, query, flinfo) &&
280
						  tinfo->f_le(query, key->upper, flinfo));
281
			break;
282
		case BTGreaterStrategyNumber:
283
			if (is_leaf)
284
				retval = tinfo->f_lt(query, key->upper, flinfo);
285
			else
286
				retval = tinfo->f_le(query, key->upper, flinfo);
287
			break;
288
		case BTGreaterEqualStrategyNumber:
289
			retval = tinfo->f_le(query, key->upper, flinfo);
290
			break;
291
		case BtreeGistNotEqualStrategyNumber:
292
			retval = (!(tinfo->f_eq(query, key->lower, flinfo) &&
293
						tinfo->f_eq(query, key->upper, flinfo)));
294
			break;
295
		default:
296
			retval = false;
297
	}
298

299
	return retval;
300
}
301

302

303
/*
304
** The GiST distance method (for KNN-Gist)
305
*/
306

307
float8
308
gbt_num_distance(const GBT_NUMKEY_R *key,
309
				 const void *query,
310
				 bool is_leaf,
311
				 const gbtree_ninfo *tinfo,
312
				 FmgrInfo *flinfo)
313
{
314
	float8		retval;
315

316
	if (tinfo->f_dist == NULL)
317
		elog(ERROR, "KNN search is not supported for btree_gist type %d",
318
			 (int) tinfo->t);
319
	if (tinfo->f_le(query, key->lower, flinfo))
320
		retval = tinfo->f_dist(query, key->lower, flinfo);
321
	else if (tinfo->f_ge(query, key->upper, flinfo))
322
		retval = tinfo->f_dist(query, key->upper, flinfo);
323
	else
324
		retval = 0.0;
325

326
	return retval;
327
}
328

329

330
GIST_SPLITVEC *
331
gbt_num_picksplit(const GistEntryVector *entryvec, GIST_SPLITVEC *v,
332
				  const gbtree_ninfo *tinfo, FmgrInfo *flinfo)
333
{
334
	OffsetNumber i,
335
				maxoff = entryvec->n - 1;
336
	Nsrt	   *arr;
337
	int			nbytes;
338

339
	arr = (Nsrt *) palloc((maxoff + 1) * sizeof(Nsrt));
340
	nbytes = (maxoff + 2) * sizeof(OffsetNumber);
341
	v->spl_left = (OffsetNumber *) palloc(nbytes);
342
	v->spl_right = (OffsetNumber *) palloc(nbytes);
343
	v->spl_ldatum = PointerGetDatum(0);
344
	v->spl_rdatum = PointerGetDatum(0);
345
	v->spl_nleft = 0;
346
	v->spl_nright = 0;
347

348
	/* Sort entries */
349

350
	for (i = FirstOffsetNumber; i <= maxoff; i = OffsetNumberNext(i))
351
	{
352
		arr[i].t = (GBT_NUMKEY *) DatumGetPointer((entryvec->vector[i].key));
353
		arr[i].i = i;
354
	}
355
	qsort_arg((void *) &arr[FirstOffsetNumber], maxoff - FirstOffsetNumber + 1, sizeof(Nsrt), (qsort_arg_comparator) tinfo->f_cmp, (void *) flinfo);
356

357
	/* We do simply create two parts */
358

359
	for (i = FirstOffsetNumber; i <= maxoff; i = OffsetNumberNext(i))
360
	{
361
		if (i <= (maxoff - FirstOffsetNumber + 1) / 2)
362
		{
363
			gbt_num_bin_union(&v->spl_ldatum, arr[i].t, tinfo, flinfo);
364
			v->spl_left[v->spl_nleft] = arr[i].i;
365
			v->spl_nleft++;
366
		}
367
		else
368
		{
369
			gbt_num_bin_union(&v->spl_rdatum, arr[i].t, tinfo, flinfo);
370
			v->spl_right[v->spl_nright] = arr[i].i;
371
			v->spl_nright++;
372
		}
373
	}
374

375
	return v;
376
}
377