impl: rows logic

Author: nmreadelf

chdb/driver/driver.go

@@ -7,7 +7,12 @@ import (
 	"database/sql/driver"
 	"fmt"
 	"reflect"
+	"time"
 
+	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/apache/arrow/go/v14/arrow/array"
+	"github.com/apache/arrow/go/v14/arrow/decimal128"
+	"github.com/apache/arrow/go/v14/arrow/decimal256"
 	wrapper "github.com/chdb-io/chdb-go/chdb"
 	"github.com/chdb-io/chdb-go/chdbstable"
 
@@ -93,6 +98,8 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, e
 type rows struct {
 	localResult *chdbstable.LocalResult
 	reader      *ipc.FileReader
+	curRecord   arrow.Record
+	curRow      int64
 }
 
 func (r *rows) Columns() (out []string) {
@@ -104,25 +111,146 @@ func (r *rows) Columns() (out []string) {
 }
 
 func (r *rows) Close() error {
+	if r.curRecord != nil {
+		r.curRecord = nil
+	}
+	// ignore reader close
+	_ = r.reader.Close()
+	r.reader = nil
+	r.localResult = nil
 	return nil
 }
 
 func (r *rows) Next(dest []driver.Value) error {
+	if r.curRecord != nil && r.curRow == r.curRecord.NumRows() {
+		r.curRecord = nil
+	}
+	for r.curRecord == nil {
+		record, err := r.reader.Read()
+		if err != nil {
+			return err
+		}
+		if record.NumRows() == 0 {
+			continue
+		}
+		r.curRecord = record
+		r.curRow = 0
+	}
+
+	for i, col := range r.curRecord.Columns() {
+		if col.IsNull(int(r.curRow)) {
+			dest[i] = nil
+			continue
+		}
+		switch col := col.(type) {
+		case *array.Boolean:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int8:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint8:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int16:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint16:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Int64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Uint64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Float32:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Float64:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.String:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.LargeString:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Binary:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.LargeBinary:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Date32:
+			dest[i] = col.Value(int(r.curRow)).ToTime()
+		case *array.Date64:
+			dest[i] = col.Value(int(r.curRow)).ToTime()
+		case *array.Time32:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.Time32Type).Unit)
+		case *array.Time64:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.Time64Type).Unit)
+		case *array.Timestamp:
+			dest[i] = col.Value(int(r.curRow)).ToTime(col.DataType().(*arrow.TimestampType).Unit)
+		case *array.Decimal128:
+			dest[i] = col.Value(int(r.curRow))
+		case *array.Decimal256:
+			dest[i] = col.Value(int(r.curRow))
+		default:
+			return fmt.Errorf(
+				"not yet implemented populating from columns of type " + col.DataType().String(),
+			)
+		}
+	}
+
+	r.curRow++
 	return nil
 }
 
 func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
-	return ""
+	return r.reader.Schema().Field(index).Type.String()
 }
 
 func (r *rows) ColumnTypeNullable(index int) (nullable, ok bool) {
-	return
+	return r.reader.Schema().Field(index).Nullable, true
 }
 
 func (r *rows) ColumnTypePrecisionScale(index int) (precision, scale int64, ok bool) {
-	return
+	typ := r.reader.Schema().Field(index).Type
+	switch dt := typ.(type) {
+	case *arrow.Decimal128Type:
+		return int64(dt.Precision), int64(dt.Scale), true
+	case *arrow.Decimal256Type:
+		return int64(dt.Precision), int64(dt.Scale), true
+	}
+	return 0, 0, false
 }
 
 func (r *rows) ColumnTypeScanType(index int) reflect.Type {
-	return reflect.TypeOf(nil)
+	switch r.reader.Schema().Field(index).Type.ID() {
+	case arrow.BOOL:
+		return reflect.TypeOf(false)
+	case arrow.INT8:
+		return reflect.TypeOf(int8(0))
+	case arrow.UINT8:
+		return reflect.TypeOf(uint8(0))
+	case arrow.INT16:
+		return reflect.TypeOf(int16(0))
+	case arrow.UINT16:
+		return reflect.TypeOf(uint16(0))
+	case arrow.INT32:
+		return reflect.TypeOf(int32(0))
+	case arrow.UINT32:
+		return reflect.TypeOf(uint32(0))
+	case arrow.INT64:
+		return reflect.TypeOf(int64(0))
+	case arrow.UINT64:
+		return reflect.TypeOf(uint64(0))
+	case arrow.FLOAT32:
+		return reflect.TypeOf(float32(0))
+	case arrow.FLOAT64:
+		return reflect.TypeOf(float64(0))
+	case arrow.DECIMAL128:
+		return reflect.TypeOf(decimal128.Num{})
+	case arrow.DECIMAL256:
+		return reflect.TypeOf(decimal256.Num{})
+	case arrow.BINARY:
+		return reflect.TypeOf([]byte{})
+	case arrow.STRING:
+		return reflect.TypeOf(string(""))
+	case arrow.TIME32, arrow.TIME64, arrow.DATE32, arrow.DATE64, arrow.TIMESTAMP:
+		return reflect.TypeOf(time.Time{})
+	}
+	return nil
 }

chdb/driver/driver_test.go

@@ -13,19 +13,6 @@ func TestDb(t *testing.T) {
 	if db.Ping() != nil {
 		t.Errorf("ping db fail")
 	}
-	{
-		rows, err := db.Query("SELECT version()")
-		if err != nil {
-			t.Errorf("run Query fail, err:%s", err)
-		}
-		cols, err := rows.Columns()
-		if err != nil {
-			t.Errorf("get result columns fail, err: %s", err)
-		}
-		if len(cols) != 1 {
-			t.Errorf("select version(), result columns length should be 1")
-		}
-	}
 	{
 		rows, err := db.Query(`SELECT 1,'abc'`)
 		if err != nil {
@@ -38,5 +25,22 @@ func TestDb(t *testing.T) {
 		if len(cols) != 2 {
 			t.Errorf("select version(), result columns length should be 1")
 		}
+		var (
+			bar int
+			foo string
+		)
+		defer rows.Close()
+		for rows.Next() {
+			err := rows.Scan(&bar, &foo)
+			if err != nil {
+				t.Errorf("scan fail, err: %s", err)
+			}
+			if bar != 1 {
+				t.Errorf("expected error")
+			}
+			if foo != "abc" {
+				t.Errorf("expected error")
+			}
+		}
 	}
 }