diff --git a/ci/conda_env_cpp.txt b/ci/conda_env_cpp.txt index 61010a3341..9d00523f0c 100644 --- a/ci/conda_env_cpp.txt +++ b/ci/conda_env_cpp.txt @@ -20,5 +20,6 @@ compilers gtest>=1.10.0 libpq libsqlite +libduckdb ninja pkg-config diff --git a/ci/scripts/go_test.sh b/ci/scripts/go_test.sh index 645a71490a..a5627ac8aa 100755 --- a/ci/scripts/go_test.sh +++ b/ci/scripts/go_test.sh @@ -28,8 +28,8 @@ main() { fi if [[ "${CGO_ENABLED}" = 1 ]]; then - export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${install_dir}/lib" - export DYLD_LIBRARY_PATH="${DYLD_LIBRARY_PATH}:${install_dir}/lib" + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${CONDA_PREFIX}/lib:${install_dir}/lib" + export DYLD_LIBRARY_PATH="${DYLD_LIBRARY_PATH}:${CONDA_PREFIX}/lib:${install_dir}/lib" fi pushd "${source_dir}/go/adbc" diff --git a/go/adbc/drivermgr/wrapper.go b/go/adbc/drivermgr/wrapper.go index de1c855bfa..3cbc86052a 100644 --- a/go/adbc/drivermgr/wrapper.go +++ b/go/adbc/drivermgr/wrapper.go @@ -464,7 +464,16 @@ func (s *stmt) Prepare(context.Context) error { } func (s *stmt) SetSubstraitPlan(plan []byte) error { - return &adbc.Error{Code: adbc.StatusNotImplemented} + var err C.struct_AdbcError + cplan := C.CBytes(plan) + defer C.free(cplan) + + length := uint64(len(plan)) + + if code := adbc.Status(C.AdbcStatementSetSubstraitPlan(s.st, (*C.uchar)(cplan), (C.ulong)(length), &err)); code != adbc.StatusOK { + return toAdbcError(code, &err) + } + return nil } func (s *stmt) Bind(_ context.Context, values arrow.Record) error { diff --git a/go/adbc/drivermgr/wrapper_duckdb_test.go b/go/adbc/drivermgr/wrapper_duckdb_test.go new file mode 100644 index 0000000000..19e16d9734 --- /dev/null +++ b/go/adbc/drivermgr/wrapper_duckdb_test.go @@ -0,0 +1,153 @@ +// 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. + +//go:build cgo + +package drivermgr_test + +import ( + "context" + "testing" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-adbc/go/adbc/drivermgr" + "github.com/apache/arrow/go/v17/arrow" + "github.com/apache/arrow/go/v17/arrow/array" + "github.com/stretchr/testify/suite" + "github.com/substrait-io/substrait-go/extensions" + "github.com/substrait-io/substrait-go/plan" + "github.com/substrait-io/substrait-go/types" + "google.golang.org/protobuf/proto" +) + +type DriverMgrDuckDBSuite struct { + suite.Suite + + ctx context.Context + drv adbc.Driver + db adbc.Database + conn adbc.Connection +} + +func (dm *DriverMgrDuckDBSuite) SetupSuite() { + dm.ctx = context.TODO() + dm.drv = drivermgr.Driver{} + var err error + dm.db, err = dm.drv.NewDatabase(map[string]string{ + "driver": "duckdb", + "entrypoint": "duckdb_adbc_init", + }) + dm.NoError(err) + + cnxn, err := dm.db.Open(dm.ctx) + dm.NoError(err) + defer cnxn.Close() + + stmt, err := cnxn.NewStatement() + dm.NoError(err) + defer stmt.Close() + + // Setup substrait extension + dm.NoError(stmt.SetSqlQuery("INSTALL substrait")) + _, err = stmt.ExecuteUpdate(dm.ctx) + dm.NoError(err) + + dm.NoError(stmt.SetSqlQuery("LOAD substrait")) + _, err = stmt.ExecuteUpdate(dm.ctx) + dm.NoError(err) + + dm.NoError(stmt.SetSqlQuery("CREATE TABLE test_table (id INTEGER PRIMARY KEY, name TEXT)")) + // Don't check nrows, duckdb currently does not support it + _, err = stmt.ExecuteUpdate(dm.ctx) + dm.NoError(err) + + dm.NoError(stmt.SetSqlQuery("INSERT INTO test_table (id, name) VALUES (1, 'test'), (2, 'test'), (3, 'test')")) + // Don't check nrows, duckdb currently does not support it + _, err = stmt.ExecuteUpdate(dm.ctx) + dm.NoError(err) +} + +func (dm *DriverMgrDuckDBSuite) TearDownSuite() { + dm.NoError(dm.db.Close()) +} + +func (dm *DriverMgrDuckDBSuite) SetupTest() { + cnxn, err := dm.db.Open(dm.ctx) + dm.Require().NoError(err) + dm.conn = cnxn +} + +func (dm *DriverMgrDuckDBSuite) TearDownTest() { + dm.NoError(dm.conn.Close()) +} + +func (dm *DriverMgrDuckDBSuite) TestSubstraitExecute() { + // Building Substrait plan equivalent to 'SELECT COUNT(*) AS count FROM test_table;' + bldr := plan.NewBuilderDefault() + scan := bldr.NamedScan( + []string{"test_table"}, + types.NamedStruct{ + Names: []string{"id", "names"}, + Struct: types.StructType{ + Nullability: types.NullabilityRequired, + Types: []types.Type{ + &types.Int32Type{Nullability: types.NullabilityRequired}, + &types.StringType{Nullability: types.NullabilityNullable}, + }, + }, + }, + ) + + countFn, err := bldr.AggregateFn(extensions.SubstraitDefaultURIPrefix+"functions_aggregate_generic.yaml", "count", nil) + dm.Require().NoError(err) + + root, err := bldr.AggregateColumns(scan, []plan.AggRelMeasure{bldr.Measure(countFn, nil)}) + dm.Require().NoError(err) + + plan, err := bldr.Plan(root, []string{"count"}) + dm.Require().NoError(err) + + planProto, err := plan.ToProto() + dm.Require().NoError(err) + + serialized, err := proto.Marshal(planProto) + dm.Require().NoError(err) + + st, err := dm.conn.NewStatement() + dm.Require().NoError(err) + defer st.Close() + + dm.Require().NoError(st.SetSubstraitPlan(serialized)) + rdr, _, err := st.ExecuteQuery(context.TODO()) + dm.Require().NoError(err) + defer rdr.Release() + + expectedSchema := arrow.NewSchema([]arrow.Field{{Name: "count", Type: arrow.PrimitiveTypes.Int64, Nullable: true}}, nil) + dm.Require().True(rdr.Next()) + rec := rdr.Record() + dm.Require().Equal(expectedSchema, rec.Schema()) + dm.Require().Equal(int64(1), rec.NumCols()) + dm.Require().Equal(int64(1), rec.NumRows()) + dm.Require().Equal(int64(3), rec.Column(0).(*array.Int64).Value(0)) + + // Only expected one record + dm.Require().False(rdr.Next()) +} + +func TestDriverMgrDuckDB(t *testing.T) { + suite.Run(t, new(DriverMgrDuckDBSuite)) +} diff --git a/go/adbc/go.mod b/go/adbc/go.mod index ea76bcef55..ec94abb757 100644 --- a/go/adbc/go.mod +++ b/go/adbc/go.mod @@ -26,6 +26,7 @@ require ( github.com/google/uuid v1.6.0 github.com/snowflakedb/gosnowflake v1.9.0 github.com/stretchr/testify v1.9.0 + github.com/substrait-io/substrait-go v0.4.2 github.com/youmark/pkcs8 v0.0.0-20201027041543-1326539a0a0a golang.org/x/exp v0.0.0-20240318143956-a85f2c67cd81 golang.org/x/sync v0.7.0 @@ -41,6 +42,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2 // indirect github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.0 // indirect github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect + github.com/alecthomas/participle/v2 v2.1.0 // indirect github.com/andybalholm/brotli v1.1.0 // indirect github.com/apache/arrow/go/v15 v15.0.0 // indirect github.com/apache/thrift v0.20.0 // indirect @@ -61,9 +63,11 @@ require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/dvsekhvalnov/jose2go v1.6.0 // indirect + github.com/fatih/color v1.15.0 // indirect github.com/form3tech-oss/jwt-go v3.2.5+incompatible // indirect github.com/gabriel-vasile/mimetype v1.4.3 // indirect github.com/goccy/go-json v0.10.2 // indirect + github.com/goccy/go-yaml v1.11.0 // indirect github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 // indirect github.com/golang/snappy v0.0.4 // indirect github.com/google/flatbuffers v24.3.25+incompatible // indirect @@ -73,6 +77,7 @@ require ( github.com/klauspost/asmfmt v1.3.2 // indirect github.com/klauspost/compress v1.17.8 // indirect github.com/klauspost/cpuid/v2 v2.2.7 // indirect + github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect diff --git a/go/adbc/go.sum b/go/adbc/go.sum index f1498113f1..f252fa2a77 100644 --- a/go/adbc/go.sum +++ b/go/adbc/go.sum @@ -16,6 +16,12 @@ github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1 h1:DzHpqpoJVaC github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI= github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c h1:RGWPOewvKIROun94nF7v2cua9qP+thov/7M50KEoeSU= github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk= +github.com/alecthomas/assert/v2 v2.3.0 h1:mAsH2wmvjsuvyBvAmCtm7zFsBlb8mIHx5ySLVdDZXL0= +github.com/alecthomas/assert/v2 v2.3.0/go.mod h1:pXcQ2Asjp247dahGEmsZ6ru0UVwnkhktn7S0bBDLxvQ= +github.com/alecthomas/participle/v2 v2.1.0 h1:z7dElHRrOEEq45F2TG5cbQihMtNTv8vwldytDj7Wrz4= +github.com/alecthomas/participle/v2 v2.1.0/go.mod h1:Y1+hAs8DHPmc3YUFzqllV+eSQ9ljPTk0ZkPMtEdAx2c= +github.com/alecthomas/repr v0.2.0 h1:HAzS41CIzNW5syS8Mf9UwXhNH1J9aix/BvDRf1Ml2Yk= +github.com/alecthomas/repr v0.2.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/andybalholm/brotli v1.1.0 h1:eLKJA0d02Lf0mVpIDgYnqXcUn0GqVmEFny3VuID1U3M= github.com/andybalholm/brotli v1.1.0/go.mod h1:sms7XGricyQI9K10gOSf56VKKWS4oLer58Q+mhRPtnY= github.com/apache/arrow/go/v15 v15.0.0 h1:1zZACWf85oEZY5/kd9dsQS7i+2G5zVQcbKTHgslqHNA= @@ -75,12 +81,22 @@ github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkp github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/dvsekhvalnov/jose2go v1.6.0 h1:Y9gnSnP4qEI0+/uQkHvFXeD2PLPJeXEL+ySMEA2EjTY= github.com/dvsekhvalnov/jose2go v1.6.0/go.mod h1:QsHjhyTlD/lAVqn/NSbVZmSCGeDehTB/mPZadG+mhXU= +github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs= +github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw= github.com/form3tech-oss/jwt-go v3.2.5+incompatible h1:/l4kBbb4/vGSsdtB5nUe8L7B9mImVMaBPw9L/0TBHU8= github.com/form3tech-oss/jwt-go v3.2.5+incompatible/go.mod h1:pbq4aXjuKjdthFRnoDwaVPLA+WlJuPGy+QneDUgJi2k= github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= +github.com/go-playground/locales v0.13.0 h1:HyWk6mgj5qFqCT5fjGBuRArbVDfE4hi8+e8ceBS/t7Q= +github.com/go-playground/locales v0.13.0/go.mod h1:taPMhCMXrRLJO55olJkUXHZBHCxTMfnGwq/HNwmWNS8= +github.com/go-playground/universal-translator v0.17.0 h1:icxd5fm+REJzpZx7ZfpaD876Lmtgy7VtROAbHHXk8no= +github.com/go-playground/universal-translator v0.17.0/go.mod h1:UkSxE5sNxxRwHyU+Scu5vgOQjsIJAF8j9muTVoKLVtA= +github.com/go-playground/validator/v10 v10.4.1 h1:pH2c5ADXtd66mxoE0Zm9SUhxE20r7aM3F26W0hOn+GE= +github.com/go-playground/validator/v10 v10.4.1/go.mod h1:nlOn6nFhuKACm19sB/8EGNn9GlaMV7XkbRSipzJ0Ii4= github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/goccy/go-yaml v1.11.0 h1:n7Z+zx8S9f9KgzG6KtQKf+kwqXZlLNR2F6018Dgau54= +github.com/goccy/go-yaml v1.11.0/go.mod h1:H+mJrWtjPTJAHvRbV09MCK9xYwODM+wRTVFFTWckfng= github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 h1:ZpnhV/YsD2/4cESfV5+Hoeu/iUR3ruzNvZ+yQfO03a0= github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2/go.mod h1:bBOAhwG1umN6/6ZUMtDFBMQR8jRg9O75tm9K00oMsK4= github.com/golang-jwt/jwt/v5 v5.2.0 h1:d/ix8ftRUorsN+5eMIlF4T6J8CAt9rch3My2winC1Jw= @@ -101,6 +117,8 @@ github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c h1:6rhixN/i8 github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c/go.mod h1:NMPJylDgVpX0MLRlPy15sqSwOFv/U1GZ2m21JhFfek0= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= @@ -119,6 +137,11 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/leodido/go-urn v1.2.0 h1:hpXL4XnriNwQ/ABnpepYM/1vCLWNDfUNts8dX3xTG6Y= +github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= @@ -152,6 +175,8 @@ github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/substrait-io/substrait-go v0.4.2 h1:buDnjsb3qAqTaNbOR7VKmNgXf4lYQxWEcnSGUWBtmN8= +github.com/substrait-io/substrait-go v0.4.2/go.mod h1:qhpnLmrcvAnlZsUyPXZRqldiHapPTXC3t7xFgDi3aQg= github.com/youmark/pkcs8 v0.0.0-20201027041543-1326539a0a0a h1:fZHgsYlfvtyqToslyjUt3VOPF4J7aK/3MPcK7xp3PDk= github.com/youmark/pkcs8 v0.0.0-20201027041543-1326539a0a0a/go.mod h1:ul22v+Nro/R083muKhosV54bj5niojjWZvU8xrevuH4= github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= @@ -174,6 +199,7 @@ golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=