modified read benchmark test and codebase related to len method

Author: bluepal-prasanthi-moparthi

test/benchmarks_test.go

@@ -87,7 +87,7 @@ func BenchmarkV1BulkInsert100KDocs(b *testing.B) {
 }
 
 func bulkRead(b *testing.B, docSize int) {
-	db, col := setup(b)
+	_, col := setup(b)
 
 	// -----------------------------
 	// Prepare and insert documents
@@ -110,31 +110,24 @@ func bulkRead(b *testing.B, docSize int) {
 	require.NoError(b, err)
 
 	// -----------------------------------------
-	// Sub-benchmark 1: Read entire collection
+	// Sub-benchmark 1: Read entire collection using ReadDocuments
 	// -----------------------------------------
 	b.Run("ReadAllDocsOnce", func(b *testing.B) {
-		query := fmt.Sprintf("FOR d IN %s RETURN d", col.Name())
+		// Prepare keys for reading
+		keys := make([]string, docSize)
+		for j := 0; j < docSize; j++ {
+			keys[j] = fmt.Sprintf("doc_%d", j)
+		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			cursor, err := db.Query(ctx, query, nil)
+			readDocs := make([]TestDoc, docSize)
+			_, _, err := col.ReadDocuments(ctx, keys, readDocs)
 			require.NoError(b, err)
 
-			count := 0
-			for {
-				var doc TestDoc
-				_, err := cursor.ReadDocument(ctx, &doc)
-				if driver.IsNoMoreDocuments(err) {
-					break
-				}
-				require.NoError(b, err)
-				count++
-			}
-			// require.Equal(b, docSize, count, "expected to read all documents")
-			_ = cursor.Close()
 			// sanity check
-			if count != docSize {
-				b.Fatalf("expected to read %d docs, got %d", docSize, count)
+			if len(readDocs) != docSize {
+				b.Fatalf("expected to read %d docs, got %d", docSize, len(readDocs))
 			}
 		}
 	})

v2/arangodb/collection_documents_create_impl.go

@@ -150,13 +150,26 @@ type collectionDocumentCreateResponseReader struct {
 	}
 	shared.ReadAllReader[CollectionDocumentCreateResponse, *collectionDocumentCreateResponseReader]
 
-	// Cache for len() method
+	// Cache for len() method - allows Read() to work after Len() is called
 	cachedResults []CollectionDocumentCreateResponse
 	cachedErrors  []error
 	cached        bool
+	readIndex     int // Track position in cache for Read() after Len()
 }
 
 func (c *collectionDocumentCreateResponseReader) Read() (CollectionDocumentCreateResponse, error) {
+	// If Len() was called, serve from cache
+	if c.cached {
+		if c.readIndex >= len(c.cachedResults) {
+			return CollectionDocumentCreateResponse{}, shared.NoMoreDocumentsError{}
+		}
+		result := c.cachedResults[c.readIndex]
+		err := c.cachedErrors[c.readIndex]
+		c.readIndex++
+		return result, err
+	}
+
+	// Normal streaming read
 	if !c.array.More() {
 		return CollectionDocumentCreateResponse{}, shared.NoMoreDocumentsError{}
 	}
@@ -191,11 +204,13 @@ func (c *collectionDocumentCreateResponseReader) Read() (CollectionDocumentCreat
 	return meta, nil
 }
 
-// Len returns the number of items in the response
+// Len returns the number of items in the response.
+// After calling Len(), you can still use Read() to iterate through items.
 func (c *collectionDocumentCreateResponseReader) Len() int {
 	if !c.cached {
 		c.cachedResults, c.cachedErrors = c.ReadAll()
 		c.cached = true
+		c.readIndex = 0 // Reset read position to allow Read() after Len()
 	}
 	return len(c.cachedResults)
 }

v2/arangodb/collection_documents_delete_impl.go

@@ -115,13 +115,27 @@ type collectionDocumentDeleteResponseReader struct {
 	array   *connection.Array
 	options *CollectionDocumentDeleteOptions
 	shared.ReadAllIntoReader[CollectionDocumentDeleteResponse, *collectionDocumentDeleteResponseReader]
-	// Cache for len() method
+	// Cache for len() method - allows Read() to work after Len() is called
 	cachedResults []CollectionDocumentDeleteResponse
 	cachedErrors  []error
 	cached        bool
+	readIndex     int // Track position in cache for Read() after Len()
 }
 
 func (c *collectionDocumentDeleteResponseReader) Read(i interface{}) (CollectionDocumentDeleteResponse, error) {
+	// If Len() was called, serve from cache
+	// Note: When serving from cache, the 'i' parameter is not populated with document data
+	if c.cached {
+		if c.readIndex >= len(c.cachedResults) {
+			return CollectionDocumentDeleteResponse{}, shared.NoMoreDocumentsError{}
+		}
+		result := c.cachedResults[c.readIndex]
+		err := c.cachedErrors[c.readIndex]
+		c.readIndex++
+		return result, err
+	}
+
+	// Normal streaming read
 	if !c.array.More() {
 		return CollectionDocumentDeleteResponse{}, shared.NoMoreDocumentsError{}
 	}
@@ -172,12 +186,15 @@ func (c *collectionDocumentDeleteResponseReader) Read(i interface{}) (Collection
 	return meta, nil
 }
 
-// Len returns the number of items in the response
+// Len returns the number of items in the response.
+// After calling Len(), you can still use Read() to iterate through items.
+// Note: When Read() serves from cache, the document data parameter is not populated.
 func (c *collectionDocumentDeleteResponseReader) Len() int {
 	if !c.cached {
 		var dummySlice []interface{}
 		c.cachedResults, c.cachedErrors = c.ReadAll(&dummySlice)
 		c.cached = true
+		c.readIndex = 0 // Reset read position to allow Read() after Len()
 	}
 	return len(c.cachedResults)
 }

v2/arangodb/collection_documents_read_impl.go

@@ -107,13 +107,27 @@ type collectionDocumentReadResponseReader struct {
 	array   *connection.Array
 	options *CollectionDocumentReadOptions
 	shared.ReadAllIntoReader[CollectionDocumentReadResponse, *collectionDocumentReadResponseReader]
-	// Cache for len() method
+	// Cache for len() method - allows Read() to work after Len() is called
 	cachedResults []CollectionDocumentReadResponse
 	cachedErrors  []error
 	cached        bool
+	readIndex     int // Track position in cache for Read() after Len()
 }
 
 func (c *collectionDocumentReadResponseReader) Read(i interface{}) (CollectionDocumentReadResponse, error) {
+	// If Len() was called, serve from cache
+	// Note: When serving from cache, the 'i' parameter is not populated with document data
+	if c.cached {
+		if c.readIndex >= len(c.cachedResults) {
+			return CollectionDocumentReadResponse{}, shared.NoMoreDocumentsError{}
+		}
+		result := c.cachedResults[c.readIndex]
+		err := c.cachedErrors[c.readIndex]
+		c.readIndex++
+		return result, err
+	}
+
+	// Normal streaming read
 	if !c.array.More() {
 		return CollectionDocumentReadResponse{}, shared.NoMoreDocumentsError{}
 	}
@@ -148,12 +162,15 @@ func (c *collectionDocumentReadResponseReader) Read(i interface{}) (CollectionDo
 	return meta, nil
 }
 
-// Len returns the number of items in the response
+// Len returns the number of items in the response.
+// After calling Len(), you can still use Read() to iterate through items.
+// Note: When Read() serves from cache, the document data parameter is not populated.
 func (c *collectionDocumentReadResponseReader) Len() int {
 	if !c.cached {
 		var dummySlice []interface{}
 		c.cachedResults, c.cachedErrors = c.ReadAll(&dummySlice)
 		c.cached = true
+		c.readIndex = 0 // Reset read position to allow Read() after Len()
 	}
 	return len(c.cachedResults)
 }

v2/arangodb/collection_documents_replace_impl.go

@@ -150,13 +150,26 @@ type collectionDocumentReplaceResponseReader struct {
 	}
 	shared.ReadAllReader[CollectionDocumentReplaceResponse, *collectionDocumentReplaceResponseReader]
 
-	// Cache for len() method
+	// Cache for len() method - allows Read() to work after Len() is called
 	cachedResults []CollectionDocumentReplaceResponse
 	cachedErrors  []error
 	cached        bool
+	readIndex     int // Track position in cache for Read() after Len()
 }
 
 func (c *collectionDocumentReplaceResponseReader) Read() (CollectionDocumentReplaceResponse, error) {
+	// If Len() was called, serve from cache
+	if c.cached {
+		if c.readIndex >= len(c.cachedResults) {
+			return CollectionDocumentReplaceResponse{}, shared.NoMoreDocumentsError{}
+		}
+		result := c.cachedResults[c.readIndex]
+		err := c.cachedErrors[c.readIndex]
+		c.readIndex++
+		return result, err
+	}
+
+	// Normal streaming read
 	if !c.array.More() {
 		return CollectionDocumentReplaceResponse{}, shared.NoMoreDocumentsError{}
 	}
@@ -198,11 +211,13 @@ func (c *collectionDocumentReplaceResponseReader) Read() (CollectionDocumentRepl
 	return meta, nil
 }
 
-// Len returns the number of items in the response
+// Len returns the number of items in the response.
+// After calling Len(), you can still use Read() to iterate through items.
 func (c *collectionDocumentReplaceResponseReader) Len() int {
 	if !c.cached {
 		c.cachedResults, c.cachedErrors = c.ReadAll()
 		c.cached = true
+		c.readIndex = 0 // Reset read position to allow Read() after Len()
 	}
 	return len(c.cachedResults)
 }

v2/arangodb/collection_documents_update_impl.go

@@ -150,13 +150,26 @@ type collectionDocumentUpdateResponseReader struct {
 	}
 	shared.ReadAllReader[CollectionDocumentUpdateResponse, *collectionDocumentUpdateResponseReader]
 
-	// Cache for len() method
+	// Cache for len() method - allows Read() to work after Len() is called
 	cachedResults []CollectionDocumentUpdateResponse
 	cachedErrors  []error
 	cached        bool
+	readIndex     int // Track position in cache for Read() after Len()
 }
 
 func (c *collectionDocumentUpdateResponseReader) Read() (CollectionDocumentUpdateResponse, error) {
+	// If Len() was called, serve from cache
+	if c.cached {
+		if c.readIndex >= len(c.cachedResults) {
+			return CollectionDocumentUpdateResponse{}, shared.NoMoreDocumentsError{}
+		}
+		result := c.cachedResults[c.readIndex]
+		err := c.cachedErrors[c.readIndex]
+		c.readIndex++
+		return result, err
+	}
+
+	// Normal streaming read
 	if !c.array.More() {
 		return CollectionDocumentUpdateResponse{}, shared.NoMoreDocumentsError{}
 	}
@@ -198,11 +211,13 @@ func (c *collectionDocumentUpdateResponseReader) Read() (CollectionDocumentUpdat
 	return meta, nil
 }
 
-// Len returns the number of items in the response
+// Len returns the number of items in the response.
+// After calling Len(), you can still use Read() to iterate through items.
 func (c *collectionDocumentUpdateResponseReader) Len() int {
 	if !c.cached {
 		c.cachedResults, c.cachedErrors = c.ReadAll()
 		c.cached = true
+		c.readIndex = 0 // Reset read position to allow Read() after Len()
 	}
 	return len(c.cachedResults)
 }

v2/tests/benchmarks_test.go

@@ -209,7 +209,7 @@ func BenchmarkV2BulkInsert100KDocs(b *testing.B) {
 }
 
 func bulkRead(b *testing.B, docSize int) {
-	db, col := setup(b)
+	_, col := setup(b)
 
 	// -----------------------------
 	// Prepare and insert documents
@@ -237,28 +237,30 @@ func bulkRead(b *testing.B, docSize int) {
 	require.NoError(b, err)
 
 	// -----------------------------------------
-	// Sub-benchmark 1: Read entire collection
+	// Sub-benchmark 1: Read entire collection using ReadDocuments
 	// -----------------------------------------
 	b.Run("ReadAllDocsOnce", func(b *testing.B) {
-		query := fmt.Sprintf("FOR d IN %s RETURN d", col.Name())
+		// Prepare keys for reading
+		keys := make([]string, docSize)
+		for j := 0; j < docSize; j++ {
+			keys[j] = fmt.Sprintf("doc_%d", j)
+		}
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			cursor, err := db.Query(ctx, query, nil)
+			resp, err := col.ReadDocuments(ctx, keys)
 			require.NoError(b, err)
 
 			count := 0
 			for {
 				var doc TestDoc
-				_, err := cursor.ReadDocument(ctx, &doc)
+				_, err := resp.Read(&doc)
 				if shared.IsNoMoreDocuments(err) {
 					break
 				}
 				require.NoError(b, err)
 				count++
 			}
-			// require.Equal(b, docSize, count, "expected to read all documents")
-			_ = cursor.Close()
 			// sanity check
 			if count != docSize {
 				b.Fatalf("expected to read %d docs, got %d", docSize, count)