├── .github
    └── workflows
    │   ├── release.yml
    │   └── tests.yml
├── .gitignore
├── Dockerfile
├── LICENSE
├── README.md
├── ROADMAP.md
├── bench_test.go
├── bkdtree.go
├── bkdtree_compact.go
├── bkdtree_erase.go
├── bkdtree_insert.go
├── bkdtree_intersect.go
├── bkdtree_test.go
├── go.mod
├── go.sum
├── kdtree.go
├── kdtree_test.go
├── point.go
├── point_test.go
└── utils.go


/.github/workflows/release.yml:
--------------------------------------------------------------------------------
  1 | name: release
  2 | 
  3 | on:
  4 |   schedule:
  5 |     - cron: '0 * * * *'  # schedule nightly build daily at midnight UTC
  6 |   # The "create tags" trigger is specifically focused on the creation of new tags, while the "push tags" trigger is activated when tags are pushed, including both new tag creations and updates to existing tags.
  7 |   create:
  8 |     tags:
  9 |       - "v*.*.*"                  # normal release
 10 |       - "nightly"                 # the only one mutable tag
 11 | 
 12 | jobs:
 13 |   release:
 14 |     runs-on: ubuntu-latest
 15 |     steps:
 16 |     - name: Ensure workspace ownership
 17 |       run: echo "chown -R $USER $GITHUB_WORKSPACE" && sudo chown -R $USER $GITHUB_WORKSPACE
 18 | 
 19 |     # https://github.com/actions/checkout/blob/v3/README.md
 20 |     - name: Check out code
 21 |       uses: actions/checkout@v3
 22 |       with:
 23 |         ssh-key: ${{ secrets.MY_DEPLOY_KEY }}
 24 | 
 25 |     - name: Prepare release body
 26 |       run: |
 27 |         if [[ $GITHUB_EVENT_NAME == 'create' ]]; then
 28 |           RELEASE_TAG=${GITHUB_REF#refs/tags/}
 29 |           if [[ $RELEASE_TAG == 'nightly' ]]; then
 30 |             PRERELEASE=true
 31 |           else
 32 |             PRERELEASE=false
 33 |           fi
 34 |           echo "Workflow triggered by create tag: $RELEASE_TAG"
 35 |         else
 36 |           RELEASE_TAG=nightly
 37 |           PRERELEASE=true
 38 |           echo "Workflow triggered by schedule"
 39 |         fi
 40 |         echo "RELEASE_TAG=$RELEASE_TAG" >> $GITHUB_ENV
 41 |         echo "PRERELEASE=$PRERELEASE" >> $GITHUB_ENV
 42 |         RELEASE_DATETIME=$(date --rfc-3339=seconds)
 43 |         cat <<EOF > release_template.md
 44 |         Release $RELEASE_TAG created from $GITHUB_SHA at $RELEASE_DATETIME
 45 |         EOF
 46 |         envsubst < release_template.md > release_body.md
 47 | 
 48 |     - name: Move the existing mutable tag
 49 |       # https://github.com/softprops/action-gh-release/issues/171
 50 |       run: |
 51 |         if [[ $GITHUB_EVENT_NAME == 'schedule' ]]; then
 52 |           # Determine if a given tag exists and matches a specific Git commit.
 53 |           # actions/checkout@v3 fetch-tags doesn't work when triggered by schedule
 54 |           git fetch --tags
 55 |           if [ "$(git rev-parse -q --verify "refs/tags/$RELEASE_TAG")" = "$GITHUB_SHA" ]; then
 56 |             echo "mutalbe tag $RELEASE_TAG exists and matchs $GITHUB_SHA"
 57 |           else
 58 |             git tag -f $RELEASE_TAG $GITHUB_SHA
 59 |             git push -f origin $RELEASE_TAG:refs/tags/$RELEASE_TAG
 60 |             echo "created/moved mutalbe tag $RELEASE_TAG to $GITHUB_SHA"
 61 |           fi
 62 |         fi
 63 | 
 64 |     # https://github.com/actions/setup-go
 65 |     - name: Setup go
 66 |       uses: actions/setup-go@v4
 67 |       with:
 68 |         go-version: '^1.21.5' # The Go version to download (if necessary) and use.
 69 | 
 70 |     - name: Build test binary
 71 |       run: go test -c .
 72 | 
 73 |     - name: Create or overwrite a release
 74 |       # https://github.com/actions/upload-release-asset has been replaced by https://github.com/softprops/action-gh-release
 75 |       uses: softprops/action-gh-release@v1
 76 |       with:
 77 |         token: ${{ secrets.MY_GITHUB_TOKEN }}  # Use the secret as an environment variable
 78 |         prerelease: ${{ env.PRERELEASE }}
 79 |         draft: false
 80 |         tag_name: ${{ env.RELEASE_TAG }}
 81 |         # The body field does not support environment variable substitution directly.
 82 |         body_path: release_body.md
 83 |         files: |
 84 |           bkdtree.test
 85 | 
 86 |     - name: Set up QEMU
 87 |       uses: docker/setup-qemu-action@v3
 88 | 
 89 |     - name: Set up Docker Buildx
 90 |       uses: docker/setup-buildx-action@v3
 91 | 
 92 |     # https://github.com/marketplace/actions/docker-login
 93 |     - name: Login to Docker Hub
 94 |       uses: docker/login-action@v3
 95 |       with:
 96 |         username: yuzhichang
 97 |         password: ${{ secrets.DOCKERHUB_TOKEN }}
 98 | 
 99 |     # https://github.com/marketplace/actions/build-and-push-docker-images
100 |     - name: Build and push
101 |       uses: docker/build-push-action@v5
102 |       with:
103 |         context: .
104 |         tags: yuzhichang/bkdtree:${{ env.RELEASE_TAG }}
105 |         file: Dockerfile
106 |         push: true
107 | 


--------------------------------------------------------------------------------
/.github/workflows/tests.yml:
--------------------------------------------------------------------------------
 1 | name: tests
 2 | 
 3 | on:
 4 |   push:
 5 |     branches: [ master ]
 6 |   pull_request:
 7 |     types: [ opened, synchronize, reopened, edited ]
 8 | 
 9 | jobs:
10 |   tests:
11 |     name: tests
12 |     runs-on: ubuntu-latest
13 |     steps:
14 |     - name: Ensure workspace ownership
15 |       run: echo "chown -R $USER $GITHUB_WORKSPACE" && sudo chown -R $USER $GITHUB_WORKSPACE
16 | 
17 |     - name: Check out code
18 |       uses: actions/checkout@v3
19 |       with:
20 |         ssh-key: ${{ secrets.MY_DEPLOY_KEY }}
21 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects)
 2 | *.o
 3 | *.a
 4 | *.so
 5 | 
 6 | # Folders
 7 | _obj
 8 | _test
 9 | .vscode
10 | 
11 | # Architecture specific extensions/prefixes
12 | *.[568vq]
13 | [568vq].out
14 | 
15 | *.cgo1.go
16 | *.cgo2.c
17 | _cgo_defun.c
18 | _cgo_gotypes.go
19 | _cgo_export.*
20 | 
21 | _testmain.go
22 | 
23 | *.exe
24 | *.test
25 | *.prof
26 | *.svg
27 | 


--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
1 | FROM debian:stable-slim
2 | 
3 | COPY bkdtree.test /usr/bin
4 | 
5 | ENTRYPOINT [ "/usr/bin/bkdtree.test" ]
6 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "{}"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright {yyyy} {name of copyright owner}
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # BKD tree
 2 | BKD tree in go.
 3 | Refers to
 4 | - O. Procopiuc, P.K. Agarwal, L. Arge, J.S. Vitter. Bkd-tree: A Dynamic Scalable kd-Tree. Proceedings of SSTD 2003.
 5 | - BKD C++ impl, https://github.com/thomasmoelhave/tpie/tree/master/apps/bkdtree.
 6 | - BKD Java impl, https://www.elastic.co/blog/this-week-in-elasticsearch-and-apache-lucene-2017-02-13.
 7 | - KD Golang impl, https://github.com/hongshibao/go-kdtree.
 8 | 
 9 | Documentation
10 | 
11 |     https://godoc.org/github.com/deepfabric/bkdtree
12 | 


--------------------------------------------------------------------------------
/ROADMAP.md:
--------------------------------------------------------------------------------
 1 | # Roadmap
 2 | 
 3 | This document describes the roadmap for BKD tree development.
 4 | 
 5 | #### KD tree (memory)
 6 | - [D] build - Kdtree in mem
 7 | - [D] intersect - Kdtree in mem
 8 | - [D] insert - Kdtree in mem
 9 | - [D] erase - Kdtree in mem
10 | 
11 | ##### BKD tree (memory + file)
12 | - [D] build
13 | - [D] insert 
14 | - [D] erase
15 | - [D] intersect
16 | - [D] compatible file format to allow multiple versions
17 | - [D] performance optimization - mmap, point encoding/decoding etc.
18 | - [D] disaster recovery - open, close
19 | - [D] concurrent access - singel writer, multiple reader
20 | - [D] concurrent access - background compact
21 | - [ ] make mmap optional
22 | - [ ] custom error type for invalid arguments, not-permitted operations
23 | - [ ] performance optimization - (*PointArrayExt).GetPoint
24 | 


--------------------------------------------------------------------------------
/bench_test.go:
--------------------------------------------------------------------------------
 1 | package bkdtree
 2 | 
 3 | import (
 4 | 	"testing"
 5 | )
 6 | 
 7 | func BenchmarkBkdInsert(b *testing.B) {
 8 | 	t0mCap := 1000
 9 | 	leafCap := 50
10 | 	intraCap := 4
11 | 	numDims := 2
12 | 	bytesPerDim := 4
13 | 	dir := "/tmp"
14 | 	prefix := "bkd"
15 | 	bkd, err := NewBkdTree(t0mCap, leafCap, intraCap, numDims, bytesPerDim, dir, prefix)
16 | 	if err != nil {
17 | 		b.Fatalf("%+v", err)
18 | 	}
19 | 	//fmt.Printf("created BkdTree %v\n", bkd)
20 | 
21 | 	b.ResetTimer()
22 | 	for i := 0; i < b.N; i++ {
23 | 		err := bkd.Insert(Point{[]uint64{uint64(i), uint64(i)}, uint64(i)})
24 | 		if err != nil {
25 | 			b.Fatalf("bkd.Insert failed, i=%v, err: %+v", i, err)
26 | 		}
27 | 	}
28 | 	return
29 | }
30 | 
31 | func BenchmarkBkdErase(b *testing.B) {
32 | 	var maxVal uint64 = 1000
33 | 	bkd, points, err := prepareBkdTree(maxVal)
34 | 	if err != nil {
35 | 		b.Fatalf("%+v", err)
36 | 	}
37 | 
38 | 	b.ResetTimer()
39 | 	for i := 0; i < b.N; i++ {
40 | 		_, err = bkd.Erase(points[i])
41 | 		if err != nil {
42 | 			b.Fatalf("%+v", err)
43 | 		}
44 | 	}
45 | }
46 | 
47 | func BenchmarkBkdIntersect(b *testing.B) {
48 | 	var maxVal uint64 = 1000
49 | 	bkd, points, err := prepareBkdTree(maxVal)
50 | 	if err != nil {
51 | 		b.Fatalf("%+v", err)
52 | 	}
53 | 	var lowPoint, highPoint Point
54 | 	var visitor *IntersectCollector
55 | 	lowPoint = points[7]
56 | 	highPoint = lowPoint
57 | 	visitor = &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
58 | 
59 | 	b.ResetTimer()
60 | 	for i := 0; i < b.N; i++ {
61 | 		err = bkd.Intersect(visitor)
62 | 		if err != nil {
63 | 			b.Fatalf("%+v", err)
64 | 		} else if len(visitor.Points) <= 0 {
65 | 			b.Errorf("found 0 matchs, however some expected")
66 | 		}
67 | 	}
68 | }
69 | 


