Get sort vector fix - DRAFT

astra/cursors/abstract_cursor.go

@@ -319,14 +319,22 @@ func (c *abstractCursorImpl[Raw]) Next(ctx context.Context) bool {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	if c.buffered() > 0 {
+	// Only consume an item if we're already positioned at one (cursor has started)
+	// On the first call (CursorStateIdle), we don't consume anything
+	if c.state == CursorStateStarted && c.buffered() > 0 {
 		*c.acs.buffer() = (*c.acs.buffer())[1:]
 	}
 
 	if c.buffered() == 0 {
 		return c.fetchIfEmpty(ctx)
 	}
 
+	// If we have items and cursor was idle, transition to started
+	// (this handles the case where GetSortVector pre-fetched the first page)
+	if c.state == CursorStateIdle {
+		c.state = CursorStateStarted
+	}
+
 	return true
 }
 

astra/cursors/find_cursor.go

@@ -83,10 +83,14 @@ func newFindCursorImpl(source findLikeCursorSource[json.RawMessage], fetcher fin
 
 // mapPage implements findLikeCursorSource.mapPage
 func (c *findCursorImpl) mapPage(resp *findResponse, targetCtx serdes.TargetDecodeCtx) *findLikePage[json.RawMessage] {
+	var sortVector *datatypes.Vector
+	if resp.Status != nil {
+		sortVector = resp.Status.SortVector
+	}
 	return &findLikePage[json.RawMessage]{
 		NextPageState: resp.Data.NextPageState,
 		Results:       resp.Data.Documents,
-		SortVector:    resp.Data.SortVector,
+		SortVector:    sortVector,
 		targetCtx:     targetCtx,
 	}
 }

astra/cursors/find_like_cursor.go

@@ -55,6 +55,7 @@ type findLikeCursorImpl[Raw any] struct {
 	fcs         findLikeCursorSource[Raw]
 	currentPage *findLikePage[Raw]
 	initialPage *findLikePage[Raw]
+	sortVector  *datatypes.Vector // Cached sort vector, preserved even when cursor is closed
 	warnings    results.Warnings
 	fetcher     findLikeCursorFetcher
 	target      serdes.Target
@@ -84,13 +85,29 @@ func (c *findLikeCursorImpl[Raw]) GetSortVector(ctx context.Context) *datatypes.
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	// If we already have the sort vector cached, return it
+	if c.sortVector != nil {
+		return c.sortVector
+	}
+
+	// If cursor hasn't started yet and sort vector is requested, fetch the first page
+	// but keep the cursor in idle state so Next() doesn't skip the first item
 	if c.state == CursorStateIdle && c.fcs.includeSortVector() {
-		c.fetchIfEmpty(ctx)
+		c.nextPage, c.err = c.acs.fetchNextPage(ctx)
+		if c.err != nil {
+			c.state = CursorStateClosed
+			c.nextPage = false
+		}
+		// Note: we intentionally do NOT change state to CursorStateStarted here
+		// so that the first Next() call will not consume the first item
 	}
-	if c.currentPage == nil {
-		return nil
+
+	// Cache the sort vector from current page if available
+	if c.currentPage != nil && c.currentPage.SortVector != nil {
+		c.sortVector = c.currentPage.SortVector
 	}
-	return c.currentPage.SortVector
+
+	return c.sortVector
 }
 
 func (c *findLikeCursorImpl[Raw]) Warnings() results.Warnings {
@@ -110,9 +127,9 @@ type findResponse struct {
 	Data struct {
 		Documents     []json.RawMessage `json:"documents"`
 		NextPageState *string           `json:"nextPageState"`
-		SortVector    *datatypes.Vector `json:"sortVector,omitempty"`
 	} `json:"data"`
 	Status *struct {
+		SortVector        *datatypes.Vector `json:"sortVector,omitempty"`
 		DocumentResponses []struct {
 			Scores map[string]float32 `json:"scores"`
 		} `json:"documentResponses"`
@@ -141,6 +158,11 @@ func (c *findLikeCursorImpl[Raw]) fetchNextPage(ctx context.Context) (bool, erro
 
 	c.currentPage = c.fcs.mapPage(&resp, schema)
 
+	// Cache the sort vector from the first page
+	if c.sortVector == nil && c.currentPage.SortVector != nil {
+		c.sortVector = c.currentPage.SortVector
+	}
+
 	return c.currentPage.NextPageState != nil, nil
 }
 
@@ -150,9 +172,11 @@ func (c *findLikeCursorImpl[Raw]) decode(raw Raw, result any) error {
 
 func (c *findLikeCursorImpl[Raw]) rewind() {
 	c.currentPage = c.initialPage
+	c.sortVector = nil
 	c.warnings = nil
 }
 
 func (c *findLikeCursorImpl[Raw]) close() {
 	c.currentPage = nil
+	// Note: sortVector is NOT cleared on close, so it remains accessible
 }

astra/cursors/rerank_cursor.go

@@ -113,10 +113,14 @@ func (c *findAndRerankCursorImpl) mapPage(resp *findResponse, targetCtx serdes.T
 		}
 	}
 
+	var sortVector *datatypes.Vector
+	if resp.Status != nil {
+		sortVector = resp.Status.SortVector
+	}
 	return &findLikePage[rawRerankedResult]{
 		NextPageState: resp.Data.NextPageState,
 		Results:       res,
-		SortVector:    resp.Data.SortVector,
+		SortVector:    sortVector,
 		targetCtx:     targetCtx,
 	}
 }

astra/cursors/rerank_cursor_test.go

@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"testing"
 
+	"github.com/datastax/astra-db-go/v2/astra/datatypes"
 	"github.com/datastax/astra-db-go/v2/astra/options"
 	"github.com/datastax/astra-db-go/v2/astra/serdes"
 )
@@ -30,6 +31,7 @@ func TestFindAndRerankCursorImpl_MapPage(t *testing.T) {
 		json.RawMessage(`{"_id": "2"}`),
 	}
 	resp.Status = &struct {
+		SortVector        *datatypes.Vector `json:"sortVector,omitempty"`
 		DocumentResponses []struct {
 			Scores map[string]float32 `json:"scores"`
 		} `json:"documentResponses"`

go.mod

@@ -2,6 +2,8 @@ module github.com/datastax/astra-db-go/v2
 
 go 1.24.0
 
+replace github.com/datastax/astra-db-go/v2 => ./
+
 // test deps
 require (
 	github.com/DeanPDX/dotconfig v1.0.2