From 08b75ab254646dd6ef8cf22f192e4eb4635f4a1f Mon Sep 17 00:00:00 2001 From: Stefano Lottini <236640031+sl-at-ibm@users.noreply.github.com> Date: Mon, 15 Jun 2026 12:18:32 +0200 Subject: [PATCH 1/4] WIP get sort vector fix --- astra/cursors/find_cursor.go | 6 +++++- astra/cursors/find_like_cursor.go | 2 +- astra/cursors/rerank_cursor.go | 6 +++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/astra/cursors/find_cursor.go b/astra/cursors/find_cursor.go index cde48c3..10daf12 100644 --- a/astra/cursors/find_cursor.go +++ b/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, } } diff --git a/astra/cursors/find_like_cursor.go b/astra/cursors/find_like_cursor.go index 65d00c8..cb40556 100644 --- a/astra/cursors/find_like_cursor.go +++ b/astra/cursors/find_like_cursor.go @@ -110,9 +110,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"` diff --git a/astra/cursors/rerank_cursor.go b/astra/cursors/rerank_cursor.go index ec11530..d4c16cf 100644 --- a/astra/cursors/rerank_cursor.go +++ b/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, } } From 90abad4f86a5864e502244b97cbd1597ace83963 Mon Sep 17 00:00:00 2001 From: Stefano Lottini <236640031+sl-at-ibm@users.noreply.github.com> Date: Tue, 16 Jun 2026 10:58:09 +0200 Subject: [PATCH 2/4] GetSortVector --- astra/cursors/abstract_cursor.go | 10 +++++++++- astra/cursors/find_like_cursor.go | 32 +++++++++++++++++++++++++++---- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/astra/cursors/abstract_cursor.go b/astra/cursors/abstract_cursor.go index b27552e..ba6b711 100644 --- a/astra/cursors/abstract_cursor.go +++ b/astra/cursors/abstract_cursor.go @@ -319,7 +319,9 @@ 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:] } @@ -327,6 +329,12 @@ func (c *abstractCursorImpl[Raw]) Next(ctx context.Context) bool { 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 } diff --git a/astra/cursors/find_like_cursor.go b/astra/cursors/find_like_cursor.go index cb40556..701ec26 100644 --- a/astra/cursors/find_like_cursor.go +++ b/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 { @@ -140,6 +157,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 } From 1b862350c4571e7c52f6098bb6af7f7649e0b528 Mon Sep 17 00:00:00 2001 From: Stefano Lottini <236640031+sl-at-ibm@users.noreply.github.com> Date: Tue, 16 Jun 2026 11:23:12 +0200 Subject: [PATCH 3/4] gofmt --- astra/cursors/find_like_cursor.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/astra/cursors/find_like_cursor.go b/astra/cursors/find_like_cursor.go index 701ec26..6015599 100644 --- a/astra/cursors/find_like_cursor.go +++ b/astra/cursors/find_like_cursor.go @@ -101,12 +101,12 @@ func (c *findLikeCursorImpl[Raw]) GetSortVector(ctx context.Context) *datatypes. // Note: we intentionally do NOT change state to CursorStateStarted here // so that the first Next() call will not consume the first item } - + // Cache the sort vector from current page if available if c.currentPage != nil && c.currentPage.SortVector != nil { c.sortVector = c.currentPage.SortVector } - + return c.sortVector } @@ -157,7 +157,7 @@ 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 From 5e051201824225d2e9b850cb6d4bf4f136e92f7d Mon Sep 17 00:00:00 2001 From: Stefano Lottini <236640031+sl-at-ibm@users.noreply.github.com> Date: Tue, 16 Jun 2026 11:30:55 +0200 Subject: [PATCH 4/4] something with tests --- astra/cursors/rerank_cursor_test.go | 2 ++ go.mod | 2 ++ 2 files changed, 4 insertions(+) diff --git a/astra/cursors/rerank_cursor_test.go b/astra/cursors/rerank_cursor_test.go index 67ec03f..0009d8d 100644 --- a/astra/cursors/rerank_cursor_test.go +++ b/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"` diff --git a/go.mod b/go.mod index a6c5841..15eb1ce 100644 --- a/go.mod +++ b/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