Line data    Source code

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