--------------------------------------------------------------------------------
/bkdtree.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 	"encoding/binary"
  6 | 	"fmt"
  7 | 	"io"
  8 | 	"os"
  9 | 	"path/filepath"
 10 | 	"sort"
 11 | 	"strconv"
 12 | 	"sync"
 13 | 
 14 | 	"github.com/pkg/errors"
 15 | )
 16 | 
 17 | type KdTreeExtNodeInfo struct {
 18 | 	Offset    uint64 //offset in file. It's a leaf if less than pointsOffEnd, otherwise an intra node.
 19 | 	NumPoints uint64 //number of points of subtree rooted at this node
 20 | }
 21 | 
 22 | const KdTreeExtNodeInfoSize int64 = 8 + 8
 23 | 
 24 | //KdTreeExtIntraNode is struct of intra node.
 25 | /**
 26 |  * invariants:
 27 |  * 1. NumStrips == 1 + len(SplitValues) == len(Children).
 28 |  * 2. values in SplitValues are in non-decreasing order.
 29 |  * 3. offset in Children are in increasing order.
 30 |  */
 31 | type KdTreeExtIntraNode struct {
 32 | 	SplitDim    uint32
 33 | 	NumStrips   uint32
 34 | 	SplitValues []uint64
 35 | 	Children    []KdTreeExtNodeInfo
 36 | }
 37 | 
 38 | // KdTreeExtMeta is persisted at the end of file.
 39 | /**
 40 |  * Some fields are redundant in order to make the file be self-descriptive.
 41 |  * Attention:
 42 |  * 1. Keep all fields exported to allow one invoke of binary.Read() to parse the whole struct.
 43 |  * 2. Keep KdTreeExtMeta be 4 bytes aligned.
 44 |  * 3. Keep formatVer one byte, and be the last member.
 45 |  * 4. Keep KdMetaSize be sizeof(KdTreeExtMeta);
 46 |  */
 47 | type KdTreeExtMeta struct {
 48 | 	PointsOffEnd uint64 //the offset end of points
 49 | 	RootOff      uint64 //the offset of root KdTreeExtIntraNode
 50 | 	NumPoints    uint64 //the current number of points. Deleting points could trigger rebuilding the tree.
 51 | 	LeafCap      uint16
 52 | 	IntraCap     uint16
 53 | 	NumDims      uint8
 54 | 	BytesPerDim  uint8
 55 | 	PointSize    uint8
 56 | 	FormatVer    uint8 //the file format version. shall be the last byte of the file.
 57 | }
 58 | 
 59 | //KdTreeExtMetaSize is sizeof(KdTreeExtMeta)
 60 | const KdTreeExtMetaSize int = 8*3 + 4 + 4
 61 | 
 62 | type BkdSubTree struct {
 63 | 	meta KdTreeExtMeta
 64 | 	f    *os.File
 65 | 	data []byte //file content via mmap
 66 | }
 67 | 
 68 | //BkdTree is a BKD tree
 69 | type BkdTree struct {
 70 | 	t0mCap      int // M in the paper, the capacity of in-memory buffer
 71 | 	leafCap     int // limit of points a leaf node can hold
 72 | 	intraCap    int // limit of children of a intra node can hold
 73 | 	NumDims     int // number of point dimensions
 74 | 	BytesPerDim int // number of bytes of each encoded dimension
 75 | 	pointSize   int
 76 | 	dir         string //directory of files which hold the persisted kdtrees
 77 | 	prefix      string //prefix of file names
 78 | 	NumPoints   int
 79 | 	t0m         BkdSubTree // T0M in the paper, in-memory buffer.
 80 | 	trees       []BkdSubTree
 81 | 	rwlock      sync.RWMutex //reader: Intersect, GetCap. writers: Insert, Erase, Open, Close, Destroy, Compact.
 82 | 	open        bool         //closed: allow Open, Close; open: allow all operations except Open.
 83 | }
 84 | 
 85 | func (n *KdTreeExtIntraNode) Read(r io.Reader) (err error) {
 86 | 	//According to https://golang.org/pkg/encoding/binary/#Read,
 87 | 	//"Data must be a pointer to a fixed-size value or a slice of fixed-size values."
 88 | 	//Slice shall be adjusted to the expected length before calling binary.Read().
 89 | 	err = binary.Read(r, binary.BigEndian, &n.SplitDim)
 90 | 	if err != nil {
 91 | 		err = errors.Wrap(err, "")
 92 | 		return
 93 | 	}
 94 | 	err = binary.Read(r, binary.BigEndian, &n.NumStrips)
 95 | 	if err != nil {
 96 | 		err = errors.Wrap(err, "")
 97 | 		return
 98 | 	}
 99 | 	n.SplitValues = make([]uint64, n.NumStrips-1)
100 | 	err = binary.Read(r, binary.BigEndian, &n.SplitValues)
101 | 	if err != nil {
102 | 		err = errors.Wrap(err, "")
103 | 		return
104 | 	}
105 | 	n.Children = make([]KdTreeExtNodeInfo, n.NumStrips)
106 | 	err = binary.Read(r, binary.BigEndian, &n.Children) //TODO: why n.children doesn't work?
107 | 	if err != nil {
108 | 		err = errors.Wrap(err, "")
109 | 		return
110 | 	}
111 | 	return
112 | }
113 | 
114 | func (n *KdTreeExtIntraNode) Write(w io.Writer) (err error) {
115 | 	//According to https://golang.org/pkg/encoding/binary/#Write,
116 | 	//"Data must be a fixed-size value or a slice of fixed-size values, or a pointer to such data."
117 | 	//Structs with slice members can not be used with binary.Write. Slice members shall be write explictly.
118 | 	err = binary.Write(w, binary.BigEndian, &n.SplitDim)
119 | 	if err != nil {
120 | 		err = errors.Wrap(err, "")
121 | 		return
122 | 	}
123 | 	err = binary.Write(w, binary.BigEndian, &n.NumStrips)
124 | 	if err != nil {
125 | 		err = errors.Wrap(err, "")
126 | 		return
127 | 	}
128 | 	err = binary.Write(w, binary.BigEndian, &n.SplitValues)
129 | 	if err != nil {
130 | 		err = errors.Wrap(err, "")
131 | 		return
132 | 	}
133 | 	err = binary.Write(w, binary.BigEndian, &n.Children)
134 | 	if err != nil {
135 | 		err = errors.Wrap(err, "")
136 | 		return
137 | 	}
138 | 	return
139 | }
140 | 
141 | //NewBkdTree creates a BKDTree. This is used for construct a BkdTree from scratch. Existing files, if any, will be removed.
142 | func NewBkdTree(t0mCap, leafCap, intraCap, numDims, bytesPerDim int, dir, prefix string) (bkd *BkdTree, err error) {
143 | 	if t0mCap <= 0 || leafCap <= 0 || leafCap >= int(^uint16(0)) || intraCap <= 2 ||
144 | 		numDims <= 0 || (bytesPerDim != 1 && bytesPerDim != 2 && bytesPerDim != 4 && bytesPerDim != 8) {
145 | 		err = errors.Errorf("invalid parameter")
146 | 		return
147 | 	}
148 | 	bkd = &BkdTree{
149 | 		t0mCap:      t0mCap,
150 | 		leafCap:     leafCap,
151 | 		intraCap:    intraCap,
152 | 		NumDims:     numDims,
153 | 		BytesPerDim: bytesPerDim,
154 | 		pointSize:   numDims*bytesPerDim + 8,
155 | 		dir:         dir,
156 | 		prefix:      prefix,
157 | 		//t0m is initialized later
158 | 		trees: make([]BkdSubTree, 0),
159 | 	}
160 | 	if err = bkd.initT0M(); err != nil {
161 | 		return
162 | 	}
163 | 	if err = rmTreeList(dir, prefix); err != nil {
164 | 		return
165 | 	}
166 | 	bkd.open = true
167 | 	return
168 | }
169 | 
170 | //Destroy close and remove all files
171 | func (bkd *BkdTree) Destroy() (err error) {
172 | 	bkd.rwlock.Lock()
173 | 	defer bkd.rwlock.Unlock()
174 | 	if err = bkd.close(); err != nil {
175 | 		return
176 | 	}
177 | 	if err = rmTreeList(bkd.dir, bkd.prefix); err != nil {
178 | 		return
179 | 	}
180 | 	if err = os.Remove(bkd.T0mPath()); err != nil {
181 | 		return
182 | 	}
183 | 	return
184 | }
185 | 
186 | //Close close unmap and all files
187 | func (bkd *BkdTree) Close() (err error) {
188 | 	bkd.rwlock.Lock()
189 | 	defer bkd.rwlock.Unlock()
190 | 	err = bkd.close()
191 | 	return
192 | }
193 | 
194 | //close close all files without holding the write lock
195 | func (bkd *BkdTree) close() (err error) {
196 | 	if !bkd.open {
197 | 		return
198 | 	}
199 | 	bkd.open = false
200 | 
201 | 	if err = FileMunmap(bkd.t0m.data); err != nil {
202 | 		return
203 | 	} else if err = bkd.t0m.f.Close(); err != nil {
204 | 		return
205 | 	}
206 | 	for i := 0; i < len(bkd.trees); i++ {
207 | 		if bkd.trees[i].meta.NumPoints == 0 {
208 | 			continue
209 | 		}
210 | 		if err = FileMunmap(bkd.trees[i].data); err != nil {
211 | 			return
212 | 		} else if err = bkd.trees[i].f.Close(); err != nil {
213 | 			return
214 | 		}
215 | 	}
216 | 	return
217 | }
218 | 
219 | //NewBkdTreeExt create a BKdTree based on exisiting files.
220 | func NewBkdTreeExt(dir, prefix string) (bkd *BkdTree, err error) {
221 | 	bkd = &BkdTree{
222 | 		dir:    dir,
223 | 		prefix: prefix,
224 | 	}
225 | 	err = bkd.Open()
226 | 	return
227 | }
228 | 
229 | //Open open existing files. Assumes dir and prefix are already poluplated.
230 | func (bkd *BkdTree) Open() (err error) {
231 | 	bkd.rwlock.Lock()
232 | 	defer bkd.rwlock.Unlock()
233 | 	if bkd.open {
234 | 		err = errors.Errorf("(*BkdTree).Open is not allowed at open state")
235 | 		return
236 | 	}
237 | 
238 | 	var nums []int
239 | 	if err = bkd.openT0M(); err != nil {
240 | 		return
241 | 	}
242 | 	bkd.NumPoints = int(bkd.t0m.meta.NumPoints)
243 | 	if nums, err = getTreeList(bkd.dir, bkd.prefix); err != nil {
244 | 		return
245 | 	}
246 | 	for _, num := range nums {
247 | 		if len(bkd.trees) <= num+1 {
248 | 			delta := num + 1 - len(bkd.trees)
249 | 			for i := 0; i < delta; i++ {
250 | 				kd := BkdSubTree{
251 | 					meta: KdTreeExtMeta{}, //NumPoints default value zero is good
252 | 				}
253 | 				bkd.trees = append(bkd.trees, kd)
254 | 			}
255 | 		}
256 | 		fp := bkd.TiPath(num)
257 | 		if err = bkd.trees[num].open(fp); err != nil {
258 | 			return
259 | 		}
260 | 		bkd.NumPoints += int(bkd.trees[num].meta.NumPoints)
261 | 	}
262 | 	bkd.open = true
263 | 	return
264 | }
265 | 
266 | //T0mPath returns T0M path
267 | func (bkd *BkdTree) T0mPath() string {
268 | 	fpT0M := filepath.Join(bkd.dir, fmt.Sprintf("%s_t0m", bkd.prefix))
269 | 	return fpT0M
270 | }
271 | 
272 | //TiPath returns Ti path
273 | func (bkd *BkdTree) TiPath(i int) string {
274 | 	fpTi := filepath.Join(bkd.dir, fmt.Sprintf("%s_t%d", bkd.prefix, i))
275 | 	return fpTi
276 | }
277 | 
278 | func (bkd *BkdTree) initT0M() (err error) {
279 | 	if err = os.MkdirAll(bkd.dir, 0700); err != nil {
280 | 		err = errors.Wrap(err, "")
281 | 		return
282 | 	}
283 | 	fT0M, err1 := os.OpenFile(bkd.T0mPath(), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
284 | 	if err1 != nil {
285 | 		err = errors.Wrap(err1, "")
286 | 		return
287 | 	}
288 | 	meta := KdTreeExtMeta{
289 | 		PointsOffEnd: uint64(bkd.pointSize * bkd.t0mCap),
290 | 		RootOff:      0, //not used in T0M
291 | 		NumPoints:    0,
292 | 		LeafCap:      uint16(bkd.leafCap),
293 | 		IntraCap:     uint16(bkd.intraCap),
294 | 		NumDims:      uint8(bkd.NumDims),
295 | 		BytesPerDim:  uint8(bkd.BytesPerDim),
296 | 		PointSize:    uint8(bkd.pointSize),
297 | 		FormatVer:    0,
298 | 	}
299 | 	buf := make([]byte, meta.PointsOffEnd)
300 | 	if _, err = fT0M.Write(buf); err != nil {
301 | 		err = errors.Wrap(err, "")
302 | 	}
303 | 	if err = binary.Write(fT0M, binary.BigEndian, &meta); err != nil {
304 | 		err = errors.Wrap(err, "")
305 | 	}
306 | 	data, err := FileMmap(fT0M)
307 | 	if err != nil {
308 | 		return
309 | 	}
310 | 	bkd.t0m = BkdSubTree{
311 | 		meta: meta,
312 | 		f:    fT0M,
313 | 		data: data,
314 | 	}
315 | 	return
316 | }
317 | 
318 | func (bkd *BkdTree) openT0M() (err error) {
319 | 	bkd.t0m = BkdSubTree{}
320 | 	if err = bkd.t0m.open(bkd.T0mPath()); err != nil {
321 | 		return
322 | 	}
323 | 	bkd.t0mCap = (len(bkd.t0m.data) - KdTreeExtMetaSize) / int(bkd.t0m.meta.PointSize)
324 | 	bkd.NumDims = int(bkd.t0m.meta.NumDims)
325 | 	bkd.BytesPerDim = int(bkd.t0m.meta.BytesPerDim)
326 | 	bkd.pointSize = int(bkd.t0m.meta.PointSize)
327 | 	bkd.leafCap = int(bkd.t0m.meta.LeafCap)
328 | 	bkd.intraCap = int(bkd.t0m.meta.IntraCap)
329 | 	return
330 | }
331 | 
332 | func (bst *BkdSubTree) open(fp string) (err error) {
333 | 	if bst.f, err = os.OpenFile(fp, os.O_RDWR, 0600); err != nil {
334 | 		err = errors.Wrap(err, "")
335 | 		return
336 | 	}
337 | 	if bst.data, err = FileMmap(bst.f); err != nil {
338 | 		return
339 | 	}
340 | 	br := bytes.NewReader(bst.data[len(bst.data)-KdTreeExtMetaSize:])
341 | 	if err = binary.Read(br, binary.BigEndian, &bst.meta); err != nil {
342 | 		err = errors.Wrap(err, "")
343 | 		return
344 | 	}
345 | 	return
346 | }
347 | 
348 | func getTreeList(dir, prefix string) (numList []int, err error) {
349 | 	var matches [][]string
350 | 	var num int
351 | 	patt := fmt.Sprintf("^%s_t(?P<num>[0-9]+)$", prefix)
352 | 	if matches, err = FilepathGlob(dir, patt); err != nil {
353 | 		return
354 | 	}
355 | 	for _, match := range matches {
356 | 		num, err = strconv.Atoi(match[1])
357 | 		if err != nil {
358 | 			err = errors.Wrap(err, "")
359 | 			return
360 | 		}
361 | 		numList = append(numList, num)
362 | 	}
363 | 	sort.Ints(numList)
364 | 	return
365 | }
366 | 
367 | func rmTreeList(dir, prefix string) (err error) {
368 | 	patt := fmt.Sprintf("^%s_t(?P<num>[0-9]+)$", prefix)
369 | 	err = FilepathGlobRm(dir, patt)
370 | 	return
371 | }
372 | 


--------------------------------------------------------------------------------
/bkdtree_compact.go:
--------------------------------------------------------------------------------
 1 | package bkdtree
 2 | 
 3 | //Compact comopact subtrees if necessary.
 4 | func (bkd *BkdTree) Compact() (err error) {
 5 | 	bkd.rwlock.RLock()
 6 | 	if !bkd.open {
 7 | 		bkd.rwlock.RUnlock()
 8 | 		return
 9 | 	}
10 | 	k := bkd.getMaxCompactPos()
11 | 	bkd.rwlock.RUnlock()
12 | 	if k < 0 {
13 | 		return
14 | 	}
15 | 
16 | 	bkd.rwlock.Lock()
17 | 	err = bkd.compactTo(k)
18 | 	bkd.rwlock.Unlock()
19 | 	return
20 | }
21 | 
22 | //caclulate the max compoint position. Returns -1 if not found.
23 | func (bkd *BkdTree) getMaxCompactPos() (pos int) {
24 | 	//find the largest index k in [0, len(trees)) at which trees[k] is empty, or its capacity is no less than the sum of size of t0m + trees[0:k+1]
25 | 	pos = -1
26 | 	sum := int(bkd.t0m.meta.NumPoints)
27 | 	for k := 0; k < len(bkd.trees); k++ {
28 | 		if bkd.trees[k].meta.NumPoints == 0 {
29 | 			pos = k
30 | 			continue
31 | 		}
32 | 		sum += int(bkd.trees[k].meta.NumPoints)
33 | 		capK := bkd.t0mCap << uint(k)
34 | 		if capK >= sum {
35 | 			pos = k
36 | 		}
37 | 	}
38 | 	return
39 | }
40 | 


--------------------------------------------------------------------------------
/bkdtree_erase.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 	"encoding/binary"
  6 | 
  7 | 	"github.com/pkg/errors"
  8 | )
  9 | 
 10 | //Erase erases given point.
 11 | func (bkd *BkdTree) Erase(point Point) (found bool, err error) {
 12 | 	bkd.rwlock.Lock()
 13 | 	defer bkd.rwlock.Unlock()
 14 | 	if !bkd.open {
 15 | 		err = errors.Errorf("(*BkdTree).Erase is not allowed at closed state")
 16 | 		return
 17 | 	}
 18 | 
 19 | 	//Query T0M with p; if found, delete it and return.
 20 | 	found = bkd.eraseT0M(point)
 21 | 	if found {
 22 | 		bkd.NumPoints--
 23 | 		return
 24 | 	}
 25 | 
 26 | 	//Query each non-empty tree in the forest with p; if found, delete it and return
 27 | 	for i := 0; i < len(bkd.trees); i++ {
 28 | 		found, err = bkd.eraseTi(point, i)
 29 | 		if err != nil {
 30 | 			return
 31 | 		} else if found {
 32 | 			bkd.NumPoints--
 33 | 			return
 34 | 		}
 35 | 	}
 36 | 	return
 37 | }
 38 | 
 39 | func (bkd *BkdTree) eraseT0M(point Point) (found bool) {
 40 | 	pae := PointArrayExt{
 41 | 		data:        bkd.t0m.data,
 42 | 		numPoints:   int(bkd.t0m.meta.NumPoints),
 43 | 		byDim:       0, //not used
 44 | 		bytesPerDim: bkd.BytesPerDim,
 45 | 		numDims:     bkd.NumDims,
 46 | 		pointSize:   bkd.pointSize,
 47 | 	}
 48 | 	found = pae.Erase(point)
 49 | 	if found {
 50 | 		bkd.t0m.meta.NumPoints--
 51 | 		writeMetaNumPoints(bkd.t0m.data, &bkd.t0m.meta)
 52 | 	}
 53 | 	return
 54 | }
 55 | 
 56 | func (bkd *BkdTree) eraseTi(point Point, idx int) (found bool, err error) {
 57 | 	if bkd.trees[idx].meta.NumPoints <= 0 {
 58 | 		return
 59 | 	}
 60 | 
 61 | 	//depth-first erasing from the root node
 62 | 	meta := &bkd.trees[idx].meta
 63 | 	found, err = bkd.eraseNode(point, bkd.trees[idx].data, meta, int(meta.RootOff))
 64 | 	if err != nil {
 65 | 		return
 66 | 	}
 67 | 	if found {
 68 | 		bkd.trees[idx].meta.NumPoints--
 69 | 		writeMetaNumPoints(bkd.trees[idx].data, &bkd.trees[idx].meta)
 70 | 		return
 71 | 	}
 72 | 	return
 73 | }
 74 | 
 75 | func (bkd *BkdTree) eraseNode(point Point, data []byte, meta *KdTreeExtMeta, nodeOffset int) (found bool, err error) {
 76 | 	var node KdTreeExtIntraNode
 77 | 	br := bytes.NewReader(data[nodeOffset:])
 78 | 	err = node.Read(br)
 79 | 	if err != nil {
 80 | 		return
 81 | 	}
 82 | 	for i, child := range node.Children {
 83 | 		if child.NumPoints <= 0 {
 84 | 			continue
 85 | 		}
 86 | 		if child.Offset < meta.PointsOffEnd {
 87 | 			//leaf node
 88 | 			pae := PointArrayExt{
 89 | 				data:        data[int(child.Offset):],
 90 | 				numPoints:   int(child.NumPoints),
 91 | 				byDim:       0, //not used
 92 | 				bytesPerDim: bkd.BytesPerDim,
 93 | 				numDims:     bkd.NumDims,
 94 | 				pointSize:   bkd.pointSize,
 95 | 			}
 96 | 			found = pae.Erase(point)
 97 | 		} else {
 98 | 			//intra node
 99 | 			found, err = bkd.eraseNode(point, data, meta, int(child.Offset))
100 | 		}
101 | 		if err != nil {
102 | 			return
103 | 		}
104 | 		if found {
105 | 			child.NumPoints--
106 | 			//Attention: offset calculation shall be synced with KdTreeExtIntraNode definion.
107 | 			off := nodeOffset + 8*int(node.NumStrips) + 16*i + 8
108 | 			binary.BigEndian.PutUint64(data[off:], child.NumPoints)
109 | 			break
110 | 		}
111 | 	}
112 | 	return
113 | }
114 | 


