-
Notifications
You must be signed in to change notification settings - Fork 97
/
arrow_test.go
117 lines (90 loc) · 2.65 KB
/
arrow_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
//go:build !no_duckdb_arrow
package duckdb
import (
"context"
"database/sql/driver"
"testing"
"github.com/stretchr/testify/require"
)
func TestArrow(t *testing.T) {
t.Parallel()
db := openDB(t)
defer db.Close()
createFooTable(db, t)
conn, err := db.Conn(context.Background())
require.NoError(t, err)
defer conn.Close()
t.Run("select series", func(t *testing.T) {
c, err := NewConnector("", nil)
require.NoError(t, err)
defer c.Close()
conn, err := c.Connect(context.Background())
require.NoError(t, err)
defer conn.Close()
ar, err := NewArrowFromConn(conn)
require.NoError(t, err)
rdr, err := ar.QueryContext(context.Background(), "SELECT * FROM generate_series(1, 10)")
require.NoError(t, err)
defer rdr.Release()
for rdr.Next() {
rec := rdr.Record()
require.Equal(t, int64(10), rec.NumRows())
require.NoError(t, err)
}
require.NoError(t, rdr.Err())
})
t.Run("select long series", func(t *testing.T) {
c, err := NewConnector("", nil)
require.NoError(t, err)
defer c.Close()
conn, err := c.Connect(context.Background())
require.NoError(t, err)
defer conn.Close()
ar, err := NewArrowFromConn(conn)
require.NoError(t, err)
rdr, err := ar.QueryContext(context.Background(), "SELECT * FROM generate_series(1, 10000)")
require.NoError(t, err)
defer rdr.Release()
var totalRows int64
for rdr.Next() {
rec := rdr.Record()
totalRows += rec.NumRows()
}
require.Equal(t, int64(10000), totalRows)
require.NoError(t, rdr.Err())
})
t.Run("query table and filter results", func(t *testing.T) {
err = conn.Raw(func(driverConn any) error {
conn, ok := driverConn.(driver.Conn)
require.True(t, ok)
ar, err := NewArrowFromConn(conn)
require.NoError(t, err)
_, err = db.ExecContext(context.Background(), "INSERT INTO foo (bar, baz) VALUES ('lala', 2), ('dada', 3)")
require.NoError(t, err)
reader, err := ar.QueryContext(context.Background(), "SELECT bar, baz FROM foo WHERE baz > ?", 1)
require.NoError(t, err)
defer reader.Release()
for reader.Next() {
rec := reader.Record()
require.Equal(t, int64(2), rec.NumRows())
require.Equal(t, "lala", rec.Column(0).ValueStr(0))
require.Equal(t, "dada", rec.Column(0).ValueStr(1))
}
require.NoError(t, reader.Err())
return nil
})
require.NoError(t, err)
})
t.Run("query error", func(t *testing.T) {
err := conn.Raw(func(driverConn any) error {
conn, ok := driverConn.(driver.Conn)
require.True(t, ok)
ar, err := NewArrowFromConn(conn)
require.NoError(t, err)
_, err = ar.QueryContext(context.Background(), "SELECT bar")
require.Error(t, err)
return nil
})
require.NoError(t, err)
})
}