Skip to content

Commit

Permalink
feat: support binary query interface
Browse files Browse the repository at this point in the history
  • Loading branch information
@huskar-t
huskar-t committed May 31, 2024
1 parent 2532a8e commit 65f96bf
Show file tree
Hide file tree
Showing 8 changed files with 290 additions and 271 deletions.
126 changes: 57 additions & 69 deletions common/parser/block.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package parser
import (
"database/sql/driver"
"math"
"unicode/utf8"
"unsafe"

"github.com/taosdata/driver-go/v3/common"
Expand Down Expand Up @@ -108,28 +109,9 @@ type rawConvertFunc func(pStart unsafe.Pointer, row int, arg ...interface{}) dri

type rawConvertVarDataFunc func(pHeader, pStart unsafe.Pointer, row int) driver.Value

var rawConvertFuncMap = map[uint8]rawConvertFunc{
uint8(common.TSDB_DATA_TYPE_BOOL): rawConvertBool,
uint8(common.TSDB_DATA_TYPE_TINYINT): rawConvertTinyint,
uint8(common.TSDB_DATA_TYPE_SMALLINT): rawConvertSmallint,
uint8(common.TSDB_DATA_TYPE_INT): rawConvertInt,
uint8(common.TSDB_DATA_TYPE_BIGINT): rawConvertBigint,
uint8(common.TSDB_DATA_TYPE_UTINYINT): rawConvertUTinyint,
uint8(common.TSDB_DATA_TYPE_USMALLINT): rawConvertUSmallint,
uint8(common.TSDB_DATA_TYPE_UINT): rawConvertUInt,
uint8(common.TSDB_DATA_TYPE_UBIGINT): rawConvertUBigint,
uint8(common.TSDB_DATA_TYPE_FLOAT): rawConvertFloat,
uint8(common.TSDB_DATA_TYPE_DOUBLE): rawConvertDouble,
uint8(common.TSDB_DATA_TYPE_TIMESTAMP): rawConvertTime,
}

var rawConvertVarDataMap = map[uint8]rawConvertVarDataFunc{
uint8(common.TSDB_DATA_TYPE_BINARY): rawConvertBinary,
uint8(common.TSDB_DATA_TYPE_NCHAR): rawConvertNchar,
uint8(common.TSDB_DATA_TYPE_JSON): rawConvertJson,
uint8(common.TSDB_DATA_TYPE_VARBINARY): rawConvertVarBinary,
uint8(common.TSDB_DATA_TYPE_GEOMETRY): rawConvertGeometry,
}
var rawConvertFuncSlice = [15]rawConvertFunc{}

var rawConvertVarDataSlice = [21]rawConvertVarDataFunc{}

func ItemIsNull(pHeader unsafe.Pointer, row int) bool {
offset := CharOffset(row)
Expand Down Expand Up @@ -196,41 +178,35 @@ func rawConvertTime(pStart unsafe.Pointer, row int, arg ...interface{}) driver.V
}

func rawConvertVarBinary(pHeader, pStart unsafe.Pointer, row int) driver.Value {
return rawGetBytes(pHeader, pStart, row)
}

func rawGetBytes(pHeader, pStart unsafe.Pointer, row int) []byte {
offset := *((*int32)(pointer.AddUintptr(pHeader, uintptr(row*4))))
if offset == -1 {
return nil
}
currentRow := pointer.AddUintptr(pStart, uintptr(offset))
clen := *((*uint16)(currentRow))
currentRow = unsafe.Pointer(uintptr(currentRow) + 2)

binaryVal := make([]byte, clen)

for index := uint16(0); index < clen; index++ {
binaryVal[index] = *((*byte)(unsafe.Pointer(uintptr(currentRow) + uintptr(index))))
if clen == 0 {
return make([]byte, 0)
}
return binaryVal[:]
currentRow = pointer.AddUintptr(currentRow, 2)
result := make([]byte, clen)
Copy(currentRow, result, 0, int(clen))
return result
}

func rawConvertGeometry(pHeader, pStart unsafe.Pointer, row int) driver.Value {
return rawConvertVarBinary(pHeader, pStart, row)
}

func rawConvertBinary(pHeader, pStart unsafe.Pointer, row int) driver.Value {
offset := *((*int32)(pointer.AddUintptr(pHeader, uintptr(row*4))))
if offset == -1 {
result := rawGetBytes(pHeader, pStart, row)
if result == nil {
return nil
}
currentRow := pointer.AddUintptr(pStart, uintptr(offset))
clen := *((*uint16)(currentRow))
currentRow = unsafe.Pointer(uintptr(currentRow) + 2)

binaryVal := make([]byte, clen)

for index := uint16(0); index < clen; index++ {
binaryVal[index] = *((*byte)(unsafe.Pointer(uintptr(currentRow) + uintptr(index))))
}
return string(binaryVal[:])
return *(*string)(unsafe.Pointer(&result))
}

func rawConvertNchar(pHeader, pStart unsafe.Pointer, row int) driver.Value {
Expand All @@ -240,31 +216,22 @@ func rawConvertNchar(pHeader, pStart unsafe.Pointer, row int) driver.Value {
}
currentRow := pointer.AddUintptr(pStart, uintptr(offset))
clen := *((*uint16)(currentRow)) / 4
if clen == 0 {
return ""
}
currentRow = unsafe.Pointer(uintptr(currentRow) + 2)

binaryVal := make([]rune, clen)

for index := uint16(0); index < clen; index++ {
binaryVal[index] = *((*rune)(unsafe.Pointer(uintptr(currentRow) + uintptr(index*4))))
utf8Bytes := make([]byte, clen*utf8.UTFMax)
index := 0
utf32Slice := (*[1 << 30]rune)(currentRow)[:clen:clen]
for _, runeValue := range utf32Slice {
index += utf8.EncodeRune(utf8Bytes[index:], runeValue)
}
return string(binaryVal)
utf8Bytes = utf8Bytes[:index]
return *(*string)(unsafe.Pointer(&utf8Bytes))
}

func rawConvertJson(pHeader, pStart unsafe.Pointer, row int) driver.Value {
offset := *((*int32)(pointer.AddUintptr(pHeader, uintptr(row*4))))
if offset == -1 {
return nil
}
currentRow := pointer.AddUintptr(pStart, uintptr(offset))
clen := *((*uint16)(currentRow))
currentRow = pointer.AddUintptr(currentRow, 2)

binaryVal := make([]byte, clen)

for index := uint16(0); index < clen; index++ {
binaryVal[index] = *((*byte)(pointer.AddUintptr(currentRow, uintptr(index))))
}
return binaryVal[:]
return rawConvertVarBinary(pHeader, pStart, row)
}

func ReadBlockSimple(block unsafe.Pointer, precision int) [][]driver.Value {
Expand All @@ -290,7 +257,7 @@ func ReadBlock(block unsafe.Pointer, blockSize int, colTypes []uint8, precision
for column := 0; column < colCount; column++ {
colLength := *((*int32)(pointer.AddUintptr(block, lengthOffset+uintptr(column)*Int32Size)))
if IsVarDataType(colTypes[column]) {
convertF := rawConvertVarDataMap[colTypes[column]]
convertF := rawConvertVarDataSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, Int32Size*uintptr(blockSize))
for row := 0; row < blockSize; row++ {
if column == 0 {
Expand All @@ -299,7 +266,7 @@ func ReadBlock(block unsafe.Pointer, blockSize int, colTypes []uint8, precision
r[row][column] = convertF(pHeader, pStart, row)
}
} else {
convertF := rawConvertFuncMap[colTypes[column]]
convertF := rawConvertFuncSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, nullBitMapOffset)
for row := 0; row < blockSize; row++ {
if column == 0 {
Expand All @@ -326,11 +293,11 @@ func ReadRow(dest []driver.Value, block unsafe.Pointer, blockSize int, row int,
for column := 0; column < colCount; column++ {
colLength := *((*int32)(pointer.AddUintptr(block, lengthOffset+uintptr(column)*Int32Size)))
if IsVarDataType(colTypes[column]) {
convertF := rawConvertVarDataMap[colTypes[column]]
convertF := rawConvertVarDataSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, Int32Size*uintptr(blockSize))
dest[column] = convertF(pHeader, pStart, row)
} else {
convertF := rawConvertFuncMap[colTypes[column]]
convertF := rawConvertFuncSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, nullBitMapOffset)
if ItemIsNull(pHeader, row) {
dest[column] = nil
Expand All @@ -352,7 +319,7 @@ func ReadBlockWithTimeFormat(block unsafe.Pointer, blockSize int, colTypes []uin
for column := 0; column < colCount; column++ {
colLength := *((*int32)(pointer.AddUintptr(block, lengthOffset+uintptr(column)*Int32Size)))
if IsVarDataType(colTypes[column]) {
convertF := rawConvertVarDataMap[colTypes[column]]
convertF := rawConvertVarDataSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, uintptr(4*blockSize))
for row := 0; row < blockSize; row++ {
if column == 0 {
Expand All @@ -361,7 +328,7 @@ func ReadBlockWithTimeFormat(block unsafe.Pointer, blockSize int, colTypes []uin
r[row][column] = convertF(pHeader, pStart, row)
}
} else {
convertF := rawConvertFuncMap[colTypes[column]]
convertF := rawConvertFuncSlice[colTypes[column]]
pStart = pointer.AddUintptr(pHeader, nullBitMapOffset)
for row := 0; row < blockSize; row++ {
if column == 0 {
Expand All @@ -381,12 +348,33 @@ func ReadBlockWithTimeFormat(block unsafe.Pointer, blockSize int, colTypes []uin

func ItemRawBlock(colType uint8, pHeader, pStart unsafe.Pointer, row int, precision int, timeFormat FormatTimeFunc) driver.Value {
if IsVarDataType(colType) {
return rawConvertVarDataMap[colType](pHeader, pStart, row)
return rawConvertVarDataSlice[colType](pHeader, pStart, row)
} else {
if ItemIsNull(pHeader, row) {
return nil
} else {
return rawConvertFuncMap[colType](pStart, row, precision, timeFormat)
return rawConvertFuncSlice[colType](pStart, row, precision, timeFormat)
}
}
}

func init() {
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_BOOL)] = rawConvertBool
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_TINYINT)] = rawConvertTinyint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_SMALLINT)] = rawConvertSmallint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_INT)] = rawConvertInt
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_BIGINT)] = rawConvertBigint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_UTINYINT)] = rawConvertUTinyint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_USMALLINT)] = rawConvertUSmallint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_UINT)] = rawConvertUInt
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_UBIGINT)] = rawConvertUBigint
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_FLOAT)] = rawConvertFloat
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_DOUBLE)] = rawConvertDouble
rawConvertFuncSlice[uint8(common.TSDB_DATA_TYPE_TIMESTAMP)] = rawConvertTime