--------------------------------------------------------------------------------
/bkdtree_insert.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 	"encoding/binary"
  6 | 	"fmt"
  7 | 	"os"
  8 | 	"unsafe"
  9 | 
 10 | 	"github.com/pkg/errors"
 11 | )
 12 | 
 13 | //Insert inserts given point. Fail if the tree is full.
 14 | func (bkd *BkdTree) Insert(point Point) (err error) {
 15 | 	bkd.rwlock.Lock()
 16 | 	defer bkd.rwlock.Unlock()
 17 | 	if !bkd.open {
 18 | 		err = errors.Errorf("(*BkdTree).Inset is not allowed at closed state")
 19 | 		return
 20 | 	}
 21 | 
 22 | 	//insert into in-memory buffer t0m. If t0m is not full, return.
 23 | 	bkd.insertT0M(point)
 24 | 	bkd.NumPoints++
 25 | 	if int(bkd.t0m.meta.NumPoints) < bkd.t0mCap {
 26 | 		return
 27 | 	}
 28 | 	//find the smallest index k in [0, len(trees)) at which trees[k] is empty, or its capacity is no less than the sum of size of t0m + trees[0:k+1]
 29 | 	k := bkd.getMinCompactPos()
 30 | 	if k == len(bkd.trees) {
 31 | 		kd := BkdSubTree{
 32 | 			meta: KdTreeExtMeta{
 33 | 				PointsOffEnd: 0,
 34 | 				RootOff:      0,
 35 | 				NumPoints:    0,
 36 | 				LeafCap:      uint16(bkd.leafCap),
 37 | 				IntraCap:     uint16(bkd.intraCap),
 38 | 				NumDims:      uint8(bkd.NumDims),
 39 | 				BytesPerDim:  uint8(bkd.BytesPerDim),
 40 | 				PointSize:    uint8(bkd.pointSize),
 41 | 				FormatVer:    0,
 42 | 			},
 43 | 		}
 44 | 		bkd.trees = append(bkd.trees, kd)
 45 | 	}
 46 | 
 47 | 	err = bkd.compactTo(k)
 48 | 	return
 49 | }
 50 | 
 51 | //caclulate the min compoint position. Returns len(bkd.trees) if not found.
 52 | func (bkd *BkdTree) getMinCompactPos() (k int) {
 53 | 	//find the smallest index k in [0, len(trees)) at which trees[k] is empty, or its capacity is no less than the sum of size of t0m + trees[0:k+1]
 54 | 	sum := int(bkd.t0m.meta.NumPoints)
 55 | 	for k = 0; k < len(bkd.trees); k++ {
 56 | 		if bkd.trees[k].meta.NumPoints == 0 {
 57 | 			return
 58 | 		}
 59 | 		sum += int(bkd.trees[k].meta.NumPoints)
 60 | 		capK := bkd.t0mCap << uint(k)
 61 | 		if capK >= sum {
 62 | 			return
 63 | 		}
 64 | 	}
 65 | 	return
 66 | }
 67 | 
 68 | //compact T0M and trees[0:k+1] into tree[k]. Assumes write lock has been acquired.
 69 | func (bkd *BkdTree) compactTo(k int) (err error) {
 70 | 	//extract all points from t0m and trees[0:k+1] into a file F
 71 | 	fpK := bkd.TiPath(k)
 72 | 	tmpFpK := fpK + ".tmp"
 73 | 	tmpFK, err := os.OpenFile(tmpFpK, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
 74 | 	if err != nil {
 75 | 		err = errors.Wrap(err, "")
 76 | 		return
 77 | 	}
 78 | 	defer tmpFK.Close()
 79 | 
 80 | 	err = bkd.extractT0M(tmpFK)
 81 | 	if err != nil {
 82 | 		return
 83 | 	}
 84 | 	for i := 0; i <= k; i++ {
 85 | 		err = bkd.extractTi(tmpFK, i)
 86 | 		if err != nil {
 87 | 			return
 88 | 		}
 89 | 	}
 90 | 	meta, err := bkd.bulkLoad(tmpFK)
 91 | 	if err != nil {
 92 | 		return
 93 | 	}
 94 | 
 95 | 	//empty T0M and Ti, 0<=i<k
 96 | 	bkd.clearT0M()
 97 | 	for i := 0; i <= k; i++ {
 98 | 		if bkd.trees[i].meta.NumPoints <= 0 {
 99 | 			continue
100 | 		} else if err = FileMunmap(bkd.trees[i].data); err != nil {
101 | 			return
102 | 		} else if err = bkd.trees[i].f.Close(); err != nil {
103 | 			err = errors.Wrap(err, "")
104 | 			return
105 | 		} else if err = os.Remove(bkd.trees[i].f.Name()); err != nil {
106 | 			err = errors.Wrap(err, "")
107 | 			return
108 | 		}
109 | 		bkd.trees[i].meta.NumPoints = 0
110 | 	}
111 | 	if err = os.Rename(tmpFpK, fpK); err != nil {
112 | 		err = errors.Wrap(err, "")
113 | 		return
114 | 	}
115 | 	fK, err := os.OpenFile(fpK, os.O_RDWR, 0600)
116 | 	if err != nil {
117 | 		err = errors.Wrap(err, "")
118 | 		return
119 | 	}
120 | 	data, err := FileMmap(fK)
121 | 	if err != nil {
122 | 		return
123 | 	}
124 | 	bkd.trees[k] = BkdSubTree{
125 | 		meta: *meta,
126 | 		f:    fK,
127 | 		data: data,
128 | 	}
129 | 	return
130 | }
131 | 
132 | func writeMetaNumPoints(data []byte, meta *KdTreeExtMeta) {
133 | 	off := len(data) - KdTreeExtMetaSize
134 | 	off += int(unsafe.Offsetof(meta.NumPoints))
135 | 	binary.BigEndian.PutUint64(data[off:], meta.NumPoints)
136 | }
137 | 
138 | func (bkd *BkdTree) insertT0M(point Point) {
139 | 	pae := PointArrayExt{
140 | 		data:        bkd.t0m.data,
141 | 		numPoints:   int(bkd.t0m.meta.NumPoints),
142 | 		byDim:       0, //not used
143 | 		bytesPerDim: bkd.BytesPerDim,
144 | 		numDims:     bkd.NumDims,
145 | 		pointSize:   bkd.pointSize,
146 | 	}
147 | 	pae.Append(point)
148 | 	bkd.t0m.meta.NumPoints++
149 | 	writeMetaNumPoints(bkd.t0m.data, &bkd.t0m.meta)
150 | }
151 | 
152 | func (bkd *BkdTree) clearT0M() {
153 | 	bkd.t0m.meta.NumPoints = 0
154 | 	writeMetaNumPoints(bkd.t0m.data, &bkd.t0m.meta)
155 | }
156 | 
157 | func (bkd *BkdTree) extractT0M(tmpF *os.File) (err error) {
158 | 	size := int(bkd.t0m.meta.NumPoints) * bkd.pointSize
159 | 	_, err = tmpF.Write(bkd.t0m.data[:size])
160 | 	if err != nil {
161 | 		err = errors.Wrap(err, "")
162 | 		return
163 | 	}
164 | 	return
165 | }
166 | 
167 | func (bkd *BkdTree) extractTi(dstF *os.File, idx int) (err error) {
168 | 	if bkd.trees[idx].meta.NumPoints <= 0 {
169 | 		return
170 | 	}
171 | 	srcF, err := os.Open(bkd.TiPath(idx))
172 | 	if err != nil {
173 | 		err = errors.Wrap(err, "")
174 | 		return
175 | 	}
176 | 	defer srcF.Close()
177 | 
178 | 	//depth-first extracting from the root node
179 | 	meta := &bkd.trees[idx].meta
180 | 	err = bkd.extractNode(dstF, bkd.trees[idx].data, meta, int(meta.RootOff))
181 | 	return
182 | }
183 | 
184 | func (bkd *BkdTree) extractNode(dstF *os.File, data []byte, meta *KdTreeExtMeta, nodeOffset int) (err error) {
185 | 	var node KdTreeExtIntraNode
186 | 	bf := bytes.NewReader(data[nodeOffset:])
187 | 	err = node.Read(bf)
188 | 	if err != nil {
189 | 		err = errors.Wrap(err, "")
190 | 		return
191 | 	}
192 | 	for _, child := range node.Children {
193 | 		if child.Offset < meta.PointsOffEnd {
194 | 			//leaf node
195 | 			length := int(child.NumPoints) * int(meta.PointSize)
196 | 			_, err = dstF.Write(data[int(child.Offset) : int(child.Offset)+length])
197 | 			if err != nil {
198 | 				err = errors.Wrap(err, "")
199 | 				return
200 | 			}
201 | 		} else {
202 | 			//intra node
203 | 			err = bkd.extractNode(dstF, data, meta, int(child.Offset))
204 | 			if err != nil {
205 | 				return
206 | 			}
207 | 		}
208 | 	}
209 | 	return
210 | }
211 | 
212 | func (bkd *BkdTree) bulkLoad(tmpF *os.File) (meta *KdTreeExtMeta, err error) {
213 | 	pointsOffEnd, err := tmpF.Seek(0, 1) //get current position
214 | 	if err != nil {
215 | 		err = errors.Wrap(err, "")
216 | 		return
217 | 	}
218 | 	var data []byte
219 | 	if data, err = FileMmap(tmpF); err != nil {
220 | 		return
221 | 	}
222 | 	defer FileMunmap(data)
223 | 
224 | 	numPoints := int(pointsOffEnd / int64(bkd.pointSize))
225 | 	rootOff, err1 := bkd.createKdTreeExt(tmpF, data, 0, numPoints, 0)
226 | 	if err1 != nil {
227 | 		err = err1
228 | 		return
229 | 	}
230 | 	//record meta info at end
231 | 	meta = &KdTreeExtMeta{
232 | 		PointsOffEnd: uint64(pointsOffEnd),
233 | 		RootOff:      uint64(rootOff),
234 | 		NumPoints:    uint64(numPoints),
235 | 		LeafCap:      uint16(bkd.leafCap),
236 | 		IntraCap:     uint16(bkd.intraCap),
237 | 		NumDims:      uint8(bkd.NumDims),
238 | 		BytesPerDim:  uint8(bkd.BytesPerDim),
239 | 		PointSize:    uint8(bkd.pointSize),
240 | 		FormatVer:    0,
241 | 	}
242 | 	err = binary.Write(tmpF, binary.BigEndian, meta)
243 | 	if err != nil {
244 | 		err = errors.Wrap(err, "")
245 | 		return
246 | 	}
247 | 	return
248 | }
249 | 
250 | func getCurrentOffset(f *os.File) (offset int64, err error) {
251 | 	offset, err = f.Seek(0, 1) //get current position
252 | 	if err != nil {
253 | 		err = errors.Wrap(err, "")
254 | 		return
255 | 	}
256 | 	return
257 | }
258 | 
259 | func (bkd *BkdTree) createKdTreeExt(tmpF *os.File, data []byte, begin, end, depth int) (offset int64, err error) {
260 | 	if begin >= end {
261 | 		err = errors.New(fmt.Sprintf("assertion begin>=end failed, begin %v, end %v", begin, end))
262 | 		return
263 | 	}
264 | 
265 | 	splitDim := depth % bkd.NumDims
266 | 	numStrips := (end - begin + bkd.leafCap - 1) / bkd.leafCap
267 | 	if numStrips > bkd.intraCap {
268 | 		numStrips = bkd.intraCap
269 | 	}
270 | 
271 | 	pae := PointArrayExt{
272 | 		data:        data[begin*bkd.pointSize:],
273 | 		numPoints:   end - begin,
274 | 		byDim:       splitDim,
275 | 		bytesPerDim: bkd.BytesPerDim,
276 | 		numDims:     bkd.NumDims,
277 | 		pointSize:   bkd.pointSize,
278 | 	}
279 | 	splitValues, splitPoses := SplitPoints(&pae, numStrips)
280 | 
281 | 	children := make([]KdTreeExtNodeInfo, 0, numStrips)
282 | 	var childOffset int64
283 | 	for strip := 0; strip < numStrips; strip++ {
284 | 		posBegin := begin
285 | 		if strip != 0 {
286 | 			posBegin = begin + splitPoses[strip-1]
287 | 		}
288 | 		posEnd := end
289 | 		if strip != numStrips-1 {
290 | 			posEnd = begin + splitPoses[strip]
291 | 		}
292 | 		if posEnd-posBegin <= bkd.leafCap {
293 | 			info := KdTreeExtNodeInfo{
294 | 				Offset:    uint64(posBegin * bkd.pointSize),
295 | 				NumPoints: uint64(posEnd - posBegin),
296 | 			}
297 | 			children = append(children, info)
298 | 		} else {
299 | 			childOffset, err = bkd.createKdTreeExt(tmpF, data, posBegin, posEnd, depth+1)
300 | 			if err != nil {
301 | 				return
302 | 			}
303 | 			info := KdTreeExtNodeInfo{
304 | 				Offset:    uint64(childOffset),
305 | 				NumPoints: uint64(posEnd - posBegin),
306 | 			}
307 | 			children = append(children, info)
308 | 		}
309 | 	}
310 | 
311 | 	offset, err = getCurrentOffset(tmpF)
312 | 	if err != nil {
313 | 		return
314 | 	}
315 | 
316 | 	node := &KdTreeExtIntraNode{
317 | 		SplitDim:    uint32(splitDim),
318 | 		NumStrips:   uint32(numStrips),
319 | 		SplitValues: splitValues,
320 | 		Children:    children,
321 | 	}
322 | 	err = node.Write(tmpF)
323 | 	if err != nil {
324 | 		err = errors.Wrap(err, "")
325 | 		return
326 | 	}
327 | 	return
328 | }
329 | 


--------------------------------------------------------------------------------
/bkdtree_intersect.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 
  6 | 	"github.com/pkg/errors"
  7 | )
  8 | 
  9 | //Intersect does window query
 10 | func (bkd *BkdTree) Intersect(visitor IntersectVisitor) (err error) {
 11 | 	bkd.rwlock.RLock()
 12 | 	defer bkd.rwlock.RUnlock()
 13 | 	if !bkd.open {
 14 | 		err = errors.Errorf("(*BkdTree).Intersect is not allowed at closed state")
 15 | 		return
 16 | 	}
 17 | 
 18 | 	bkd.intersectT0M(visitor)
 19 | 	for i := 0; i < len(bkd.trees); i++ {
 20 | 		err = bkd.intersectTi(visitor, i)
 21 | 		if err != nil {
 22 | 			return
 23 | 		}
 24 | 	}
 25 | 	return
 26 | }
 27 | 
 28 | func (bkd *BkdTree) intersectT0M(visitor IntersectVisitor) {
 29 | 	lowP := visitor.GetLowPoint()
 30 | 	highP := visitor.GetHighPoint()
 31 | 	pae := PointArrayExt{
 32 | 		data:        bkd.t0m.data,
 33 | 		numPoints:   int(bkd.t0m.meta.NumPoints),
 34 | 		byDim:       0, //not used
 35 | 		bytesPerDim: bkd.BytesPerDim,
 36 | 		numDims:     bkd.NumDims,
 37 | 		pointSize:   bkd.pointSize,
 38 | 	}
 39 | 	for i := 0; i < pae.numPoints; i++ {
 40 | 		point := pae.GetPoint(i)
 41 | 		if point.Inside(lowP, highP) {
 42 | 			visitor.VisitPoint(point)
 43 | 		}
 44 | 	}
 45 | 	return
 46 | }
 47 | 
 48 | func (bkd *BkdTree) intersectTi(visitor IntersectVisitor, idx int) (err error) {
 49 | 	if bkd.trees[idx].meta.NumPoints <= 0 {
 50 | 		return
 51 | 	}
 52 | 	//depth-first visiting from the root node
 53 | 	meta := &bkd.trees[idx].meta
 54 | 	err = bkd.intersectNode(visitor, bkd.trees[idx].data, meta, int(meta.RootOff))
 55 | 	return
 56 | }
 57 | 
 58 | func (bkd *BkdTree) intersectNode(visitor IntersectVisitor, data []byte,
 59 | 	meta *KdTreeExtMeta, nodeOffset int) (err error) {
 60 | 	lowP := visitor.GetLowPoint()
 61 | 	highP := visitor.GetHighPoint()
 62 | 	var node KdTreeExtIntraNode
 63 | 	bf := bytes.NewReader(data[nodeOffset:])
 64 | 	err = node.Read(bf)
 65 | 	if err != nil {
 66 | 		return
 67 | 	}
 68 | 	lowVal := lowP.Vals[node.SplitDim]
 69 | 	highVal := highP.Vals[node.SplitDim]
 70 | 	for i, child := range node.Children {
 71 | 		if child.NumPoints <= 0 {
 72 | 			continue
 73 | 		}
 74 | 		if i < int(node.NumStrips)-1 && node.SplitValues[i] < lowVal {
 75 | 			continue
 76 | 		}
 77 | 		if i != 0 && node.SplitValues[i-1] > highVal {
 78 | 			continue
 79 | 		}
 80 | 		if child.Offset < meta.PointsOffEnd {
 81 | 			//leaf node
 82 | 			pae := PointArrayExt{
 83 | 				data:        data[int(child.Offset):],
 84 | 				numPoints:   int(child.NumPoints),
 85 | 				byDim:       0, //not used
 86 | 				bytesPerDim: bkd.BytesPerDim,
 87 | 				numDims:     bkd.NumDims,
 88 | 				pointSize:   bkd.pointSize,
 89 | 			}
 90 | 			for i := 0; i < pae.numPoints; i++ {
 91 | 				point := pae.GetPoint(i)
 92 | 				if point.Inside(lowP, highP) {
 93 | 					visitor.VisitPoint(point)
 94 | 				}
 95 | 			}
 96 | 		} else {
 97 | 			//intra node
 98 | 			err = bkd.intersectNode(visitor, data, meta, int(child.Offset))
 99 | 		}
100 | 		if err != nil {
101 | 			return
102 | 		}
103 | 	}
104 | 	return
105 | }
106 | 


