Skip to content

Commit 8436489

Browse files
tglsfdctglsfdc
committed
Add KNNGIST support to contrib/btree_gist.
This extends GiST's support for nearest-neighbor searches to many of the standard data types. Teodor Sigaev
1 parent 6094c24 commit 8436489

40 files changed

+1546
-59
lines changed

‎contrib/btree_gist/btree_cash.c

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,15 @@ PG_FUNCTION_INFO_V1(gbt_cash_compress);
1818
PG_FUNCTION_INFO_V1(gbt_cash_union);
1919
PG_FUNCTION_INFO_V1(gbt_cash_picksplit);
2020
PG_FUNCTION_INFO_V1(gbt_cash_consistent);
21+
PG_FUNCTION_INFO_V1(gbt_cash_distance);
2122
PG_FUNCTION_INFO_V1(gbt_cash_penalty);
2223
PG_FUNCTION_INFO_V1(gbt_cash_same);
2324

2425
Datum gbt_cash_compress(PG_FUNCTION_ARGS);
2526
Datum gbt_cash_union(PG_FUNCTION_ARGS);
2627
Datum gbt_cash_picksplit(PG_FUNCTION_ARGS);
2728
Datum gbt_cash_consistent(PG_FUNCTION_ARGS);
29+
Datum gbt_cash_distance(PG_FUNCTION_ARGS);
2830
Datum gbt_cash_penalty(PG_FUNCTION_ARGS);
2931
Datum gbt_cash_same(PG_FUNCTION_ARGS);
3032

@@ -71,6 +73,12 @@ gbt_cashkey_cmp(const void *a, const void *b)
7173
return (ia->lower > ib->lower) ? 1 : -1;
7274
}
7375

76+
static float8
77+
gbt_cash_dist(const void *a, const void *b)
78+
{
79+
return GET_FLOAT_DISTANCE(Cash, a, b);
80+
}
81+
7482

7583
static const gbtree_ninfo tinfo =
7684
{
@@ -81,10 +89,33 @@ static const gbtree_ninfo tinfo =
8189
gbt_casheq,
8290
gbt_cashle,
8391
gbt_cashlt,
84-
gbt_cashkey_cmp
92+
gbt_cashkey_cmp,
93+
gbt_cash_dist
8594
};
8695

8796

97+
PG_FUNCTION_INFO_V1(cash_dist);
98+
Datum cash_dist(PG_FUNCTION_ARGS);
99+
Datum
100+
cash_dist(PG_FUNCTION_ARGS)
101+
{
102+
Cash a = PG_GETARG_CASH(0);
103+
Cash b = PG_GETARG_CASH(1);
104+
Cash r;
105+
Cash ra;
106+
107+
r = a - b;
108+
ra = Abs(r);
109+
110+
/* Overflow check. */
111+
if (ra < 0 || (!SAMESIGN(a, b) && !SAMESIGN(r, a)))
112+
ereport(ERROR,
113+
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
114+
errmsg("money out of range")));
115+
116+
PG_RETURN_CASH(ra);
117+
}
118+
88119
/**************************************************
89120
* Cash ops
90121
**************************************************/
@@ -124,6 +155,25 @@ gbt_cash_consistent(PG_FUNCTION_ARGS)
124155
}
125156

126157

158+
Datum
159+
gbt_cash_distance(PG_FUNCTION_ARGS)
160+
{
161+
GISTENTRY *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
162+
Cash query = PG_GETARG_CASH(1);
163+
164+
/* Oid subtype = PG_GETARG_OID(3); */
165+
cashKEY *kkk = (cashKEY *) DatumGetPointer(entry->key);
166+
GBT_NUMKEY_R key;
167+
168+
key.lower = (GBT_NUMKEY *) &kkk->lower;
169+
key.upper = (GBT_NUMKEY *) &kkk->upper;
170+
171+
PG_RETURN_FLOAT8(
172+
gbt_num_distance(&key, (void *) &query, GIST_LEAF(entry), &tinfo)
173+
);
174+
}
175+
176+
127177
Datum
128178
gbt_cash_union(PG_FUNCTION_ARGS)
129179
{

‎contrib/btree_gist/btree_date.c

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,15 @@ PG_FUNCTION_INFO_V1(gbt_date_compress);
1818
PG_FUNCTION_INFO_V1(gbt_date_union);
1919
PG_FUNCTION_INFO_V1(gbt_date_picksplit);
2020
PG_FUNCTION_INFO_V1(gbt_date_consistent);
21+
PG_FUNCTION_INFO_V1(gbt_date_distance);
2122
PG_FUNCTION_INFO_V1(gbt_date_penalty);
2223
PG_FUNCTION_INFO_V1(gbt_date_same);
2324

2425
Datum gbt_date_compress(PG_FUNCTION_ARGS);
2526
Datum gbt_date_union(PG_FUNCTION_ARGS);
2627
Datum gbt_date_picksplit(PG_FUNCTION_ARGS);
2728
Datum gbt_date_consistent(PG_FUNCTION_ARGS);
29+
Datum gbt_date_distance(PG_FUNCTION_ARGS);
2830
Datum gbt_date_penalty(PG_FUNCTION_ARGS);
2931
Datum gbt_date_same(PG_FUNCTION_ARGS);
3032

@@ -84,6 +86,17 @@ gbt_datekey_cmp(const void *a, const void *b)
8486
return res;
8587
}
8688

