mirror of
https://github.com/lionsoul2014/ip2region.git
synced 2025-12-08 19:25:22 +00:00
add vector index and xdb content load util func
This commit is contained in:
parent
db4aa26547
commit
059446caa7
@ -54,52 +54,6 @@ func (s *Searcher) Close() {
|
||||
}
|
||||
}
|
||||
|
||||
// LoadVectorIndex load and cache the vector index for search speedup.
|
||||
// this will take up VectorIndexRows x VectorIndexCols x VectorIndexSize bytes memory.
|
||||
func (s *Searcher) LoadVectorIndex() error {
|
||||
// loaded already
|
||||
if s.vectorIndex != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// load all the vector index block
|
||||
_, err := s.handle.Seek(HeaderInfoLength, 0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("seek to vector index: %w", err)
|
||||
}
|
||||
|
||||
var buff = make([]byte, VectorIndexRows*VectorIndexCols*VectorIndexSize)
|
||||
rLen, err := s.handle.Read(buff)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if rLen != len(buff) {
|
||||
return fmt.Errorf("incomplete read: readed bytes should be %d", len(buff))
|
||||
}
|
||||
|
||||
// decode the vector index blocks
|
||||
var vectorIndex = make([][]*VectorIndexBlock, VectorIndexRows)
|
||||
for r := 0; r < VectorIndexRows; r++ {
|
||||
vectorIndex[r] = make([]*VectorIndexBlock, VectorIndexCols)
|
||||
for c := 0; c < VectorIndexCols; c++ {
|
||||
offset := r*VectorIndexCols*VectorIndexSize + c*VectorIndexSize
|
||||
vectorIndex[r][c], err = VectorIndexBlockDecode(buff[offset:])
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode vector index at [%d][%d]: %w", r, c, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
s.vectorIndex = vectorIndex
|
||||
return nil
|
||||
}
|
||||
|
||||
// ClearVectorIndex clear preloaded vector index cache
|
||||
func (s *Searcher) ClearVectorIndex() {
|
||||
s.vectorIndex = nil
|
||||
}
|
||||
|
||||
// SearchByStr find the region for the specified ip string
|
||||
func (s *Searcher) SearchByStr(str string) (string, error) {
|
||||
ip, err := CheckIP(str)
|
||||
|
||||
@ -3,6 +3,7 @@ package xdb
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
@ -43,3 +44,86 @@ func Long2IP(ip uint32) string {
|
||||
func MidIP(sip uint32, eip uint32) uint32 {
|
||||
return uint32((uint64(sip) + uint64(eip)) >> 1)
|
||||
}
|
||||
|
||||
// LoadVectorIndex util function to load the vector index from the specified file handle
|
||||
func LoadVectorIndex(handle *os.File) ([][]*VectorIndexBlock, error) {
|
||||
// load all the vector index block
|
||||
_, err := handle.Seek(HeaderInfoLength, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("seek to vector index: %w", err)
|
||||
}
|
||||
|
||||
var buff = make([]byte, VectorIndexRows*VectorIndexCols*VectorIndexSize)
|
||||
rLen, err := handle.Read(buff)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if rLen != len(buff) {
|
||||
return nil, fmt.Errorf("incomplete read: readed bytes should be %d", len(buff))
|
||||
}
|
||||
|
||||
// decode the vector index blocks
|
||||
var vectorIndex = make([][]*VectorIndexBlock, VectorIndexRows)
|
||||
for r := 0; r < VectorIndexRows; r++ {
|
||||
vectorIndex[r] = make([]*VectorIndexBlock, VectorIndexCols)
|
||||
for c := 0; c < VectorIndexCols; c++ {
|
||||
offset := r*VectorIndexCols*VectorIndexSize + c*VectorIndexSize
|
||||
vectorIndex[r][c], err = VectorIndexBlockDecode(buff[offset:])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode vector index at [%d][%d]: %w", r, c, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return vectorIndex, nil
|
||||
}
|
||||
|
||||
// LoadVectorIndexFromFile load vector index from a specified file path
|
||||
func LoadVectorIndexFromFile(dbFile string) ([][]*VectorIndexBlock, error) {
|
||||
handle, err := os.OpenFile(dbFile, os.O_RDONLY, 0600)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open xdb file `%s`: %w", dbFile, err)
|
||||
}
|
||||
|
||||
return LoadVectorIndex(handle)
|
||||
}
|
||||
|
||||
// LoadContent load the whole xdb content from the specified file handle
|
||||
func LoadContent(handle *os.File) ([]byte, error) {
|
||||
// get file size
|
||||
fi, err := handle.Stat()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("stat: %w", err)
|
||||
}
|
||||
|
||||
size := fi.Size()
|
||||
|
||||
// seek to the head of the file
|
||||
_, err = handle.Seek(0, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("seek to get xdb file length: %w", err)
|
||||
}
|
||||
|
||||
var buff = make([]byte, size)
|
||||
rLen, err := handle.Read(buff)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if rLen != len(buff) {
|
||||
return nil, fmt.Errorf("incomplete read: readed bytes should be %d", len(buff))
|
||||
}
|
||||
|
||||
return buff, nil
|
||||
}
|
||||
|
||||
// LoadContentFromFile load the whole xdb content from the specified db file path
|
||||
func LoadContentFromFile(dbFile string) ([]byte, error) {
|
||||
handle, err := os.OpenFile(dbFile, os.O_RDONLY, 0600)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open xdb file `%s`: %w", dbFile, err)
|
||||
}
|
||||
|
||||
return LoadContent(handle)
|
||||
}
|
||||
|
||||
29
binding/golang/xdb/util_test.go
Normal file
29
binding/golang/xdb/util_test.go
Normal file
@ -0,0 +1,29 @@
|
||||
// Copyright 2022 The Ip2Region Authors. All rights reserved.
|
||||
// Use of this source code is governed by a Apache2.0-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package xdb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadVectorIndex(t *testing.T) {
|
||||
vIndex, err := LoadVectorIndexFromFile("../../../data/ip2region.xdb")
|
||||
if err != nil {
|
||||
fmt.Printf("failed to load vector index: %s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("vIndex length: %d\n", len(vIndex))
|
||||
}
|
||||
|
||||
func TestLoadContent(t *testing.T) {
|
||||
buff, err := LoadContentFromFile("../../../data/ip2region.xdb")
|
||||
if err != nil {
|
||||
fmt.Printf("failed to load xdb content: %s\n", err)
|
||||
}
|
||||
|
||||
fmt.Printf("buff length: %d\n", len(buff))
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user