--------------------------------------------------------------------------------
/bkdtree_test.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 	"encoding/binary"
  6 | 	"fmt"
  7 | 	"math/rand"
  8 | 	"os"
  9 | 	"testing"
 10 | 	"time"
 11 | 
 12 | 	"github.com/juju/testing/checkers"
 13 | 	"github.com/pkg/errors"
 14 | )
 15 | 
 16 | func (n *KdTreeExtIntraNode) equal(n2 *KdTreeExtIntraNode) (res bool) {
 17 | 	if n.SplitDim != n2.SplitDim || n.NumStrips != n2.NumStrips ||
 18 | 		len(n.SplitValues) != len(n2.SplitValues) ||
 19 | 		len(n.Children) != len(n2.Children) {
 20 | 		res = false
 21 | 		return
 22 | 	}
 23 | 	for i := 0; i < len(n.SplitValues); i++ {
 24 | 		if n.SplitValues[i] != n2.SplitValues[i] {
 25 | 			res = false
 26 | 			return
 27 | 		}
 28 | 	}
 29 | 	for i := 0; i < len(n.Children); i++ {
 30 | 		if n.Children[i] != n2.Children[i] {
 31 | 			res = false
 32 | 			return
 33 | 		}
 34 | 	}
 35 | 	res = true
 36 | 	return
 37 | }
 38 | 
 39 | func TestIntraNodeReadWrite(t *testing.T) {
 40 | 	n := KdTreeExtIntraNode{
 41 | 		SplitDim:    1,
 42 | 		NumStrips:   4,
 43 | 		SplitValues: []uint64{3, 5, 7},
 44 | 		Children: []KdTreeExtNodeInfo{
 45 | 			{Offset: 0, NumPoints: 7},
 46 | 			{Offset: 10, NumPoints: 9},
 47 | 			{Offset: 20, NumPoints: 137},
 48 | 			{Offset: 180, NumPoints: 999},
 49 | 		},
 50 | 	}
 51 | 	bf := new(bytes.Buffer)
 52 | 	if err := n.Write(bf); err != nil {
 53 | 		t.Fatalf("%+v", err)
 54 | 	}
 55 | 
 56 | 	var n2 KdTreeExtIntraNode
 57 | 	if err := n2.Read(bf); err != nil {
 58 | 		t.Fatalf("%+v", err)
 59 | 	}
 60 | 
 61 | 	if !n.equal(&n2) {
 62 | 		t.Fatalf("KdTreeExtIntraNode changes after encode and decode: %v, %v", n, n2)
 63 | 	}
 64 | }
 65 | func TestBkdInsert(t *testing.T) {
 66 | 	t0mCap := 1000
 67 | 	treesCap := 5
 68 | 	bkdCap := t0mCap<<uint(treesCap) - 1
 69 | 	leafCap := 50
 70 | 	intraCap := 4
 71 | 	numDims := 2
 72 | 	bytesPerDim := 4
 73 | 	dir := "/tmp"
 74 | 	prefix := "bkd"
 75 | 	bkd, err := NewBkdTree(t0mCap, leafCap, intraCap, numDims, bytesPerDim, dir, prefix)
 76 | 	if err != nil {
 77 | 		t.Fatalf("%+v", err)
 78 | 	}
 79 | 
 80 | 	maxVal := uint64(1000)
 81 | 	size := bkdCap + 1
 82 | 	points := NewRandPoints(numDims, maxVal, size)
 83 | 
 84 | 	//insert points
 85 | 	for i := 0; i < bkdCap; i++ {
 86 | 		err := bkd.Insert(points[i])
 87 | 		if err != nil {
 88 | 			t.Fatalf("bkd.Insert failed, i=%v, err: %+v", i, err)
 89 | 		}
 90 | 		if bkd.NumPoints != i+1 {
 91 | 			t.Fatalf("incorrect numPoints. numPoints=%v, i=%v", bkd.NumPoints, i)
 92 | 		}
 93 | 		remained := bkd.NumPoints % bkd.t0mCap
 94 | 		quotient := bkd.NumPoints / bkd.t0mCap
 95 | 		if int(bkd.t0m.meta.NumPoints) != remained {
 96 | 			t.Fatalf("bkd.numPoints %d, bkd.t0m %d is incorect, want %d", bkd.NumPoints, int(bkd.t0m.meta.NumPoints), remained)
 97 | 		}
 98 | 		for i := 0; i < len(bkd.trees); i++ {
 99 | 			tiCap := bkd.t0mCap << uint(i)
100 | 			want := tiCap * (quotient % 2)
101 | 			if bkd.trees[i].meta.NumPoints != uint64(want) {
102 | 				t.Fatalf("bkd.numPoints %d, bkd.tree[%d].numPoints %d is incorrect, want %d", bkd.NumPoints, i, bkd.trees[i].meta.NumPoints, want)
103 | 			}
104 | 			quotient >>= 1
105 | 		}
106 | 	}
107 | }
108 | 
109 | func prepareBkdTree(maxVal uint64) (bkd *BkdTree, points []Point, err error) {
110 | 	t0mCap := 1000
111 | 	treesCap := 5
112 | 	bkdCap := t0mCap<<uint(treesCap) - 1
113 | 	leafCap := 50
114 | 	intraCap := 4
115 | 	numDims := 2
116 | 	bytesPerDim := 4
117 | 	dir := "/tmp"
118 | 	prefix := "bkd"
119 | 	bkd, err = NewBkdTree(t0mCap, leafCap, intraCap, numDims, bytesPerDim, dir, prefix)
120 | 	if err != nil {
121 | 		return
122 | 	}
123 | 	//fmt.Printf("created BkdTree %v\n", bkd)
124 | 
125 | 	size := bkdCap
126 | 	points = NewRandPoints(numDims, maxVal, size)
127 | 	for i := 0; i < bkdCap; i++ {
128 | 		err = bkd.Insert(points[i])
129 | 		if err != nil {
130 | 			err = errors.Errorf("bkd.Insert failed, i=%v, err: %+v", i, err)
131 | 		}
132 | 	}
133 | 	return
134 | }
135 | 
136 | func TestBkdIntersect(t *testing.T) {
137 | 	var maxVal uint64 = 1000
138 | 	bkd, points, err := prepareBkdTree(maxVal)
139 | 	if err != nil {
140 | 		t.Fatalf("%+v", err)
141 | 	}
142 | 	var lowPoint, highPoint Point
143 | 	var visitor *IntersectCollector
144 | 
145 | 	//some intersect
146 | 	lowPoint = points[7]
147 | 	highPoint = lowPoint
148 | 	visitor = &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
149 | 	err = bkd.Intersect(visitor)
150 | 	if err != nil {
151 | 		t.Fatalf("%+v", err)
152 | 	}
153 | 	if len(visitor.Points) <= 0 {
154 | 		t.Errorf("found 0 matchs, however some expected")
155 | 	}
156 | 	var matched int
157 | 	for _, point := range visitor.Points {
158 | 		isInside := point.Inside(lowPoint, highPoint)
159 | 		if !isInside {
160 | 			t.Errorf("point %v is ouside of range", point)
161 | 		}
162 | 		if point.Equal(lowPoint) {
163 | 			matched++
164 | 		}
165 | 	}
166 | 	if matched != 1 {
167 | 		t.Errorf("found %d points equal to %v", matched, lowPoint)
168 | 	}
169 | 
170 | 	//all intersect
171 | 	lowPoint = Point{[]uint64{0, 0, 0}, 0}
172 | 	highPoint = Point{[]uint64{maxVal, maxVal, maxVal}, 0}
173 | 	visitor = &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
174 | 	err = bkd.Intersect(visitor)
175 | 	if err != nil {
176 | 		t.Fatalf("%+v", err)
177 | 	}
178 | 	if len(visitor.Points) != len(points) {
179 | 		t.Errorf("found %d matchs, want %d", len(visitor.Points), len(points))
180 | 	}
181 | }
182 | 
183 | func verifyBkdMeta(bkd *BkdTree) (err error) {
184 | 	cnt := int(bkd.t0m.meta.NumPoints)
185 | 	var f *os.File
186 | 	for i := 0; i < len(bkd.trees); i++ {
187 | 		if bkd.trees[i].meta.NumPoints <= 0 {
188 | 			continue
189 | 		}
190 | 		f, err = os.OpenFile(bkd.TiPath(i), os.O_RDONLY, 0)
191 | 		if err != nil {
192 | 			return
193 | 		}
194 | 		_, err = f.Seek(-int64(KdTreeExtMetaSize), 2)
195 | 		if err != nil {
196 | 			return
197 | 		}
198 | 		var meta KdTreeExtMeta
199 | 		err = binary.Read(f, binary.BigEndian, &meta)
200 | 		if err != nil {
201 | 			return
202 | 		}
203 | 		if meta != bkd.trees[i].meta {
204 | 			err = errors.Errorf("bkd.trees[%d].meta does not match file content, has %v, want %v", i, bkd.trees[i].meta, meta)
205 | 			return
206 | 		}
207 | 		cnt += int(meta.NumPoints)
208 | 	}
209 | 	if cnt != bkd.NumPoints {
210 | 		err = errors.Errorf("bkd.numPoints does not match file content, has %v, want %v", bkd.NumPoints, cnt)
211 | 		return
212 | 	}
213 | 	return
214 | }
215 | 
216 | func countPoint(bkd *BkdTree, point Point) (cnt int, err error) {
217 | 	lowPoint, highPoint := point, point
218 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
219 | 	err = bkd.Intersect(visitor)
220 | 	if err != nil {
221 | 		return
222 | 	}
223 | 	for _, p := range visitor.Points {
224 | 		if p.Equal(point) {
225 | 			cnt++
226 | 		}
227 | 	}
228 | 	return
229 | }
230 | 
231 | func TestBkdErase(t *testing.T) {
232 | 	var maxVal uint64 = 1000
233 | 	bkd, points, err := prepareBkdTree(maxVal)
234 | 	if err != nil {
235 | 		t.Fatalf("%+v", err)
236 | 	}
237 | 	var target Point
238 | 	var cnt int
239 | 	var found bool
240 | 
241 | 	//erase an non-existing point
242 | 	target = points[17]
243 | 	target.UserData = uint64(len(points))
244 | 	found, err = bkd.Erase(target)
245 | 	if err != nil {
246 | 		t.Fatalf("%v", err)
247 | 	} else if found {
248 | 		t.Fatalf("point %v found, want non-existing", target)
249 | 	}
250 | 
251 | 	//erase an existing point, verify the point really erased.
252 | 	target = points[13]
253 | 	found, err = bkd.Erase(target)
254 | 	if err != nil {
255 | 		t.Fatalf("%+v", err)
256 | 	} else if !found {
257 | 		t.Fatalf("point %v not found", target)
258 | 	} else if bkd.NumPoints != len(points)-1 {
259 | 		t.Fatalf("incorrect bkd.numPoints %d, want %d", bkd.NumPoints, len(points)-1)
260 | 	} else if err = verifyBkdMeta(bkd); err != nil {
261 | 		t.Fatalf("%+v", err)
262 | 	}
263 | 
264 | 	cnt, err = countPoint(bkd, target)
265 | 	if err != nil {
266 | 		t.Fatalf("%+v", err)
267 | 	} else if cnt != 0 {
268 | 		t.Errorf("point %v still exists", target)
269 | 	}
270 | 	//there's room for insertion
271 | 	err = bkd.Insert(target)
272 | 	if err != nil {
273 | 		t.Fatalf("bkd.Insert failed, err: %+v", err)
274 | 	} else if bkd.NumPoints != len(points) {
275 | 		t.Fatalf("incorrect bkd.numPoints %d, want %d", bkd.NumPoints, len(points))
276 | 	} else if err = verifyBkdMeta(bkd); err != nil {
277 | 		t.Fatalf("%+v", err)
278 | 	}
279 | 
280 | 	cnt, err = countPoint(bkd, target)
281 | 	if err != nil {
282 | 		t.Fatalf("%+v", err)
283 | 	} else if cnt != 1 {
284 | 		t.Errorf("point %v still exists", target)
285 | 	}
286 | }
287 | 
288 | func TestBkdDestroy(t *testing.T) {
289 | 	var bkd *BkdTree
290 | 	var err error
291 | 	var maxVal uint64 = 1000
292 | 	bkd, _, err = prepareBkdTree(maxVal)
293 | 	if err != nil {
294 | 		t.Fatalf("%+v", err)
295 | 	}
296 | 
297 | 	if err = bkd.Destroy(); err != nil {
298 | 		t.Fatalf("%+v", err)
299 | 	}
300 | 	//Verify t0m has been removed
301 | 	if _, err = NewBkdTreeExt(bkd.dir, bkd.prefix); err == nil {
302 | 		t.Fatalf("%s is still there", bkd.t0m.f.Name())
303 | 	}
304 | }
305 | 
306 | func TestBkdOpenClose(t *testing.T) {
307 | 	var bkd, bkd2 *BkdTree
308 | 	var err error
309 | 	var maxVal uint64 = 1000
310 | 	bkd, _, err = prepareBkdTree(maxVal)
311 | 	if err != nil {
312 | 		t.Fatalf("%+v", err)
313 | 	}
314 | 
315 | 	if err = bkd.Close(); err != nil {
316 | 		t.Fatalf("%+v", err)
317 | 	}
318 | 
319 | 	if bkd2, err = NewBkdTreeExt(bkd.dir, bkd.prefix); err != nil {
320 | 		t.Fatalf("%+v", err)
321 | 	}
322 | 	if err = bkd2.Close(); err != nil {
323 | 		t.Fatalf("%+v", err)
324 | 	}
325 | 	//Compare two structs recursively and record the difference.
326 | 	//TODO: How to ignore specific fields effciently?
327 | 	bkd2.t0m.data, bkd.t0m.data = make([]byte, 0), make([]byte, 0)
328 | 	for i := 0; i < len(bkd.trees); i++ {
329 | 		bkd2.trees[i].data, bkd.trees[i].data = make([]byte, 0), make([]byte, 0)
330 | 	}
331 | 	isEqual, err := checkers.DeepEqual(bkd, bkd2)
332 | 	if !isEqual {
333 | 		t.Fatalf("bkd changed with close and open. %+v", err)
334 | 	}
335 | }
336 | 
337 | func TestBkdCompact(t *testing.T) {
338 | 	var bkd *BkdTree
339 | 	var points []Point
340 | 	var err error
341 | 	var maxVal uint64 = 1000
342 | 	bkd, points, err = prepareBkdTree(maxVal)
343 | 	if err != nil {
344 | 		t.Fatalf("%+v", err)
345 | 	}
346 | 
347 | 	/*
348 | 		points distribution of bkd: T0M(1), 1, 2, 4, 8, 16
349 | 		test case:
350 | 		1. remove all points of trees[len-2] and trees[len-4].
351 | 		2. compact
352 | 		3. verify points of T0M and trees 0..len-2(inclusive) are moved to trees[len-2].
353 | 	*/
354 | 	treeLen := len(bkd.trees)
355 | 	tree1End := bkd.t0mCap << uint(treeLen-1)
356 | 	tree2End := tree1End + bkd.t0mCap<<uint(treeLen-2)
357 | 	tree3End := tree2End + bkd.t0mCap<<uint(treeLen-3)
358 | 	tree4End := tree3End + bkd.t0mCap<<uint(treeLen-4)
359 | 
360 | 	sum := int(bkd.t0m.meta.NumPoints)
361 | 	for i := 0; i < treeLen; i++ {
362 | 		if i == treeLen-4 || i == treeLen-2 || i == treeLen-1 {
363 | 			continue
364 | 		}
365 | 		sum += int(bkd.trees[i].meta.NumPoints)
366 | 	}
367 | 
368 | 	for i := tree1End; i < tree2End; i++ {
369 | 		found, err := bkd.Erase(points[i])
370 | 		if err != nil {
371 | 			t.Fatalf("%+v", err)
372 | 		} else if !found {
373 | 			t.Fatalf("point %v not found", points[i])
374 | 		}
375 | 	}
376 | 
377 | 	for i := tree3End; i < tree4End; i++ {
378 | 		found, err := bkd.Erase(points[i])
379 | 		if err != nil {
380 | 			t.Fatalf("%+v", err)
381 | 		} else if !found {
382 | 			t.Fatalf("point %v not found", points[i])
383 | 		}
384 | 	}
385 | 
386 | 	if err = bkd.Compact(); err != nil {
387 | 		t.Fatalf("%+v", err)
388 | 	}
389 | 
390 | 	if bkd.t0m.meta.NumPoints != 0 {
391 | 		t.Fatalf("bkd.t0m.NumPoints is incorrect, is %d, want %d", bkd.t0m.meta.NumPoints, 0)
392 | 	}
393 | 	for i := 0; i < treeLen-2; i++ {
394 | 		if bkd.trees[i].meta.NumPoints != 0 {
395 | 			t.Fatalf("bkd.trees[%d].NumPoints is incorrect, is %d, want %d", i, bkd.trees[i].meta.NumPoints, 0)
396 | 		}
397 | 	}
398 | 	if int(bkd.trees[treeLen-2].meta.NumPoints) != sum {
399 | 		t.Fatalf("bkd.trees[%d].NumPoints is incorrect, is %d, want %d", treeLen-2, int(bkd.trees[treeLen-2].meta.NumPoints), sum)
400 | 	}
401 | }
402 | 
403 | func bkdCloser(abort chan interface{}, bkd *BkdTree) {
404 | 	var interval time.Duration = 5 * time.Second
405 | FOR_LOOP:
406 | 	for {
407 | 		select {
408 | 		case <-abort:
409 | 			break FOR_LOOP //break for
410 | 		default: // adding default makes select nonblocking
411 | 		}
412 | 
413 | 		if err := bkd.Close(); err != nil {
414 | 			panic(err)
415 | 		}
416 | 		if err := bkd.Open(); err != nil {
417 | 			panic(err)
418 | 		}
419 | 		time.Sleep(interval)
420 | 	}
421 | }
422 | 
423 | func bkdWriter(abort chan interface{}, bkd *BkdTree, points []Point) {
424 | FOR_LOOP:
425 | 	for {
426 | 		select {
427 | 		case <-abort:
428 | 			break FOR_LOOP
429 | 		default:
430 | 		}
431 | 		idx := rand.Intn(len(points))
432 | 		//TODO: custom error type
433 | 		_, _ = bkd.Erase(points[idx])
434 | 		_ = bkd.Insert(points[idx])
435 | 	}
436 | }
437 | 
438 | func bkdReader(abort chan interface{}, bkd *BkdTree, points []Point) {
439 | FOR_LOOP:
440 | 	for {
441 | 		select {
442 | 		case <-abort:
443 | 			break FOR_LOOP
444 | 		default:
445 | 		}
446 | 
447 | 		idx1 := rand.Intn(len(points))
448 | 		idx2 := rand.Intn(len(points))
449 | 		visitor := &IntersectCollector{points[idx1], points[idx2], make([]Point, 0)}
450 | 		//TODO: custom error type
451 | 		_ = bkd.Intersect(visitor)
452 | 	}
453 | }
454 | 
455 | func TestBkdConcurrentOps(t *testing.T) {
456 | 	var bkd *BkdTree
457 | 	var points []Point
458 | 	var err error
459 | 	var maxVal uint64 = 1000
460 | 	bkd, points, err = prepareBkdTree(maxVal)
461 | 	if err != nil {
462 | 		t.Fatalf("%+v", err)
463 | 	}
464 | 	ch := make(chan interface{})
465 | 	for i := 0; i < 1; i++ {
466 | 		go bkdCloser(ch, bkd)
467 | 	}
468 | 	for i := 0; i < 3; i++ {
469 | 		go bkdWriter(ch, bkd, points)
470 | 	}
471 | 	for i := 0; i < 5; i++ {
472 | 		go bkdReader(ch, bkd, points)
473 | 	}
474 | 	//sleep a while, send message to abort readers and writers
475 | 	time.Sleep(20 * time.Second)
476 | 	close(ch)
477 | 	time.Sleep(1 * time.Second)
478 | 	fmt.Println("children shall all have quited")
479 | }
480 | 