89+
static float8
90+
gdb_date_dist(const void *a, const void *b)
91+
{
92+
/* we assume the difference can't overflow */
93+
Datum diff = DirectFunctionCall2(date_mi,
94+
DateADTGetDatum(*((const DateADT *) a)),
95+
DateADTGetDatum(*((const DateADT *) b)));
96+
97+
return (float8) Abs(DatumGetInt32(diff));
98+
}
99+
87100

88101
static const gbtree_ninfo tinfo =
89102
{
@@ -94,10 +107,25 @@ static const gbtree_ninfo tinfo =
94107
gbt_dateeq,
95108
gbt_datele,
96109
gbt_datelt,
97-
gbt_datekey_cmp
110+
gbt_datekey_cmp,
111+
gdb_date_dist
98112
};
99113

100114

115+
PG_FUNCTION_INFO_V1(date_dist);
116+
Datum date_dist(PG_FUNCTION_ARGS);
117+
Datum
118+
date_dist(PG_FUNCTION_ARGS)
119+
{
120+
/* we assume the difference can't overflow */
121+
Datum diff = DirectFunctionCall2(date_mi,
122+
PG_GETARG_DATUM(0),
123+
PG_GETARG_DATUM(1));
124+
125+
PG_RETURN_INT32(Abs(DatumGetInt32(diff)));
126+
}
127+
128+
101129
/**************************************************
102130
* date ops
103131
**************************************************/
@@ -139,6 +167,25 @@ gbt_date_consistent(PG_FUNCTION_ARGS)
139167
}
140168

141169

170+
Datum
171+
gbt_date_distance(PG_FUNCTION_ARGS)
172+
{
173+
GISTENTRY *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
174+
DateADT query = PG_GETARG_DATEADT(1);
175+
176+
/* Oid subtype = PG_GETARG_OID(3); */
177+
dateKEY *kkk = (dateKEY *) DatumGetPointer(entry->key);
178+
GBT_NUMKEY_R key;
179+
180+
key.lower = (GBT_NUMKEY *) &kkk->lower;
181+
key.upper = (GBT_NUMKEY *) &kkk->upper;
182+
183+
PG_RETURN_FLOAT8(
184+
gbt_num_distance(&key, (void *) &query, GIST_LEAF(entry), &tinfo)
185+
);
186+
}
187+
188+
142189
Datum
143190
gbt_date_union(PG_FUNCTION_ARGS)
144191
{

‎contrib/btree_gist/btree_float4.c

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@ PG_FUNCTION_INFO_V1(gbt_float4_compress);
1717
PG_FUNCTION_INFO_V1(gbt_float4_union);
1818
PG_FUNCTION_INFO_V1(gbt_float4_picksplit);
1919
PG_FUNCTION_INFO_V1(gbt_float4_consistent);
20+
PG_FUNCTION_INFO_V1(gbt_float4_distance);
2021
PG_FUNCTION_INFO_V1(gbt_float4_penalty);
2122
PG_FUNCTION_INFO_V1(gbt_float4_same);
2223

2324
Datum gbt_float4_compress(PG_FUNCTION_ARGS);
2425
Datum gbt_float4_union(PG_FUNCTION_ARGS);
2526
Datum gbt_float4_picksplit(PG_FUNCTION_ARGS);
2627
Datum gbt_float4_consistent(PG_FUNCTION_ARGS);
28+
Datum gbt_float4_distance(PG_FUNCTION_ARGS);
2729
Datum gbt_float4_penalty(PG_FUNCTION_ARGS);
2830
Datum gbt_float4_same(PG_FUNCTION_ARGS);
2931

@@ -70,6 +72,12 @@ gbt_float4key_cmp(const void *a, const void *b)
7072
return (ia->lower > ib->lower) ? 1 : -1;
7173
}
7274

75+
static float8
76+
gbt_float4_dist(const void *a, const void *b)
77+
{
78+
return GET_FLOAT_DISTANCE(float4, a, b);
79+
}
80+
7381

7482
static const gbtree_ninfo tinfo =
7583
{
@@ -80,10 +88,27 @@ static const gbtree_ninfo tinfo =
8088
gbt_float4eq,
8189
gbt_float4le,
8290
gbt_float4lt,
83-
gbt_float4key_cmp
91+
gbt_float4key_cmp,
92+
gbt_float4_dist
8493
};
8594

8695

96+
PG_FUNCTION_INFO_V1(float4_dist);
97+
Datum float4_dist(PG_FUNCTION_ARGS);
98+
Datum
99+
float4_dist(PG_FUNCTION_ARGS)
100+
{
101+
float4 a = PG_GETARG_FLOAT4(0);
102+
float4 b = PG_GETARG_FLOAT4(1);
103+
float4 r;
104+
105+
r = a - b;
106+
CHECKFLOATVAL(r, isinf(a) || isinf(b), true);
107+
108+
PG_RETURN_FLOAT4( Abs(r) );
109+
}
110+
111+
87112
/**************************************************
88113
* float4 ops
89114
**************************************************/
@@ -123,6 +148,25 @@ gbt_float4_consistent(PG_FUNCTION_ARGS)
123148
}
124149