rawConvertVarDataSlice[uint8(common.TSDB_DATA_TYPE_BINARY)] = rawConvertBinary
rawConvertVarDataSlice[uint8(common.TSDB_DATA_TYPE_NCHAR)] = rawConvertNchar
rawConvertVarDataSlice[uint8(common.TSDB_DATA_TYPE_JSON)] = rawConvertJson
rawConvertVarDataSlice[uint8(common.TSDB_DATA_TYPE_VARBINARY)] = rawConvertVarBinary
rawConvertVarDataSlice[uint8(common.TSDB_DATA_TYPE_GEOMETRY)] = rawConvertGeometry
}
12 changes: 12 additions & 0 deletions common/parser/mem.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package parser

import "unsafe"

//go:noescape
func memmove(to, from unsafe.Pointer, n uintptr)

//go:linkname memmove runtime.memmove

func Copy(source unsafe.Pointer, data []byte, index int, length int) {
memmove(unsafe.Pointer(&data[index]), source, uintptr(length))
}
Empty file added common/parser/mem.s
View file
Empty file.
20 changes: 20 additions & 0 deletions common/parser/mem_test.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
package parser

import (
"testing"
"unsafe"

"github.com/stretchr/testify/assert"
)

func TestCopy(t *testing.T) {
data := []byte("World")
data1 := make([]byte, 10)
data1[0] = 'H'
data1[1] = 'e'
data1[2] = 'l'
data1[3] = 'l'
data1[4] = 'o'
Copy(unsafe.Pointer(&data[0]), data1, 5, 5)
assert.Equal(t, "HelloWorld", string(data1))
}
2 changes: 0 additions & 2 deletions taosSql/rows.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ type rows struct {
block unsafe.Pointer
blockOffset int
blockSize int
lengthList []int
result unsafe.Pointer
precision int
isStmt bool
Expand Down Expand Up @@ -107,7 +106,6 @@ func (rs *rows) taosFetchBlock() error {
}
rs.blockSize = result.N
rs.block = wrapper.TaosGetRawBlock(result.Res)
rs.lengthList = wrapper.FetchLengths(rs.result, len(rs.rowsHeader.ColLength))
rs.blockOffset = 0
return nil
}
Expand Down
Loading

0 comments on commit 65f96bf

Please sign in to comment.