--------------------------------------------------------------------------------
/go.mod:
--------------------------------------------------------------------------------
 1 | module github.com/deepfabric/bkdtree
 2 | 
 3 | go 1.21.5
 4 | 
 5 | require (
 6 | 	github.com/deepfabric/go-datastructures v0.0.0-20170927014437-f6355768d70e
 7 | 	github.com/juju/testing v1.1.0
 8 | 	github.com/keegancsmith/nth v0.0.0-20160926112203-ee21de2f07b8
 9 | 	github.com/pkg/errors v0.9.1
10 | )
11 | 
12 | require (
13 | 	github.com/juju/loggo v1.0.0 // indirect
14 | 	github.com/kr/pretty v0.3.1 // indirect
15 | 	github.com/kr/text v0.2.0 // indirect
16 | 	github.com/rogpeppe/go-internal v1.9.0 // indirect
17 | 	github.com/stretchr/testify v1.8.4 // indirect
18 | 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
19 | 	gopkg.in/yaml.v2 v2.4.0 // indirect
20 | )
21 | 


--------------------------------------------------------------------------------
/go.sum:
--------------------------------------------------------------------------------
 1 | github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 4 | github.com/deepfabric/go-datastructures v0.0.0-20170927014437-f6355768d70e h1:znsjfbR2/zGoKYztUmGElpU1jTCf4GB+ibRdBEAeCeE=
 5 | github.com/deepfabric/go-datastructures v0.0.0-20170927014437-f6355768d70e/go.mod h1:XzfNeeU3SD84cKQARZs5MODWAzwO3v5SbNrVAVpN2pY=
 6 | github.com/juju/ansiterm v0.0.0-20180109212912-720a0952cc2a/go.mod h1:UJSiEoRfvx3hP73CvoARgeLjaIOjybY9vj8PUPPFGeU=
 7 | github.com/juju/errors v1.0.0 h1:yiq7kjCLll1BiaRuNY53MGI0+EQ3rF6GB+wvboZDefM=
 8 | github.com/juju/errors v1.0.0/go.mod h1:B5x9thDqx0wIMH3+aLIMP9HjItInYWObRovoCFM5Qe8=
 9 | github.com/juju/loggo v1.0.0 h1:Y6ZMQOGR9Aj3BGkiWx7HBbIx6zNwNkxhVNOHU2i1bl0=