125150

151+
Datum
152+
gbt_float4_distance(PG_FUNCTION_ARGS)
153+
{
154+
GISTENTRY *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
155+
float4 query = PG_GETARG_FLOAT4(1);
156+
157+
/* Oid subtype = PG_GETARG_OID(3); */
158+
float4KEY *kkk = (float4KEY *) DatumGetPointer(entry->key);
159+
GBT_NUMKEY_R key;
160+
161+
key.lower = (GBT_NUMKEY *) &kkk->lower;
162+
key.upper = (GBT_NUMKEY *) &kkk->upper;
163+
164+
PG_RETURN_FLOAT8(
165+
gbt_num_distance(&key, (void *) &query, GIST_LEAF(entry), &tinfo)
166+
);
167+
}
168+
169+
126170
Datum
127171
gbt_float4_union(PG_FUNCTION_ARGS)
128172
{

‎contrib/btree_gist/btree_float8.c

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@ PG_FUNCTION_INFO_V1(gbt_float8_compress);
1717
PG_FUNCTION_INFO_V1(gbt_float8_union);
1818
PG_FUNCTION_INFO_V1(gbt_float8_picksplit);
1919
PG_FUNCTION_INFO_V1(gbt_float8_consistent);
20+
PG_FUNCTION_INFO_V1(gbt_float8_distance);
2021
PG_FUNCTION_INFO_V1(gbt_float8_penalty);
2122
PG_FUNCTION_INFO_V1(gbt_float8_same);
2223

2324
Datum gbt_float8_compress(PG_FUNCTION_ARGS);
2425
Datum gbt_float8_union(PG_FUNCTION_ARGS);
2526
Datum gbt_float8_picksplit(PG_FUNCTION_ARGS);
2627
Datum gbt_float8_consistent(PG_FUNCTION_ARGS);
28+
Datum gbt_float8_distance(PG_FUNCTION_ARGS);
2729
Datum gbt_float8_penalty(PG_FUNCTION_ARGS);
2830
Datum gbt_float8_same(PG_FUNCTION_ARGS);
2931

@@ -71,6 +73,19 @@ gbt_float8key_cmp(const void *a, const void *b)
7173
return (ia->lower > ib->lower) ? 1 : -1;
7274
}
7375

76+
static float8
77+
gbt_float8_dist(const void *a, const void *b)
78+
{
79+
float8 arg1 = *(const float8 *)a;
80+
float8 arg2 = *(const float8 *)b;
81+
float8 r;
82+
83+
r = arg1 - arg2;
84+
CHECKFLOATVAL(r, isinf(arg1) || isinf(arg2), true);
85+
86+
return Abs(r);
87+
}
88+
7489

7590
static const gbtree_ninfo tinfo =
7691
{
@@ -81,10 +96,26 @@ static const gbtree_ninfo tinfo =
8196
gbt_float8eq,
8297
gbt_float8le,
8398
gbt_float8lt,
84-
gbt_float8key_cmp
99+
gbt_float8key_cmp,
100+
gbt_float8_dist
85101
};
86102

87103

104+
PG_FUNCTION_INFO_V1(float8_dist);
105+
Datum float8_dist(PG_FUNCTION_ARGS);
106+
Datum
107+
float8_dist(PG_FUNCTION_ARGS)
108+
{
109+
float8 a = PG_GETARG_FLOAT8(0);
110+
float8 b = PG_GETARG_FLOAT8(1);
111+
float8 r;
112+
113+
r = a - b;
114+
CHECKFLOATVAL(r, isinf(a) || isinf(b), true);
115+
116+
PG_RETURN_FLOAT8( Abs(r) );
117+
}
118+
88119
/**************************************************
89120
* float8 ops
90121
**************************************************/
@@ -124,6 +155,25 @@ gbt_float8_consistent(PG_FUNCTION_ARGS)
124155
}
125156

126157

158+
Datum
159+
gbt_float8_distance(PG_FUNCTION_ARGS)
160+
{
161+
GISTENTRY *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
162+
float8 query = PG_GETARG_FLOAT8(1);
163+
164+
/* Oid subtype = PG_GETARG_OID(3); */
165+
float8KEY *kkk = (float8KEY *) DatumGetPointer(entry->key);
166+
GBT_NUMKEY_R key;
167+
168+
key.lower = (GBT_NUMKEY *) &kkk->lower;
169+
key.upper = (GBT_NUMKEY *) &kkk->upper;
170+
171+
PG_RETURN_FLOAT8(
172+
gbt_num_distance(&key, (void *) &query, GIST_LEAF(entry), &tinfo)
173+
);
174+
}
175+
176+
127177
Datum
128178
gbt_float8_union(PG_FUNCTION_ARGS)
129179
{

0 commit comments

Comments
 (0)