diff --git a/go/adbc/driver/mocks/mock_builder.go b/go/adbc/driver/mocks/mock_builder.go new file mode 100644 index 0000000000..f34224787d --- /dev/null +++ b/go/adbc/driver/mocks/mock_builder.go @@ -0,0 +1,497 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package mocks + +import ( + "math" + + "github.com/apache/arrow/go/v16/arrow" + "github.com/apache/arrow/go/v16/arrow/array" + "github.com/apache/arrow/go/v16/arrow/float16" + "github.com/apache/arrow/go/v16/arrow/memory" + "golang.org/x/exp/constraints" +) + +func mockStruct(mem memory.Allocator, rows int, structType *arrow.StructType) arrow.Array { + structBuilder := array.NewStructBuilder(mem, structType) + defer structBuilder.Release() + structBuilder.Reserve(structType.NumFields()) + for i := 0; i < structType.NumFields(); i++ { + fieldBuilder := structBuilder.FieldBuilder(i) + fieldBuilder.Resize(rows) + valid := make([]bool, rows) + for i := 0; i < rows; i++ { + valid[i] = true + } + structBuilder.AppendValues(valid) + switch fieldBuilder.Type() { + case arrow.PrimitiveTypes.Int8: + fillValue(fieldBuilder.(*array.Int8Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int16: + fillValue(fieldBuilder.(*array.Int16Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int32: + fillValue(fieldBuilder.(*array.Int32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int64: + fillValue(fieldBuilder.(*array.Int64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint8: + fillValue(fieldBuilder.(*array.Uint8Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint16: + fillValue(fieldBuilder.(*array.Uint16Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint32: + fillValue(fieldBuilder.(*array.Uint32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint64: + fillValue(fieldBuilder.(*array.Uint64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Float32: + fillValue(fieldBuilder.(*array.Float32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Float64: + fillValue(fieldBuilder.(*array.Float64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Date32, arrow.FixedWidthTypes.Date32: + fillValue(fieldBuilder.(*array.Date32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Date64, arrow.FixedWidthTypes.Date64: + fillValue(fieldBuilder.(*array.Date64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time32s: + fillValue(fieldBuilder.(*array.Time32Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time32ms: + fillValue(fieldBuilder.(*array.Time32Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time64us: + fillValue(fieldBuilder.(*array.Time64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time64ns: + fillValue(fieldBuilder.(*array.Time64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_s: + fillValue(fieldBuilder.(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_ms: + fillValue(fieldBuilder.(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_us: + fillValue(fieldBuilder.(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_ns: + fillValue(fieldBuilder.(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_s: + fillValue(fieldBuilder.(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_ms: + fillValue(fieldBuilder.(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_us: + fillValue(fieldBuilder.(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_ns: + fillValue(fieldBuilder.(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.MonthInterval: + fillValue(fieldBuilder.(*array.MonthIntervalBuilder).AppendValues, rows, 0) + + case arrow.FixedWidthTypes.Boolean: + fillBoolValue(fieldBuilder.(*array.BooleanBuilder).AppendValues, rows) + case arrow.FixedWidthTypes.DayTimeInterval: + fillDayTimeIntervalValue(fieldBuilder.(*array.DayTimeIntervalBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.MonthDayNanoInterval: + fillMonthDayNanoIntervalValue(fieldBuilder.(*array.MonthDayNanoIntervalBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Float16: + fillFloat16Value(fieldBuilder.(*array.Float16Builder).AppendValues, rows, 0) + case arrow.BinaryTypes.Binary: + fillBinaryValue(fieldBuilder.(*array.BinaryBuilder).AppendValues, rows, 0) + case arrow.BinaryTypes.String: + fillStringValue(fieldBuilder.(*array.StringBuilder).AppendValues, rows, 0) + } + } + return structBuilder.NewArray() +} + +func mockArray(mem memory.Allocator, rows int, elemType arrow.DataType) arrow.Array { + listBuilder := array.NewListBuilder(mem, elemType) + defer listBuilder.Release() + listBuilder.Append(true) + + switch elemType { + case arrow.PrimitiveTypes.Int8: + fillValue(listBuilder.ValueBuilder().(*array.Int8Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int16: + fillValue(listBuilder.ValueBuilder().(*array.Int16Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int32: + fillValue(listBuilder.ValueBuilder().(*array.Int32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Int64: + fillValue(listBuilder.ValueBuilder().(*array.Int64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint8: + fillValue(listBuilder.ValueBuilder().(*array.Uint8Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint16: + fillValue(listBuilder.ValueBuilder().(*array.Uint16Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint32: + fillValue(listBuilder.ValueBuilder().(*array.Uint32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Uint64: + fillValue(listBuilder.ValueBuilder().(*array.Uint64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Float32: + fillValue(listBuilder.ValueBuilder().(*array.Float32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Float64: + fillValue(listBuilder.ValueBuilder().(*array.Float64Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Date32, arrow.FixedWidthTypes.Date32: + fillValue(listBuilder.ValueBuilder().(*array.Date32Builder).AppendValues, rows, 0) + case arrow.PrimitiveTypes.Date64, arrow.FixedWidthTypes.Date64: + fillValue(listBuilder.ValueBuilder().(*array.Date64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time32s: + fillValue(listBuilder.ValueBuilder().(*array.Time32Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time32ms: + fillValue(listBuilder.ValueBuilder().(*array.Time32Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time64us: + fillValue(listBuilder.ValueBuilder().(*array.Time64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Time64ns: + fillValue(listBuilder.ValueBuilder().(*array.Time64Builder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_s: + fillValue(listBuilder.ValueBuilder().(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_ms: + fillValue(listBuilder.ValueBuilder().(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_us: + fillValue(listBuilder.ValueBuilder().(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Timestamp_ns: + fillValue(listBuilder.ValueBuilder().(*array.TimestampBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_s: + fillValue(listBuilder.ValueBuilder().(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_ms: + fillValue(listBuilder.ValueBuilder().(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_us: + fillValue(listBuilder.ValueBuilder().(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Duration_ns: + fillValue(listBuilder.ValueBuilder().(*array.DurationBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.MonthInterval: + fillValue(listBuilder.ValueBuilder().(*array.MonthIntervalBuilder).AppendValues, rows, 0) + + case arrow.FixedWidthTypes.Boolean: + fillBoolValue(listBuilder.ValueBuilder().(*array.BooleanBuilder).AppendValues, rows) + case arrow.FixedWidthTypes.DayTimeInterval: + fillDayTimeIntervalValue(listBuilder.ValueBuilder().(*array.DayTimeIntervalBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.MonthDayNanoInterval: + fillMonthDayNanoIntervalValue(listBuilder.ValueBuilder().(*array.MonthDayNanoIntervalBuilder).AppendValues, rows, 0) + case arrow.FixedWidthTypes.Float16: + fillFloat16Value(listBuilder.ValueBuilder().(*array.Float16Builder).AppendValues, rows, 0) + case arrow.BinaryTypes.Binary: + fillBinaryValue(listBuilder.ValueBuilder().(*array.BinaryBuilder).AppendValues, rows, 0) + case arrow.BinaryTypes.String: + fillStringValue(listBuilder.ValueBuilder().(*array.StringBuilder).AppendValues, rows, 0) + } + + return listBuilder.NewArray() +} + +func fillValue[T constraints.Integer | constraints.Float](append func(value []T, valid []bool), rows int, start int) { + append(getSlice(rows, T(start)), nil) +} + +func fillBoolValue(append func(value []bool, valid []bool), rows int) { + append(getBoolSlice(rows), nil) +} + +func fillDayTimeIntervalValue(append func(value []arrow.DayTimeInterval, valid []bool), rows int, start int) { + append(getDayTimeIntervalSlice(rows, start), nil) +} + +func fillMonthDayNanoIntervalValue(append func(value []arrow.MonthDayNanoInterval, valid []bool), rows int, start int) { + append(getMonthDayNanoIntervalSlice(rows, start), nil) +} + +func fillFloat16Value(append func(value []float16.Num, valid []bool), rows int, start int) { + append(getFloat16Slice(rows, float32(start)), nil) +} + +func fillBinaryValue(append func(value [][]byte, valid []bool), rows int, start int) { + append(getBinarySlice(rows, start), nil) +} + +func fillStringValue(append func(value []string, valid []bool), rows int, start int) { + append(getStringSlice(rows, start), nil) +} + +func getSlice[T constraints.Integer | constraints.Float](rows int, start T) []T { + slice := make([]T, rows) + for i := int64(0); i < int64(rows); i++ { + slice[i] = T(i) + start + } + return slice +} + +func getBoolSlice(rows int) []bool { + slice := make([]bool, rows) + for i := 0; i < rows; i++ { + slice[i] = i%2 == 0 + } + return slice +} + +func getDayTimeIntervalSlice(rows int, start int) []arrow.DayTimeInterval { + slice := make([]arrow.DayTimeInterval, rows) + for i := start; i < start+rows; i++ { + slice[i] = arrow.DayTimeInterval{Days: int32(i), Milliseconds: int32(i)} + } + return slice +} + +func getMonthDayNanoIntervalSlice(rows int, start int) []arrow.MonthDayNanoInterval { + slice := make([]arrow.MonthDayNanoInterval, rows) + for i := start; i < start+rows; i++ { + slice[i] = arrow.MonthDayNanoInterval{Months: int32(i), Days: int32(i), Nanoseconds: int64(i)} + } + return slice +} + +func getFloat16Slice(rows int, start float32) []float16.Num { + slice := make([]float16.Num, rows) + for i := 0; i < rows; i++ { + slice[i] = float16.New(start + float32(i)) + } + return slice +} + +func getBinarySlice(rows int, start int) [][]byte { + slice := make([][]byte, rows) + for i := 0; i < rows; i++ { + slice[i] = []byte{byte(i + start), byte(i + start + 1), byte(i + start + 2), byte(i + start + 3)} + } + return slice +} + +var stringData = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + +func getStringSlice(rows int, start int) []string { + slice := make([]string, rows) + for i := 0; i < rows; i++ { + pos := (i + start) % len(stringData) + slice[i] = stringData[pos : pos+10] + } + return slice +} + +func mockInt8(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewInt8Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewInt8Array() +} + +func mockInt16(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewInt16Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewInt16Array() +} + +func mockInt32(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewInt32Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewInt32Array() +} + +func mockInt64(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewInt64Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewInt64Array() +} + +func mockUint8(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewUint8Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewUint8Array() +} + +func mockUint16(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewUint16Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewUint16Array() +} + +func mockUint32(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewUint32Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewUint32Array() +} + +func mockUint64(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewUint64Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewUint64Array() +} + +func mockFloat32(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewFloat32Builder(mem) + defer ib.Release() + values := make([]float32, 0, rows) + if rows > 3 { + values = append(values, math.SmallestNonzeroFloat32, -1, 0, math.MaxFloat32) + } + for i := len(values); i < rows; i++ { + values = append(values, float32(i)) + } + ib.AppendValues(values, nil) + return ib.NewFloat32Array() +} + +func mockFloat64(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewFloat64Builder(mem) + defer ib.Release() + values := make([]float64, 0, rows) + if rows > 3 { + values = append(values, math.SmallestNonzeroFloat64, -1, 0, math.MaxFloat64) + } + for i := len(values); i < rows; i++ { + values = append(values, float64(i)) + } + ib.AppendValues(values, nil) + return ib.NewFloat64Array() +} + +func mockDate32(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDate32Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDate32Array() +} + +func mockDate64(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDate64Builder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDate64Array() +} + +func mockBinary(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary) + defer ib.Release() + fillBinaryValue(ib.AppendValues, rows, 0) + return ib.NewBinaryArray() +} + +func mockString(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewStringBuilder(mem) + defer ib.Release() + fillStringValue(ib.AppendValues, rows, 0) + return ib.NewStringArray() +} + +func mockDayTimeInterval(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDayTimeIntervalBuilder(mem) + defer ib.Release() + fillDayTimeIntervalValue(ib.AppendValues, rows, 0) + return ib.NewDayTimeIntervalArray() +} + +func mockDuration_s(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Second}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDurationArray() +} + +func mockDuration_ms(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Millisecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDurationArray() +} + +func mockDuration_us(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Microsecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDurationArray() +} + +func mockDuration_ns(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Nanosecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewDurationArray() +} + +func mockFloat16(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewFloat16Builder(mem) + defer ib.Release() + values := make([]float16.Num, rows) + for i := 0; i < rows; i++ { + values[i] = float16.New(float32(i + 1)) + } + ib.AppendValues(values, nil) + return ib.NewFloat16Array() +} + +func mockMonthInterval(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewMonthIntervalBuilder(mem) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewMonthIntervalArray() +} + +func mockTime32s(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTime32Builder(mem, &arrow.Time32Type{Unit: arrow.Second}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTime32Array() +} + +func mockTime32ms(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTime32Builder(mem, &arrow.Time32Type{Unit: arrow.Millisecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTime32Array() +} + +func mockTime64us(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTime64Builder(mem, &arrow.Time64Type{Unit: arrow.Microsecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTime64Array() +} + +func mockTime64ns(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTime64Builder(mem, &arrow.Time64Type{Unit: arrow.Nanosecond}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTime64Array() +} + +func mockTimestamp_s(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Second, TimeZone: "UTC"}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTimestampArray() +} + +func mockTimestamp_ms(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Millisecond, TimeZone: "UTC"}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTimestampArray() +} + +func mockTimestamp_us(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTimestampArray() +} + +func mockTimestamp_ns(mem memory.Allocator, rows int) arrow.Array { + ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}) + defer ib.Release() + fillValue(ib.AppendValues, rows, 0) + return ib.NewTimestampArray() +} + +// func mockList(mem memory.Allocator, rows, length int64, innerList arrow.Array) arrow.Array { +// lb := array.NewListBuilder(mem, arrow.ListOf(innerList.DataType())) +// defer lb.Release() +// // vb := lb.ValueBuilder().(*array.ListBuilder).Append() +// } diff --git a/go/adbc/driver/mocks/mock_reader.go b/go/adbc/driver/mocks/mock_reader.go index 730256462e..a437d83373 100644 --- a/go/adbc/driver/mocks/mock_reader.go +++ b/go/adbc/driver/mocks/mock_reader.go @@ -19,7 +19,6 @@ package mocks import ( "fmt" - "math" "regexp" "strconv" "strings" @@ -28,7 +27,6 @@ import ( "github.com/apache/arrow-adbc/go/adbc" "github.com/apache/arrow/go/v16/arrow" "github.com/apache/arrow/go/v16/arrow/array" - "github.com/apache/arrow/go/v16/arrow/float16" "github.com/apache/arrow/go/v16/arrow/memory" ) @@ -41,11 +39,13 @@ type mockReader struct { type typeBuilder struct { field arrow.Field - builder func(memory.Allocator,int64) arrow.Array + builder func(memory.Allocator, int) arrow.Array } var ( rowsRegex = regexp.MustCompile(`^(?P\d+)[\s]*:`) + listRegex = regexp.MustCompile(`^list<(?P\d*):?(?P.+)>`) + structRegex = regexp.MustCompile(`^struct\<(?P.+)\>`) availableTypes = map[string]typeBuilder{ "int8": typeBuilder{ field: arrow.Field{Name: "int8", Type: arrow.PrimitiveTypes.Int8}, @@ -166,15 +166,14 @@ var ( } ) -func parseQuery(query string) ([]arrow.Field, []arrow.Array, int64, error) { - query = strings.TrimSpace(query) +func parseQuery(query string, innerRows int) ([]arrow.Field, []arrow.Array, int, error) { matches := rowsRegex.FindStringSubmatch(query) - rows := int64(1) + rows := innerRows rowsString := "" if len(matches) == 2 { var err error rowsString = matches[rowsRegex.SubexpIndex("rows")] - if rows, err = strconv.ParseInt(rowsString, 10, 64); err != nil { + if rows, err = strconv.Atoi(rowsString); err != nil { return nil, nil, -1, err } } @@ -182,45 +181,77 @@ func parseQuery(query string) ([]arrow.Field, []arrow.Array, int64, error) { fields := make([]arrow.Field, 0) fieldValues := make([]arrow.Array, 0) - query = strings.TrimSpace(query[len(rowsString):]) + query = query[len(rowsString):] if query[0] == ':' { query = query[1:] } var t string for { - var requestTypes []string - if strings.HasPrefix(query, "list") { - // TODO: check if this is a valid list type - // query = strings.TrimSpace(query[4:]) - // if query[0] == '<' && query[len(query)-1] == '>' { - // query = query[1 : len(query)-1] - // } - - break - } - requestTypes = strings.SplitN(query, ",", 2) + requestTypes := strings.SplitN(query, ",", 2) if len(requestTypes) == 2 { - t = strings.TrimSpace(requestTypes[0]) - query = strings.TrimSpace(requestTypes[1]) + t = requestTypes[0] + query = requestTypes[1] } else if len(requestTypes) == 1 { - t = strings.TrimSpace(requestTypes[0]) + t = requestTypes[0] query = "" } if builder, ok := availableTypes[t]; ok { fields = append(fields, builder.field) fieldValues = append(fieldValues, builder.builder(mem, rows)) - } else if t == "list" { - fields = append(fields, arrow.Field{Name: "list", Type: arrow.ListOf(arrow.PrimitiveTypes.Int8)}) - fieldValues = append(fieldValues, mockListInt8(mem)) - } else { - return nil, nil, -1, adbc.Error{ - Code: adbc.StatusInvalidArgument, - Msg: fmt.Sprintf("unknown type %s", t), + } else if strings.HasPrefix(t, "list") { + if t == "list" { + fields = append(fields, arrow.Field{Name: "list", Type: arrow.ListOf(arrow.PrimitiveTypes.Int8)}) + fieldValues = append(fieldValues, mockArray(mem, rows, arrow.PrimitiveTypes.Int8)) + } else if listMatch := listRegex.FindStringSubmatch(t); len(listMatch) > 1 { + listLen := 1 + if lenStr := listMatch[listRegex.SubexpIndex("len")]; len(lenStr) > 0 { + listLen, _ = strconv.Atoi(lenStr) + } + elmTypeStr := listMatch[listRegex.SubexpIndex("typename")] + if innerType, ok := availableTypes[elmTypeStr]; ok { + fields = append(fields, arrow.Field{Name: "list", Type: arrow.ListOf(innerType.field.Type)}) + fieldValues = append(fieldValues, mockArray(mem, listLen, innerType.field.Type)) + } else { + // array.NewListBuilder(mem,) + // innerFields, innerValues, _, err := parseQuery(elmTypeStr,listLen) + // if err != nil { + // return nil, nil, -1, err + // } + // fields = append(fields, arrow.Field{Name: "list", Type: arrow.ListOf(innerFields[0].Type)}) + // fieldValues = append(fieldValues, mockList(innerValues[0])) + return nil, nil, -1, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("unknown type %s", t), + } + } + } else { + return nil, nil, -1, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("unknown type %s", t), + } + } + } else if strings.HasPrefix(t, "struct") { + fmt.Println("oh it's a struct!") + structMatch := structRegex.FindStringSubmatch(t) + fmt.Printf("structMatch: %v\n", structMatch) + if structMatch := structRegex.FindStringSubmatch(t); structMatch != nil { + structString := strings.Split(structMatch[structRegex.SubexpIndex("struct")], ",") + fmt.Printf("structString: %v\n", structString) + innerTypeFields := make([]arrow.Field, 0) + for _, innerType := range structString { + innerTypeFields = append(innerTypeFields, availableTypes[innerType].field) + } + fields = append(fields, arrow.Field{Name: "struct", Type: arrow.StructOf(innerTypeFields...)}) + fieldValues = append(fieldValues, mockStruct(mem, rows, arrow.StructOf(innerTypeFields...))) + } else { + return nil, nil, -1, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("unknown type %s", t), + } } } - if len(query) == 0 { break } @@ -229,16 +260,23 @@ func parseQuery(query string) ([]arrow.Field, []arrow.Array, int64, error) { return fields, fieldValues, rows, nil } +func parseJsonQuery(query string) ([]arrow.Field, []arrow.Array, int, error) { + +} + // Create a mockReader according to the query // The query should be a list of types separated by commas // The returned mockReader will have the types in the same order func NewMockReader(query string) (*mockReader, error) { - fields, fieldValues, rows, err := parseQuery(query) + query = strings.ReplaceAll(query, "\t", "") + query = strings.ReplaceAll(query, " ", "") + query = strings.ReplaceAll(query, "\n", "") + fields, fieldValues, rows, err := parseQuery(query, 1) if err != nil { return nil, err } schema := arrow.NewSchema(fields, nil) - rec := array.NewRecord(schema, fieldValues, rows) + rec := array.NewRecord(schema, fieldValues, int64(rows)) rec.Retain() return &mockReader{ refCount: 1, @@ -279,355 +317,3 @@ func (r *mockReader) Next() bool { } return false } - -func mockListInt8(mem memory.Allocator) arrow.Array { - listBuilder := array.NewListBuilder(mem, arrow.PrimitiveTypes.Int8) - defer listBuilder.Release() - - int8Builder := listBuilder.ValueBuilder().(*array.Int8Builder) - - // Create a list of int8 values. For example, create 3 lists with varying lengths. - lists := [][]int8{ - {1, 2, 3}, - {4, 5}, - {6}, - {}, - } - - for _, vals := range lists { - listBuilder.Append(true) // Append true to indicate a valid list item - for _, val := range vals { - int8Builder.Append(val) - } - } - - return listBuilder.NewArray() -} - -func mockInt8(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewInt8Builder(mem) - defer ib.Release() - values := make([]int8, rows) - for i := int64(0); i < rows; i++ { - values[i] = int8(i) - } - ib.AppendValues(values, nil) - return ib.NewInt8Array() -} - -func mockInt16(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewInt16Builder(mem) - defer ib.Release() - values := make([]int16, rows) - for i := int64(0); i < rows; i++ { - values[i] = int16(i) - } - ib.AppendValues(values, nil) - return ib.NewInt16Array() -} - - -func mockInt32(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewInt32Builder(mem) - defer ib.Release() - values := make([]int32, rows) - for i := int64(0); i < rows; i++ { - values[i] = int32(i) - } - ib.AppendValues(values, nil) - return ib.NewInt32Array() -} - -func mockInt64(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewInt64Builder(mem) - defer ib.Release() - values := make([]int64, rows) - for i := int64(0); i < rows; i++ { - values[i] = int64(i) - } - ib.AppendValues(values, nil) - return ib.NewInt64Array() -} - -func mockUint8(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewUint8Builder(mem) - defer ib.Release() - values := make([]uint8, rows) - for i := int64(0); i < rows; i++ { - values[i] = uint8(i) - } - ib.AppendValues(values, nil) - return ib.NewUint8Array() -} - -func mockUint16(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewUint16Builder(mem) - defer ib.Release() - values := make([]uint16, rows) - for i := int64(0); i < rows; i++ { - values[i] = uint16(i) - } - ib.AppendValues(values, nil) - return ib.NewUint16Array() -} - -func mockUint32(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewUint32Builder(mem) - defer ib.Release() - values := make([]uint32, rows) - for i := int64(0); i < rows; i++ { - values[i] = uint32(i) - } - ib.AppendValues(values, nil) - return ib.NewUint32Array() -} - -func mockUint64(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewUint64Builder(mem) - defer ib.Release() - values := make([]uint64, rows) - for i := int64(0); i < rows; i++ { - values[i] = uint64(i) - } - ib.AppendValues(values, nil) - return ib.NewUint64Array() -} - -func mockFloat32(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewFloat32Builder(mem) - defer ib.Release() - values := make([]float32, 0, rows) - if rows > 3 { - values = append(values, math.SmallestNonzeroFloat32, -1, 0, math.MaxFloat32) - } - for i := int64(len(values)); i < rows; i++ { - values = append(values, float32(i)) - } - ib.AppendValues(values, nil) - return ib.NewFloat32Array() -} - -func mockFloat64(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewFloat64Builder(mem) - defer ib.Release() - values := make([]float64, 0, rows) - if rows > 3 { - values = append(values, math.SmallestNonzeroFloat64, -1, 0, math.MaxFloat64) - } - for i := int64(len(values)); i < rows; i++ { - values = append(values, float64(i)) - } - ib.AppendValues(values, nil) - return ib.NewFloat64Array() -} - -func mockDate32(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDate32Builder(mem) - defer ib.Release() - values := make([]arrow.Date32, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Date32(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDate32Array() -} - -func mockDate64(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDate64Builder(mem) - defer ib.Release() - values := make([]arrow.Date64, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Date64(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDate64Array() -} - -func mockBinary(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary) - defer ib.Release() - values := make([][]byte, rows) - for i := int64(0); i < rows; i++ { - values[i] = []byte{byte(i)} - } - ib.AppendValues(values, nil) - return ib.NewBinaryArray() -} - -func mockString(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewStringBuilder(mem) - defer ib.Release() - values := make([]string, rows) - for i := int64(0); i < rows; i++ { - values[i] = string(65 + int(i) % 26) - } - ib.AppendValues(values, nil) - return ib.NewStringArray() -} - - -func mockDayTimeInterval(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDayTimeIntervalBuilder(mem) - defer ib.Release() - values := make([]arrow.DayTimeInterval, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.DayTimeInterval{Days: int32(i), Milliseconds: int32(i + 1)} - } - ib.AppendValues(values, nil) - return ib.NewDayTimeIntervalArray() -} - -func mockDuration_s(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Second}) - defer ib.Release() - values := make([]arrow.Duration, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Duration(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDurationArray() -} - -func mockDuration_ms(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Millisecond}) - defer ib.Release() - values := make([]arrow.Duration, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Duration(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDurationArray() -} - -func mockDuration_us(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Microsecond}) - defer ib.Release() - values := make([]arrow.Duration, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Duration(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDurationArray() -} - -func mockDuration_ns(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewDurationBuilder(mem, &arrow.DurationType{Unit: arrow.Nanosecond}) - defer ib.Release() - values := make([]arrow.Duration, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Duration(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewDurationArray() -} - -func mockFloat16(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewFloat16Builder(mem) - defer ib.Release() - values := make([]float16.Num, rows) - for i := int64(0); i < rows; i++ { - values[i] = float16.New(float32(i + 1)) - } - ib.AppendValues(values, nil) - return ib.NewFloat16Array() -} - -func mockMonthInterval(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewMonthIntervalBuilder(mem) - defer ib.Release() - values := make([]arrow.MonthInterval, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.MonthInterval(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewMonthIntervalArray() -} - -func mockTime32s(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTime32Builder(mem, &arrow.Time32Type{Unit: arrow.Second}) - defer ib.Release() - values := make([]arrow.Time32, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Time32(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewTime32Array() -} - -func mockTime32ms(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTime32Builder(mem, &arrow.Time32Type{Unit: arrow.Millisecond}) - defer ib.Release() - values := make([]arrow.Time32, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Time32(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewTime32Array() -} - -func mockTime64us(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTime64Builder(mem, &arrow.Time64Type{Unit: arrow.Microsecond}) - defer ib.Release() - values := make([]arrow.Time64, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Time64(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewTime64Array() -} - -func mockTime64ns(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTime64Builder(mem, &arrow.Time64Type{Unit: arrow.Nanosecond}) - defer ib.Release() - values := make([]arrow.Time64, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Time64(i + 1) - } - ib.AppendValues(values, nil) - return ib.NewTime64Array() -} - - -func mockTimestamp_s(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Second, TimeZone: "UTC"}) - defer ib.Release() - values := make([]arrow.Timestamp, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Timestamp(i) - } - ib.AppendValues(values, nil) - return ib.NewTimestampArray() -} - -func mockTimestamp_ms(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Millisecond, TimeZone: "UTC"}) - defer ib.Release() - values := make([]arrow.Timestamp, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Timestamp(i) - } - ib.AppendValues(values, nil) - return ib.NewTimestampArray() -} - -func mockTimestamp_us(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}) - defer ib.Release() - values := make([]arrow.Timestamp, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Timestamp(i) - } - ib.AppendValues(values, nil) - return ib.NewTimestampArray() -} - -func mockTimestamp_ns(mem memory.Allocator, rows int64) arrow.Array { - ib := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"}) - defer ib.Release() - values := make([]arrow.Timestamp, rows) - for i := int64(0); i < rows; i++ { - values[i] = arrow.Timestamp(i) - } - ib.AppendValues(values, nil) - return ib.NewTimestampArray() -}