10 | github.com/juju/loggo v1.0.0/go.mod h1:NIXFioti1SmKAlKNuUwbMenNdef59IF52+ZzuOmHYkg=
11 | github.com/juju/testing v1.1.0 h1:+WWez0vCu6dtnpLIzfuuo3bN3x62LBIyMDCfvMYP+Qg=
12 | github.com/juju/testing v1.1.0/go.mod h1:1XQGptw6JWFvRWb3ewilUdTBG0oGcoI2kdX9Z1VEzhU=
13 | github.com/keegancsmith/nth v0.0.0-20160926112203-ee21de2f07b8 h1:QCdRsV8WNLoLAUwAm6G4FZF3OpCkRtU0QDjGkp5/1r4=
14 | github.com/keegancsmith/nth v0.0.0-20160926112203-ee21de2f07b8/go.mod h1:GGie0rcBbRSciSNcfYjdzMR0jL66DbnsHDOhwIJ4L7g=
15 | github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
16 | github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
17 | github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
18 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
19 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
20 | github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
21 | github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
22 | github.com/lunixbochs/vtclean v0.0.0-20160125035106-4fbf7632a2c6/go.mod h1:pHhQNgMf3btfWnGBVipUOjRYhoOsdGqdm/+2c2E2WMI=
23 | github.com/mattn/go-colorable v0.0.6/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU=
24 | github.com/mattn/go-isatty v0.0.0-20160806122752-66b8e73f3f5c/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4=
25 | github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
26 | github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
27 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
28 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
29 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
30 | github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
31 | github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
32 | github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
33 | github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
34 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
35 | gopkg.in/check.v1 v1.0.0-20160105164936-4f90aeace3a2/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
36 | gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
37 | gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
38 | gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
39 | gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
40 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
41 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
42 | 


--------------------------------------------------------------------------------
/kdtree.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"math/rand"
  5 | 	"sort"
  6 | )
  7 | 
  8 | type U64Slice []uint64
  9 | 
 10 | func (a U64Slice) Len() int           { return len(a) }
 11 | func (a U64Slice) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
 12 | func (a U64Slice) Less(i, j int) bool { return a[i] < a[j] }
 13 | 
 14 | type KdTreeNode interface {
 15 | 	intersect(visitor IntersectVisitor, numDims int)
 16 | 	insert(point Point, numDims int)
 17 | 	erase(point Point, numDims int) bool
 18 | }
 19 | 
 20 | type KdTreeIntraNode struct {
 21 | 	splitDim    int
 22 | 	splitValues []uint64
 23 | 	children    []KdTreeNode
 24 | }
 25 | 
 26 | type KdTreeLeafNode struct {
 27 | 	points []Point
 28 | }
 29 | 
 30 | type IntersectVisitor interface {
 31 | 	GetLowPoint() Point
 32 | 	GetHighPoint() Point
 33 | 	VisitPoint(point Point)
 34 | }
 35 | 
 36 | type IntersectCollector struct {
 37 | 	LowPoint  Point
 38 | 	HighPoint Point
 39 | 	Points    []Point
 40 | }
 41 | 
 42 | func (d *IntersectCollector) GetLowPoint() Point     { return d.LowPoint }
 43 | func (d *IntersectCollector) GetHighPoint() Point    { return d.HighPoint }
 44 | func (d *IntersectCollector) VisitPoint(point Point) { d.Points = append(d.Points, point) }
 45 | 
 46 | type KdTree struct {
 47 | 	root     KdTreeNode
 48 | 	NumDims  int
 49 | 	leafCap  int // limit of points a leaf node can hold
 50 | 	intraCap int // limit of children of a intra node can hold
 51 | }
 52 | 
 53 | func NewKdTree(points []Point, numDims, leafCap, intraCap int) (kd *KdTree) {
 54 | 	if len(points) == 0 || numDims <= 0 ||
 55 | 		leafCap <= 0 || leafCap >= int(^uint16(0)) || intraCap <= 2 || intraCap >= int(^uint16(0)) {
 56 | 		return
 57 | 	}
 58 | 	kd = &KdTree{
 59 | 		root:     createKdTree(points, 0, numDims, leafCap, intraCap),
 60 | 		NumDims:  numDims,
 61 | 		leafCap:  leafCap,
 62 | 		intraCap: intraCap,
 63 | 	}
 64 | 	return
 65 | }
 66 | 
 67 | func createKdTree(points []Point, depth, numDims, leafCap, intraCap int) KdTreeNode {
 68 | 	if len(points) == 0 {
 69 | 		return nil
 70 | 	}
 71 | 	if len(points) <= leafCap {
 72 | 		pointsCopy := make([]Point, len(points))
 73 | 		copy(pointsCopy, points)
 74 | 		ret := &KdTreeLeafNode{
 75 | 			points: pointsCopy,
 76 | 		}
 77 | 		return ret
 78 | 	}
 79 | 
 80 | 	splitDim := depth % numDims
 81 | 	numStrips := (len(points) + leafCap - 1) / leafCap
 82 | 	if numStrips > intraCap {
 83 | 		numStrips = intraCap
 84 | 	}
 85 | 
 86 | 	pam := PointArrayMem{
 87 | 		points: points,
 88 | 		byDim:  splitDim,
 89 | 	}
 90 | 
 91 | 	splitValues, splitPoses := SplitPoints(&pam, numStrips)
 92 | 
 93 | 	children := make([]KdTreeNode, 0, numStrips)
 94 | 	for strip := 0; strip < numStrips; strip++ {
 95 | 		posBegin := 0
 96 | 		if strip != 0 {
 97 | 			posBegin = splitPoses[strip-1]
 98 | 		}
 99 | 		posEnd := len(points)
100 | 		if strip != numStrips-1 {
101 | 			posEnd = splitPoses[strip]
102 | 		}
103 | 		child := createKdTree(points[posBegin:posEnd], depth+1, numDims, leafCap, intraCap)
104 | 		children = append(children, child)
105 | 	}
106 | 	ret := &KdTreeIntraNode{
107 | 		splitDim:    splitDim,
108 | 		splitValues: splitValues,
109 | 		children:    children,
110 | 	}
111 | 	return ret
112 | }
113 | 
114 | func (n *KdTreeIntraNode) intersect(visitor IntersectVisitor, numDims int) {
115 | 	lowVal := visitor.GetLowPoint().Vals[n.splitDim]
116 | 	highVal := visitor.GetHighPoint().Vals[n.splitDim]
117 | 	numSplits := len(n.splitValues)
118 | 	//calculate children[begin:end) need to visit
119 | 	end := sort.Search(numSplits, func(i int) bool { return n.splitValues[i] > highVal })
120 | 	begin := sort.Search(end, func(i int) bool { return n.splitValues[i] >= lowVal })
121 | 	end++
122 | 	for strip := begin; strip < end; strip++ {
123 | 		n.children[strip].intersect(visitor, numDims)
124 | 	}
125 | }
126 | 
127 | func (n *KdTreeLeafNode) intersect(visitor IntersectVisitor, numDims int) {
128 | 	lowPoint := visitor.GetLowPoint()
129 | 	highPoint := visitor.GetHighPoint()
130 | 	for _, point := range n.points {
131 | 		isInside := point.Inside(lowPoint, highPoint)
132 | 		if isInside {
133 | 			visitor.VisitPoint(point)
134 | 		}
135 | 	}
136 | }
137 | 
138 | func (t *KdTree) Intersect(visitor IntersectVisitor) {
139 | 	t.root.intersect(visitor, t.NumDims)
140 | }
141 | 
142 | func (n *KdTreeIntraNode) insert(point Point, numDims int) {
143 | 	lowVal := point.Vals[n.splitDim]
144 | 	highVal := lowVal
145 | 	numSplits := len(n.splitValues)
146 | 	//calculate children[begin:end) need to visit
147 | 	end := sort.Search(numSplits, func(i int) bool { return n.splitValues[i] > highVal })
148 | 	begin := sort.Search(end, func(i int) bool { return n.splitValues[i] >= lowVal })
149 | 	end++
150 | 	//if multiple strips could cover the point, select one randomly.
151 | 	strip := begin + rand.Intn(end-begin)
152 | 	n.children[strip].insert(point, numDims)
153 | }
154 | 
155 | func (n *KdTreeLeafNode) insert(point Point, numDims int) {
156 | 	//append blindly, no rebalance
157 | 	n.points = append(n.points, point)
158 | }
159 | 
160 | func (t *KdTree) Insert(point Point) {
161 | 	t.root.insert(point, t.NumDims)
162 | }
163 | 
164 | func (n *KdTreeIntraNode) erase(point Point, numDims int) (found bool) {
165 | 	lowVal := point.Vals[n.splitDim]
166 | 	highVal := lowVal
167 | 	numSplits := len(n.splitValues)
168 | 	//calculate children[begin:end) need to visit
169 | 	end := sort.Search(numSplits, func(i int) bool { return n.splitValues[i] > highVal })
170 | 	begin := sort.Search(end, func(i int) bool { return n.splitValues[i] >= lowVal })
171 | 	end++
172 | 	//if multiple strips could cover the point, iterate them. And stop iteration if found at middle way.
173 | 	for strip := begin; strip < end; strip++ {
174 | 		found = n.children[strip].erase(point, numDims)
175 | 		if found {
176 | 			break
177 | 		}
178 | 	}
179 | 	return
180 | }
181 | 
182 | func (n *KdTreeLeafNode) erase(point Point, numDims int) (found bool) {
183 | 	found = false
184 | 	idx := len(n.points)
185 | 	for i, point2 := range n.points {
186 | 		//assumes each point's userData is unique
187 | 		if point.Equal(point2) {
188 | 			idx = i
189 | 			break
190 | 		}
191 | 	}
192 | 	if idx < len(n.points) {
193 | 		n.points = append(n.points[:idx], n.points[idx+1:]...)
194 | 		found = true
195 | 	}
196 | 	return
197 | }
198 | 
199 | func (t *KdTree) Erase(point Point) {
200 | 	t.root.erase(point, t.NumDims)
201 | }
202 | 


--------------------------------------------------------------------------------
/kdtree_test.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"testing"
  5 | )
  6 | 
  7 | func TestKdIntersectSome(t *testing.T) {
  8 | 	numDims := 3
  9 | 	maxVal := uint64(1000)
 10 | 	size := 1000
 11 | 	leafCap := 50
 12 | 	intraCap := 4
 13 | 	points := NewRandPoints(numDims, maxVal, size)
 14 | 	kdt := NewKdTree(points, numDims, leafCap, intraCap)
 15 | 
 16 | 	lowPoint := points[0]
 17 | 	highPoint := points[0]
 18 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0, 1)}
 19 | 	kdt.Intersect(visitor)
 20 | 
 21 | 	//fmt.Printf("%v\n", visitor.points)
 22 | 	if len(visitor.Points) <= 0 {
 23 | 		t.Errorf("found 0 matchs, however some expected")
 24 | 	}
 25 | 	for _, point := range visitor.Points {
 26 | 		isInside := point.Inside(lowPoint, highPoint)
 27 | 		if !isInside {
 28 | 			t.Errorf("point %v is ouside of range", point)
 29 | 		}
 30 | 	}
 31 | }
 32 | func TestKdIntersectAll(t *testing.T) {
 33 | 	numDims := 3
 34 | 	maxVal := uint64(1000)
 35 | 	size := 1000
 36 | 	leafCap := 50
 37 | 	intraCap := 4
 38 | 	points := NewRandPoints(numDims, maxVal, size)
 39 | 	kdt := NewKdTree(points, numDims, leafCap, intraCap)
 40 | 
 41 | 	lowPoint := Point{[]uint64{0, 0, 0}, 0}
 42 | 	highPoint := Point{[]uint64{maxVal, maxVal, maxVal}, 0}
 43 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0, size)}
 44 | 	kdt.Intersect(visitor)
 45 | 	if len(visitor.Points) != size {
 46 | 		t.Errorf("found %v matchs, however %v expected", len(visitor.Points), size)
 47 | 	}
 48 | }
 49 | 
 50 | func TestKdIntersect(t *testing.T) {
 51 | 	numDims := 3
 52 | 	maxVal := uint64(1000)
 53 | 	size := 100000
 54 | 	leafCap := 50
 55 | 	intraCap := 4
 56 | 	points := NewRandPoints(numDims, maxVal, size)
 57 | 	kdt := NewKdTree(points, numDims, leafCap, intraCap)
 58 | 
 59 | 	lowPoint := Point{[]uint64{20, 30, 40}, 0}
 60 | 	highPoint := Point{[]uint64{maxVal, maxVal, maxVal}, 0}
 61 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
 62 | 	kdt.Intersect(visitor)
 63 | 
 64 | 	//fmt.Printf("%v\n", visitor.points)
 65 | 	for _, point := range visitor.Points {
 66 | 		isInside := point.Inside(lowPoint, highPoint)
 67 | 		if !isInside {
 68 | 			t.Errorf("point %v is ouside of range", point)
 69 | 		}
 70 | 	}
 71 | }
 72 | 
 73 | func TestKdInsert(t *testing.T) {
 74 | 	numDims := 3
 75 | 	maxVal := uint64(1000)
 76 | 	size := 1000
 77 | 	leafCap := 50
 78 | 	intraCap := 4
 79 | 	points := NewRandPoints(numDims, maxVal, size)
 80 | 	kdt := NewKdTree(points, numDims, leafCap, intraCap)
 81 | 
 82 | 	newPoint := Point{[]uint64{40, 30, 20}, maxVal} //use unique userData
 83 | 	kdt.Insert(newPoint)
 84 | 
 85 | 	lowPoint := newPoint
 86 | 	highPoint := newPoint
 87 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
 88 | 	kdt.Intersect(visitor)
 89 | 
 90 | 	//fmt.Printf("%v\n", visitor.points)
 91 | 	if len(visitor.Points) <= 0 {
 92 | 		t.Errorf("found 0 matchs, however some expected")
 93 | 	}
 94 | 	numMatchs := 0
 95 | 	for _, point := range visitor.Points {
 96 | 		isInside := point.Inside(lowPoint, highPoint)
 97 | 		if !isInside {
 98 | 			t.Errorf("point %v is ouside of range", point)
 99 | 		}
100 | 		if point.UserData == newPoint.UserData {
101 | 			numMatchs++
102 | 		}
103 | 	}
104 | 	if numMatchs != 1 {
105 | 		t.Errorf("found %v matchs, however 1 expected", numMatchs)
106 | 	}
107 | }
108 | 
109 | func TestKdErase(t *testing.T) {
110 | 	numDims := 3
111 | 	maxVal := uint64(1000)
112 | 	size := 1000
113 | 	leafCap := 50
114 | 	intraCap := 4
115 | 	points := NewRandPoints(numDims, maxVal, size)
116 | 	kdt := NewKdTree(points, numDims, leafCap, intraCap)
117 | 
118 | 	kdt.Erase(points[0])
119 | 
120 | 	lowPoint := points[0]
121 | 	highPoint := points[0]
122 | 	visitor := &IntersectCollector{lowPoint, highPoint, make([]Point, 0)}
123 | 	kdt.Intersect(visitor)
124 | 
125 | 	//fmt.Printf("%v\n", visitor.points)
126 | 	if len(visitor.Points) != 0 {
127 | 		t.Errorf("found %v matchs, however 0 expected", len(visitor.Points))
128 | 	}
129 | }
130 | 


