Sa/support query plan

.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

.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 ..

.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

.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 ..

.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 ..

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
+}

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",
+	})
+}

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
+}