diff --git a/.github/scripts/install_tiledb_linux.sh b/.github/scripts/install_tiledb_linux.sh index 550ee061..a67bfb4d 100755 --- a/.github/scripts/install_tiledb_linux.sh +++ b/.github/scripts/install_tiledb_linux.sh @@ -1,4 +1,4 @@ set -e -x -curl --location -o tiledb.tar.gz https://github.com/TileDB-Inc/TileDB/releases/download/2.19.1/tiledb-linux-x86_64-2.19.1-29ceb3e7.tar.gz \ +curl --location -o tiledb.tar.gz https://github.com/TileDB-Inc/TileDB/releases/download/2.20.0-rc2/tiledb-linux-x86_64-2.20.0-rc2-40552aa.tar.gz \ && sudo tar -C /usr/local -xf tiledb.tar.gz sudo ldconfig /usr/local/lib diff --git a/.github/scripts/install_tiledb_linux_debug.sh b/.github/scripts/install_tiledb_linux_debug.sh index 019d6c9e..2a395143 100755 --- a/.github/scripts/install_tiledb_linux_debug.sh +++ b/.github/scripts/install_tiledb_linux_debug.sh @@ -1,5 +1,5 @@ set -e -x -git clone https://github.com/TileDB-Inc/TileDB.git -b 2.19.1 +git clone https://github.com/TileDB-Inc/TileDB.git -b 2.20.0-rc2 cd TileDB mkdir build && cd build cmake -DTILEDB_WERROR=OFF -DTILEDB_VCPKG=ON -DSANITIZER=leak -DTILEDB_VERBOSE=OFF -DTILEDB_S3=ON -DTILEDB_SERIALIZATION=ON -DCMAKE_BUILD_TYPE=Debug -DCMAKE_INSTALL_PREFIX=/usr/local .. diff --git a/.github/scripts/install_tiledb_macos.sh b/.github/scripts/install_tiledb_macos.sh index 5bf7d1c5..2d231c1a 100755 --- a/.github/scripts/install_tiledb_macos.sh +++ b/.github/scripts/install_tiledb_macos.sh @@ -1,3 +1,3 @@ set -e -x -curl --location -o tiledb.tar.gz https://github.com/TileDB-Inc/TileDB/releases/download/2.19.1/tiledb-macos-x86_64-2.19.1-29ceb3e7.tar.gz \ +curl --location -o tiledb.tar.gz https://github.com/TileDB-Inc/TileDB/releases/download/2.20.0-rc2/tiledb-macos-x86_64-2.20.0-rc2-40552aa.tar.gz \ && sudo tar -C /usr/local -xf tiledb.tar.gz diff --git a/.github/scripts/install_tiledb_source_linux.sh b/.github/scripts/install_tiledb_source_linux.sh index 16c910eb..0a81d929 100755 --- a/.github/scripts/install_tiledb_source_linux.sh +++ b/.github/scripts/install_tiledb_source_linux.sh @@ -1,5 +1,5 @@ set -e -x -git clone https://github.com/TileDB-Inc/TileDB.git -b 2.19.1 +git clone https://github.com/TileDB-Inc/TileDB.git -b 2.20.0-rc2 cd TileDB mkdir build && cd build cmake -DTILEDB_WERROR=OFF -DTILEDB_VCPKG=ON -DTILEDB_VERBOSE=OFF -DTILEDB_S3=ON -DTILEDB_SERIALIZATION=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/usr/local .. diff --git a/.github/scripts/install_tiledb_source_macos.sh b/.github/scripts/install_tiledb_source_macos.sh index 072d401e..3b7886fc 100755 --- a/.github/scripts/install_tiledb_source_macos.sh +++ b/.github/scripts/install_tiledb_source_macos.sh @@ -1,5 +1,5 @@ set -e -x -git clone https://github.com/TileDB-Inc/TileDB.git -b 2.19.1 +git clone https://github.com/TileDB-Inc/TileDB.git -b 2.20.0-rc2 cd TileDB mkdir build && cd build cmake -DTILEDB_WERROR=OFF -DTILEDB_VCPKG=ON -DTILEDB_VERBOSE=OFF -DTILEDB_S3=ON -DTILEDB_SERIALIZATION=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/usr/local .. diff --git a/query_experimental.go b/query_experimental.go index 0e411080..738df7c9 100644 --- a/query_experimental.go +++ b/query_experimental.go @@ -41,3 +41,43 @@ func (q *Query) StatusDetails() (QueryStatusDetails, error) { details.IncompleteReason = QueryStatusDetailsReason(cDetails.incomplete_reason) return details, nil } + +// GetPlan returns a json encoding of the query plan for the query. +// Example: +// +// { +// "TileDB Query Plan": { +// "Array.Type": "sparse", +// "Array.URI": "file:///tmp/TestHandleQueryPlanRequest732268097/001/t-testhandlequeryplanrequest-b757271e", +// "Query.Attributes": [ +// "a1", +// "a2", +// "a3", +// "a4", +// "a5" +// ], +// "Query.Dimensions": [ +// "dim1" +// ], +// "Query.Layout": "unordered", +// "Query.Strategy.Name": "UnorderedWriter", +// "VFS.Backend": "file" +// } +// } +func (q *Query) GetPlan() (string, error) { + var plan *C.tiledb_string_t + + ret := C.tiledb_query_get_plan(q.context.tiledbContext, q.tiledbQuery, &plan) + if ret != C.TILEDB_OK { + return "", fmt.Errorf("Error getting query plan: %s", q.context.LastError()) + } + + var sPlan *C.char + var sPlanSize C.size_t + ret = C.tiledb_string_view(plan, &sPlan, &sPlanSize) + if ret != C.TILEDB_OK { + return "", fmt.Errorf("Error extracting query query: %s", q.context.LastError()) + } + + return C.GoStringN(sPlan, C.int(sPlanSize)), nil +} diff --git a/query_experimental_test.go b/query_experimental_test.go index d0907222..ca5a8c75 100644 --- a/query_experimental_test.go +++ b/query_experimental_test.go @@ -3,6 +3,10 @@ package tiledb import ( + "bytes" + "encoding/json" + "html/template" + "reflect" "testing" "github.com/stretchr/testify/assert" @@ -150,3 +154,149 @@ func TestQueryStatusDetails(t *testing.T) { err = array.Close() require.NoError(t, err) } + +// templateQueryPlan is the query plan expected for the arrays of the test. +// It is parameterized with the array URI which is different every time. +// The other fields should not change except there are changes in the array schema +// or the core query plan implementation. +const templateQueryPlan = `{ + "TileDB Query Plan": { + "Array.Type": "dense", + "Array.URI": "{{.uri}}", + "Query.Attributes": [ + "v" + ], + "Query.Dimensions": [ + "x" + ], + "Query.Layout": "{{.layout}}", + "Query.Strategy.Name": "{{.strategy}}", + "VFS.Backend": "file" + } +}` + +// requirePlanAsExpected checks if a query plan conforms to the query plan template +func requirePlanAsExpected(t *testing.T, arrayPath, actualPlan string, diffs map[string]interface{}) { + var expectedPlan bytes.Buffer + require.NoError(t, template.Must(template.New("plan").Parse(templateQueryPlan)).Execute(&expectedPlan, diffs)) + + m1 := map[string]any{} + require.NoError(t, json.Unmarshal(expectedPlan.Bytes(), &m1)) + m2 := map[string]any{} + require.NoError(t, json.Unmarshal([]byte(actualPlan), &m2)) + + require.True(t, reflect.DeepEqual(m1, m2)) +} + +func TestQueryPlan(t *testing.T) { + // Create an 1d array + + // Create configuration + config, err := NewConfig() + require.NoError(t, err) + + // Create context with config + context, err := NewContext(config) + require.NoError(t, err) + + // Create dimension + dimension, err := NewDimension(context, "x", TILEDB_INT8, []int8{0, 9}, int8(5)) + require.NoError(t, err) + assert.NotNil(t, dimension) + + // Create domain + domain, err := NewDomain(context) + require.NoError(t, err) + assert.NotNil(t, domain) + + // Add dimension + require.NoError(t, domain.AddDimensions(dimension)) + + // Create array schema + arraySchema, err := NewArraySchema(context, TILEDB_DENSE) + require.NoError(t, err) + assert.NotNil(t, arraySchema) + + // Create attribute to add to schema + attribute, err := NewAttribute(context, "v", TILEDB_INT32) + require.NoError(t, err) + assert.NotNil(t, attribute) + + // Add attribute to schema + err = arraySchema.AddAttributes(attribute) + require.NoError(t, err) + + // Set Domain + err = arraySchema.SetDomain(domain) + require.NoError(t, err) + + // Validate Schema + err = arraySchema.Check() + require.NoError(t, err) + + // Create array on disk + tmpArrayPath := t.TempDir() + array, err := NewArray(context, tmpArrayPath) + require.NoError(t, err) + assert.NotNil(t, array) + err = array.Create(arraySchema) + require.NoError(t, err) + + // Write to array + + // Open array for writing + err = array.Open(TILEDB_WRITE) + require.NoError(t, err) + + // Create write query + query, err := NewQuery(context, array) + require.NoError(t, err) + assert.NotNil(t, query) + err = query.SetSubArray([]int8{0, 9}) + require.NoError(t, err) + + // Initialize the data buffer + bufferV := []int32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9} + _, err = query.SetDataBuffer("v", bufferV) + require.NoError(t, err) + + // Submit write query + err = query.Submit() + require.NoError(t, err) + + // Validate status, since query was used this is should be complete + status, err := query.Status() + require.NoError(t, err) + assert.Equal(t, TILEDB_COMPLETED, status) + + // close array + err = array.Close() + require.NoError(t, err) + + // Read from the array. We will test an incomplete query + // Open array for reading + err = array.Open(TILEDB_READ) + require.NoError(t, err) + + // Create read query + query, err = NewQuery(context, array) + require.NoError(t, err) + assert.NotNil(t, query) + err = query.SetSubArray([]int8{0, 9}) // we want to read the whole array, 2 full tiles + require.NoError(t, err) + + // Initialize the data buffer + bufferV = make([]int32, 10) + _, err = query.SetDataBuffer("v", bufferV) + require.NoError(t, err) + + // Get query plan + actualPlan, err := query.GetPlan() + require.NoError(t, err) + + requirePlanAsExpected(t, tmpArrayPath, actualPlan, map[string]interface{}{ + "uri": "file://" + tmpArrayPath, + "layout": "row-major", + "strategy": "DenseReader", + }) +} diff --git a/serialize.go b/serialize.go index 1fa6bd50..b7d36f1b 100644 --- a/serialize.go +++ b/serialize.go @@ -534,3 +534,25 @@ func HandleArrayDeleteFragmentsListRequest(context *Context, array *Array, buffe runtime.KeepAlive(buffer) return nil } + +// HandleQueryPlanRequest handles a request for a query plan. This is used by TileDB-Cloud +// It returns a buffer with the serialized response. The caller should free the buffer after use. +func HandleQueryPlanRequest(array *Array, serializationType SerializationType, request *Buffer) (*Buffer, error) { + opContext := array.context + + response, err := NewBuffer(opContext) + if err != nil { + return nil, fmt.Errorf("Error allocating tiledb buffer: %s", opContext.LastError()) + } + + ret := C.tiledb_handle_query_plan_request(opContext.tiledbContext, array.tiledbArray, C.tiledb_serialization_type_t(serializationType), + request.tiledbBuffer, response.tiledbBuffer) + if ret != C.TILEDB_OK { + return nil, fmt.Errorf("Error handling query plan request: %s", opContext.LastError()) + } + + runtime.KeepAlive(request) + runtime.KeepAlive(array) + + return response, nil +}