--------------------------------------------------------------------------------
/point.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"encoding/binary"
  5 | 	"sort"
  6 | 
  7 | 	datastructures "github.com/deepfabric/go-datastructures"
  8 | 	"github.com/keegancsmith/nth"
  9 | )
 10 | 
 11 | type Point struct {
 12 | 	Vals     []uint64
 13 | 	UserData uint64
 14 | }
 15 | 
 16 | type PointArray interface {
 17 | 	sort.Interface
 18 | 	GetPoint(idx int) Point
 19 | 	GetValue(idx int) uint64
 20 | 	SubArray(begin, end int) PointArray
 21 | 	Erase(point Point) bool
 22 | 	Append(point Point)
 23 | }
 24 | 
 25 | type PointArrayMem struct {
 26 | 	points []Point
 27 | 	byDim  int
 28 | }
 29 | 
 30 | type PointArrayExt struct {
 31 | 	data        []byte
 32 | 	numPoints   int
 33 | 	byDim       int
 34 | 	bytesPerDim int
 35 | 	numDims     int
 36 | 	pointSize   int
 37 | }
 38 | 
 39 | // Compare is part of datastructures.Comparable interface
 40 | func (p Point) Compare(other datastructures.Comparable) int {
 41 | 	rhs := other.(Point)
 42 | 	for dim := 0; dim < len(p.Vals); dim++ {
 43 | 		if p.Vals[dim] != rhs.Vals[dim] {
 44 | 			return int(p.Vals[dim] - rhs.Vals[dim])
 45 | 		}
 46 | 	}
 47 | 	return int(p.UserData - rhs.UserData)
 48 | }
 49 | 
 50 | func (p *Point) Inside(lowPoint, highPoint Point) (isInside bool) {
 51 | 	for dim := 0; dim < len(p.Vals); dim++ {
 52 | 		if p.Vals[dim] < lowPoint.Vals[dim] || p.Vals[dim] > highPoint.Vals[dim] {
 53 | 			return
 54 | 		}
 55 | 	}
 56 | 	isInside = true
 57 | 	return
 58 | }
 59 | 
 60 | func (p Point) LessThan(rhs Point) (res bool) {
 61 | 	for dim := 0; dim < len(p.Vals); dim++ {
 62 | 		if p.Vals[dim] != rhs.Vals[dim] {
 63 | 			return p.Vals[dim] < rhs.Vals[dim]
 64 | 		}
 65 | 	}
 66 | 	return p.UserData < rhs.UserData
 67 | }
 68 | 
 69 | func (p *Point) Equal(rhs Point) (res bool) {
 70 | 	if p.UserData != rhs.UserData || len(p.Vals) != len(rhs.Vals) {
 71 | 		return
 72 | 	}
 73 | 	for dim := 0; dim < len(p.Vals); dim++ {
 74 | 		if p.Vals[dim] != rhs.Vals[dim] {
 75 | 			return
 76 | 		}
 77 | 	}
 78 | 	res = true
 79 | 	return
 80 | }
 81 | 
 82 | //Encode encode in place. Refers to binary.Write impl in standard library.
 83 | //len(b) shall be no less than bytesPerDim*numDims+8
 84 | func (p *Point) Encode(b []byte, bytesPerDim int) {
 85 | 	numDims := len(p.Vals)
 86 | 	for i := 0; i < numDims; i++ {
 87 | 		switch bytesPerDim {
 88 | 		case 1:
 89 | 			b[i] = byte(p.Vals[i])
 90 | 		case 2:
 91 | 			binary.BigEndian.PutUint16(b[2*i:], uint16(p.Vals[i]))
 92 | 		case 4:
 93 | 			binary.BigEndian.PutUint32(b[4*i:], uint32(p.Vals[i]))
 94 | 		case 8:
 95 | 			binary.BigEndian.PutUint64(b[8*i:], p.Vals[i])
 96 | 		}
 97 | 	}
 98 | 	binary.BigEndian.PutUint64(b[numDims*bytesPerDim:], p.UserData)
 99 | 	return
100 | }
101 | 
102 | func (p *Point) Decode(b []byte, numDims int, bytesPerDim int) {
103 | 	p.Vals = make([]uint64, numDims)
104 | 	for i := 0; i < numDims; i++ {
105 | 		switch bytesPerDim {
106 | 		case 1:
107 | 			p.Vals[i] = uint64(b[i])
108 | 		case 2:
109 | 			p.Vals[i] = uint64(binary.BigEndian.Uint16(b[2*i:]))
110 | 		case 4:
111 | 			p.Vals[i] = uint64(binary.BigEndian.Uint32(b[4*i:]))
112 | 		case 8:
113 | 			p.Vals[i] = binary.BigEndian.Uint64(b[8*i:])
114 | 		}
115 | 	}
116 | 	p.UserData = binary.BigEndian.Uint64(b[numDims*bytesPerDim:])
117 | 	return
118 | }
119 | 
120 | // Len is part of sort.Interface.
121 | func (s *PointArrayMem) Len() int {
122 | 	return len(s.points)
123 | }
124 | 
125 | // Swap is part of sort.Interface.
126 | func (s *PointArrayMem) Swap(i, j int) {
127 | 	s.points[i], s.points[j] = s.points[j], s.points[i]
128 | }
129 | 
130 | // Less is part of sort.Interface.
131 | func (s *PointArrayMem) Less(i, j int) bool {
132 | 	return s.points[i].Vals[s.byDim] < s.points[j].Vals[s.byDim]
133 | }
134 | 
135 | func (s *PointArrayMem) GetPoint(idx int) (point Point) {
136 | 	point = s.points[idx]
137 | 	return
138 | }
139 | 
140 | func (s *PointArrayMem) GetValue(idx int) (val uint64) {
141 | 	val = s.points[idx].Vals[s.byDim]
142 | 	return
143 | }
144 | 
145 | func (s *PointArrayMem) SubArray(begin, end int) (sub PointArray) {
146 | 	sub = &PointArrayMem{
147 | 		points: s.points[begin:end],
148 | 		byDim:  s.byDim,
149 | 	}
150 | 	return
151 | }
152 | 
153 | func (s *PointArrayMem) Erase(point Point) (found bool) {
154 | 	idx := 0
155 | 	for i, point2 := range s.points {
156 | 		//assumes each point's userData is unique
157 | 		if point.Equal(point2) {
158 | 			idx = i
159 | 			found = true
160 | 			break
161 | 		}
162 | 	}
163 | 	if found {
164 | 		s.points = append(s.points[:idx], s.points[idx+1:]...)
165 | 	}
166 | 	return
167 | }
168 | 
169 | func (s *PointArrayMem) Append(point Point) {
170 | 	s.points = append(s.points, point)
171 | }
172 | 
173 | func (s *PointArrayMem) ToExt(bytesPerDim int) (pae *PointArrayExt) {
174 | 	numDims := len(s.points[0].Vals)
175 | 	pointSize := numDims*bytesPerDim + 8
176 | 	size := len(s.points) * pointSize
177 | 	data := make([]byte, size)
178 | 	pae = &PointArrayExt{
179 | 		data:        data,
180 | 		numPoints:   len(s.points),
181 | 		byDim:       s.byDim,
182 | 		numDims:     numDims,
183 | 		bytesPerDim: bytesPerDim,
184 | 		pointSize:   pointSize,
185 | 	}
186 | 	for i, point := range s.points {
187 | 		point.Encode(data[i*pointSize:], bytesPerDim)
188 | 	}
189 | 	return
190 | }
191 | 
192 | // Len is part of sort.Interface.
193 | func (s *PointArrayExt) Len() int {
194 | 	return s.numPoints
195 | }
196 | 
197 | // Swap is part of sort.Interface.
198 | func (s *PointArrayExt) Swap(i, j int) {
199 | 	offI := i * s.pointSize
200 | 	offJ := j * s.pointSize
201 | 	for idx := 0; idx < s.pointSize; idx++ {
202 | 		s.data[offI+idx], s.data[offJ+idx] = s.data[offJ+idx], s.data[offI+idx]
203 | 	}
204 | }
205 | 
206 | // Less is part of sort.Interface.
207 | func (s *PointArrayExt) Less(i, j int) bool {
208 | 	valI := s.GetValue(i)
209 | 	valJ := s.GetValue(j)
210 | 	return valI < valJ
211 | }
212 | 
213 | func (s *PointArrayExt) GetPoint(i int) (point Point) {
214 | 	point.Decode(s.data[i*s.pointSize:], s.numDims, s.bytesPerDim)
215 | 	return
216 | }
217 | 
218 | func (s *PointArrayExt) GetValue(i int) (val uint64) {
219 | 	offI := i * s.pointSize
220 | 	switch s.bytesPerDim {
221 | 	case 1:
222 | 		val = uint64(s.data[offI+s.byDim])
223 | 	case 2:
224 | 		val = uint64(binary.BigEndian.Uint16(s.data[offI+2*s.byDim:]))
225 | 	case 4:
226 | 		val = uint64(binary.BigEndian.Uint32(s.data[offI+4*s.byDim:]))
227 | 	case 8:
228 | 		val = binary.BigEndian.Uint64(s.data[offI+8*s.byDim:])
229 | 	}
230 | 	return
231 | }
232 | 
233 | func (s *PointArrayExt) SubArray(begin, end int) (sub PointArray) {
234 | 	sub = &PointArrayExt{
235 | 		data:        s.data[begin*s.pointSize : end*s.pointSize],
236 | 		numPoints:   end - begin,
237 | 		byDim:       s.byDim,
238 | 		bytesPerDim: s.bytesPerDim,
239 | 		numDims:     s.numDims,
240 | 		pointSize:   s.pointSize,
241 | 	}
242 | 	return
243 | }
244 | 
245 | func (s *PointArrayExt) Erase(point Point) (found bool) {
246 | 	var i int
247 | 	for i = 0; i < s.numPoints; i++ {
248 | 		pI := s.GetPoint(i)
249 | 		//assumes each point's userData is unique
250 | 		found = point.Equal(pI)
251 | 		if found {
252 | 			break
253 | 		}
254 | 	}
255 | 	if found {
256 | 		//replace the matched point with the last point and decrease the array length
257 | 		offI := i * s.pointSize
258 | 		offJ := (s.numPoints - 1) * s.pointSize
259 | 		for idx := 0; idx < s.pointSize; idx++ {
260 | 			s.data[offI+idx] = s.data[offJ+idx]
261 | 			s.data[offJ+idx] = 0
262 | 		}
263 | 		s.numPoints--
264 | 	}
265 | 	return
266 | }
267 | 
268 | func (s *PointArrayExt) Append(point Point) {
269 | 	off := s.numPoints * s.pointSize
270 | 	point.Encode(s.data[off:], s.bytesPerDim)
271 | }
272 | 
273 | func (s *PointArrayExt) ToMem() (pam *PointArrayMem) {
274 | 	points := make([]Point, s.numPoints)
275 | 	for i := 0; i < s.numPoints; i++ {
276 | 		p := s.GetPoint(i)
277 | 		points[i] = p
278 | 	}
279 | 	pam = &PointArrayMem{
280 | 		points: points,
281 | 		byDim:  s.byDim,
282 | 	}
283 | 	return
284 | }
285 | 
286 | // SplitPoints splits points per byDim
287 | func SplitPoints(points PointArray, numStrips int) (splitValues []uint64, splitPoses []int) {
288 | 	if numStrips <= 1 {
289 | 		return
290 | 	}
291 | 	splitPos := points.Len() / 2
292 | 	nth.Element(points, splitPos)
293 | 	splitValue := points.GetValue(splitPos)
294 | 
295 | 	numStrips1 := (numStrips + 1) / 2
296 | 	numStrips2 := numStrips - numStrips1
297 | 	splitValues1, splitPoses1 := SplitPoints(points.SubArray(0, splitPos), numStrips1)
298 | 	splitValues = append(splitValues, splitValues1...)
299 | 	splitPoses = append(splitPoses, splitPoses1...)
300 | 	splitValues = append(splitValues, splitValue)
301 | 	splitPoses = append(splitPoses, splitPos)
302 | 	splitValues2, splitPoses2 := SplitPoints(points.SubArray(splitPos, points.Len()), numStrips2)
303 | 	splitValues = append(splitValues, splitValues2...)
304 | 	for i := 0; i < len(splitPoses2); i++ {
305 | 		splitPoses = append(splitPoses, splitPos+splitPoses2[i])
306 | 	}
307 | 	return
308 | }
309 | 


