| 1 | /*------------------------------------------------------------------------- |
| 2 | * |
| 3 | * spgkdtreeproc.c |
| 4 | * implementation of k-d tree over points for SP-GiST |
| 5 | * |
| 6 | * |
| 7 | * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group |
| 8 | * Portions Copyright (c) 1994, Regents of the University of California |
| 9 | * |
| 10 | * IDENTIFICATION |
| 11 | * src/backend/access/spgist/spgkdtreeproc.c |
| 12 | * |
| 13 | *------------------------------------------------------------------------- |
| 14 | */ |
| 15 | |
| 16 | #include "postgres.h" |
| 17 | |
| 18 | #include "access/spgist.h" |
| 19 | #include "access/spgist_private.h" |
| 20 | #include "access/stratnum.h" |
| 21 | #include "catalog/pg_type.h" |
| 22 | #include "utils/builtins.h" |
| 23 | #include "utils/float.h" |
| 24 | #include "utils/geo_decls.h" |
| 25 | |
| 26 | |
| 27 | Datum |
| 28 | spg_kd_config(PG_FUNCTION_ARGS) |
| 29 | { |
| 30 | /* spgConfigIn *cfgin = (spgConfigIn *) PG_GETARG_POINTER(0); */ |
| 31 | spgConfigOut *cfg = (spgConfigOut *) PG_GETARG_POINTER(1); |
| 32 | |
| 33 | cfg->prefixType = FLOAT8OID; |
| 34 | cfg->labelType = VOIDOID; /* we don't need node labels */ |
| 35 | cfg->canReturnData = true; |
| 36 | cfg->longValuesOK = false; |
| 37 | PG_RETURN_VOID(); |
| 38 | } |
| 39 | |
| 40 | static int |
| 41 | getSide(double coord, bool isX, Point *tst) |
| 42 | { |
| 43 | double tstcoord = (isX) ? tst->x : tst->y; |
| 44 | |
| 45 | if (coord == tstcoord) |
| 46 | return 0; |
| 47 | else if (coord > tstcoord) |
| 48 | return 1; |
| 49 | else |
| 50 | return -1; |
| 51 | } |
| 52 | |
| 53 | Datum |
| 54 | spg_kd_choose(PG_FUNCTION_ARGS) |
| 55 | { |
| 56 | spgChooseIn *in = (spgChooseIn *) PG_GETARG_POINTER(0); |
| 57 | spgChooseOut *out = (spgChooseOut *) PG_GETARG_POINTER(1); |
| 58 | Point *inPoint = DatumGetPointP(in->datum); |
| 59 | double coord; |
| 60 | |
| 61 | if (in->allTheSame) |
| 62 | elog(ERROR, "allTheSame should not occur for k-d trees" ); |
| 63 | |
| 64 | Assert(in->hasPrefix); |
| 65 | coord = DatumGetFloat8(in->prefixDatum); |
| 66 | |
| 67 | Assert(in->nNodes == 2); |
| 68 | |
| 69 | out->resultType = spgMatchNode; |
| 70 | out->result.matchNode.nodeN = |
| 71 | (getSide(coord, in->level % 2, inPoint) > 0) ? 0 : 1; |
| 72 | out->result.matchNode.levelAdd = 1; |
| 73 | out->result.matchNode.restDatum = PointPGetDatum(inPoint); |
| 74 | |
| 75 | PG_RETURN_VOID(); |
| 76 | } |
| 77 | |
| 78 | typedef struct SortedPoint |
| 79 | { |
| 80 | Point *p; |
| 81 | int i; |
| 82 | } SortedPoint; |
| 83 | |
| 84 | static int |
| 85 | x_cmp(const void *a, const void *b) |
| 86 | { |
| 87 | SortedPoint *pa = (SortedPoint *) a; |
| 88 | SortedPoint *pb = (SortedPoint *) b; |
| 89 | |
| 90 | if (pa->p->x == pb->p->x) |
| 91 | return 0; |
| 92 | return (pa->p->x > pb->p->x) ? 1 : -1; |
| 93 | } |
| 94 | |
| 95 | static int |
| 96 | y_cmp(const void *a, const void *b) |
| 97 | { |
| 98 | SortedPoint *pa = (SortedPoint *) a; |
| 99 | SortedPoint *pb = (SortedPoint *) b; |
| 100 | |
| 101 | if (pa->p->y == pb->p->y) |
| 102 | return 0; |
| 103 | return (pa->p->y > pb->p->y) ? 1 : -1; |
| 104 | } |
| 105 | |
| 106 | |
| 107 | Datum |
| 108 | spg_kd_picksplit(PG_FUNCTION_ARGS) |
| 109 | { |
| 110 | spgPickSplitIn *in = (spgPickSplitIn *) PG_GETARG_POINTER(0); |
| 111 | spgPickSplitOut *out = (spgPickSplitOut *) PG_GETARG_POINTER(1); |
| 112 | int i; |
| 113 | int middle; |
| 114 | SortedPoint *sorted; |
| 115 | double coord; |
| 116 | |
| 117 | sorted = palloc(sizeof(*sorted) * in->nTuples); |
| 118 | for (i = 0; i < in->nTuples; i++) |
| 119 | { |
| 120 | sorted[i].p = DatumGetPointP(in->datums[i]); |
| 121 | sorted[i].i = i; |
| 122 | } |
| 123 | |
| 124 | qsort(sorted, in->nTuples, sizeof(*sorted), |
| 125 | (in->level % 2) ? x_cmp : y_cmp); |
| 126 | middle = in->nTuples >> 1; |
| 127 | coord = (in->level % 2) ? sorted[middle].p->x : sorted[middle].p->y; |
| 128 | |
| 129 | out->hasPrefix = true; |
| 130 | out->prefixDatum = Float8GetDatum(coord); |
| 131 | |
| 132 | out->nNodes = 2; |
| 133 | out->nodeLabels = NULL; /* we don't need node labels */ |
| 134 | |
| 135 | out->mapTuplesToNodes = palloc(sizeof(int) * in->nTuples); |
| 136 | out->leafTupleDatums = palloc(sizeof(Datum) * in->nTuples); |
| 137 | |
| 138 | /* |
| 139 | * Note: points that have coordinates exactly equal to coord may get |
| 140 | * classified into either node, depending on where they happen to fall in |
| 141 | * the sorted list. This is okay as long as the inner_consistent function |
| 142 | * descends into both sides for such cases. This is better than the |
| 143 | * alternative of trying to have an exact boundary, because it keeps the |
| 144 | * tree balanced even when we have many instances of the same point value. |
| 145 | * So we should never trigger the allTheSame logic. |
| 146 | */ |
| 147 | for (i = 0; i < in->nTuples; i++) |
| 148 | { |
| 149 | Point *p = sorted[i].p; |
| 150 | int n = sorted[i].i; |
| 151 | |
| 152 | out->mapTuplesToNodes[n] = (i < middle) ? 0 : 1; |
| 153 | out->leafTupleDatums[n] = PointPGetDatum(p); |
| 154 | } |
| 155 | |
| 156 | PG_RETURN_VOID(); |
| 157 | } |
| 158 | |
| 159 | Datum |
| 160 | spg_kd_inner_consistent(PG_FUNCTION_ARGS) |
| 161 | { |
| 162 | spgInnerConsistentIn *in = (spgInnerConsistentIn *) PG_GETARG_POINTER(0); |
| 163 | spgInnerConsistentOut *out = (spgInnerConsistentOut *) PG_GETARG_POINTER(1); |
| 164 | double coord; |
| 165 | int which; |
| 166 | int i; |
| 167 | BOX bboxes[2]; |
| 168 | |
| 169 | Assert(in->hasPrefix); |
| 170 | coord = DatumGetFloat8(in->prefixDatum); |
| 171 | |
| 172 | if (in->allTheSame) |
| 173 | elog(ERROR, "allTheSame should not occur for k-d trees" ); |
| 174 | |
| 175 | Assert(in->nNodes == 2); |
| 176 | |
| 177 | /* "which" is a bitmask of children that satisfy all constraints */ |
| 178 | which = (1 << 1) | (1 << 2); |
| 179 | |
| 180 | for (i = 0; i < in->nkeys; i++) |
| 181 | { |
| 182 | Point *query = DatumGetPointP(in->scankeys[i].sk_argument); |
| 183 | BOX *boxQuery; |
| 184 | |
| 185 | switch (in->scankeys[i].sk_strategy) |
| 186 | { |
| 187 | case RTLeftStrategyNumber: |
| 188 | if ((in->level % 2) != 0 && FPlt(query->x, coord)) |
| 189 | which &= (1 << 1); |
| 190 | break; |
| 191 | case RTRightStrategyNumber: |
| 192 | if ((in->level % 2) != 0 && FPgt(query->x, coord)) |
| 193 | which &= (1 << 2); |
| 194 | break; |
| 195 | case RTSameStrategyNumber: |
| 196 | if ((in->level % 2) != 0) |
| 197 | { |
| 198 | if (FPlt(query->x, coord)) |
| 199 | which &= (1 << 1); |
| 200 | else if (FPgt(query->x, coord)) |
| 201 | which &= (1 << 2); |
| 202 | } |
| 203 | else |
| 204 | { |
| 205 | if (FPlt(query->y, coord)) |
| 206 | which &= (1 << 1); |
| 207 | else if (FPgt(query->y, coord)) |
| 208 | which &= (1 << 2); |
| 209 | } |
| 210 | break; |
| 211 | case RTBelowStrategyNumber: |
| 212 | if ((in->level % 2) == 0 && FPlt(query->y, coord)) |
| 213 | which &= (1 << 1); |
| 214 | break; |
| 215 | case RTAboveStrategyNumber: |
| 216 | if ((in->level % 2) == 0 && FPgt(query->y, coord)) |
| 217 | which &= (1 << 2); |
| 218 | break; |
| 219 | case RTContainedByStrategyNumber: |
| 220 | |
| 221 | /* |
| 222 | * For this operator, the query is a box not a point. We |
| 223 | * cheat to the extent of assuming that DatumGetPointP won't |
| 224 | * do anything that would be bad for a pointer-to-box. |
| 225 | */ |
| 226 | boxQuery = DatumGetBoxP(in->scankeys[i].sk_argument); |
| 227 | |
| 228 | if ((in->level % 2) != 0) |
| 229 | { |
| 230 | if (FPlt(boxQuery->high.x, coord)) |
| 231 | which &= (1 << 1); |
| 232 | else if (FPgt(boxQuery->low.x, coord)) |
| 233 | which &= (1 << 2); |
| 234 | } |
| 235 | else |
| 236 | { |
| 237 | if (FPlt(boxQuery->high.y, coord)) |
| 238 | which &= (1 << 1); |
| 239 | else if (FPgt(boxQuery->low.y, coord)) |
| 240 | which &= (1 << 2); |
| 241 | } |
| 242 | break; |
| 243 | default: |
| 244 | elog(ERROR, "unrecognized strategy number: %d" , |
| 245 | in->scankeys[i].sk_strategy); |
| 246 | break; |
| 247 | } |
| 248 | |
| 249 | if (which == 0) |
| 250 | break; /* no need to consider remaining conditions */ |
| 251 | } |
| 252 | |
| 253 | /* We must descend into the children identified by which */ |
| 254 | out->nNodes = 0; |
| 255 | |
| 256 | /* Fast-path for no matching children */ |
| 257 | if (!which) |
| 258 | PG_RETURN_VOID(); |
| 259 | |
| 260 | out->nodeNumbers = (int *) palloc(sizeof(int) * 2); |
| 261 | |
| 262 | /* |
| 263 | * When ordering scan keys are specified, we've to calculate distance for |
| 264 | * them. In order to do that, we need calculate bounding boxes for both |
| 265 | * children nodes. Calculation of those bounding boxes on non-zero level |
| 266 | * require knowledge of bounding box of upper node. So, we save bounding |
| 267 | * boxes to traversalValues. |
| 268 | */ |
| 269 | if (in->norderbys > 0) |
| 270 | { |
| 271 | BOX infArea; |
| 272 | BOX *area; |
| 273 | |
| 274 | out->distances = (double **) palloc(sizeof(double *) * in->nNodes); |
| 275 | out->traversalValues = (void **) palloc(sizeof(void *) * in->nNodes); |
| 276 | |
| 277 | if (in->level == 0) |
| 278 | { |
| 279 | float8 inf = get_float8_infinity(); |
| 280 | |
| 281 | infArea.high.x = inf; |
| 282 | infArea.high.y = inf; |
| 283 | infArea.low.x = -inf; |
| 284 | infArea.low.y = -inf; |
| 285 | area = &infArea; |
| 286 | } |
| 287 | else |
| 288 | { |
| 289 | area = (BOX *) in->traversalValue; |
| 290 | Assert(area); |
| 291 | } |
| 292 | |
| 293 | bboxes[0].low = area->low; |
| 294 | bboxes[1].high = area->high; |
| 295 | |
| 296 | if (in->level % 2) |
| 297 | { |
| 298 | /* split box by x */ |
| 299 | bboxes[0].high.x = bboxes[1].low.x = coord; |
| 300 | bboxes[0].high.y = area->high.y; |
| 301 | bboxes[1].low.y = area->low.y; |
| 302 | } |
| 303 | else |
| 304 | { |
| 305 | /* split box by y */ |
| 306 | bboxes[0].high.y = bboxes[1].low.y = coord; |
| 307 | bboxes[0].high.x = area->high.x; |
| 308 | bboxes[1].low.x = area->low.x; |
| 309 | } |
| 310 | } |
| 311 | |
| 312 | for (i = 1; i <= 2; i++) |
| 313 | { |
| 314 | if (which & (1 << i)) |
| 315 | { |
| 316 | out->nodeNumbers[out->nNodes] = i - 1; |
| 317 | |
| 318 | if (in->norderbys > 0) |
| 319 | { |
| 320 | MemoryContext oldCtx = MemoryContextSwitchTo(in->traversalMemoryContext); |
| 321 | BOX *box = box_copy(&bboxes[i - 1]); |
| 322 | |
| 323 | MemoryContextSwitchTo(oldCtx); |
| 324 | |
| 325 | out->traversalValues[out->nNodes] = box; |
| 326 | |
| 327 | out->distances[out->nNodes] = spg_key_orderbys_distances(BoxPGetDatum(box), false, |
| 328 | in->orderbys, in->norderbys); |
| 329 | } |
| 330 | |
| 331 | out->nNodes++; |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | /* Set up level increments, too */ |
| 336 | out->levelAdds = (int *) palloc(sizeof(int) * 2); |
| 337 | out->levelAdds[0] = 1; |
| 338 | out->levelAdds[1] = 1; |
| 339 | |
| 340 | PG_RETURN_VOID(); |
| 341 | } |
| 342 | |
| 343 | /* |
| 344 | * spg_kd_leaf_consistent() is the same as spg_quad_leaf_consistent(), |
| 345 | * since we support the same operators and the same leaf data type. |
| 346 | * So we just borrow that function. |
| 347 | */ |
| 348 | |