--------------------------------------------------------------------------------
/point_test.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"bytes"
  5 | 	"math/rand"
  6 | 	"os"
  7 | 	"testing"
  8 | )
  9 | 
 10 | type CaseInside struct {
 11 | 	point, lowPoint, highPoint Point
 12 | 	numDims                    int
 13 | 	isInside                   bool
 14 | }
 15 | 
 16 | func TestIsInside(t *testing.T) {
 17 | 	cases := []CaseInside{
 18 | 		{
 19 | 			Point{[]uint64{30, 80, 40}, 0},
 20 | 			Point{[]uint64{30, 80, 40}, 0},
 21 | 			Point{[]uint64{50, 90, 50}, 0},
 22 | 			3,
 23 | 			true,
 24 | 		},
 25 | 		{
 26 | 			Point{[]uint64{30, 79, 40}, 0},
 27 | 			Point{[]uint64{30, 80, 40}, 0},
 28 | 			Point{[]uint64{50, 90, 50}, 0},
 29 | 			3,
 30 | 			false,
 31 | 		},
 32 | 		{ //invalid range
 33 | 			Point{[]uint64{30, 80, 40}, 0},
 34 | 			Point{[]uint64{30, 80, 40}, 0},
 35 | 			Point{[]uint64{50, 90, 39}, 0},
 36 | 			3,
 37 | 			false,
 38 | 		},
 39 | 	}
 40 | 
 41 | 	for i, tc := range cases {
 42 | 		res := tc.point.Inside(tc.lowPoint, tc.highPoint)
 43 | 		if res != tc.isInside {
 44 | 			t.Fatalf("case %v failed\n", i)
 45 | 		}
 46 | 	}
 47 | }
 48 | 
 49 | type CaseCodec struct {
 50 | 	point       Point
 51 | 	numDims     int
 52 | 	bytesPerDim int
 53 | 	bytesP      []byte
 54 | }
 55 | 
 56 | func TestPointCodec(t *testing.T) {
 57 | 	cases := []CaseCodec{
 58 | 		{
 59 | 			Point{[]uint64{6, 92, 68}, 8},
 60 | 			3,
 61 | 			4,
 62 | 			[]byte{0x0, 0x0, 0x0, 0x6, 0x0, 0x0, 0x0, 0x5c, 0x0, 0x0, 0x0, 0x44, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x8},
 63 | 		},
 64 | 		{
 65 | 			Point{[]uint64{6, 92, 68}, 256},
 66 | 			3,
 67 | 			4,
 68 | 			[]byte{0x0, 0x0, 0x0, 0x6, 0x0, 0x0, 0x0, 0x5c, 0x0, 0x0, 0x0, 0x44, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x0},
 69 | 		},
 70 | 	}
 71 | 
 72 | 	var point Point
 73 | 	for i, tc := range cases {
 74 | 		b := make([]byte, tc.numDims*tc.bytesPerDim+8)
 75 | 		tc.point.Encode(b, tc.bytesPerDim)
 76 | 		if !bytes.Equal(b, tc.bytesP) {
 77 | 			t.Fatalf("point %d Encode as %v", i, b)
 78 | 		}
 79 | 		point.Decode(tc.bytesP, tc.numDims, tc.bytesPerDim)
 80 | 		if !point.Equal(tc.point) {
 81 | 			t.Fatalf("point %d Decode as %v", i, point)
 82 | 		}
 83 | 	}
 84 | }
 85 | 
 86 | func NewRandPoints(numDims int, maxVal uint64, size int) (points []Point) {
 87 | 	for i := 0; i < size; i++ {
 88 | 		vals := make([]uint64, 0, numDims)
 89 | 		for j := 0; j < numDims; j++ {
 90 | 			vals = append(vals, rand.Uint64()%maxVal)
 91 | 		}
 92 | 		point := Point{vals, uint64(i)}
 93 | 		points = append(points, point)
 94 | 	}
 95 | 	return
 96 | }
 97 | 
 98 | func TestPointArrayExt_ToMem(t *testing.T) {
 99 | 	numDims := 3
100 | 	maxVal := uint64(100)
101 | 	size := 10000
102 | 	points := NewRandPoints(numDims, maxVal, size)
103 | 	pam := PointArrayMem{
104 | 		points: points,
105 | 		byDim:  1,
106 | 	}
107 | 
108 | 	bytesPerDim := 4
109 | 	pae := pam.ToExt(bytesPerDim)
110 | 	pam2 := pae.ToMem()
111 | 	if pam.byDim != pam2.byDim {
112 | 		t.Fatalf("point array meta info changes after convertion")
113 | 	}
114 | 	if len(pam.points) != len(pam2.points) {
115 | 		t.Fatalf("point array length changes after convertion: %d %d", len(pam.points), len(pam2.points))
116 | 	}
117 | 	for i := 0; i < len(pam.points); i++ {
118 | 		p1, p2 := pam.points[i], pam2.points[i]
119 | 		if !p1.Equal(p2) {
120 | 			t.Fatalf("point %d changes after convertion: %v %v", i, p1, p2)
121 | 		}
122 | 	}
123 | }
124 | 
125 | //verify if lhs and rhs contains the same points. order doesn't matter.
126 | func areSmaePoints(lhs, rhs []Point, numDims int) (res bool) {
127 | 	if len(lhs) != len(rhs) {
128 | 		return
129 | 	}
130 | 	numPoints := len(lhs)
131 | 	mapLhs := make(map[uint64]Point, numPoints)
132 | 	mapRhs := make(map[uint64]Point, numPoints)
133 | 	for i := 0; i < numPoints; i++ {
134 | 		mapLhs[lhs[i].UserData] = lhs[i]
135 | 		mapRhs[rhs[i].UserData] = rhs[i]
136 | 	}
137 | 	for k, v := range mapLhs {
138 | 		v2, found := mapRhs[k]
139 | 		if !found || !v.Equal(v2) {
140 | 			return
141 | 		}
142 | 	}
143 | 	res = true
144 | 	return
145 | }
146 | 
147 | func verifySplit(t *testing.T, pam *PointArrayMem, numStrips int, splitValues []uint64, splitPoses []int) {
148 | 	//fmt.Printf("points: %v\nsplitValues: %v\nsplitPoses:%v\n", points, splitValues, splitPoses)
149 | 	if len(splitValues) != numStrips-1 || len(splitValues) != len(splitPoses) {
150 | 		t.Fatalf("incorrect size of splitValues or splitPoses\n")
151 | 	}
152 | 	for i := 0; i < len(splitValues)-1; i++ {
153 | 		if splitValues[i] > splitValues[i+1] {
154 | 			t.Fatalf("incorrect splitValues\n")
155 | 		}
156 | 		if splitPoses[i] >= splitPoses[i+1] {
157 | 			t.Fatalf("incorrect splitPoses\n")
158 | 		}
159 | 	}
160 | 	numSplits := len(splitValues)
161 | 	for strip := 0; strip < numStrips; strip++ {
162 | 		posBegin := 0
163 | 		minValue := uint64(0)
164 | 		if strip != 0 {
165 | 			posBegin = splitPoses[strip-1]
166 | 			minValue = splitValues[strip-1]
167 | 		}
168 | 		posEnd := len(pam.points)
169 | 		maxValue := ^uint64(0)
170 | 		if strip != numSplits {
171 | 			posEnd = splitPoses[strip]
172 | 			maxValue = splitValues[strip]
173 | 		}
174 | 
175 | 		for pos := posBegin; pos < posEnd; pos++ {
176 | 			val := pam.points[pos].Vals[pam.byDim]
177 | 			if val < minValue || val > maxValue {
178 | 				t.Fatalf("points[%v] dim %v val %v is not in range %v-%v", pos, pam.byDim, val, minValue, maxValue)
179 | 			}
180 | 		}
181 | 	}
182 | 	return
183 | }
184 | 
185 | func TestSplitPoints(t *testing.T) {
186 | 	//TODO: use suite setup to initialize points
187 | 	numDims := 3
188 | 	maxVal := uint64(100)
189 | 	size := 10000
190 | 	numStrips := 4
191 | 	points := NewRandPoints(numDims, maxVal, size)
192 | 	pointsSaved := make([]Point, size)
193 | 	copy(pointsSaved, points)
194 | 	//test SplitPoints(PointArrayMem)
195 | 	for dim := 0; dim < numDims; dim++ {
196 | 		pam := &PointArrayMem{
197 | 			points: points,
198 | 			byDim:  dim,
199 | 		}
200 | 		splitValues, splitPoses := SplitPoints(pam, numStrips)
201 | 		verifySplit(t, pam, numStrips, splitValues, splitPoses)
202 | 		if !areSmaePoints(pointsSaved, pam.points, numDims) {
203 | 			t.Fatalf("point set changes after split")
204 | 		}
205 | 	}
206 | 
207 | 	//test SplitPoints(PointArrayExt)
208 | 	bytesPerDim := 4
209 | 	pam := &PointArrayMem{
210 | 		points: points,
211 | 		byDim:  0,
212 | 	}
213 | 	pae := pam.ToExt(bytesPerDim)
214 | 	for dim := 0; dim < numDims; dim++ {
215 | 		pae.byDim = dim
216 | 		splitValues, splitPoses := SplitPoints(pae, numStrips)
217 | 		pam2 := pae.ToMem()
218 | 		verifySplit(t, pam2, numStrips, splitValues, splitPoses)
219 | 		if !areSmaePoints(pam.points, pam2.points, numDims) {
220 | 			t.Fatalf("point set changes after split")
221 | 		}
222 | 	}
223 | 
224 | 	//test SplitPoints(PointArrayExt) on external temp file
225 | 	tmpF, err := os.OpenFile("/tmp/point_test", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
226 | 	if err != nil {
227 | 		t.Fatalf("%+v", err)
228 | 	}
229 | 	defer tmpF.Close()
230 | 	_, err = tmpF.Write(pae.data)
231 | 	if err != nil {
232 | 		t.Fatalf("%+v", err)
233 | 	}
234 | 	data, err := FileMmap(tmpF)
235 | 	if err != nil {
236 | 		t.Fatalf("%+v", err)
237 | 	}
238 | 	defer FileMunmap(data)
239 | 	pae.data = data
240 | 	for dim := 0; dim < numDims; dim++ {
241 | 		pae.byDim = dim
242 | 		splitValues, splitPoses := SplitPoints(pae, numStrips)
243 | 		pam2 := pae.ToMem()
244 | 		verifySplit(t, pam2, numStrips, splitValues, splitPoses)
245 | 		if !areSmaePoints(pam.points, pam2.points, numDims) {
246 | 			t.Fatalf("point set changes after split")
247 | 		}
248 | 	}
249 | }
250 | 


--------------------------------------------------------------------------------
/utils.go:
--------------------------------------------------------------------------------
  1 | package bkdtree
  2 | 
  3 | import (
  4 | 	"encoding/json"
  5 | 	"os"
  6 | 	"path/filepath"
  7 | 	"regexp"
  8 | 	"sort"
  9 | 	"syscall"
 10 | 
 11 | 	"github.com/pkg/errors"
 12 | )
 13 | 
 14 | //FileMmap mmaps the given file.
 15 | //https://medium.com/@arpith/adventures-with-mmap-463b33405223
 16 | func FileMmap(f *os.File) (data []byte, err error) {
 17 | 	info, err1 := f.Stat()
 18 | 	if err1 != nil {
 19 | 		err = errors.Wrap(err1, "")
 20 | 		return
 21 | 	}
 22 | 	prots := []int{syscall.PROT_WRITE | syscall.PROT_READ, syscall.PROT_READ}
 23 | 	for _, prot := range prots {
 24 | 		data, err = syscall.Mmap(int(f.Fd()), 0, int(info.Size()), prot, syscall.MAP_SHARED)
 25 | 		if err == nil {
 26 | 			break
 27 | 		}
 28 | 	}
 29 | 	if err != nil {
 30 | 		err = errors.Wrap(err, "")
 31 | 		return
 32 | 	}
 33 | 	return
 34 | }
 35 | 
 36 | //FileMunmap unmaps the given file.
 37 | func FileMunmap(data []byte) (err error) {
 38 | 	err = syscall.Munmap(data)
 39 | 	if err != nil {
 40 | 		err = errors.Wrap(err, "")
 41 | 		return
 42 | 	}
 43 | 	return
 44 | }
 45 | 
 46 | //FileUnmarshal unmarshals the given file to object.
 47 | func FileUnmarshal(fp string, v interface{}) (err error) {
 48 | 	var f *os.File
 49 | 	var data []byte
 50 | 	if f, err = os.Open(fp); err != nil {
 51 | 		return
 52 | 	}
 53 | 	defer f.Close()
 54 | 	if data, err = FileMmap(f); err != nil {
 55 | 		return
 56 | 	}
 57 | 	defer FileMunmap(data)
 58 | 	err = json.Unmarshal(data, v)
 59 | 	return
 60 | }
 61 | 
 62 | //FileMarshal marshals the given object to file.
 63 | func FileMarshal(fp string, v interface{}) (err error) {
 64 | 	var f *os.File
 65 | 	var data []byte
 66 | 	var count int
 67 | 	if data, err = json.Marshal(v); err != nil {
 68 | 		return
 69 | 	}
 70 | 	if f, err = os.OpenFile(fp, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600); err != nil {
 71 | 		err = errors.Wrap(err, "")
 72 | 		return
 73 | 	}
 74 | 	defer f.Close()
 75 | 	if count, err = f.Write(data); err != nil {
 76 | 		err = errors.Wrap(err, "")
 77 | 		return
 78 | 	}
 79 | 	if count != len(data) {
 80 | 		err = errors.Errorf("%s partial wirte %d, want %d", fp, count, len(data))
 81 | 		return
 82 | 	}
 83 | 	return
 84 | }
 85 | 
 86 | //FilepathGlob is enhanced version of standard path.filepath::Glob()
 87 | func FilepathGlob(dir, patt string) (matches [][]string, err error) {
 88 | 	var d *os.File
 89 | 	var fns []string
 90 | 	d, err = os.Open(dir)
 91 | 	if os.IsNotExist(err) {
 92 | 		err = nil
 93 | 		return
 94 | 	}
 95 | 	defer d.Close()
 96 | 	fns, err = d.Readdirnames(0)
 97 | 	if err != nil {
 98 | 		err = errors.Wrap(err, "")
 99 | 		return
100 | 	}
101 | 	sort.Strings(fns)
102 | 	re := regexp.MustCompile(patt)
103 | 	for _, fn := range fns {
104 | 		subs := re.FindStringSubmatch(fn)
105 | 		if subs == nil {
106 | 			continue
107 | 		}
108 | 		matches = append(matches, subs)
109 | 	}
110 | 	return
111 | }
112 | 
113 | //FilepathGlobRm remove given files under the given directory.
114 | func FilepathGlobRm(dir, patt string) (err error) {
115 | 	var matches [][]string
116 | 	if matches, err = FilepathGlob(dir, patt); err != nil {
117 | 		return
118 | 	}
119 | 	for _, match := range matches {
120 | 		fp := filepath.Join(dir, match[0])
121 | 		if err = os.Remove(fp); err != nil {
122 | 			err = errors.Wrap(err, "")
123 | 		}
124 | 	}
125 | 	return
126 | }
127 | 


--------------------------------------------------------------------------------